From 3f7e7725fdc93f2c453700bc0f80666ac9cdf642 Mon Sep 17 00:00:00 2001 From: Jakub Dobry Date: Wed, 8 Jul 2026 14:47:04 -0700 Subject: [PATCH] refactor(storage): shared compressed-upload framework + storage limiter Generalize the per-provider compressed-upload path into a shared package-level storeFileCompressed driven by a partUploader interface (GCS XML multipart implements it), replacing the duplicated per-object methods. Add clampCloudMinPartSize for the shared S3/GCS multipart limits, io_wrappers multi-slice part bodies, and the part-uploader contract test harness. Also generalize the GCS-only upload limiter (limit/gcloud.go) into a provider-agnostic storage upload limiter and thread it through the AWS and GCS providers. AWS still rejects compressed uploads here; the AWS compressed provider is added in a follow-up. Pure refactor of existing GCS behavior; no functional change. --- packages/shared/pkg/featureflags/flags.go | 7 +- packages/shared/pkg/limit/gcloud.go | 18 ---- packages/shared/pkg/limit/limiter.go | 12 +-- packages/shared/pkg/limit/storage.go | 33 +++++++ packages/shared/pkg/limit/storage_test.go | 41 ++++++++ packages/shared/pkg/limit/upload.go | 4 +- .../shared/pkg/storage/compress_upload.go | 61 ++++++++++++ .../pkg/storage/compress_upload_test.go | 16 +++ packages/shared/pkg/storage/gcp_multipart.go | 40 +------- .../shared/pkg/storage/gcp_multipart_test.go | 31 ++++++ packages/shared/pkg/storage/io_wrappers.go | 46 +++++++++ .../shared/pkg/storage/part_uploader_test.go | 98 +++++++++++++++++++ packages/shared/pkg/storage/paths.go | 2 +- packages/shared/pkg/storage/storage.go | 19 ++-- packages/shared/pkg/storage/storage_aws.go | 16 ++- packages/shared/pkg/storage/storage_google.go | 84 ++++------------ .../pkg/storage/storageopts/storageopts.go | 23 +++++ 17 files changed, 406 insertions(+), 145 deletions(-) delete mode 100644 packages/shared/pkg/limit/gcloud.go create mode 100644 packages/shared/pkg/limit/storage.go create mode 100644 packages/shared/pkg/limit/storage_test.go create mode 100644 packages/shared/pkg/storage/part_uploader_test.go diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index 7bb5d40b01..2373a115be 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -268,9 +268,10 @@ func NewIntFlag(name string, fallback int) IntFlag { } var ( - MaxSandboxesPerNode = NewIntFlag("max-sandboxes-per-node", 200) - GcloudConcurrentUploadLimit = NewIntFlag("gcloud-concurrent-upload-limit", 8) - GcloudMaxTasks = NewIntFlag("gcloud-max-tasks", 16) + MaxSandboxesPerNode = NewIntFlag("max-sandboxes-per-node", 200) + // The LD keys keep the legacy "gcloud-" prefix, but the limits apply to uploads on all storage providers. + StorageConcurrentUploadLimit = NewIntFlag("gcloud-concurrent-upload-limit", 8) + StorageMaxUploadTasks = NewIntFlag("gcloud-max-tasks", 16) ClickhouseBatcherMaxBatchSize = NewIntFlag("clickhouse-batcher-max-batch-size", 100) ClickhouseBatcherMaxDelay = NewIntFlag("clickhouse-batcher-max-delay", 1000) // 1s in milliseconds ClickhouseBatcherQueueSize = NewIntFlag("clickhouse-batcher-queue-size", 1000) diff --git a/packages/shared/pkg/limit/gcloud.go b/packages/shared/pkg/limit/gcloud.go deleted file mode 100644 index db28d0d4d3..0000000000 --- a/packages/shared/pkg/limit/gcloud.go +++ /dev/null @@ -1,18 +0,0 @@ -package limit - -import ( - "context" - - "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" - "github.com/e2b-dev/infra/packages/shared/pkg/utils" -) - -func (l *Limiter) GCloudUploadLimiter() *utils.AdjustableSemaphore { - return l.gCloudUploadLimiter -} - -func (l *Limiter) GCloudMaxTasks(ctx context.Context) int { - maxTasks := l.featureFlags.IntFlag(ctx, featureflags.GcloudMaxTasks) - - return maxTasks -} diff --git a/packages/shared/pkg/limit/limiter.go b/packages/shared/pkg/limit/limiter.go index 12ce8d448b..75d8f50c2f 100644 --- a/packages/shared/pkg/limit/limiter.go +++ b/packages/shared/pkg/limit/limiter.go @@ -9,23 +9,23 @@ import ( ) type Limiter struct { - gCloudUploadLimiter *utils.AdjustableSemaphore - featureFlags *featureflags.Client + storageUploadLimiter *utils.AdjustableSemaphore + featureFlags *featureflags.Client done chan struct{} closeOnce sync.Once } func New(ctx context.Context, featureFlags *featureflags.Client) (*Limiter, error) { - uploadLimiter, err := utils.NewAdjustableSemaphore(int64(featureflags.GcloudConcurrentUploadLimit.Fallback())) + uploadLimiter, err := utils.NewAdjustableSemaphore(int64(featureflags.StorageConcurrentUploadLimit.Fallback())) if err != nil { return nil, err } l := &Limiter{ - gCloudUploadLimiter: uploadLimiter, - featureFlags: featureFlags, - done: make(chan struct{}), + storageUploadLimiter: uploadLimiter, + featureFlags: featureFlags, + done: make(chan struct{}), } go l.UpdateUploadLimitSemaphore(ctx) diff --git a/packages/shared/pkg/limit/storage.go b/packages/shared/pkg/limit/storage.go new file mode 100644 index 0000000000..2822344731 --- /dev/null +++ b/packages/shared/pkg/limit/storage.go @@ -0,0 +1,33 @@ +package limit + +import ( + "context" + "fmt" + + "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" +) + +// AcquireUploadSlot reserves a slot in the shared storage-upload semaphore and +// returns a func releasing it. Safe on a nil receiver: without a limiter the +// upload is not throttled and release is a no-op. +func (l *Limiter) AcquireUploadSlot(ctx context.Context) (release func(), err error) { + if l == nil { + return func() {}, nil + } + + if err := l.storageUploadLimiter.Acquire(ctx, 1); err != nil { + return nil, fmt.Errorf("failed to acquire semaphore: %w", err) + } + + return func() { l.storageUploadLimiter.Release(1) }, nil +} + +// MaxUploadTasks returns the per-upload concurrency limit. Safe on a nil +// receiver: without a limiter it returns the flag's fallback value. +func (l *Limiter) MaxUploadTasks(ctx context.Context) int { + if l == nil { + return featureflags.StorageMaxUploadTasks.Fallback() + } + + return l.featureFlags.IntFlag(ctx, featureflags.StorageMaxUploadTasks) +} diff --git a/packages/shared/pkg/limit/storage_test.go b/packages/shared/pkg/limit/storage_test.go new file mode 100644 index 0000000000..ebdfdbd5b5 --- /dev/null +++ b/packages/shared/pkg/limit/storage_test.go @@ -0,0 +1,41 @@ +package limit + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/shared/pkg/utils" +) + +func TestAcquireUploadSlotNilLimiter(t *testing.T) { + t.Parallel() + + var l *Limiter + release, err := l.AcquireUploadSlot(t.Context()) + require.NoError(t, err) + release() +} + +func TestAcquireUploadSlotBlocksWhenFull(t *testing.T) { + t.Parallel() + + sem, err := utils.NewAdjustableSemaphore(1) + require.NoError(t, err) + l := &Limiter{storageUploadLimiter: sem} + + release, err := l.AcquireUploadSlot(t.Context()) + require.NoError(t, err) + + ctx, cancel := context.WithTimeout(t.Context(), 20*time.Millisecond) + defer cancel() + _, err = l.AcquireUploadSlot(ctx) + require.ErrorIs(t, err, context.DeadlineExceeded) + + release() + release, err = l.AcquireUploadSlot(t.Context()) + require.NoError(t, err) + release() +} diff --git a/packages/shared/pkg/limit/upload.go b/packages/shared/pkg/limit/upload.go index afbd289f53..a0a09f2f0e 100644 --- a/packages/shared/pkg/limit/upload.go +++ b/packages/shared/pkg/limit/upload.go @@ -17,10 +17,10 @@ func (l *Limiter) UpdateUploadLimitSemaphore(ctx context.Context) { for { select { case <-ticker.C: - uploadLimitFlag := l.featureFlags.IntFlag(ctx, featureflags.GcloudConcurrentUploadLimit) + uploadLimitFlag := l.featureFlags.IntFlag(ctx, featureflags.StorageConcurrentUploadLimit) // Update the semaphore with the new value - if err := l.gCloudUploadLimiter.SetLimit(int64(uploadLimitFlag)); err != nil { + if err := l.storageUploadLimiter.SetLimit(int64(uploadLimitFlag)); err != nil { logger.L().Error(ctx, "failed to adjust upload semaphore", zap.Error(err)) } case <-l.done: diff --git a/packages/shared/pkg/storage/compress_upload.go b/packages/shared/pkg/storage/compress_upload.go index f296945d4d..2e3e5452af 100644 --- a/packages/shared/pkg/storage/compress_upload.go +++ b/packages/shared/pkg/storage/compress_upload.go @@ -7,11 +7,14 @@ import ( "errors" "fmt" "io" + "os" "slices" "sync" "sync/atomic" "golang.org/x/sync/errgroup" + + "github.com/e2b-dev/infra/packages/shared/pkg/units" ) type partUploader interface { @@ -21,6 +24,55 @@ type partUploader interface { Close() error } +const ( + // cloudMinPartSizeMB is the smallest non-final multipart part both S3 and + // the GCS XML API accept (5 MiB). Smaller configured values would fail with + // EntityTooSmall at CompleteMultipartUpload, after all bytes are shipped. + cloudMinPartSizeMB = 5 + // cloudMaxParts bounds the part count; S3 and the GCS XML API cap multipart + // uploads at 10,000 parts. Kept below the hard cap because parts are sized + // by *compressed* bytes and incompressible frames can expand slightly past + // the uncompressed file size this bound is computed from. + cloudMaxParts = 9000 +) + +// clampCloudMinPartSize resolves the configured minimum part size against the +// multipart limits shared by S3 and the GCS XML API: at least 5 MiB per +// non-final part, and large enough that even an incompressible file of +// fileSize bytes stays under the 10,000-part cap. +func clampCloudMinPartSize(cfg CompressConfig, fileSize int64) CompressConfig { + minMB := units.BytesToMB(cfg.MinPartSize()) // resolves the <= 0 default (50 MB) + cfg.MinPartSizeMB = int(max(minMB, cloudMinPartSizeMB, units.BytesToMB(fileSize)/cloudMaxParts+1)) + + return cfg +} + +// storeFileCompressed streams localPath through compressStream into a +// provider-specific multipart upload. It owns the shared recipe — open, stat, +// stamp the uncompressed size into the object metadata (multipart APIs only +// accept metadata at initiate time) — so providers supply just the uploader. +func storeFileCompressed(ctx context.Context, localPath string, cfg CompressConfig, maxUploadConcurrency int, putOpts PutOptions, newUploader func(metadata ObjectMetadata) (partUploader, error)) (*FullFrameTable, [32]byte, error) { + file, err := os.Open(localPath) + if err != nil { + return nil, [32]byte{}, fmt.Errorf("failed to open local file %s: %w", localPath, err) + } + defer file.Close() + + fi, err := file.Stat() + if err != nil { + return nil, [32]byte{}, fmt.Errorf("failed to stat local file %s: %w", localPath, err) + } + + cfg = clampCloudMinPartSize(cfg, fi.Size()) + + uploader, err := newUploader(putOpts.Metadata.WithUncompressedSize(fi.Size())) + if err != nil { + return nil, [32]byte{}, fmt.Errorf("failed to create multipart uploader: %w", err) + } + + return compressStream(ctx, file, cfg, uploader, maxUploadConcurrency, putOpts.FrameSink) +} + type memPartUploader struct { mu sync.Mutex parts map[int][]byte @@ -175,6 +227,15 @@ func compressStream(ctx context.Context, in io.Reader, cfg CompressConfig, uploa return nil, [32]byte{}, err } + // Zero-byte input produces no parts, but S3 and the GCS XML API both + // refuse to complete a multipart upload with zero parts — ship a single + // empty final part so empty files still store successfully. + if len(frameSizes) == 0 { + if err := uploader.UploadPart(ctx, 1); err != nil { + return nil, [32]byte{}, fmt.Errorf("upload empty part: %w", err) + } + } + if err := uploader.Complete(ctx); err != nil { return nil, [32]byte{}, fmt.Errorf("complete upload: %w", err) } diff --git a/packages/shared/pkg/storage/compress_upload_test.go b/packages/shared/pkg/storage/compress_upload_test.go index 8efe3f2568..9f71fdbac5 100644 --- a/packages/shared/pkg/storage/compress_upload_test.go +++ b/packages/shared/pkg/storage/compress_upload_test.go @@ -210,6 +210,22 @@ func TestCompressStreamContextCancel(t *testing.T) { require.ErrorIs(t, err, context.Canceled) } +func TestClampCloudMinPartSize(t *testing.T) { + t.Parallel() + + // Below S3/GCS XML API's 5 MiB non-final part minimum → clamped. + require.Equal(t, 5, clampCloudMinPartSize(CompressConfig{MinPartSizeMB: 1}, 0).MinPartSizeMB) + // Unset resolves to the 50 MB default from MinPartSize(). + require.Equal(t, 50, clampCloudMinPartSize(CompressConfig{}, 0).MinPartSizeMB) + require.Equal(t, 50, clampCloudMinPartSize(CompressConfig{MinPartSizeMB: 50}, 0).MinPartSizeMB) + + // Large files raise the part size to stay under the 10,000-part cap: + // 900 GiB at 50 MB parts would need ~18,400 parts. + clamped := clampCloudMinPartSize(CompressConfig{}, 900<<30) + require.Equal(t, 103, clamped.MinPartSizeMB) + require.LessOrEqual(t, int64(900<<30)/clamped.MinPartSize(), int64(cloudMaxParts)) +} + func TestCompressStreamPartSizeMinimum(t *testing.T) { t.Parallel() diff --git a/packages/shared/pkg/storage/gcp_multipart.go b/packages/shared/pkg/storage/gcp_multipart.go index d9945be96a..def9d037c8 100644 --- a/packages/shared/pkg/storage/gcp_multipart.go +++ b/packages/shared/pkg/storage/gcp_multipart.go @@ -287,44 +287,6 @@ func (m *MultipartUploader) uploadPart(ctx context.Context, uploadID string, par return etag, nil } -type multiSliceReader struct { - slices [][]byte - idx int - off int -} - -func (r *multiSliceReader) Read(p []byte) (int, error) { - if len(p) == 0 { - if r.idx >= len(r.slices) { - return 0, io.EOF - } - - return 0, nil - } - - var n int - for len(p) > 0 && r.idx < len(r.slices) { - current := r.slices[r.idx] - if r.off >= len(current) { - r.idx++ - r.off = 0 - - continue - } - - copied := copy(p, current[r.off:]) - n += copied - r.off += copied - p = p[copied:] - } - - if n > 0 { - return n, nil - } - - return 0, io.EOF -} - // uploadPartSlices uploads a part from multiple byte slices without concatenating them. func (m *MultipartUploader) uploadPartSlices(ctx context.Context, uploadID string, partNumber int, slices [][]byte) (string, error) { totalLen := 0 @@ -337,7 +299,7 @@ func (m *MultipartUploader) uploadPartSlices(ctx context.Context, uploadID strin // Use a ReaderFunc so the retryable client can replay the body on retries bodyFn := func() (io.Reader, error) { - return &multiSliceReader{slices: slices}, nil + return newMultiSliceReader(slices), nil } req, err := retryablehttp.NewRequestWithContext(ctx, "PUT", url, retryablehttp.ReaderFunc(bodyFn)) diff --git a/packages/shared/pkg/storage/gcp_multipart_test.go b/packages/shared/pkg/storage/gcp_multipart_test.go index 18eb0d1c73..5964a4ecbc 100644 --- a/packages/shared/pkg/storage/gcp_multipart_test.go +++ b/packages/shared/pkg/storage/gcp_multipart_test.go @@ -60,6 +60,37 @@ func createTestMultipartUploader(t *testing.T, handler http.HandlerFunc, retryCo return uploader } +func TestMultipartUploader_PartUploaderContract(t *testing.T) { + t.Parallel() + + testPartUploaderContract(t, partUploaderTestAdapter{ + new: func(t *testing.T, recorder *partUploaderRecorder) partUploader { + t.Helper() + + return createTestMultipartUploader(t, func(w http.ResponseWriter, r *http.Request) { + switch { + case r.Method == http.MethodPost && r.URL.RawQuery == uploadsPath: + recorder.started = true + response := InitiateMultipartUploadResult{Bucket: testBucketName, Key: testObjectName, UploadID: "contract-upload-id"} + xmlData, _ := xml.Marshal(response) + w.WriteHeader(http.StatusOK) + w.Write(xmlData) + case r.Method == http.MethodPut: + recordUploadedPart(t, recorder, w, r) + case r.Method == http.MethodPost && strings.Contains(r.URL.RawQuery, "uploadId=contract-upload-id"): + recorder.completed = true + w.WriteHeader(http.StatusOK) + case r.Method == http.MethodDelete && strings.Contains(r.URL.RawQuery, "uploadId=contract-upload-id"): + recorder.aborted = true + w.WriteHeader(http.StatusNoContent) + default: + t.Fatalf("unexpected GCP multipart request: %s %s", r.Method, r.URL.String()) + } + }) + }, + }) +} + func TestMultipartUploader_InitiateUpload_Success(t *testing.T) { t.Parallel() expectedUploadID := "test-upload-id-123" diff --git a/packages/shared/pkg/storage/io_wrappers.go b/packages/shared/pkg/storage/io_wrappers.go index 413128ad99..61043ce027 100644 --- a/packages/shared/pkg/storage/io_wrappers.go +++ b/packages/shared/pkg/storage/io_wrappers.go @@ -183,3 +183,49 @@ func (r *spanReader) Close(ctx context.Context) (*ReadStats, error) { return stats, closeErr } + +// sliceReaderAt is a stateless io.ReaderAt over multiple byte slices; see +// newMultiSliceReader. +type sliceReaderAt struct { + slices [][]byte +} + +func (r sliceReaderAt) ReadAt(p []byte, off int64) (int, error) { + // io.ReaderAt requires a non-negative offset; return an error rather than + // panicking on the slice reslice below if used directly (io.SectionReader, + // the only current caller, already guards this). + if off < 0 { + return 0, errors.New("storage: sliceReaderAt.ReadAt: negative offset") + } + + var n int + for _, s := range r.slices { + if off >= int64(len(s)) { + off -= int64(len(s)) + + continue + } + + n += copy(p[n:], s[off:]) + off = 0 + if n == len(p) { + return n, nil + } + } + + return n, io.EOF +} + +// newMultiSliceReader streams multiple byte slices as one seekable body +// without concatenating them. Used as the multipart-part request body by both +// the GCP XML uploader (recreated per retry via ReaderFunc) and the AWS part +// uploader, where the SDK seeks to compute the payload hash/length and to +// rewind on retries. +func newMultiSliceReader(slices [][]byte) *io.SectionReader { + var size int64 + for _, s := range slices { + size += int64(len(s)) + } + + return io.NewSectionReader(sliceReaderAt{slices: slices}, 0, size) +} diff --git a/packages/shared/pkg/storage/part_uploader_test.go b/packages/shared/pkg/storage/part_uploader_test.go new file mode 100644 index 0000000000..79ec2f45c1 --- /dev/null +++ b/packages/shared/pkg/storage/part_uploader_test.go @@ -0,0 +1,98 @@ +package storage + +import ( + "cmp" + "fmt" + "io" + "net/http" + "strconv" + "testing" + + "github.com/stretchr/testify/require" +) + +type partUploaderTestAdapter struct { + abortsOnClose bool + new func(t *testing.T, recorder *partUploaderRecorder) partUploader +} + +type partUploaderRecorder struct { + started bool + completed bool + aborted bool + parts []recordedPart +} + +type recordedPart struct { + number int + body string +} + +func testPartUploaderContract(t *testing.T, adapter partUploaderTestAdapter) { + t.Helper() + + t.Run("uploads slices and completes in part order", func(t *testing.T) { + t.Parallel() + + recorder := &partUploaderRecorder{} + uploader := adapter.new(t, recorder) + + require.NoError(t, uploader.Start(t.Context())) + require.NoError(t, uploader.UploadPart(t.Context(), 2, []byte("two"))) + require.NoError(t, uploader.UploadPart(t.Context(), 1, []byte("one"), []byte("-split"))) + require.NoError(t, uploader.Complete(t.Context())) + require.NoError(t, uploader.Close()) + + require.True(t, recorder.started) + require.True(t, recorder.completed) + require.False(t, recorder.aborted) + require.Equal(t, []recordedPart{ + {number: 2, body: "two"}, + {number: 1, body: "one-split"}, + }, recorder.parts) + }) + + if !adapter.abortsOnClose { + return + } + + t.Run("close aborts unfinished upload", func(t *testing.T) { + t.Parallel() + + recorder := &partUploaderRecorder{} + uploader := adapter.new(t, recorder) + + require.NoError(t, uploader.Start(t.Context())) + require.NoError(t, uploader.Close()) + + require.True(t, recorder.started) + require.False(t, recorder.completed) + require.True(t, recorder.aborted) + }) +} + +// recordUploadedPart handles a multipart PUT in a fake server: it records the +// part number and body and replies with a per-part ETag, echoing the CRC32 +// checksum back like S3 does (the body must be read first — the SDK may send +// the checksum as an HTTP trailer). +func recordUploadedPart(t *testing.T, recorder *partUploaderRecorder, w http.ResponseWriter, r *http.Request) { + t.Helper() + + n, err := strconv.Atoi(r.URL.Query().Get("partNumber")) + require.NoError(t, err) + recorder.parts = append(recorder.parts, recordedPart{number: n, body: readAllString(t, r.Body)}) + if c := cmp.Or(r.Header.Get("x-amz-checksum-crc32"), r.Trailer.Get("x-amz-checksum-crc32")); c != "" { + w.Header().Set("x-amz-checksum-crc32", c) + } + w.Header().Set("ETag", fmt.Sprintf(`"etag%d"`, n)) + w.WriteHeader(http.StatusOK) +} + +func readAllString(t *testing.T, r io.Reader) string { + t.Helper() + + b, err := io.ReadAll(r) + require.NoError(t, err) + + return string(b) +} diff --git a/packages/shared/pkg/storage/paths.go b/packages/shared/pkg/storage/paths.go index d3086e0eb0..b9e6e861a2 100644 --- a/packages/shared/pkg/storage/paths.go +++ b/packages/shared/pkg/storage/paths.go @@ -105,7 +105,7 @@ func StripCompression(name string) string { // "/data/memfile.zstd.uncompressed-size"). Used by the FS backend where // GCS-style object metadata is unavailable. func SizeSidecar(objectPath string) string { - return objectPath + "." + MetadataKeyUncompressedSize + return objectPath + "." + ObjectMetadataUncompressedSize } // seekableObjectType derives the metric file_type and codec from a data-file diff --git a/packages/shared/pkg/storage/storage.go b/packages/shared/pkg/storage/storage.go index 8eaff56c34..01dcd814f2 100644 --- a/packages/shared/pkg/storage/storage.go +++ b/packages/shared/pkg/storage/storage.go @@ -54,10 +54,6 @@ const ( // MemoryChunkSize must always be bigger or equal to the block size. MemoryChunkSize = 4 * 1024 * 1024 // 4 MB - - // MetadataKeyUncompressedSize stores the original size so that Size() - // returns the uncompressed size for compressed objects. - MetadataKeyUncompressedSize = "uncompressed-size" ) // GetProviderType returns the configured storage provider type from the @@ -121,12 +117,13 @@ type ( ) const ( - ObjectMetadataTeamID = storageopts.ObjectMetadataTeamID - ObjectMetadataTemplateID = storageopts.ObjectMetadataTemplateID - ObjectMetadataBuildOrigin = storageopts.ObjectMetadataBuildOrigin - ObjectMetadataLogicalSize = storageopts.ObjectMetadataLogicalSize - ObjectMetadataMappedSize = storageopts.ObjectMetadataMappedSize - ObjectMetadataDiffSize = storageopts.ObjectMetadataDiffSize + ObjectMetadataTeamID = storageopts.ObjectMetadataTeamID + ObjectMetadataTemplateID = storageopts.ObjectMetadataTemplateID + ObjectMetadataBuildOrigin = storageopts.ObjectMetadataBuildOrigin + ObjectMetadataUncompressedSize = storageopts.ObjectMetadataUncompressedSize + ObjectMetadataLogicalSize = storageopts.ObjectMetadataLogicalSize + ObjectMetadataMappedSize = storageopts.ObjectMetadataMappedSize + ObjectMetadataDiffSize = storageopts.ObjectMetadataDiffSize ObjectOriginPause = storageopts.ObjectOriginPause ObjectOriginTemplateBuild = storageopts.ObjectOriginTemplateBuild @@ -317,7 +314,7 @@ func GetStorageProvider(ctx context.Context, cfg StorageConfig) (StorageProvider // cloud bucket-based storage switch provider { case AWSStorageProvider: - return newAWSStorage(ctx, bucketName) + return newAWSStorage(ctx, bucketName, cfg.limiter) case GCPStorageProvider: return NewGCP(ctx, bucketName, cfg.limiter) } diff --git a/packages/shared/pkg/storage/storage_aws.go b/packages/shared/pkg/storage/storage_aws.go index 191f96507c..cccbf6f2fb 100644 --- a/packages/shared/pkg/storage/storage_aws.go +++ b/packages/shared/pkg/storage/storage_aws.go @@ -18,6 +18,7 @@ import ( "go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-sdk-go-v2/otelaws" "go.uber.org/zap" + "github.com/e2b-dev/infra/packages/shared/pkg/limit" "github.com/e2b-dev/infra/packages/shared/pkg/logger" ) @@ -31,6 +32,7 @@ type awsStorage struct { client *s3.Client presignClient *s3.PresignClient bucketName string + limiter *limit.Limiter } var _ StorageProvider = (*awsStorage)(nil) @@ -39,6 +41,7 @@ type awsObject struct { client *s3.Client path string bucketName string + limiter *limit.Limiter } var ( @@ -46,7 +49,7 @@ var ( _ Blob = (*awsObject)(nil) ) -func newAWSStorage(ctx context.Context, bucketName string) (*awsStorage, error) { +func newAWSStorage(ctx context.Context, bucketName string, limiter *limit.Limiter) (*awsStorage, error) { cfg, err := config.LoadDefaultConfig(ctx) if err != nil { return nil, err @@ -71,6 +74,7 @@ func newAWSStorage(ctx context.Context, bucketName string) (*awsStorage, error) client: client, presignClient: presignClient, bucketName: bucketName, + limiter: limiter, }, nil } @@ -145,6 +149,7 @@ func (s *awsStorage) OpenSeekable(_ context.Context, path string) (Seekable, err client: s.client, bucketName: s.bucketName, path: path, + limiter: s.limiter, }, nil } @@ -153,6 +158,7 @@ func (s *awsStorage) OpenBlob(_ context.Context, path string) (Blob, error) { client: s.client, bucketName: s.bucketName, path: path, + limiter: s.limiter, }, nil } @@ -186,6 +192,12 @@ func (o *awsObject) StoreFile(ctx context.Context, path string, opts ...PutOptio return nil, [32]byte{}, errors.New("compressed uploads are not supported on AWS (builds target GCP only)") } + release, err := o.limiter.AcquireUploadSlot(ctx) + if err != nil { + return nil, [32]byte{}, err + } + defer release() + // Inherit the caller's context for the multipart upload. The AWS SDK's // manager.Uploader reuses the same ctx for CreateMultipartUpload, every // UploadPart (Concurrency=8, PartSize=10MB), and the final Complete/Abort — @@ -204,7 +216,7 @@ func (o *awsObject) StoreFile(ctx context.Context, path string, opts ...PutOptio o.client, func(u *manager.Uploader) { u.PartSize = 10 * 1024 * 1024 // 10 MB - u.Concurrency = 8 // eight parts in flight + u.Concurrency = o.limiter.MaxUploadTasks(ctx) }, ) diff --git a/packages/shared/pkg/storage/storage_google.go b/packages/shared/pkg/storage/storage_google.go index 4a1d20ce91..24e1e78f74 100644 --- a/packages/shared/pkg/storage/storage_google.go +++ b/packages/shared/pkg/storage/storage_google.go @@ -10,10 +10,8 @@ import ( "fmt" "hash" "io" - "maps" "net/http" "os" - "strconv" "time" "cloud.google.com/go/storage" @@ -34,16 +32,15 @@ import ( ) const ( - googleReadTimeout = 10 * time.Second - googleOperationTimeout = 5 * time.Second - googleBufferSize = 4 << 20 // 4 MiB - googleInitialBackoff = 10 * time.Millisecond - googleMaxBackoff = 10 * time.Second - googleBackoffMultiplier = 2 - googleMaxAttempts = 10 - defaultGRPCConnectionPoolSize = 8 - defaultGCSEnableDirectPath = false - gcloudDefaultUploadConcurrency = 16 + googleReadTimeout = 10 * time.Second + googleOperationTimeout = 5 * time.Second + googleBufferSize = 4 << 20 // 4 MiB + googleInitialBackoff = 10 * time.Millisecond + googleMaxBackoff = 10 * time.Second + googleBackoffMultiplier = 2 + googleMaxAttempts = 10 + defaultGRPCConnectionPoolSize = 8 + defaultGCSEnableDirectPath = false gcsOperationAttr = "operation" gcsOperationAttrWrite = "Write" @@ -227,11 +224,8 @@ func (o *gcpObject) Size(ctx context.Context) (_ int64, err error) { return 0, fmt.Errorf("failed to get GCS object (%q) attributes: %w", o.path, err) } - if v, ok := attrs.Metadata[MetadataKeyUncompressedSize]; ok { - parsed, parseErr := strconv.ParseInt(v, 10, 64) - if parseErr == nil { - return parsed, nil - } + if size, ok := ObjectMetadata(attrs.Metadata).UncompressedSize(); ok { + return size, nil } return attrs.Size, nil @@ -418,27 +412,23 @@ func (o *gcpObject) StoreFile(ctx context.Context, path string, opts ...PutOptio attribute.String("compression.type", compressionMetricType), ) - maxConcurrency := gcloudDefaultUploadConcurrency - if o.limiter != nil { - uploadLimiter := o.limiter.GCloudUploadLimiter() - if uploadLimiter != nil { - semaphoreErr := uploadLimiter.Acquire(ctx, 1) - if semaphoreErr != nil { - timer.Failure(ctx, 0) - - return nil, [32]byte{}, fmt.Errorf("failed to acquire semaphore: %w", semaphoreErr) - } - defer uploadLimiter.Release(1) - } + release, err := o.limiter.AcquireUploadSlot(ctx) + if err != nil { + timer.Failure(ctx, 0) - maxConcurrency = o.limiter.GCloudMaxTasks(ctx) + return nil, [32]byte{}, err } + defer release() + + maxConcurrency := o.limiter.MaxUploadTasks(ctx) // Compressed uploads always go through the multipart compressed path, // regardless of file size. if cfg.IsCompressionEnabled() { start := time.Now() - ft, checksum, err := o.storeFileCompressed(ctx, path, cfg, maxConcurrency, putOpts) + ft, checksum, err := storeFileCompressed(ctx, path, cfg, maxConcurrency, putOpts, func(metadata ObjectMetadata) (partUploader, error) { + return NewMultipartUploaderWithRetryConfig(ctx, bucketName, objectName, DefaultRetryConfig(), metadata) + }) if err != nil { timer.Failure(ctx, fileInfo.Size()) logger.L().Error(ctx, "Failed to upload file to GCS", @@ -545,38 +535,6 @@ func (o *gcpObject) StoreFile(ctx context.Context, path string, opts ...PutOptio return nil, sum256(hasher), e } -func (o *gcpObject) storeFileCompressed(ctx context.Context, localPath string, cfg CompressConfig, maxConcurrency int, putOpts PutOptions) (*FullFrameTable, [32]byte, error) { - file, err := os.Open(localPath) - if err != nil { - return nil, [32]byte{}, fmt.Errorf("failed to open local file %s: %w", localPath, err) - } - defer file.Close() - - fi, err := file.Stat() - if err != nil { - return nil, [32]byte{}, fmt.Errorf("failed to stat local file %s: %w", localPath, err) - } - - // Merge caller metadata (e.g. team_id) with our internal uncompressed-size - // bookkeeping. Internal key wins on collision. - metadata := make(map[string]string, len(putOpts.Metadata)+1) - maps.Copy(metadata, putOpts.Metadata) - metadata[MetadataKeyUncompressedSize] = strconv.FormatInt(fi.Size(), 10) - - uploader, err := NewMultipartUploaderWithRetryConfig( - ctx, - o.storage.bucket.BucketName(), - o.path, - DefaultRetryConfig(), - metadata, - ) - if err != nil { - return nil, [32]byte{}, fmt.Errorf("failed to create multipart uploader: %w", err) - } - - return compressStream(ctx, file, cfg, uploader, maxConcurrency, putOpts.FrameSink) -} - type gcpServiceToken struct { ClientEmail string `json:"client_email"` PrivateKey string `json:"private_key"` diff --git a/packages/shared/pkg/storage/storageopts/storageopts.go b/packages/shared/pkg/storage/storageopts/storageopts.go index d749369cad..60e2418923 100644 --- a/packages/shared/pkg/storage/storageopts/storageopts.go +++ b/packages/shared/pkg/storage/storageopts/storageopts.go @@ -5,6 +5,7 @@ package storageopts import ( "context" "maps" + "strconv" ) type ObjectMetadata map[string]string @@ -14,8 +15,30 @@ const ( ObjectMetadataTeamID = "team_id" ObjectMetadataTemplateID = "template_id" ObjectMetadataBuildOrigin = "build_origin" + + // ObjectMetadataUncompressedSize stores the original size of a compressed + // object so that Size() can report it without fetching the frame table. + ObjectMetadataUncompressedSize = "uncompressed-size" ) +// WithUncompressedSize returns a copy of the metadata with the original +// (uncompressed) size recorded. Internal key wins on collision. +func (m ObjectMetadata) WithUncompressedSize(size int64) ObjectMetadata { + out := make(ObjectMetadata, len(m)+1) + maps.Copy(out, m) + out[ObjectMetadataUncompressedSize] = strconv.FormatInt(size, 10) + + return out +} + +// UncompressedSize reports the original size of a compressed object, or false +// if absent (uncompressed object) or malformed. +func (m ObjectMetadata) UncompressedSize() (int64, bool) { + n, err := strconv.ParseInt(m[ObjectMetadataUncompressedSize], 10, 64) + + return n, err == nil +} + // ObjectOrigin is the immutable operation that created a build, stored as the // ObjectMetadataBuildOrigin value. type ObjectOrigin string