Skip to content

Commit 0bf4e57

Browse files
snajpae2b-bot[bot]
authored andcommitted
fix(storage): normalize upload concurrency limits
GitOrigin-RevId: 9be9346295f9b3311ae66d944804191b0841f9b4
1 parent b999fd6 commit 0bf4e57

5 files changed

Lines changed: 205 additions & 2 deletions

File tree

‎packages/shared/pkg/storage/compress_upload_test.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,22 @@ func TestCompressStreamContextCancel(t *testing.T) {
210210
require.ErrorIs(t, err, context.Canceled)
211211
}
212212

213+
func TestCompressStreamNonPositiveUploadConcurrency(t *testing.T) {
214+
t.Parallel()
215+
216+
for _, maxConcurrency := range []int{-1, 0} {
217+
t.Run(fmt.Sprintf("max_%d", maxConcurrency), func(t *testing.T) {
218+
t.Parallel()
219+
220+
up := new(memPartUploader)
221+
cfg := defaultCfg(CompressionZstd, 1, 128*1024)
222+
_, _, err := compressStream(t.Context(), bytes.NewReader([]byte("payload")), cfg, up, maxConcurrency, nil)
223+
require.NoError(t, err)
224+
require.NotEmpty(t, up.parts)
225+
})
226+
}
227+
}
228+
213229
func TestClampCloudMinPartSize(t *testing.T) {
214230
t.Parallel()
215231

‎packages/shared/pkg/storage/gcp_multipart.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -583,7 +583,9 @@ func (m *MultipartUploader) UploadFileInParallel(ctx context.Context, filePath s
583583

584584
func (m *MultipartUploader) uploadParts(ctx context.Context, maxConcurrency int, numParts int, fileSize int64, file *os.File, uploadID string) ([]Part, error) {
585585
g, ctx := errgroup.WithContext(ctx) // Context ONLY for waitgroup goroutines; canceled after errgroup finishes
586-
g.SetLimit(maxConcurrency) // Limit concurrent goroutines
586+
// SetLimit(0) blocks every Go call, while a negative limit disables the
587+
// bound. Always retain at least one admitted uploader.
588+
g.SetLimit(max(maxConcurrency, 1))
587589

588590
// Thread-safe map to collect parts
589591
var partsMu sync.Mutex

‎packages/shared/pkg/storage/gcp_multipart_test.go‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import (
1919
"sync"
2020
"sync/atomic"
2121
"testing"
22+
"testing/synctest"
2223
"time"
2324

2425
"github.com/aws/aws-sdk-go-v2/aws"
@@ -36,6 +37,12 @@ const (
3637
uploadsPath = "uploads"
3738
)
3839

40+
type multipartRoundTripFunc func(*http.Request) (*http.Response, error)
41+
42+
func (f multipartRoundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
43+
return f(req)
44+
}
45+
3946
// createTestMultipartUploader creates a test uploader with a mock HTTP client
4047
func createTestMultipartUploader(t *testing.T, handler http.HandlerFunc, retryConfig ...RetryConfig) *MultipartUploader {
4148
t.Helper()
@@ -145,6 +152,103 @@ func TestGCPUploadFileInParallelAbortsOnPartFailure(t *testing.T) {
145152
require.True(t, aborted.Load(), "failed parallel upload should abort the multipart upload")
146153
}
147154

155+
func TestMultipartUploaderUploadPartsZeroConcurrencyDoesNotDeadlock(t *testing.T) {
156+
t.Parallel()
157+
158+
inputPath := writeTempFile(t, []byte("x"))
159+
input, err := os.Open(inputPath)
160+
require.NoError(t, err)
161+
t.Cleanup(func() { require.NoError(t, input.Close()) })
162+
163+
uploader := createTestMultipartUploader(t, func(w http.ResponseWriter, r *http.Request) {
164+
if r.Method != http.MethodPut {
165+
t.Errorf("unexpected request method: %s", r.Method)
166+
167+
return
168+
}
169+
if _, err := io.Copy(io.Discard, r.Body); err != nil {
170+
t.Errorf("read request body: %v", err)
171+
172+
return
173+
}
174+
w.Header().Set("ETag", `"etag"`)
175+
w.WriteHeader(http.StatusOK)
176+
})
177+
178+
done := make(chan error, 1)
179+
go func() {
180+
_, err := uploader.uploadParts(t.Context(), 0, 1, 1, input, "upload-id")
181+
done <- err
182+
}()
183+
184+
select {
185+
case err := <-done:
186+
require.NoError(t, err)
187+
case <-time.After(time.Second):
188+
t.Fatal("zero upload concurrency blocked part admission")
189+
}
190+
}
191+
192+
func TestMultipartUploaderUploadPartsNegativeConcurrencyUsesOneWorker(t *testing.T) {
193+
t.Parallel()
194+
195+
inputPath := filepath.Join(t.TempDir(), "multipart.bin")
196+
input, err := os.OpenFile(inputPath, os.O_CREATE|os.O_RDWR, 0o600)
197+
require.NoError(t, err)
198+
t.Cleanup(func() { require.NoError(t, input.Close()) })
199+
require.NoError(t, input.Truncate(gcpMultipartUploadChunkSize+1))
200+
201+
synctest.Test(t, func(t *testing.T) {
202+
var active atomic.Int32
203+
var maxActive atomic.Int32
204+
release := make(chan struct{})
205+
206+
client := retryablehttp.NewClient()
207+
client.RetryMax = 0
208+
client.Logger = nil
209+
client.HTTPClient.Transport = multipartRoundTripFunc(func(req *http.Request) (*http.Response, error) {
210+
current := active.Add(1)
211+
defer active.Add(-1)
212+
for {
213+
previous := maxActive.Load()
214+
if current <= previous || maxActive.CompareAndSwap(previous, current) {
215+
break
216+
}
217+
}
218+
219+
<-release
220+
221+
response := new(http.Response)
222+
response.StatusCode = http.StatusOK
223+
response.Header = make(http.Header)
224+
response.Header.Set("ETag", `"etag"`)
225+
response.Body = http.NoBody
226+
response.Request = req
227+
228+
return response, nil
229+
})
230+
231+
uploader := new(MultipartUploader)
232+
uploader.objectName = "object"
233+
uploader.token = "token"
234+
uploader.client = client
235+
uploader.baseURL = "http://storage.test"
236+
237+
done := make(chan error, 1)
238+
go func() {
239+
_, err := uploader.uploadParts(t.Context(), -1, 2, gcpMultipartUploadChunkSize+1, input, "upload-id")
240+
done <- err
241+
}()
242+
243+
synctest.Wait()
244+
gotMax := maxActive.Load()
245+
close(release)
246+
require.NoError(t, <-done)
247+
require.Equal(t, int32(1), gotMax, "negative concurrency admitted multiple GCS parts")
248+
require.Zero(t, active.Load())
249+
})
250+
}
251+
148252
func TestMultipartUploader_InitiateUpload_Success(t *testing.T) {
149253
t.Parallel()
150254
expectedUploadID := "test-upload-id-123"

‎packages/shared/pkg/storage/storage_aws.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,10 @@ const (
3131
awsMultipartUploadPartSize = 10 * 1024 * 1024
3232
)
3333

34+
// normalizeAWSUploadConcurrency preserves the SDK's zero-as-default contract
35+
// while preventing negative values from reaching multipart channel capacities.
36+
func normalizeAWSUploadConcurrency(tasks int) int { return max(tasks, 0) }
37+
3438
type awsStorage struct {
3539
client *s3.Client
3640
presignClient *s3.PresignClient
@@ -297,7 +301,9 @@ func (o *awsObject) StoreFile(ctx context.Context, path string, opts ...PutOptio
297301
o.client,
298302
func(u *manager.Uploader) {
299303
u.PartSize = awsMultipartUploadPartSize
300-
u.Concurrency = o.limiter.MaxUploadTasks(ctx)
304+
// The SDK maps zero to its default; negative values survive into
305+
// multipart channel capacities and panic before useful I/O.
306+
u.Concurrency = normalizeAWSUploadConcurrency(o.limiter.MaxUploadTasks(ctx))
301307
},
302308
)
303309

‎packages/shared/pkg/storage/storage_aws_sse_test.go‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,25 @@
11
package storage
22

33
import (
4+
"context"
5+
"io"
46
"net/http"
57
"net/url"
8+
"os"
9+
"path/filepath"
610
"strings"
711
"sync"
812
"testing"
913
"time"
1014

1115
"github.com/aws/aws-sdk-go-v2/service/s3"
16+
"github.com/launchdarkly/go-sdk-common/v3/ldvalue"
17+
"github.com/launchdarkly/go-server-sdk/v7/testhelpers/ldtestdata"
1218
"github.com/stretchr/testify/assert"
1319
"github.com/stretchr/testify/require"
20+
21+
"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
22+
"github.com/e2b-dev/infra/packages/shared/pkg/limit"
1423
)
1524

1625
const (
@@ -82,6 +91,11 @@ func newSSETestStorage(t *testing.T, spec Spec, rec *sseRecorder) *awsStorage {
8291
w.WriteHeader(http.StatusOK)
8392
w.Write([]byte(`<InitiateMultipartUploadResult><UploadId>sse-upload-id</UploadId></InitiateMultipartUploadResult>`))
8493
case r.Method == http.MethodPut && strings.Contains(r.URL.RawQuery, "partNumber="):
94+
if _, err := io.Copy(io.Discard, r.Body); err != nil {
95+
t.Errorf("read multipart body: %v", err)
96+
97+
return
98+
}
8599
w.Header().Set("ETag", `"etag1"`)
86100
w.WriteHeader(http.StatusOK)
87101
case r.Method == http.MethodPost && strings.Contains(r.URL.RawQuery, "uploadId=sse-upload-id"):
@@ -106,6 +120,67 @@ func newSSETestStorage(t *testing.T, spec Spec, rec *sseRecorder) *awsStorage {
106120
}
107121
}
108122

123+
func newTestUploadLimiter(t *testing.T, maxUploadTasks int) *limit.Limiter {
124+
t.Helper()
125+
126+
source := ldtestdata.DataSource()
127+
source.Update(source.Flag(featureflags.StorageMaxUploadTasks.Key()).ValueForAll(ldvalue.Int(maxUploadTasks)))
128+
ff, err := featureflags.NewClientWithDatasource(source)
129+
require.NoError(t, err)
130+
t.Cleanup(func() { require.NoError(t, ff.Close(context.WithoutCancel(t.Context()))) })
131+
132+
limiter, err := limit.New(t.Context(), ff)
133+
require.NoError(t, err)
134+
t.Cleanup(func() { require.NoError(t, limiter.Close(context.WithoutCancel(t.Context()))) })
135+
136+
return limiter
137+
}
138+
139+
func TestNormalizeAWSUploadConcurrency(t *testing.T) {
140+
t.Parallel()
141+
142+
tests := []struct {
143+
name string
144+
tasks int
145+
want int
146+
}{
147+
{name: "negative uses SDK default", tasks: -1, want: 0},
148+
{name: "zero keeps SDK default", tasks: 0, want: 0},
149+
{name: "positive passes through", tasks: 16, want: 16},
150+
}
151+
152+
for _, tt := range tests {
153+
t.Run(tt.name, func(t *testing.T) {
154+
t.Parallel()
155+
require.Equal(t, tt.want, normalizeAWSUploadConcurrency(tt.tasks))
156+
})
157+
}
158+
}
159+
160+
func TestAWSUncompressedStoreFileNegativeConcurrencyUsesSDKDefault(t *testing.T) {
161+
t.Parallel()
162+
163+
var spec Spec
164+
provider := newSSETestStorage(t, spec, new(sseRecorder))
165+
provider.limiter = newTestUploadLimiter(t, -1)
166+
167+
obj, err := provider.OpenSeekable(t.Context(), testObjectName)
168+
require.NoError(t, err)
169+
170+
inputPath := filepath.Join(t.TempDir(), "multipart.bin")
171+
input, err := os.Create(inputPath)
172+
require.NoError(t, err)
173+
t.Cleanup(func() { _ = input.Close() })
174+
require.NoError(t, input.Truncate(awsMultipartUploadPartSize+1))
175+
require.NoError(t, input.Close())
176+
177+
var storeErr error
178+
require.NotPanics(t, func() {
179+
_, _, storeErr = obj.StoreFile(t.Context(), inputPath)
180+
})
181+
require.NoError(t, storeErr)
182+
}
183+
109184
func TestAWSPutRequestsServerSideEncryption(t *testing.T) {
110185
t.Parallel()
111186

0 commit comments

Comments
 (0)