feat(client): parallel git snapshot download with bounded in-memory streaming

client/chunk_sink.go

@@ -0,0 +1,51 @@
+package client
+
+import (
+	"context"
+	"io"
+
+	"github.com/alecthomas/errors"
+)
+
+// ChunkSink is the destination ParallelGet places fetched chunks into. Place
+// is called once per chunk, concurrently, with the chunk's absolute offset and
+// a body of exactly length bytes (length < 0 means "read the whole body", the
+// single-stream fallback). Place must close body and may block to bound
+// memory, but must abort when ctx is cancelled.
+type ChunkSink interface {
+	Place(ctx context.Context, off, length int64, body io.ReadCloser) error
+}
+
+// DiskSink is a ChunkSink that writes each chunk straight to its offset in an
+// io.WriterAt such as an *os.File. Unlike StreamSink it needs no concurrent
+// reader, so ParallelGet may run to completion synchronously. On error the
+// destination is left partially written and must be discarded by the caller.
+type DiskSink struct{ W io.WriterAt }
+
+// Place streams the chunk straight to its offset in the underlying WriterAt.
+func (d DiskSink) Place(_ context.Context, off, length int64, body io.ReadCloser) error {
+	dst := io.NewOffsetWriter(d.W, off)
+	if length < 0 {
+		_, err := io.Copy(dst, body)
+		return errors.Join(errors.Wrap(err, "write chunk"), body.Close())
+	}
+	n, err := io.Copy(dst, io.LimitReader(body, length))
+	if err != nil {
+		return errors.Join(errors.Errorf("write chunk at offset %d: %w", off, err), body.Close())
+	}
+	if n != length {
+		return errors.Join(errors.Errorf("chunk at offset %d: wrote %d of %d bytes", off, n, length), body.Close())
+	}
+	if overlong(body) {
+		return errors.Join(errors.Errorf("chunk at offset %d: read more than the expected %d bytes", off, length), body.Close())
+	}
+	return errors.WithStack(body.Close())
+}
+
+// overlong reports whether r has bytes left, detecting a body longer than the
+// requested chunk without buffering the excess.
+func overlong(r io.Reader) bool {
+	var probe [1]byte
+	n, _ := io.ReadFull(r, probe[:]) //nolint:errcheck // any byte past the chunk is overlong, regardless of the error
+	return n > 0
+}

client/git.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"strconv"
 	"strings"
 	"sync"
 
@@ -184,43 +185,83 @@ func (c *Client) openGitArtifact(ctx context.Context, repoURL, suffix string) (*
 	}, nil
 }
 
