fix (cache): Migrate cache S3 manager to AWS transfermanager #4118
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4,25 +4,23 @@ import ( | |
| "context" | ||
| "errors" | ||
| "fmt" | ||
| "log/slog" | ||
| "net/http" | ||
| "net/url" | ||
| "os" | ||
| "path" | ||
| "strconv" | ||
| "strings" | ||
| "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/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" | ||
| "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 *transfermanager.Client used by downloadWithRetry | ||
| type objectDownloader interface { | ||
| 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 *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() }() | ||
|
|
||
| // 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 = 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 *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 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 | ||
|
buildsworth-bk-app[bot] marked this conversation as resolved.
|
||
| 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 := 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 | ||
|
zhming0 marked this conversation as resolved.
|
||
| 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. | ||
| // The concurrency and part_size_mb URL overrides restore | ||
| // path only; save always uses the fixed upload defaults. | ||
| func resolveTransferSettings(opts *Options) transferSettings { | ||
| // 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 { | ||
| downloadConcurrency = opts.Concurrency | ||
| } | ||
| if opts.PartSizeMB > 0 { | ||
|
Comment on lines
303
to
-284
Contributor
There was a problem hiding this comment. 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:.
Contributor
Author
There was a problem hiding this comment. save and restore have differing memory behaviours:
I very much agree with the generic names for concurrency/part_size_mb, are misleading, so I propose we could do one of these:
Contributor
There was a problem hiding this comment.
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?
Contributor
Author
There was a problem hiding this comment. 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? |
||
| 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 | ||
| } | ||
|
Comment on lines
+328
to
+356
Contributor
There was a problem hiding this comment. [blocking] I believe the arithmetic is based on some internal logic of transfer function, this isn't ideal because we don't want to be coupled to something that will change anytime. |
||
|
|
||
| // 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", uploadConcurrency, | ||
| ) | ||
|
|
||
| // 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 := 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", 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", 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: uploadConcurrency, | ||
| }, nil | ||
| } | ||
|
|
||
|
|
@@ -401,30 +468,16 @@ func (b *S3Blob) Download(ctx context.Context, key, destPath string) (*TransferI | |
| "concurrency", b.downloadConcurrency, | ||
| ) | ||
|
|
||
| // 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, &transfermanager.DownloadObjectInput{ | ||
| Bucket: aws.String(b.bucketName), | ||
| Key: aws.String(fullKey), | ||
| }) | ||
| 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) | ||
| } | ||
|
|
||
| // 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) | ||
|
|
||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just a question for my understanding, what was this memory budget before this change?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
So there was no explicit "budget" before - manager was inherently bounded as it was streaming rather than buffering (which is waht transferManager does).
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I see what this is trying to achieve now. But I am not sure if hard coding to 256MB is the right path:
At this point, I think it'd be fair to pause to consider if transfer manager is mature enough to justify this change.