refactor(storage): shared compressed-upload framework + storage limiter (#3232)

Author: dobrac

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)

packages/shared/pkg/limit/gcloud.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)

packages/shared/pkg/limit/limiter.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)
+}

packages/shared/pkg/limit/storage.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()
+}

packages/shared/pkg/limit/storage_test.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:

packages/shared/pkg/limit/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)
 	}

packages/shared/pkg/storage/compress_upload.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()
 

packages/shared/pkg/storage/compress_upload_test.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))

packages/shared/pkg/storage/gcp_multipart.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"