fix(snapshot): bound S3 parallelGet and make snapshot serving range-aware

The S3 full-object download path (parallelGet) let workers read every chunk
into memory ahead of a slow consumer, so a single large snapshot served to a
slow client could buffer the whole object in the Go heap. Bound it with a
window of in-flight chunks (peak now numWorkers x chunkSize).

The git snapshot endpoint also ignored client Range headers, so it always
served the full object via the unbounded path. Forward Range/If-Range to the
cache on the cache-hit path, return 206/Content-Range (416 when not
satisfiable), and advertise Accept-Ranges, so clients can fetch with bounded
parallel range requests (client.ParallelGet) instead of one full GET.

Co-authored-by: Amp <amp@ampcode.com>
Amp-Thread-ID: https://ampcode.com/threads/T-019ef6fe-55c6-71e8-96dc-ca3ef4301d36

Author: worstell

internal/cache/s3_parallel_get.go

@@ -69,9 +69,15 @@ func (c *cancelReadCloser) Close() error {
 
 // parallelGet downloads an S3 object in parallel chunks and writes them in
 // order to w. Each worker downloads its chunk into memory so the TCP
-// connection stays active at full speed. Peak memory: numWorkers × chunkSize.
-// All chunk requests are pinned to the given etag to ensure consistency.
-// An errgroup cancels all workers on the first error from any goroutine.
+// connection stays active at full speed. All chunk requests are pinned to the
+// given etag to ensure consistency. An errgroup cancels all workers on the
+// first error from any goroutine.
+//
+// A window token is acquired before each chunk is fetched and released by the
+// writer once that chunk has been consumed, bounding the number of
+// downloaded-but-unwritten chunks to numWorkers. Without it, workers race ahead
+// of a slow consumer and buffer the whole object in memory. Peak memory is
+// therefore numWorkers × chunkSize.
 func (s *S3) parallelGet(ctx context.Context, bucket, objectName string, size int64, etag string, w io.Writer) error {
 	chunkSize := int64(s.config.DownloadPartSizeMB) << 20 // #nosec G115 -- DownloadPartSizeMB is a small operator-supplied tuning value.
 	numChunks := int((size + chunkSize - 1) / chunkSize)
@@ -83,6 +89,10 @@ func (s *S3) parallelGet(ctx context.Context, bucket, objectName string, size in
 		results[i] = make(chan []byte, 1)
 	}
 
+	// Bounds downloaded-but-unwritten chunks to numWorkers: a token is held from
+	// before a chunk is fetched until the writer has consumed it.
+	window := make(chan struct{}, numWorkers)
+
 	// Work queue of chunk indices.
 	work := make(chan int, numChunks)
 	for i := range numChunks {
@@ -97,7 +107,11 @@ func (s *S3) parallelGet(ctx context.Context, bucket, objectName string, size in
 	for range numWorkers {
 		eg.Go(func() error {
 			for seq := range work {
-				if egCtx.Err() != nil {
+				// Block until the writer has caught up enough to admit another
+				// in-flight chunk, or the group is cancelled.
+				select {
+				case window <- struct{}{}:
+				case <-egCtx.Done():
 					return egCtx.Err()
 				}
 
@@ -142,6 +156,7 @@ func (s *S3) parallelGet(ctx context.Context, bucket, objectName string, size in
 				if _, err := w.Write(data); err != nil {
 					return errors.Wrap(err, "write chunk")
 				}
+				<-window // Admit another in-flight chunk now this one is consumed.
 			case <-egCtx.Done():
 				return egCtx.Err()
 			}

internal/cache/s3_test.go

@@ -1,7 +1,9 @@
 package cache_test
 
 import (
+	"bytes"
 	"context"
+	"io"
 	"log/slog"
 	"os"
 	"testing"
@@ -71,6 +73,50 @@ func TestS3ContextCancellationAbortsUpload(t *testing.T) {
 	assert.IsError(t, err, os.ErrNotExist)
 }
 
+// TestS3ParallelGetMultiChunk exercises the multi-chunk parallel download path
+// (small part size + concurrency over a multi-part object) and verifies the
+// reassembled bytes are correct and in order. This guards the bounded-window
+// reordering in parallelGet against off-by-one or interleaving regressions.
+func TestS3ParallelGetMultiChunk(t *testing.T) {
+	bucket := s3clienttest.Start(t)
+	s3clienttest.CleanBucket(t, bucket)
+	_, ctx := logging.Configure(t.Context(), logging.Config{})
+
+	clientProvider := s3client.NewClientProvider(ctx, s3client.Config{Endpoint: s3clienttest.Addr, UseSSL: false})
+	c, err := cache.NewS3(ctx, cache.S3Config{
+		Bucket:              bucket,
+		MaxTTL:              time.Hour,
+		UploadPartSizeMB:    5,
+		UploadConcurrency:   1, // single-threaded upload; concurrent upload exercises a minio-go crc64 race unrelated to this test
+		DownloadPartSizeMB:  5, // 5 MiB chunks
+		DownloadConcurrency: 3, // fewer workers than chunks, so the window throttles
+	}, clientProvider)
+	assert.NoError(t, err)
+	defer c.Close()
+
+	// 17 MiB => 4 chunks at 5 MiB, with a deterministic pseudo-random pattern so
+	// any mis-ordered or duplicated chunk is detected on read-back.
+	want := make([]byte, 17<<20)
+	for i := range want {
+		want[i] = byte(i*2654435761 + i>>3) //nolint:gosec // arbitrary deterministic fill
+	}
+
+	key := cache.NewKey("parallel-get-multichunk")
+	w, err := c.Create(ctx, key, nil, time.Hour)
+	assert.NoError(t, err)
+	_, err = w.Write(want)
+	assert.NoError(t, err)
+	assert.NoError(t, w.Close())
+
+	rc, _, err := c.Open(ctx, key)
+	assert.NoError(t, err)
+	got, err := io.ReadAll(rc)
+	assert.NoError(t, err)
+	assert.NoError(t, rc.Close())
+	assert.Equal(t, len(want), len(got))
+	assert.True(t, bytes.Equal(want, got), "reassembled object must match the original bytes")
+}
+
 func TestS3CacheSoak(t *testing.T) {
 	if os.Getenv("SOAK_TEST") == "" {
 		t.Skip("Skipping soak test; set SOAK_TEST=1 to run")

internal/httputil/conditional.go

@@ -32,6 +32,23 @@ func ConditionalOptions(r *http.Request) []client.RequestOption {
 	return opts
 }
 
+// RangeOptions extracts only the Range/If-Range options from r, for callers
+// that evaluate If-Match/If-None-Match separately (e.g. via CheckConditionals)
+// and so must not double-handle them by forwarding the full set to Open.
+// If-Range is only meaningful alongside Range, so it is dropped when Range is
+// absent.
+func RangeOptions(r *http.Request) []client.RequestOption {
+	v := r.Header.Get("Range")
+	if v == "" {
+		return nil
+	}
+	opts := []client.RequestOption{func(o *client.RequestOptions) { o.Range = v }}
+	if ir := r.Header.Get("If-Range"); ir != "" {
+		opts = append(opts, client.IfRange(ir))
+	}
+	return opts
+}
+
 // CheckConditionals evaluates RFC 7232 If-Match and If-None-Match precondition
 // headers on r against etag. It returns 0 when all preconditions pass,
 // otherwise the HTTP status the caller should send: 412 Precondition Failed for

internal/strategy/git/snapshot.go

@@ -320,7 +320,18 @@ func (s *Strategy) handleSnapshotRequest(w http.ResponseWriter, r *http.Request,
 	}
 	s.maybeBackgroundFetch(repo)
 
-	reader, headers, err := s.cache.Open(ctx, cacheKey)
+	// Forward only Range/If-Range here; If-Match/If-None-Match are evaluated
+	// against the served body below via CheckConditionals.
+	reader, headers, err := s.cache.Open(ctx, cacheKey, httputil.RangeOptions(r)...)
+	if errors.Is(err, cache.ErrRangeNotSatisfiable) {
+		w.Header().Set("Content-Type", "application/zstd")
+		w.Header().Set("Accept-Ranges", "bytes")
+		if cr := headers.Get("Content-Range"); cr != "" {
+			w.Header().Set("Content-Range", cr)
+		}
+		w.WriteHeader(http.StatusRequestedRangeNotSatisfiable)
+		return
+	}
 	if err != nil && !errors.Is(err, os.ErrNotExist) {
 		logger.ErrorContext(ctx, "Failed to open snapshot from cache", "upstream", upstreamURL, "error", err)
 		http.Error(w, "Internal server error", http.StatusInternalServerError)
@@ -517,23 +528,29 @@ func (s *Strategy) handleBundleRequest(w http.ResponseWriter, r *http.Request, h
 }
 
 func (s *Strategy) serveSnapshotWithBundle(ctx context.Context, w http.ResponseWriter, r *http.Request, reader io.ReadCloser, headers http.Header, repo *gitclone.Repository, upstreamURL, repoName string, start time.Time) error {
-	snapshotCommit := headers.Get("X-Cachew-Snapshot-Commit")
-	mirrorHead := s.getMirrorHead(ctx, repo)
-
-	// Forward the snapshot commit to the client so it knows whether the
-	// snapshot is fresh (no bundle URL = already at HEAD, skip freshen).
-	if snapshotCommit != "" {
-		w.Header().Set("X-Cachew-Snapshot-Commit", snapshotCommit)
+	// A satisfied byte range is served as-is. Bundle negotiation and full-body
+	// conditionals apply only to whole-snapshot downloads, so a partial read
+	// (e.g. a client.ParallelGet chunk) skips them and returns 206 directly.
+	if cr := headers.Get("Content-Range"); cr != "" {
+		applySnapshotCacheHeaders(w, headers)
+		w.Header().Set("Accept-Ranges", "bytes")
+		w.Header().Set("Content-Range", cr)
+		// Ranged clients (client.ParallelGet) read the freshen metadata from the
+		// discovery chunk, so it must be present on partial responses too.
+		s.setSnapshotMetadataHeaders(ctx, w, headers, repo, upstreamURL)
+		w.WriteHeader(http.StatusPartialContent)
+		n, err := io.Copy(w, reader)
+		s.metrics.recordSnapshotServe(ctx, "cache_range", repoName, n, time.Since(start))
+		if span := trace.SpanFromContext(ctx); span.SpanContext().IsValid() {
+			span.SetAttributes(attribute.String("cachew.source", "cache_range"), attribute.Int64("cachew.bytes", n))
+		}
+		return errors.Wrap(err, "stream snapshot range")
 	}
 
-	if snapshotCommit != "" && mirrorHead != "" && snapshotCommit != mirrorHead {
-		repoPath, err := gitclone.RepoPathFromURL(upstreamURL)
-		if err == nil {
-			bundleURL := fmt.Sprintf("/git/%s/snapshot.bundle?base=%s", repoPath, snapshotCommit)
-			w.Header().Set("X-Cachew-Bundle-Url", bundleURL)
-		}
+	snapshotCommit, bundleURL := s.setSnapshotMetadataHeaders(ctx, w, headers, repo, upstreamURL)
 
-		// Proactively generate and cache the bundle so any pod can serve it.
+	// Proactively generate and cache the advertised bundle so any pod can serve it.
+	if bundleURL != "" {
 		go func() {
 			bgCtx := context.WithoutCancel(ctx)
 			logger := logging.FromContext(bgCtx)
@@ -550,6 +567,7 @@ func (s *Strategy) serveSnapshotWithBundle(ctx context.Context, w http.ResponseW
 	}
 
 	applySnapshotCacheHeaders(w, headers)
+	w.Header().Set("Accept-Ranges", "bytes")
 
 	// Honour conditional GETs against the advertised ETag. ServeContent does this
 	// natively for *os.File readers, but cache backends returning non-file readers
@@ -569,6 +587,33 @@ func (s *Strategy) serveSnapshotWithBundle(ctx context.Context, w http.ResponseW
 	return errors.Wrap(err, "stream snapshot")
 }
 
+// setSnapshotMetadataHeaders advertises the snapshot's commit and, when the
+// snapshot trails the mirror's HEAD, the delta-bundle URL clients use to
+// fast-forward. Shared by the full and ranged serve paths so ranged clients
+// (client.ParallelGet) receive the same freshen metadata on the discovery
+// chunk. It returns the snapshot commit and the bundle URL it set (empty when
+// the snapshot is already at HEAD), so callers can decide whether to
+// pre-generate the bundle.
+func (s *Strategy) setSnapshotMetadataHeaders(ctx context.Context, w http.ResponseWriter, headers http.Header, repo *gitclone.Repository, upstreamURL string) (snapshotCommit, bundleURL string) {
+	snapshotCommit = headers.Get("X-Cachew-Snapshot-Commit")
+	if snapshotCommit == "" {
+		return "", ""
+	}
+	w.Header().Set("X-Cachew-Snapshot-Commit", snapshotCommit)
+
+	mirrorHead := s.getMirrorHead(ctx, repo)
+	if mirrorHead == "" || snapshotCommit == mirrorHead {
+		return snapshotCommit, ""
+	}
+	repoPath, err := gitclone.RepoPathFromURL(upstreamURL)
+	if err != nil {
+		return snapshotCommit, ""
+	}
+	bundleURL = fmt.Sprintf("/git/%s/snapshot.bundle?base=%s", repoPath, snapshotCommit)
+	w.Header().Set("X-Cachew-Bundle-Url", bundleURL)
+	return snapshotCommit, bundleURL
+}
+
 // applySnapshotCacheHeaders forwards the cached snapshot's validators so clients
 // can revalidate (ETag) and size the transfer (Content-Length). Content-Type is
 // fixed for snapshots regardless of what the cache backend recorded.

internal/strategy/git/snapshot_test.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -930,3 +931,75 @@ func TestSnapshotRemoteURLUsesUpstreamURL(t *testing.T) {
 	assert.NoError(t, err, string(output))
 	assert.Equal(t, upstreamURL+"\n", string(output))
 }
+
+func TestSnapshotGetHonorsRange(t *testing.T) {
+	if _, err := exec.LookPath("git"); err != nil {
+		t.Skip("git not found in PATH")
+	}
+
+	_, ctx := logging.Configure(context.Background(), logging.Config{})
+	tmpDir := t.TempDir()
+	mirrorRoot := filepath.Join(tmpDir, "mirrors")
+	upstreamURL := "https://github.com/org/repo"
+	mirrorPath := filepath.Join(mirrorRoot, "github.com", "org", "repo")
+	createTestMirrorRepo(t, mirrorPath)
+
+	// The memory cache returns a non-*os.File reader, so range handling must come
+	// from the strategy forwarding Range to Open rather than http.ServeContent.
+	memCache, err := cache.NewMemory(ctx, cache.MemoryConfig{MaxTTL: time.Hour})
+	assert.NoError(t, err)
+	mux := newTestMux()
+
+	cm := gitclone.NewManagerProvider(ctx, gitclone.Config{MirrorRoot: mirrorRoot}, nil)
+	s, err := git.New(ctx, git.Config{}, newTestScheduler(ctx, t), memCache, mux, cm, func() (*githubapp.TokenManager, error) { return nil, nil }) //nolint:nilnil
+	assert.NoError(t, err)
+
+	manager, err := cm()
+	assert.NoError(t, err)
+	repo, err := manager.GetOrCreate(ctx, upstreamURL)
+	assert.NoError(t, err)
+
+	waitForReady(t, s)
+	err = s.GenerateAndUploadSnapshot(ctx, repo)
+	assert.NoError(t, err)
+
+	handler := mux.handlers["GET /git/{host}/{path...}"]
+	assert.NotZero(t, handler)
+
+	get := func(rangeHeader string) *httptest.ResponseRecorder {
+		req := httptest.NewRequest(http.MethodGet, "/git/github.com/org/repo/snapshot.tar.zst", nil)
+		req = req.WithContext(ctx)
+		req.SetPathValue("host", "github.com")
+		req.SetPathValue("path", "org/repo/snapshot.tar.zst")
+		if rangeHeader != "" {
+			req.Header.Set("Range", rangeHeader)
+		}
+		w := httptest.NewRecorder()
+		handler.ServeHTTP(w, req)
+		return w
+	}
+
+	// Full GET advertises range support and yields the whole body.
+	full := get("")
+	assert.Equal(t, 200, full.Code)
+	assert.Equal(t, "bytes", full.Header().Get("Accept-Ranges"))
+	body := full.Body.Bytes()
+	assert.True(t, len(body) > 4, "snapshot body should be larger than the test range")
+	commit := full.Header().Get("X-Cachew-Snapshot-Commit")
+	assert.NotZero(t, commit, "full GET should advertise the snapshot commit")
+
+	// A satisfiable range returns 206 with the matching bytes and Content-Range.
+	partial := get("bytes=0-3")
+	assert.Equal(t, http.StatusPartialContent, partial.Code)
+	assert.Equal(t, "bytes", partial.Header().Get("Accept-Ranges"))
+	assert.Equal(t, "bytes 0-3/"+strconv.Itoa(len(body)), partial.Header().Get("Content-Range"))
+	assert.Equal(t, body[:4], partial.Body.Bytes())
+	// Ranged clients must receive the same freshen metadata as a full GET so
+	// they can apply a delta bundle after a parallel download.
+	assert.Equal(t, commit, partial.Header().Get("X-Cachew-Snapshot-Commit"))
+
+	// A range beyond the object is not satisfiable.
+	tooBig := get("bytes=" + strconv.Itoa(len(body)+10) + "-" + strconv.Itoa(len(body)+20))
+	assert.Equal(t, http.StatusRequestedRangeNotSatisfiable, tooBig.Code)
+	assert.Equal(t, "bytes */"+strconv.Itoa(len(body)), tooBig.Header().Get("Content-Range"))
+}