Optimize S3 parallel download pipeline

Tune the HTTP transport and download parameters based on benchmarking
against CI hosts (m6id.4xlarge, 16 vCPU):

- Configure HTTP transport with connection pooling (MaxIdleConnsPerHost)
  and 128KB read/write buffers for better connection reuse across
  parallel range workers.
- Default to max(16, NumCPU) download workers with 32MB chunks. This
  configuration consistently achieved ~245 MB/s on fast hosts vs ~100
  MB/s for single-stream downloads (the per-connection S3 cap).
- Add 8MB bufio buffer between S3 download pipe and pzstd to decouple
  network I/O from decompression stalls.
- Log full commit hash on cache hit for easier debugging.

Author: joshfriend

cmd/gradle-cache/main.go

@@ -2,7 +2,6 @@
 //
 // Base bundles are stored at s3://{bucket}/{commit}/{cache-key}/{bundle-file},
 // where bundle-file is the cache key with colons replaced by dashes + ".tar.zst".
-// This format is compatible with the bundled-cache-manager Ruby script.
 //
 // On restore, the tool walks the local git history (counting distinct-author
 // "blocks") to find the most recent S3 hit, downloads it, and extracts it
@@ -109,7 +108,7 @@ type RestoreCmd struct {
 	MaxBlocks      int      `help:"Number of distinct-author commit blocks to search." default:"20"`
 	GradleUserHome string   `help:"Path to GRADLE_USER_HOME." env:"GRADLE_USER_HOME"`
 	IncludedBuilds []string `help:"Included build directories whose build/ output to restore (relative to project root). Use 'dir/*' to restore build/ for all subdirectories. May be repeated." name:"included-build"`
-	Branch         string   `help:"Branch name to also apply a delta bundle for (typically $$BRANCH_NAME). The delta download runs concurrently with base extraction." optional:""`
+	Branch         string   `help:"Branch name to also apply a delta bundle for (typically $BRANCH_NAME). The delta download runs concurrently with base extraction." optional:""`
 }
 
 func (c *RestoreCmd) AfterApply() error {
@@ -169,7 +168,7 @@ func (c *RestoreCmd) Run(ctx context.Context, metrics metricsClient) error {
 		slog.Info("no cache bundle found in history")
 		return nil
 	}
-	slog.Info("cache hit", "commit", hitCommit[:min(8, len(hitCommit))], "cache-key", c.CacheKey)
+	slog.Info("cache hit", "commit", hitCommit, "cache-key", c.CacheKey)
 
 	// ── Delta pre-fetch (concurrent with base extraction) ─────────────────────
 	// If --branch is set, kick off a goroutine that stats + downloads the delta
@@ -346,7 +345,10 @@ type extractRule struct {
 // merged rather than replaced.
 func extractBundleZstd(ctx context.Context, r io.Reader, rules []extractRule, defaultDir string) error {
 	zstdCmd := zstdDecompressCmd(ctx)
-	zstdCmd.Stdin = r
+	// Buffer between S3 download and pzstd to decouple network I/O from
+	// decompression. Without this, any momentary pause in pzstd (context
+	// switch, hard block) stalls the S3 read on the synchronous pipe.
+	zstdCmd.Stdin = bufio.NewReaderSize(r, 8<<20)
 
 	var zstdStderr bytes.Buffer
 	zstdCmd.Stderr = &zstdStderr
@@ -745,7 +747,7 @@ func setupLogger(level string) {
 	slog.SetDefault(slog.New(handler))
 }
 
-// bundleFilename converts a cache key to its S3 filename, matching the Ruby bundled-cache-manager.
+// bundleFilename converts a cache key to its S3 filename.
 func bundleFilename(cacheKey string) string {
 	return strings.ReplaceAll(cacheKey, ":", "-") + ".tar.zst"
 }
@@ -843,7 +845,7 @@ func (c *countingBody) Read(p []byte) (int, error) {
 
 func extractTarZstd(ctx context.Context, r io.Reader, dir string) error {
 	zstdCmd := zstdDecompressCmd(ctx)
-	zstdCmd.Stdin = r
+	zstdCmd.Stdin = bufio.NewReaderSize(r, 8<<20)
 
 	var zstdStderr bytes.Buffer
 	zstdCmd.Stderr = &zstdStderr

cmd/gradle-cache/s3.go

@@ -10,6 +10,7 @@ import (
 	"encoding/xml"
 	"fmt"
 	"io"
+	"log/slog"
 	"net/http"
 	"net/url"
 	"os"
@@ -33,20 +34,34 @@ type awsCreds struct {
 
 // s3Client is a minimal AWS S3 client supporting HeadObject, GetObject, and PutObject.
 type s3Client struct {
-	region string
-	creds  awsCreds
-	http   *http.Client
+	region    string
+	creds     awsCreds
+	http      *http.Client
+	chunkSize int64
+	dlWorkers int
 }
 
 func newS3Client(region string) (*s3Client, error) {
 	creds, err := resolveAWSCredentials(region)
 	if err != nil {
 		return nil, errors.Wrap(err, "resolve AWS credentials")
 	}
+	workers := max(defaultDownloadWorkers, runtime.NumCPU())
+
+	transport := &http.Transport{
+		MaxIdleConnsPerHost: workers,
+		WriteBufferSize:     128 << 10,
+		ReadBufferSize:      128 << 10,
+	}
+
+	slog.Debug("s3 client config", "workers", workers, "chunk_mb", defaultDownloadChunkSize>>20)
+
 	return &s3Client{
-		region: region,
-		creds:  creds,
-		http:   &http.Client{},
+		region:    region,
+		creds:     creds,
+		http:      &http.Client{Transport: transport},
+		chunkSize: defaultDownloadChunkSize,
+		dlWorkers: workers,
 	}, nil
 }
 
@@ -71,14 +86,8 @@ func (c *s3Client) stat(ctx context.Context, bucket, key string) (int64, error)
 }
 
 const (
-	// downloadChunkSize is the size of each parallel range request.
-	// 32 MiB gives ~8 in-flight buffers = 256 MiB peak memory, matching the
-	// AWS S3 Transfer Manager default.
-	downloadChunkSize = 32 << 20
-	// downloadWorkers is the number of concurrent range requests.
-	// max(8, NumCPU) saturates S3 bandwidth on CI instances where a single
-	// TCP flow is throttled well below the available network capacity.
-	downloadWorkers = 8
+	defaultDownloadChunkSize = 32 << 20
+	defaultDownloadWorkers   = 16
 )
 
 // get downloads an object and returns its body as a streaming ReadCloser.
@@ -89,7 +98,7 @@ const (
 // The caller must close the returned reader.
 func (c *s3Client) get(ctx context.Context, bucket, key string, size int64) (io.ReadCloser, error) {
 	// Small object or unknown size: single-stream GET.
-	if size <= downloadChunkSize {
+	if size <= c.chunkSize {
 		req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.objectURL(bucket, key), nil)
 		if err != nil {
 			return nil, err
@@ -121,8 +130,8 @@ func (c *s3Client) get(ctx context.Context, bucket, key string, size int64) (io.
 // stays active at full speed. All workers run concurrently, saturating the
 // available S3 bandwidth. Peak memory is numWorkers × downloadChunkSize.
 func (c *s3Client) parallelGet(ctx context.Context, bucket, key string, size int64, w io.Writer) error {
-	numChunks := int((size + downloadChunkSize - 1) / downloadChunkSize)
-	numWorkers := max(downloadWorkers, runtime.NumCPU())
+	numChunks := int((size + c.chunkSize - 1) / c.chunkSize)
+	numWorkers := max(c.dlWorkers, runtime.NumCPU())
 
 	type chunkResult struct {
 		data []byte
@@ -149,8 +158,8 @@ func (c *s3Client) parallelGet(ctx context.Context, bucket, key string, size int
 		go func() {
 			defer wg.Done()
 			for seq := range work {
-				start := int64(seq) * downloadChunkSize
-				end := min(start+downloadChunkSize-1, size-1)
+				start := int64(seq) * c.chunkSize
+				end := min(start+c.chunkSize-1, size-1)
 
 				req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.objectURL(bucket, key), nil)
 				if err != nil {