-// GitSnapshotMetadata carries the freshen metadata returned alongside a
-// parallel snapshot download. Commit is the mirror's HEAD SHA at snapshot time
-// (empty for cold serves); BundleURL, when non-empty, points at a delta bundle
-// that brings the snapshot up to the mirror's current HEAD.
-type GitSnapshotMetadata struct {
-	Commit    string
-	BundleURL string
-}
-
-// DownloadGitSnapshot fetches the working-tree snapshot for repoURL into dst,
-// using up to concurrency concurrent range requests of chunkSize bytes each.
-// When concurrency is 1, or the server does not support ranges, it transparently
-// falls back to a single full download. dst is written at non-overlapping
-// offsets via WriteAt (e.g. an *os.File) and the caller owns its lifecycle. It
-// returns the snapshot's freshen metadata, read from the discovery response.
-// Returns os.ErrNotExist when the server has no snapshot available.
-func (c *Client) DownloadGitSnapshot(ctx context.Context, repoURL string, dst io.WriterAt, chunkSize int64, concurrency int) (GitSnapshotMetadata, error) {
+// OpenGitSnapshotParallel downloads the working-tree snapshot for repoURL with
+// up to concurrency concurrent range requests of chunkSize bytes each. It
+// returns as soon as the freshen metadata (Commit, BundleURL) is available,
+// with the download continuing in the background as the caller reads Body, so
+// extraction overlaps the transfer.
+//
+// The caller must Close the returned GitSnapshot, which cancels any in-flight
+// download. A concurrency of 1, or a server without range support, falls back
+// to a single full download. Returns os.ErrNotExist when the server has no
+// snapshot.
+func (c *Client) OpenGitSnapshotParallel(ctx context.Context, repoURL string, chunkSize int64, concurrency int) (*GitSnapshot, error) {
 	endpoint, err := gitEndpointURL(c.baseURL, repoURL, "snapshot.tar.zst")
 	if err != nil {
-		return GitSnapshotMetadata{}, err
+		return nil, err
 	}
-	reader := &gitArtifactRangeReader{client: c, endpoint: endpoint}
-	if err := ParallelGet(ctx, reader, NewKey(repoURL), dst, chunkSize, concurrency); err != nil {
-		return GitSnapshotMetadata{}, errors.Wrap(err, "download snapshot")
+	reader := &gitArtifactRangeReader{client: c, endpoint: endpoint, discovered: make(chan struct{})}
+
+	ctx, cancel := context.WithCancel(ctx)
+
+	sink := NewStreamSink(chunkSize, concurrency)
+	done := make(chan error, 1)
+	go func() {
+		err := ParallelGet(ctx, reader, NewKey(repoURL), sink, chunkSize, concurrency)
+		sink.Done(err)
+		done <- errors.Wrap(err, "download snapshot")
+	}()
+
+	// A small object can finish before discovered is observed, leaving both
+	// channels ready and select picking at random; treat done as "no snapshot"
+	// only when discovery never happened, otherwise let Body drain the buffered
+	// bytes or surface the download error.
+	select {
+	case <-reader.discovered:
+	case err := <-done:
+		if !reader.didDiscover() {
+			cancel()
+			_ = sink.Close() //nolint:errcheck
+			if err == nil {
+				return nil, errors.WithStack(os.ErrNotExist)
+			}
+			return nil, err
+		}
 	}
-	return reader.metadata(), nil
+	headers := reader.discoveryHeaders()
+	return &GitSnapshot{
+		Body:      &cancelReadCloser{ReadCloser: sink, cancel: cancel},
+		Headers:   headers,
+		Commit:    headers.Get(SnapshotCommitHeader),
+		BundleURL: headers.Get(BundleURLHeader),
+	}, nil
+}
+
+// cancelReadCloser cancels the supplied context when Closed, in addition to
+// closing the wrapped reader, so closing a streaming download stops its
+// background goroutine promptly.
+type cancelReadCloser struct {
+	io.ReadCloser
+	cancel context.CancelFunc
+}
+
+func (c *cancelReadCloser) Close() error {
+	c.cancel()
+	return errors.WithStack(c.ReadCloser.Close())
 }
 
 // gitArtifactRangeReader adapts a git artifact endpoint to the RangeReader
-// interface so ParallelGet can fetch it with concurrent range requests. The
-// object's identity is the endpoint URL, so the Key argument is ignored. It
-// records the first response's headers, which carry the snapshot's freshen
-// metadata (delivered on the discovery chunk) that ParallelGet does not surface.
+// interface. The object's identity is the endpoint URL, so the Key argument is
+// ignored. It records the first response's headers, which carry the snapshot's
+// freshen metadata that ParallelGet does not surface.
 type gitArtifactRangeReader struct {
 	client   *Client
 	endpoint string
 
+	// discovered, when non-nil, is closed once the first response's headers
+	// are recorded.
+	discovered chan struct{}
+
 	mu        sync.Mutex
 	discovery http.Header
 }
@@ -253,23 +294,39 @@ func (g *gitArtifactRangeReader) Open(ctx context.Context, _ Key, opts ...Reques
 	}
 }
 
+// didDiscover reports whether the first response's headers have been recorded.
+func (g *gitArtifactRangeReader) didDiscover() bool {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	return g.discovery != nil
+}
+
 // recordDiscovery stores the first response's headers so the freshen metadata
 // they carry survives after the bodies are consumed.
 func (g *gitArtifactRangeReader) recordDiscovery(h http.Header) {
 	g.mu.Lock()
 	defer g.mu.Unlock()
 	if g.discovery == nil {
 		g.discovery = h.Clone()
+		if g.discovered != nil {
+			close(g.discovered)
+		}
 	}
 }
 
-func (g *gitArtifactRangeReader) metadata() GitSnapshotMetadata {
+// discoveryHeaders returns the first response's headers with transport-layer
+// headers stripped. The discovery response is a 206 for only the first chunk,
+// so Content-Range/Content-Length are rewritten to describe the full
+// reassembled object streamed on GitSnapshot.Body.
+func (g *gitArtifactRangeReader) discoveryHeaders() http.Header {
 	g.mu.Lock()
 	defer g.mu.Unlock()
-	return GitSnapshotMetadata{
-		Commit:    g.discovery.Get(SnapshotCommitHeader),
-		BundleURL: g.discovery.Get(BundleURLHeader),
+	headers := filterHeaders(g.discovery, transportHeaders...)
+	if total, ok := parseContentRangeTotal(headers.Get("Content-Range")); ok {
+		headers.Del("Content-Range")
+		headers.Set("Content-Length", strconv.FormatInt(total, 10))
 	}
+	return headers
 }
 
 // gitEndpointURL builds a /git/{host}/{repoPath}/{suffix} URL from a cachew

client/git_test.go

@@ -8,6 +8,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"os"
+	"strconv"
 	"strings"
 	"sync/atomic"
 	"testing"
@@ -148,7 +149,6 @@ func TestOpenGitBundle(t *testing.T) {
 
 	api := client.NewWithHTTPClient(srv.URL, srv.Client())
 
-	// Root-relative URL (as returned in X-Cachew-Bundle-Url).
 	body, err := api.OpenGitBundle(context.Background(),
 		"/git/github.com/org/repo/snapshot.bundle?base=abc")
 	assert.NoError(t, err)
@@ -157,7 +157,6 @@ func TestOpenGitBundle(t *testing.T) {
 	assert.NoError(t, body.Close())
 	assert.Equal(t, "bundle-bytes", string(data))
 
-	// Absolute URL.
 	body, err = api.OpenGitBundle(context.Background(),
 		srv.URL+"/git/github.com/org/repo/snapshot.bundle?base=abc")
 	assert.NoError(t, err)
@@ -175,7 +174,7 @@ func TestOpenGitBundleNotFound(t *testing.T) {
 	assert.True(t, errors.Is(err, os.ErrNotExist))
 }
 
-func TestDownloadGitSnapshotParallel(t *testing.T) {
+func TestOpenGitSnapshotParallel(t *testing.T) {
 	body := make([]byte, 1000)
 	for i := range body {
 		body[i] = byte(i % 251)
@@ -190,51 +189,101 @@ func TestDownloadGitSnapshotParallel(t *testing.T) {
 		w.Header().Set("ETag", etag)
 		w.Header().Set(client.SnapshotCommitHeader, "deadbeef")
 		w.Header().Set(client.BundleURLHeader, "/git/github.com/org/repo/snapshot.bundle?base=deadbeef")
-		// ServeContent honours Range/If-Range against the ETag set above, so it
-		// returns 206 + Content-Range for the chunked requests ParallelGet makes.
 		http.ServeContent(w, r, "snapshot.tar.zst", time.Time{}, bytes.NewReader(body))
 	}))
 	defer srv.Close()
 
 	api := client.NewWithHTTPClient(srv.URL, srv.Client())
-	var dst bufferAt
-	// A 128-byte chunk over a 1000-byte body forces multiple chunks, exercising
-	// concurrent range reassembly.
-	meta, err := api.DownloadGitSnapshot(context.Background(), "https://github.com/org/repo", &dst, 128, 4)
+	snap, err := api.OpenGitSnapshotParallel(context.Background(), "https://github.com/org/repo", 128, 4)
 	assert.NoError(t, err)
-	assert.Equal(t, body, dst.buf)
-	assert.Equal(t, "deadbeef", meta.Commit)
-	assert.Equal(t, "/git/github.com/org/repo/snapshot.bundle?base=deadbeef", meta.BundleURL)
+	defer snap.Close()
+
+	assert.Equal(t, "deadbeef", snap.Commit)
+	assert.Equal(t, "/git/github.com/org/repo/snapshot.bundle?base=deadbeef", snap.BundleURL)
+
+	assert.Equal(t, "", snap.Headers.Get("Content-Range"))
+	assert.Equal(t, strconv.Itoa(len(body)), snap.Headers.Get("Content-Length"))
+
+	got, err := io.ReadAll(snap.Body)
+	assert.NoError(t, err)
+	assert.Equal(t, body, got)
 	assert.True(t, requests.Load() > 1, "expected multiple range requests, got %d", requests.Load())
 }
 
-func TestDownloadGitSnapshotFallsBackWithoutRange(t *testing.T) {
+func TestOpenGitSnapshotParallelSmallObject(t *testing.T) {
+	body := []byte("a tiny snapshot that fits in the first chunk")
+	const etag = `"snap-small"`
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/zstd")
+		w.Header().Set("ETag", etag)
+		w.Header().Set(client.SnapshotCommitHeader, "beef")
+		http.ServeContent(w, r, "snapshot.tar.zst", time.Time{}, bytes.NewReader(body))
+	}))
+	defer srv.Close()
+
+	api := client.NewWithHTTPClient(srv.URL, srv.Client())
+	for range 200 {
+		snap, err := api.OpenGitSnapshotParallel(context.Background(), "https://github.com/org/repo", 4096, 4)
+		assert.NoError(t, err)
+		assert.Equal(t, "beef", snap.Commit)
+		got, err := io.ReadAll(snap.Body)
+		assert.NoError(t, err)
+		assert.Equal(t, body, got)
+		assert.NoError(t, snap.Close())
+	}
+}
+
+func TestOpenGitSnapshotParallelFallsBackWithoutRange(t *testing.T) {
 	body := []byte("full body, server ignores ranges")
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
-		// No ETag and no range handling: always answer the full object with 200,
-		// mimicking an older server. ParallelGet must fall back to a single read.
 		w.Header().Set("Content-Type", "application/zstd")
 		w.Header().Set(client.SnapshotCommitHeader, "cafe")
 		_, _ = w.Write(body) //nolint:errcheck
 	}))
 	defer srv.Close()
 
 	api := client.NewWithHTTPClient(srv.URL, srv.Client())
-	var dst bufferAt
-	meta, err := api.DownloadGitSnapshot(context.Background(), "https://github.com/org/repo", &dst, 8, 4)
+	snap, err := api.OpenGitSnapshotParallel(context.Background(), "https://github.com/org/repo", 8, 4)
 	assert.NoError(t, err)
-	assert.Equal(t, body, dst.buf)
-	assert.Equal(t, "cafe", meta.Commit)
+	defer snap.Close()
+
+	assert.Equal(t, "cafe", snap.Commit)
+	got, err := io.ReadAll(snap.Body)
+	assert.NoError(t, err)
+	assert.Equal(t, body, got)
 }
 
