Skip to content

fix (cache): Migrate cache S3 manager to AWS transfermanager #4118

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ss1909 wants to merge 6 commits into
base: main
Choose a base branch
from
Open

fix (cache): Migrate cache S3 manager to AWS transfermanager #4118

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0712816
fix (cache): Migrate cache S3 transfers to AWS transfermanager
ss1909 Jul 23, 2026
a8dd30d
Merge branch 'main' into sneha/A-1470
ss1909 Jul 23, 2026
b969269
fix: decouple save from restore - cap the effective upload part size …
ss1909 Jul 23, 2026
488f804
Merge branch 'main' into sneha/A-1470
ss1909 Jul 24, 2026
56e0d3b
Merge branch 'main' into sneha/A-1470
ss1909 Jul 28, 2026
0441f48
refactor
ss1909 Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 42 additions & 43 deletions internal/cache/store/s3.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,25 +4,23 @@ import (
"context"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"net/url"
"os"
"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"
Expand All @@ -35,6 +33,8 @@ import (
const (
defaultDownloadConcurrency = 32
defaultDownloadPartSizeMB = 32
defaultUploadConcurrency = 5
defaultUploadPartSizeBytes = 5 * 1024 * 1024
)

// Options holds configuration for S3Blob and can be constructed from an S3 URL in a similar way to gocloud.dev
Expand Down Expand Up @@ -106,10 +106,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.
Expand Down Expand Up @@ -143,7 +142,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)
Expand All @@ -153,7 +152,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",
Expand All @@ -164,7 +167,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
Expand All @@ -173,12 +176,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
Expand Down Expand Up @@ -227,14 +231,20 @@ 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) {
o.Concurrency = settings.uploadConcurrency
o.PartSizeBytes = settings.uploadPartSize
Comment thread
buildsworth-bk-app[bot] marked this conversation as resolved.
})
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
Comment thread
zhming0 marked this conversation as resolved.
o.GetObjectType = tmtypes.GetObjectRanges
})

slog.Debug("configured S3 transfer manager",
Expand All @@ -253,6 +263,7 @@ func NewS3Blob(ctx context.Context, s3url string) (*S3Blob, error) {
prefix: opts.Prefix,
uploadConcurrency: settings.uploadConcurrency,
downloadConcurrency: settings.downloadConcurrency,
downloadPartSize: settings.downloadPartSize,
}, nil
}

Expand All @@ -271,14 +282,14 @@ type transferSettings struct {
// downloads; otherwise uploads use the SDK's upload defaults and downloads use
// the download-tuned defaults.
func resolveTransferSettings(opts *Options) transferSettings {
uploadConcurrency := manager.DefaultUploadConcurrency
uploadConcurrency := defaultUploadConcurrency
downloadConcurrency := defaultDownloadConcurrency
if opts.Concurrency > 0 {
uploadConcurrency = opts.Concurrency
downloadConcurrency = opts.Concurrency
}

uploadPartSize := manager.DefaultUploadPartSize
uploadPartSize := int64(defaultUploadPartSizeBytes)
downloadPartSize := int64(defaultDownloadPartSizeMB) * 1024 * 1024
if opts.PartSizeMB > 0 {
uploadPartSize = int64(opts.PartSizeMB) * 1024 * 1024
Comment on lines 303 to -284

@zhming0 zhming0 Jul 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, can you help me understand again (using your words) why we need to set different concurrency and part size for upload and download? I did read the buildsworth thread, that doesn't quite make sense to my brain :g_thinking:.

The thing is, our current url scheme allows a generic currency and part size setting, if we were to make this change, our URL params will become misleading.

@ss1909 ss1909 Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

save and restore have differing memory behaviours:

  • Save buffers each in-flight part fully in RAM
  • Restore streams each range through a small (32 KiB) buffer
    Save/Upload buffers (memory grows with part size); Restore/download streams (memory doesn't). That mismatch is why this need for different concurrency/part-size handling arises.

I very much agree with the generic names for concurrency/part_size_mb, are misleading, so I propose we could do one of these:

  1. Document them as restore/download only tuning
  2. Rename to download_concurrency / download_part_size_mb

@zhming0 zhming0 Jul 27, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Save buffers each in-flight part fully in RAM

I can see that the buffer behave a bit differently, but why does the the configuration need to differ I wonder? I read the other thread it's a bit avoiding excessive memory? But does the upload memory budget setting cap the memory?

@ss1909 ss1909 Jul 27, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the reason why the configs need to be different is because the numbers that make restore fast are exactly the ones that make save allocate big chunks (gbs), so a single shared value can't serve both! it's either too small (slow restore) or too large (save OOM).

Restore streams, so memory is flat regardless of part size/concurrency - the best config is high concurrency + large parts (c32/p32) purely for speed. Save buffers, so memory = (concurrency+1) × partSize - the best config is low concurrency + small parts to stay bounded. More on this here - https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html and https://github.com/aws/aws-sdk-go-v2/tree/50fd8e78e0780bc66f360bc50d59aa314d164d54/feature/s3/transfermanager

does that make sense?

Expand Down Expand Up @@ -327,8 +338,8 @@ func (b *S3Blob) Upload(ctx context.Context, filePath, key string) (*TransferInf
"concurrency", b.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,
Expand All @@ -345,7 +356,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)
Expand Down Expand Up @@ -401,30 +412,16 @@ 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(
roko.WithMaxAttempts(3),
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) {
Expand All @@ -433,10 +430,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)
Expand Down
68 changes: 54 additions & 14 deletions internal/cache/store/s3_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,16 @@ import (
"context"
"errors"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"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"
Expand Down Expand Up @@ -386,6 +385,44 @@ 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")
}
if blob.uploadConcurrency != 8 {
t.Errorf("uploadConcurrency = %d, want 8", blob.uploadConcurrency)
}
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 TestResolveTransferSettings(t *testing.T) {
const mb = int64(1024 * 1024)

Expand All @@ -398,8 +435,8 @@ 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,
Expand All @@ -410,7 +447,7 @@ func TestResolveTransferSettings(t *testing.T) {
opts: &Options{Concurrency: 50},
want: transferSettings{
uploadConcurrency: 50,
uploadPartSize: manager.DefaultUploadPartSize,
uploadPartSize: int64(defaultUploadPartSizeBytes),
downloadConcurrency: 50,
downloadPartSize: int64(defaultDownloadPartSizeMB) * mb,
maxIdleConnsPerHost: 50,
Expand All @@ -420,7 +457,7 @@ func TestResolveTransferSettings(t *testing.T) {
name: "part size override applies to both",
opts: &Options{PartSizeMB: 64},
want: transferSettings{
uploadConcurrency: manager.DefaultUploadConcurrency,
uploadConcurrency: defaultUploadConcurrency,
uploadPartSize: 64 * mb,
downloadConcurrency: defaultDownloadConcurrency,
downloadPartSize: 64 * mb,
Expand Down Expand Up @@ -474,14 +511,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
Expand All @@ -504,7 +544,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)
}
Expand All @@ -529,7 +569,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")
}
Expand All @@ -547,7 +587,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")
}
Expand Down
Loading