diff --git a/internal/cache/store/s3.go b/internal/cache/store/s3.go index 39de870f61..3973067f17 100644 --- a/internal/cache/store/s3.go +++ b/internal/cache/store/s3.go @@ -4,7 +4,6 @@ import ( "context" "errors" "fmt" - "io" "log/slog" "net/http" "net/url" @@ -12,17 +11,16 @@ import ( "path" "strconv" "strings" - "sync/atomic" "time" "github.com/aws/aws-sdk-go-v2/aws" awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http" "github.com/aws/aws-sdk-go-v2/config" - "github.com/aws/aws-sdk-go-v2/feature/s3/manager" + "github.com/aws/aws-sdk-go-v2/feature/s3/transfermanager" + tmtypes "github.com/aws/aws-sdk-go-v2/feature/s3/transfermanager/types" "github.com/aws/aws-sdk-go-v2/service/s3" "github.com/aws/aws-sdk-go-v2/service/s3/types" smithy "github.com/aws/smithy-go" - smithymiddleware "github.com/aws/smithy-go/middleware" "github.com/buildkite/agent/v3/internal/cache/internal/trace" "github.com/buildkite/roko" "go.opentelemetry.io/otel/attribute" @@ -35,6 +33,25 @@ import ( const ( defaultDownloadConcurrency = 32 defaultDownloadPartSizeMB = 32 + // Upload defaults are fixed, never URL-tunable. defaultUploadConcurrency is + // the maximum parallelism; uploadConcurrencyForSize lowers it for large + // objects so UploadObject's peak part buffers — (concurrency+2) × partSize + // (the concurrency+1 eager pool plus a separately-held first chunk) — stay + // within uploadMemoryBudget. + defaultUploadConcurrency = 5 + defaultUploadPartSizeBytes = 5 * 1024 * 1024 + // defaultUploadMultipartThreshold is pinned rather than inherited from the + // SDK so upload behaviour is explicit and stable across SDK versions. + defaultUploadMultipartThreshold = 16 * 1024 * 1024 + // uploadMemoryBudget is the target ceiling for a single upload's part-buffer + // pool. It holds for all normal object sizes; multi-TiB objects bottom out at + // concurrency 1, where one part buffer (size/uploadMaxParts) is unavoidable + // and may exceed this. + uploadMemoryBudget = 256 * 1024 * 1024 + // uploadMaxParts mirrors S3's hard limit on parts per multipart upload. It is + // pinned on the client so uploadConcurrencyForSize predicts the SDK's part + // sizing exactly. + uploadMaxParts = 10000 ) // Options holds configuration for S3Blob and can be constructed from an S3 URL in a similar way to gocloud.dev @@ -106,10 +123,9 @@ func OptionsFromURL(s3url string) (*Options, error) { return opts, nil } -// objectDownloader is the subset of manager.Downloader used by downloadWithRetry, -// declared so the retry loop can be tested with a fake. +// objectDownloader is the subset of *transfermanager.Client used by downloadWithRetry type objectDownloader interface { - Download(ctx context.Context, w io.WriterAt, input *s3.GetObjectInput, options ...func(*manager.Downloader)) (int64, error) //nolint:staticcheck // SA1019: pending migration to transfermanager + DownloadObject(ctx context.Context, input *transfermanager.DownloadObjectInput, opts ...func(*transfermanager.Options)) (*transfermanager.DownloadObjectOutput, error) } // isPreconditionFailed returns true when an error is an S3 412 PreconditionFailed. @@ -143,7 +159,7 @@ func isNotFound(err error) bool { // downloadWithRetry runs the multipart download, retrying on S3 412 // PreconditionFailed (a concurrent restore's TTL-refresh CopyObject changed the // object's ETag and invalidated the SDK's If-Match guard). Returns bytes written. -func downloadWithRetry(ctx context.Context, r *roko.Retrier, d objectDownloader, destPath string, in *s3.GetObjectInput, opts ...func(*manager.Downloader)) (int64, error) { //nolint:staticcheck // SA1019: pending migration to transfermanager +func downloadWithRetry(ctx context.Context, r *roko.Retrier, d objectDownloader, destPath string, in *transfermanager.DownloadObjectInput) (int64, error) { var bytesWritten int64 err := r.DoWithContext(ctx, func(r *roko.Retrier) error { destFile, err := os.Create(destPath) @@ -153,7 +169,11 @@ func downloadWithRetry(ctx context.Context, r *roko.Retrier, d objectDownloader, } defer func() { _ = destFile.Close() }() - n, err := d.Download(ctx, destFile, in, opts...) + // Parts arrive out of order via io.WriterAt, so point the transfer at a + // freshly-truncated file on each attempt. + in.WriterAt = destFile + + out, err := d.DownloadObject(ctx, in) if err != nil { if isPreconditionFailed(err) { slog.Warn("cache download hit 412 (concurrent ETag change), retrying", @@ -164,7 +184,7 @@ func downloadWithRetry(ctx context.Context, r *roko.Retrier, d objectDownloader, return err } - bytesWritten = n + bytesWritten = aws.ToInt64(out.ContentLength) return nil }) return bytesWritten, err @@ -173,12 +193,13 @@ func downloadWithRetry(ctx context.Context, r *roko.Retrier, d objectDownloader, // S3Blob implements the Blob interface using AWS S3 type S3Blob struct { client *s3.Client - uploader *manager.Uploader //nolint:staticcheck // SA1019: pending migration to transfermanager - downloader *manager.Downloader //nolint:staticcheck // SA1019: pending migration to transfermanager + uploader *transfermanager.Client + downloader *transfermanager.Client bucketName string prefix string uploadConcurrency int downloadConcurrency int + downloadPartSize int64 } // NewS3Blob creates a new S3Blob instance using an S3 URL and prefix @@ -227,14 +248,26 @@ func NewS3Blob(ctx context.Context, s3url string) (*S3Blob, error) { } }) - // Create the uploader and downloader with their resolved settings - uploader := manager.NewUploader(client, func(u *manager.Uploader) { //nolint:staticcheck // SA1019: pending migration to transfermanager - u.Concurrency = settings.uploadConcurrency - u.PartSize = settings.uploadPartSize + // Create the transfer-manager clients with their resolved settings. Uploads + // and downloads are tuned differently, so each gets its own client. + uploader := transfermanager.New(client, func(o *transfermanager.Options) { + // Concurrency is the per-upload maximum; Upload lowers it per object via + // uploadConcurrencyForSize to keep the buffer pool bounded. + o.Concurrency = settings.uploadConcurrency + o.PartSizeBytes = settings.uploadPartSize + o.MaxUploadParts = uploadMaxParts + // Pin the multipart threshold instead of inheriting the SDK default so + // upload behaviour is explicit and stable across SDK versions. + o.MultipartUploadThreshold = defaultUploadMultipartThreshold }) - downloader := manager.NewDownloader(client, func(d *manager.Downloader) { //nolint:staticcheck // SA1019: pending migration to transfermanager - d.Concurrency = settings.downloadConcurrency - d.PartSize = settings.downloadPartSize + downloader := transfermanager.New(client, func(o *transfermanager.Options) { + o.Concurrency = settings.downloadConcurrency + o.PartSizeBytes = settings.downloadPartSize + // Use Range-based fan-out rather than partNumber-based. The SDK default + // (PART) only fans out for objects uploaded as multipart; Ranges work on any object + // regardless of how it was uploaded, so restore parallelism + // is determined by our config (Concurrency × PartSizeBytes), not by upload history + o.GetObjectType = tmtypes.GetObjectRanges }) slog.Debug("configured S3 transfer manager", @@ -253,6 +286,7 @@ func NewS3Blob(ctx context.Context, s3url string) (*S3Blob, error) { prefix: opts.Prefix, uploadConcurrency: settings.uploadConcurrency, downloadConcurrency: settings.downloadConcurrency, + downloadPartSize: settings.downloadPartSize, }, nil } @@ -267,21 +301,18 @@ type transferSettings struct { } // resolveTransferSettings turns parsed Options into concrete transfer settings. -// An explicit concurrency or part_size_mb override applies to both uploads and -// downloads; otherwise uploads use the SDK's upload defaults and downloads use -// the download-tuned defaults. +// The concurrency and part_size_mb URL overrides restore +// path only; save always uses the fixed upload defaults. func resolveTransferSettings(opts *Options) transferSettings { - uploadConcurrency := manager.DefaultUploadConcurrency + // Upload settings are fixed and intentionally not derived from the URL + uploadConcurrency := defaultUploadConcurrency + uploadPartSize := int64(defaultUploadPartSizeBytes) downloadConcurrency := defaultDownloadConcurrency + downloadPartSize := int64(defaultDownloadPartSizeMB) * 1024 * 1024 if opts.Concurrency > 0 { - uploadConcurrency = opts.Concurrency downloadConcurrency = opts.Concurrency } - - uploadPartSize := manager.DefaultUploadPartSize - downloadPartSize := int64(defaultDownloadPartSizeMB) * 1024 * 1024 if opts.PartSizeMB > 0 { - uploadPartSize = int64(opts.PartSizeMB) * 1024 * 1024 downloadPartSize = int64(opts.PartSizeMB) * 1024 * 1024 } @@ -294,6 +325,36 @@ func resolveTransferSettings(opts *Options) transferSettings { } } +// uploadConcurrencyForSize picks how many part buffers UploadObject may hold in +// parallel for an object of the given size, so its peak allocation stays within +// uploadMemoryBudget. That peak is (concurrency+2) × partSize: the eager pool +// holds concurrency+1 buffers, and UploadObject reads the first chunk into a +// separate buffer outside the pool. +// +// transfermanager raises the part size to size/uploadMaxParts once an object +// would exceed S3's 10,000-part limit, so for large objects the part buffers +// grow with the object and we trade away concurrency to compensate. Multi-TiB +// objects bottom out at concurrency 1, where three part buffers are still +// unavoidable, so the budget is a ceiling for normal sizes, not a hard cap at +// every size. +func uploadConcurrencyForSize(size int64, maxConcurrency int) int { + partSize := int64(defaultUploadPartSizeBytes) + if forced := size/uploadMaxParts + 1; forced > partSize { + partSize = forced + } + + // Peak = (concurrency+1) pool buffers + 1 separately-held first chunk, so + // solve (concurrency+2) × partSize <= budget. + concurrency := int(uploadMemoryBudget/partSize) - 2 + if concurrency > maxConcurrency { + concurrency = maxConcurrency + } + if concurrency < 1 { + concurrency = 1 + } + return concurrency +} + // Upload uploads a file to S3 using multipart upload for parallel transfers func (b *S3Blob) Upload(ctx context.Context, filePath, key string) (*TransferInfo, error) { ctx, span := trace.Start(ctx, "S3Blob.Upload") @@ -321,17 +382,23 @@ func (b *S3Blob) Upload(ctx context.Context, filePath, key string) (*TransferInf bytesWritten := fileInfo.Size() + // Choose upload concurrency from the object size so the part-buffer pool + // stays within uploadMemoryBudget (see uploadConcurrencyForSize). + uploadConcurrency := uploadConcurrencyForSize(bytesWritten, b.uploadConcurrency) + slog.Debug("starting S3 upload", "key", fullKey, "file_size", bytesWritten, - "concurrency", b.uploadConcurrency, + "concurrency", uploadConcurrency, ) - // Upload the file to S3 using the multipart uploader - result, err := b.uploader.Upload(ctx, &s3.PutObjectInput{ //nolint:staticcheck // SA1019: pending migration to transfermanager + // Upload the file to S3 using the multipart transfer manager + result, err := b.uploader.UploadObject(ctx, &transfermanager.UploadObjectInput{ Bucket: aws.String(b.bucketName), Key: aws.String(fullKey), Body: file, + }, func(o *transfermanager.Options) { + o.Concurrency = uploadConcurrency }) if err != nil { return nil, fmt.Errorf("failed to upload file to S3: %w", err) @@ -345,7 +412,7 @@ func (b *S3Blob) Upload(ctx context.Context, filePath, key string) (*TransferInf } // Extract request ID from the upload result (only set for multipart uploads) - requestID := result.UploadID + requestID := aws.ToString(result.UploadID) duration := time.Since(start) averageSpeed := calculateTransferSpeedMBps(bytesWritten, duration) @@ -354,7 +421,7 @@ func (b *S3Blob) Upload(ctx context.Context, filePath, key string) (*TransferInf "key", fullKey, "bytes_transferred", bytesWritten, "parts_uploaded", partCount, - "concurrency", b.uploadConcurrency, + "concurrency", uploadConcurrency, "duration", duration, "transfer_speed_mbps", fmt.Sprintf("%.2f", averageSpeed), ) @@ -364,7 +431,7 @@ func (b *S3Blob) Upload(ctx context.Context, filePath, key string) (*TransferInf attribute.String("transfer_speed", fmt.Sprintf("%.2fMB/s", averageSpeed)), attribute.String("request_id", requestID), attribute.Int("part_count", partCount), - attribute.Int("concurrency", b.uploadConcurrency), + attribute.Int("concurrency", uploadConcurrency), ) return &TransferInfo{ @@ -373,7 +440,7 @@ func (b *S3Blob) Upload(ctx context.Context, filePath, key string) (*TransferInf RequestID: requestID, Duration: duration, PartCount: partCount, - Concurrency: b.uploadConcurrency, + Concurrency: uploadConcurrency, }, nil } @@ -401,8 +468,6 @@ func (b *S3Blob) Download(ctx context.Context, key, destPath string) (*TransferI "concurrency", b.downloadConcurrency, ) - var partCount atomic.Int32 - // Download the file from S3 using parallel range requests, retrying on a // 412 PreconditionFailed caused by a concurrent restore's ETag change. retrier := roko.NewRetrier( @@ -410,21 +475,9 @@ func (b *S3Blob) Download(ctx context.Context, key, destPath string) (*TransferI roko.WithStrategy(roko.ExponentialSubsecond(200*time.Millisecond)), roko.WithJitterRange(0, 250*time.Millisecond), ) - bytesWritten, err := downloadWithRetry(ctx, retrier, b.downloader, destPath, &s3.GetObjectInput{ + bytesWritten, err := downloadWithRetry(ctx, retrier, b.downloader, destPath, &transfermanager.DownloadObjectInput{ Bucket: aws.String(b.bucketName), Key: aws.String(fullKey), - }, func(d *manager.Downloader) { //nolint:staticcheck // SA1019: pending migration to transfermanager - d.ClientOptions = append(d.ClientOptions, func(o *s3.Options) { - o.APIOptions = append(o.APIOptions, func(stack *smithymiddleware.Stack) error { - return stack.Initialize.Add(smithymiddleware.InitializeMiddlewareFunc( - "PartCounter", - func(ctx context.Context, in smithymiddleware.InitializeInput, next smithymiddleware.InitializeHandler) (smithymiddleware.InitializeOutput, smithymiddleware.Metadata, error) { - partCount.Add(1) - return next.HandleInitialize(ctx, in) - }, - ), smithymiddleware.Before) - }) - }) }) if err != nil { if isNotFound(err) { @@ -433,10 +486,12 @@ func (b *S3Blob) Download(ctx context.Context, key, destPath string) (*TransferI return nil, fmt.Errorf("failed to download file from S3: %w", err) } - // Get actual part count from interceptor - actualPartCount := int(partCount.Load()) - if actualPartCount == 0 { - actualPartCount = 1 + // transfermanager doesn't report how many ranges it fetched, and an object's + // PartsCount reflects its original multipart upload, not the restore fan-out. + // For range-based downloads the range count is ceil(bytes / partSize). + actualPartCount := 1 + if b.downloadPartSize > 0 && bytesWritten > 0 { + actualPartCount = int((bytesWritten + b.downloadPartSize - 1) / b.downloadPartSize) } duration := time.Since(start) diff --git a/internal/cache/store/s3_test.go b/internal/cache/store/s3_test.go index 1c91216601..4a9a1ce260 100644 --- a/internal/cache/store/s3_test.go +++ b/internal/cache/store/s3_test.go @@ -4,7 +4,6 @@ import ( "context" "errors" "fmt" - "io" "net/http" "os" "path/filepath" @@ -12,9 +11,9 @@ import ( "testing" "time" + "github.com/aws/aws-sdk-go-v2/aws" awshttp "github.com/aws/aws-sdk-go-v2/aws/transport/http" - "github.com/aws/aws-sdk-go-v2/feature/s3/manager" - "github.com/aws/aws-sdk-go-v2/service/s3" + "github.com/aws/aws-sdk-go-v2/feature/s3/transfermanager" "github.com/aws/aws-sdk-go-v2/service/s3/types" smithy "github.com/aws/smithy-go" smithyhttp "github.com/aws/smithy-go/transport/http" @@ -386,6 +385,84 @@ func TestGetFullKey(t *testing.T) { } } +// TestNewS3Blob exercises construction against a custom endpoint with +// path-style access and explicit transfer tuning — the URL options that +// self-hosted / S3-compatible backends rely on. Construction makes no network +// calls, so this verifies the transfermanager clients and settings are wired +// through from the URL without needing a live S3. +func TestNewS3Blob(t *testing.T) { + blob, err := NewS3Blob(t.Context(), + "s3://my-bucket/cache/prefix?region=us-west-2&endpoint=http://localhost:9000&use_path_style=true&concurrency=8&part_size_mb=16") + if err != nil { + t.Fatalf("NewS3Blob: %v", err) + } + + if blob.client == nil { + t.Error("client is nil") + } + if blob.uploader == nil { + t.Error("uploader is nil") + } + if blob.downloader == nil { + t.Error("downloader is nil") + } + if blob.bucketName != "my-bucket" { + t.Errorf("bucketName = %q, want %q", blob.bucketName, "my-bucket") + } + if blob.prefix != "cache/prefix" { + t.Errorf("prefix = %q, want %q", blob.prefix, "cache/prefix") + } + // The URL's concurrency/part_size_mb tune downloads only; uploads stay on + // the fixed defaults regardless of what the URL asks for. + if blob.uploadConcurrency != defaultUploadConcurrency { + t.Errorf("uploadConcurrency = %d, want %d (URL knobs must not affect uploads)", blob.uploadConcurrency, defaultUploadConcurrency) + } + if blob.downloadConcurrency != 8 { + t.Errorf("downloadConcurrency = %d, want 8", blob.downloadConcurrency) + } + if want := int64(16 * 1024 * 1024); blob.downloadPartSize != want { + t.Errorf("downloadPartSize = %d, want %d", blob.downloadPartSize, want) + } +} + +func TestUploadConcurrencyForSize(t *testing.T) { + const ( + mib = int64(1024 * 1024) + gib = 1024 * mib + tib = 1024 * gib + ) + + tests := []struct { + name string + size int64 + want int + }{ + {name: "small file keeps full concurrency", size: 10 * mib, want: defaultUploadConcurrency}, + {name: "hundreds of GiB still full", size: 300 * gib, want: defaultUploadConcurrency}, + {name: "large object throttles to stay under budget", size: 600 * gib, want: 2}, + {name: "multi-TiB bottoms out at 1", size: 2 * tib, want: 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := uploadConcurrencyForSize(tt.size, defaultUploadConcurrency) + if got != tt.want { + t.Errorf("uploadConcurrencyForSize(%d) = %d, want %d", tt.size, got, tt.want) + } + // Peak allocation is (concurrency+2) × partSize: the concurrency+1 + // pool plus the separately-held first chunk. + partSize := int64(defaultUploadPartSizeBytes) + if forced := tt.size/uploadMaxParts + 1; forced > partSize { + partSize = forced + } + peak := int64(got+2) * partSize + if peak > uploadMemoryBudget && got != 1 { + t.Errorf("peak %d exceeds budget %d at concurrency %d", peak, int64(uploadMemoryBudget), got) + } + }) + } +} + func TestResolveTransferSettings(t *testing.T) { const mb = int64(1024 * 1024) @@ -398,41 +475,41 @@ func TestResolveTransferSettings(t *testing.T) { name: "defaults differ between upload and download", opts: &Options{}, want: transferSettings{ - uploadConcurrency: manager.DefaultUploadConcurrency, - uploadPartSize: manager.DefaultUploadPartSize, + uploadConcurrency: defaultUploadConcurrency, + uploadPartSize: int64(defaultUploadPartSizeBytes), downloadConcurrency: defaultDownloadConcurrency, downloadPartSize: int64(defaultDownloadPartSizeMB) * mb, maxIdleConnsPerHost: defaultDownloadConcurrency, }, }, { - name: "concurrency override applies to both", + name: "concurrency override applies to download only", opts: &Options{Concurrency: 50}, want: transferSettings{ - uploadConcurrency: 50, - uploadPartSize: manager.DefaultUploadPartSize, + uploadConcurrency: defaultUploadConcurrency, + uploadPartSize: int64(defaultUploadPartSizeBytes), downloadConcurrency: 50, downloadPartSize: int64(defaultDownloadPartSizeMB) * mb, maxIdleConnsPerHost: 50, }, }, { - name: "part size override applies to both", + name: "part size override applies to download only", opts: &Options{PartSizeMB: 64}, want: transferSettings{ - uploadConcurrency: manager.DefaultUploadConcurrency, - uploadPartSize: 64 * mb, + uploadConcurrency: defaultUploadConcurrency, + uploadPartSize: int64(defaultUploadPartSizeBytes), downloadConcurrency: defaultDownloadConcurrency, downloadPartSize: 64 * mb, maxIdleConnsPerHost: defaultDownloadConcurrency, }, }, { - name: "both overrides applied", + name: "overrides never affect upload settings", opts: &Options{Concurrency: 8, PartSizeMB: 16}, want: transferSettings{ - uploadConcurrency: 8, - uploadPartSize: 16 * mb, + uploadConcurrency: defaultUploadConcurrency, + uploadPartSize: int64(defaultUploadPartSizeBytes), downloadConcurrency: 8, downloadPartSize: 16 * mb, maxIdleConnsPerHost: 8, @@ -474,14 +551,17 @@ type fakeDownloadResult struct { payload []byte } -func (f *fakeDownloader) Download(_ context.Context, w io.WriterAt, _ *s3.GetObjectInput, _ ...func(*manager.Downloader)) (int64, error) { //nolint:staticcheck // SA1019: pending migration to transfermanager +func (f *fakeDownloader) DownloadObject(_ context.Context, in *transfermanager.DownloadObjectInput, _ ...func(*transfermanager.Options)) (*transfermanager.DownloadObjectOutput, error) { res := f.results[min(f.calls, len(f.results)-1)] f.calls++ if res.err != nil { - return 0, res.err + return nil, res.err + } + n, err := in.WriterAt.WriteAt(res.payload, 0) + if err != nil { + return nil, err } - n, err := w.WriteAt(res.payload, 0) - return int64(n), err + return &transfermanager.DownloadObjectOutput{ContentLength: aws.Int64(int64(n))}, nil } // testRetrier builds a retrier that runs instantly (no real sleeps) so the @@ -504,7 +584,7 @@ func TestDownloadWithRetry(t *testing.T) { {payload: payload}, }} - n, err := downloadWithRetry(t.Context(), testRetrier(), fake, destPath, &s3.GetObjectInput{}) + n, err := downloadWithRetry(t.Context(), testRetrier(), fake, destPath, &transfermanager.DownloadObjectInput{}) if err != nil { t.Fatalf("downloadWithRetry: unexpected error: %v", err) } @@ -529,7 +609,7 @@ func TestDownloadWithRetry(t *testing.T) { {err: responseErrorWithStatus(http.StatusPreconditionFailed)}, }} - _, err := downloadWithRetry(t.Context(), testRetrier(), fake, destPath, &s3.GetObjectInput{}) + _, err := downloadWithRetry(t.Context(), testRetrier(), fake, destPath, &transfermanager.DownloadObjectInput{}) if err == nil { t.Fatal("downloadWithRetry: expected error, got nil") } @@ -547,7 +627,7 @@ func TestDownloadWithRetry(t *testing.T) { {err: responseErrorWithStatus(http.StatusInternalServerError)}, }} - _, err := downloadWithRetry(t.Context(), testRetrier(), fake, destPath, &s3.GetObjectInput{}) + _, err := downloadWithRetry(t.Context(), testRetrier(), fake, destPath, &transfermanager.DownloadObjectInput{}) if err == nil { t.Fatal("downloadWithRetry: expected error, got nil") }