-func TestDownloadGitSnapshotNotFound(t *testing.T) {
+func TestOpenGitSnapshotParallelNotFound(t *testing.T) {
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		http.NotFound(w, r)
 	}))
 	defer srv.Close()
 
 	api := client.NewWithHTTPClient(srv.URL, srv.Client())
-	var dst bufferAt
-	_, err := api.DownloadGitSnapshot(context.Background(), "https://github.com/org/repo", &dst, 8, 4)
+	_, err := api.OpenGitSnapshotParallel(context.Background(), "https://github.com/org/repo", 8, 4)
 	assert.True(t, errors.Is(err, os.ErrNotExist))
 }
+
+func TestOpenGitSnapshotParallelCloseStopsDownload(t *testing.T) {
+	body := make([]byte, 1<<20)
+	for i := range body {
+		body[i] = byte(i % 251)
+	}
+	const etag = `"snap-v1"`
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/zstd")
+		w.Header().Set("ETag", etag)
+		http.ServeContent(w, r, "snapshot.tar.zst", time.Time{}, bytes.NewReader(body))
+	}))
+	defer srv.Close()
+
+	api := client.NewWithHTTPClient(srv.URL, srv.Client())
+	snap, err := api.OpenGitSnapshotParallel(context.Background(), "https://github.com/org/repo", 4096, 4)
+	assert.NoError(t, err)
+
+	buf := make([]byte, 16)
+	_, err = io.ReadFull(snap.Body, buf)
+	assert.NoError(t, err)
+	assert.NoError(t, snap.Close())
+}