From 310ef7c416a61b9396c1e64aa72416a75a062d94 Mon Sep 17 00:00:00 2001 From: Elizabeth Worstell Date: Fri, 10 Jul 2026 16:16:06 -0700 Subject: [PATCH] feat(client): parallel git snapshot download with bounded in-memory streaming Download git snapshots with up to N concurrent range requests, reassembled in order by a StreamSink and piped straight into zstd/tar extraction so transfer and extraction overlap. Peak buffering is capped at 2*concurrency*chunkSize regardless of snapshot size: fetchers block once they are a full window ahead of the consumer. Chunks are pinned to the discovery ETag so a mid-download rewrite is rejected rather than spliced; range-ignoring backends, missing ETags, and small objects fall back to a single streaming request. Chunk fetches retry transient per-replica failures, and extraction goes to a staging directory renamed into place only on success. --- client/chunk_sink.go | 51 +++++ client/git.go | 117 +++++++--- client/git_test.go | 93 ++++++-- client/parallel_get.go | 215 +++++++++++------- client/parallel_get_test.go | 330 ++++++++++++++++++++++++---- client/stream_sink.go | 208 ++++++++++++++++++ cmd/cachew/git.go | 118 +++------- internal/cache/parallel_get.go | 9 +- internal/cache/parallel_get_test.go | 52 ++--- 9 files changed, 899 insertions(+), 294 deletions(-) create mode 100644 client/chunk_sink.go create mode 100644 client/stream_sink.go diff --git a/client/chunk_sink.go b/client/chunk_sink.go new file mode 100644 index 00000000..74f735ed --- /dev/null +++ b/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 +} diff --git a/client/git.go b/client/git.go index 999509d9..553ea8e2 100644 --- a/client/git.go +++ b/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,6 +294,13 @@ 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) { @@ -260,16 +308,25 @@ func (g *gitArtifactRangeReader) recordDiscovery(h http.Header) { 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 diff --git a/client/git_test.go b/client/git_test.go index 0e2ac435..2b0858b2 100644 --- a/client/git_test.go +++ b/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,29 +189,53 @@ 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 @@ -220,21 +243,47 @@ func TestDownloadGitSnapshotFallsBackWithoutRange(t *testing.T) { 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()) +} diff --git a/client/parallel_get.go b/client/parallel_get.go index 9db919b7..9a2f1aa5 100644 --- a/client/parallel_get.go +++ b/client/parallel_get.go @@ -6,127 +6,180 @@ import ( "net/http" "strconv" "strings" + "sync/atomic" "github.com/alecthomas/errors" "golang.org/x/sync/errgroup" ) +// discoveryAttempts bounds retries when a response ignores the requested +// range, e.g. a load-balanced replica serving a streaming miss. +const discoveryAttempts = 4 + // RangeReader is the subset of cache operations ParallelGet needs: a -// conditional, Range-capable Open. Both *Client and the cache.Cache interface -// satisfy it, so ParallelGet can drive either a remote cache server or a local -// backend. +// conditional, Range-capable Open. Both *Client and cache.Cache satisfy it. type RangeReader interface { Open(ctx context.Context, key Key, opts ...RequestOption) (io.ReadCloser, http.Header, error) } -// ParallelGet downloads an object from any Range-capable RangeReader into dst, -// fetching it in chunkSize-byte chunks concurrently (up to concurrency requests -// in flight) and writing each chunk at its offset via dst.WriteAt. Latency-bound -// backends such as a remote cache can saturate bandwidth with overlapping reads. -// -// The first chunk is fetched with a ranged Open, whose response yields both the -// total size (from Content-Range) and the object's ETag; every remaining chunk -// is then requested with IfMatch pinned to that ETag. If the object changes -// mid-download, the chunk is rejected with a bodiless ErrPreconditionFailed -// (412) and ParallelGet returns an error rather than splicing bytes from two -// revisions; a server that ignores If-Match is caught by verifying each chunk's -// response ETag. A missing or truncated chunk is likewise reported as an error, -// so a partially written dst must be discarded by the caller on failure. An -// object with no ETag to pin to cannot be kept revision-safe across chunks, so -// it falls back to a single full read instead of parallelising. A concurrency of -// 1 likewise reads the whole object in one request, since chunking a single -// worker would only serialise ranged GETs for no benefit. +// ParallelGet downloads an object in chunkSize-byte chunks with up to +// concurrency requests in flight, handing each chunk to sink. Pass a +// StreamSink to reassemble the in-order byte stream (read it concurrently +// while ParallelGet runs), or a DiskSink to scatter chunks into an +// io.WriterAt. // -// dst is written via concurrent WriteAt calls at non-overlapping offsets; the -// caller owns dst's lifecycle (open, close, cleanup) and need not pre-size it, -// as WriteAt extends the destination. -func ParallelGet(ctx context.Context, c RangeReader, key Key, dst io.WriterAt, chunkSize int64, concurrency int) error { +// Chunks are pinned to the discovery response's ETag so a mid-download rewrite +// is rejected rather than spliced. No ETag, a range-ignoring backend, an +// object that fits in the first chunk, or concurrency 1 all fall back to a +// single full read. +func ParallelGet(ctx context.Context, c RangeReader, key Key, sink ChunkSink, chunkSize int64, concurrency int) error { + if max(concurrency, 1) == 1 { + return fullRead(ctx, c, key, sink) + } if chunkSize <= 0 { return errors.Errorf("parallel get: chunk size must be positive, got %d", chunkSize) } - concurrency = max(concurrency, 1) - - if concurrency == 1 { - return fullRead(ctx, c, key, dst) - } - rc, headers, err := c.Open(ctx, key, Range(0, chunkSize)) - if errors.Is(err, ErrRangeNotSatisfiable) { - return nil - } - if err != nil { - return errors.Wrap(err, "parallel get: open first chunk") + // The first ranged Open delivers chunk zero and reveals the total size and + // ETag that pin the rest. Retry when the range is ignored: each attempt may + // route to a different replica. + var ( + rc io.ReadCloser + headers http.Header + err error + ) + etag := "" + total, hasRange := int64(0), false + for attempt := range discoveryAttempts { + rc, headers, err = c.Open(ctx, key, Range(0, chunkSize)) + if errors.Is(err, ErrRangeNotSatisfiable) { + return nil // Empty object: nothing to place. + } + if err != nil { + return errors.Wrap(err, "parallel get: open first chunk") + } + etag = headers.Get(ETagKey) + total, hasRange = parseContentRangeTotal(headers.Get("Content-Range")) + if hasRange || attempt == discoveryAttempts-1 { + break + } + if err := rc.Close(); err != nil { + return errors.Wrap(err, "parallel get: close range-ignoring discovery reader") + } } - etag := headers.Get(ETagKey) - total, hasRange := parseContentRangeTotal(headers.Get("Content-Range")) - - firstLen := min(chunkSize, total) + // Range ignored or the object fits in the first chunk: this response + // already carries the whole object (negative length = unknown size). if !hasRange { - firstLen = -1 + return errors.Wrap(sink.Place(ctx, 0, -1, rc), "parallel get") } - if !hasRange || total <= chunkSize { - return errors.Wrap(writeChunkAt(dst, 0, firstLen, rc), "parallel get") + if total <= chunkSize { + return errors.Wrap(sink.Place(ctx, 0, total, rc), "parallel get") } + // Without an ETag there is nothing to pin chunks to, so a rewrite could be + // spliced undetected; fall back to a single revision-consistent read. if etag == "" { if err := rc.Close(); err != nil { return errors.Wrap(err, "parallel get: close discovery reader") } - return fullRead(ctx, c, key, dst) + return fullRead(ctx, c, key, sink) } - numChunks := int((total + chunkSize - 1) / chunkSize) + numChunks := (total + chunkSize - 1) / chunkSize eg, egCtx := errgroup.WithContext(ctx) - eg.SetLimit(concurrency) - eg.Go(func() error { return writeChunkAt(dst, 0, firstLen, rc) }) - for seq := 1; seq < numChunks; seq++ { - if egCtx.Err() != nil { - break + + var nextSeq atomic.Int64 + nextSeq.Store(1) + worker := func() error { + for { + if egCtx.Err() != nil { + return errors.WithStack(egCtx.Err()) + } + seq := nextSeq.Add(1) - 1 + if seq >= numChunks { + return nil + } + start := seq * chunkSize + end := min(start+chunkSize, total) + body, err := openChunk(egCtx, c, key, etag, start, end) + if err != nil { + return err + } + if err := sink.Place(egCtx, start, end-start, body); err != nil { + return errors.WithStack(err) + } } - start := int64(seq) * chunkSize - end := min(start+chunkSize, total) - eg.Go(func() error { return fetchChunk(egCtx, c, key, dst, start, end, etag) }) } - return errors.Wrap(eg.Wait(), "parallel get") -} -func fullRead(ctx context.Context, c RangeReader, key Key, dst io.WriterAt) error { - rc, _, err := c.Open(ctx, key) - if err != nil { - return errors.Wrap(err, "parallel get: full read") + // The first worker places chunk zero from the already-open discovery body, + // keeping in-flight requests at concurrency rather than concurrency+1. rc + // was opened with the parent ctx, so close it on egCtx cancellation or a + // failed sibling would leave Place blocked; double Close is harmless. + eg.Go(func() error { + stop := context.AfterFunc(egCtx, func() { _ = rc.Close() }) //nolint:errcheck + err := sink.Place(egCtx, 0, min(chunkSize, total), rc) + stop() + if err != nil { + return errors.WithStack(err) + } + return worker() + }) + for range min(int64(concurrency), numChunks) - 1 { + eg.Go(worker) } - return errors.Wrap(writeChunkAt(dst, 0, -1, rc), "parallel get") + return errors.Wrap(eg.Wait(), "parallel get") } -func fetchChunk(ctx context.Context, c RangeReader, key Key, dst io.WriterAt, start, end int64, etag string) error { - rc, headers, err := c.Open(ctx, key, Range(start, end), IfMatch(etag)) - if errors.Is(err, ErrPreconditionFailed) { - return errors.Errorf("open range %d-%d: object changed during read: %w", start, end, err) - } - if err != nil { - return errors.Errorf("open range %d-%d: %w", start, end, err) - } - if got := headers.Get(ETagKey); got != etag { - return errors.Join( - errors.Errorf("object changed during read at offset %d: etag %q != %q", start, got, etag), - rc.Close(), - ) +// openChunk fetches one chunk's range pinned to the discovery ETag. +// Range-ignoring responses are retried, as each attempt may route to a +// different replica; an ETag mismatch or precondition failure is fatal since +// the object changed mid-download. +func openChunk(ctx context.Context, c RangeReader, key Key, etag string, start, end int64) (io.ReadCloser, error) { + for attempt := 0; ; attempt++ { + body, h, err := c.Open(ctx, key, Range(start, end), IfMatch(etag)) + if errors.Is(err, ErrPreconditionFailed) { + return nil, errors.Errorf("parallel get: open range %d-%d: object changed during read: %w", start, end, err) + } + if err != nil { + return nil, errors.Errorf("parallel get: open range %d-%d: %w", start, end, err) + } + // Check Content-Range before trusting the ETag: a range-ignoring + // response often carries no ETag. + if cr := h.Get("Content-Range"); !strings.HasPrefix(cr, "bytes "+strconv.FormatInt(start, 10)+"-") { + if attempt < discoveryAttempts-1 { + if err := body.Close(); err != nil { + return nil, errors.Wrap(err, "parallel get: close range-ignoring chunk reader") + } + continue + } + return nil, errors.Join( + errors.Errorf("parallel get: server ignored range %d-%d (Content-Range %q)", start, end, cr), + body.Close(), + ) + } + if got := h.Get(ETagKey); got != etag { + return nil, errors.Join( + errors.Errorf("parallel get: object changed during read at offset %d: etag %q != %q", start, got, etag), + body.Close(), + ) + } + return body, nil } - return writeChunkAt(dst, start, end-start, rc) } -func writeChunkAt(dst io.WriterAt, off, want int64, src io.ReadCloser) error { - n, copyErr := io.Copy(io.NewOffsetWriter(dst, off), src) - if err := errors.Join(copyErr, src.Close()); err != nil { - return errors.Errorf("write chunk at offset %d: %w", off, err) - } - if want >= 0 && n != want { - return errors.Errorf("short chunk at offset %d: wrote %d of %d bytes", off, n, want) +// fullRead downloads the object in a single request and places it with a +// negative length, since a full-body response's size is unknown up front. +func fullRead(ctx context.Context, c RangeReader, key Key, sink ChunkSink) error { + rc, _, err := c.Open(ctx, key) + if err != nil { + return errors.Wrap(err, "parallel get: full read") } - return nil + return errors.Wrap(sink.Place(ctx, 0, -1, rc), "parallel get") } +// parseContentRangeTotal extracts the total from "bytes start-end/total", +// returning ok=false when the header is absent or unparseable. func parseContentRangeTotal(contentRange string) (total int64, ok bool) { _, size, found := strings.Cut(contentRange, "/") if !found { diff --git a/client/parallel_get_test.go b/client/parallel_get_test.go index 3d74ddfb..f93097e1 100644 --- a/client/parallel_get_test.go +++ b/client/parallel_get_test.go @@ -11,30 +11,54 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "testing" + "time" "github.com/alecthomas/assert/v2" + "github.com/alecthomas/errors" "github.com/block/cachew/client" ) +// Compile-time assertion that the concrete Client satisfies the narrow +// interface ParallelGet drives. var _ client.RangeReader = (*client.Client)(nil) -type bufferAt struct { - mu sync.Mutex - buf []byte +func patternBytes(n int) []byte { + data := make([]byte, n) + for i := range data { + data[i] = byte(i % 251) + } + return data } -func (b *bufferAt) WriteAt(p []byte, off int64) (int, error) { - b.mu.Lock() - defer b.mu.Unlock() - if end := int(off) + len(p); end > len(b.buf) { - b.buf = append(b.buf, make([]byte, end-len(b.buf))...) +// collect runs ParallelGet into a StreamSink and returns the reassembled +// bytes, reading the sink concurrently as the engine requires. The download +// error (if any) takes precedence over the read error. +func collect(c client.RangeReader, key client.Key, chunkSize int64, concurrency int) ([]byte, error) { + sink := client.NewStreamSink(chunkSize, concurrency) + defer sink.Close() //nolint:errcheck + type result struct { + data []byte + err error } - copy(b.buf[off:], p) - return len(p), nil + rc := make(chan result, 1) + go func() { + data, err := io.ReadAll(sink) + rc <- result{data: data, err: err} + }() + err := client.ParallelGet(context.Background(), c, key, sink, chunkSize, concurrency) + sink.Done(err) + res := <-rc + if err != nil { + return res.data, err + } + return res.data, res.err } +// rangeFlipReader serves correct byte ranges but reports a different ETag for +// any chunk past the first, simulating an object rewritten mid-download. type rangeFlipReader struct { data []byte firstETag string @@ -64,12 +88,14 @@ func (f *rangeFlipReader) Open(_ context.Context, _ client.Key, opts ...client.R func TestParallelGetETagMismatch(t *testing.T) { c := &rangeFlipReader{data: make([]byte, 1000), firstETag: `"v1"`, restETag: `"v2"`} - var dst bufferAt - err := client.ParallelGet(context.Background(), c, client.NewKey("k"), &dst, 100, 4) + _, err := collect(c, client.NewKey("k"), 100, 4) assert.Error(t, err) assert.Contains(t, err.Error(), "object changed during read") } +// ifMatchFlipReader honours If-Match against firstETag on the discovery request +// but reports restETag (and a 412) for any pinned chunk, modelling a server that +// enforces If-Match on an object rewritten mid-download. type ifMatchFlipReader struct { data []byte firstETag string @@ -102,12 +128,14 @@ func (f *ifMatchFlipReader) Open(_ context.Context, _ client.Key, opts ...client func TestParallelGetPreconditionFailedOnRewrite(t *testing.T) { c := &ifMatchFlipReader{data: make([]byte, 1000), firstETag: `"v1"`, restETag: `"v2"`} - var dst bufferAt - err := client.ParallelGet(context.Background(), c, client.NewKey("k"), &dst, 100, 4) + dst := &bufferAt{} + err := client.ParallelGet(context.Background(), c, client.NewKey("k"), client.DiskSink{W: dst}, 100, 4) assert.IsError(t, err, client.ErrPreconditionFailed) assert.Contains(t, err.Error(), "object changed during read") } +// noETagReader serves byte ranges but never sets an ETag, modelling a legacy +// entry or a RangeReader implementation that omits it. type noETagReader struct { data []byte } @@ -128,26 +156,24 @@ func (n *noETagReader) Open(_ context.Context, _ client.Key, opts ...client.Requ } func TestParallelGetNoETagMultiChunk(t *testing.T) { - data := make([]byte, 1000) - for i := range data { - data[i] = byte(i % 251) - } + data := patternBytes(1000) c := &noETagReader{data: data} - var dst bufferAt - err := client.ParallelGet(context.Background(), c, client.NewKey("k"), &dst, 100, 4) + got, err := collect(c, client.NewKey("k"), 100, 4) assert.NoError(t, err) - assert.Equal(t, data, dst.buf) + assert.Equal(t, data, got) } func TestParallelGetNoETagSingleChunk(t *testing.T) { data := []byte("0123456789") c := &noETagReader{data: data} - var dst bufferAt - err := client.ParallelGet(context.Background(), c, client.NewKey("k"), &dst, 100, 4) + got, err := collect(c, client.NewKey("k"), 100, 4) assert.NoError(t, err) - assert.Equal(t, data, dst.buf) + assert.Equal(t, data, got) } +// changingSizeReader serves a multi-chunk body with no ETag on the ranged +// discovery request, then a differently sized body on the subsequent full +// (non-range) read, modelling an object rewritten between the two requests. type changingSizeReader struct { discovery []byte rewritten []byte @@ -173,11 +199,21 @@ func (c *changingSizeReader) Open(_ context.Context, _ client.Key, opts ...clien return io.NopCloser(bytes.NewReader(c.discovery[start : start+length])), headers, nil } +func TestParallelGetNoETagSizeChangedBetweenRequests(t *testing.T) { + c := &changingSizeReader{discovery: make([]byte, 1000), rewritten: []byte("changed")} + got, err := collect(c, client.NewKey("k"), 100, 4) + assert.NoError(t, err) + assert.Equal(t, c.rewritten, got) +} + type openRecord struct { Range string IfMatch string } +// recordingReader serves byte ranges and records the Range and If-Match options +// of every Open call (both "" for a full, non-ranged read), so tests can assert +// how the object was fetched and how chunks were pinned. type recordingReader struct { data []byte etag string @@ -209,27 +245,28 @@ func (r *recordingReader) Open(_ context.Context, _ client.Key, opts ...client.R return io.NopCloser(bytes.NewReader(r.data[start : start+length])), headers, nil } +func TestParallelGetReassembly(t *testing.T) { + data := patternBytes(10_000) + c := &recordingReader{data: data, etag: `"v1"`} + got, err := collect(c, client.NewKey("k"), 1000, 4) + assert.NoError(t, err) + assert.Equal(t, data, got) +} + func TestParallelGetSingleWorkerFullRead(t *testing.T) { - data := make([]byte, 1000) - for i := range data { - data[i] = byte(i % 251) - } + data := patternBytes(1000) c := &recordingReader{data: data, etag: `"v1"`} - var dst bufferAt - err := client.ParallelGet(context.Background(), c, client.NewKey("k"), &dst, 100, 1) + got, err := collect(c, client.NewKey("k"), 100, 1) assert.NoError(t, err) - assert.Equal(t, data, dst.buf) + assert.Equal(t, data, got) assert.Equal(t, []openRecord{{}}, c.opens) } func TestParallelGetPinsChunksWithIfMatch(t *testing.T) { - data := make([]byte, 1000) - for i := range data { - data[i] = byte(i % 251) - } + data := patternBytes(1000) c := &recordingReader{data: data, etag: `"v1"`} - var dst bufferAt - err := client.ParallelGet(context.Background(), c, client.NewKey("k"), &dst, 100, 4) + dst := &bufferAt{} + err := client.ParallelGet(context.Background(), c, client.NewKey("k"), client.DiskSink{W: dst}, 100, 4) assert.NoError(t, err) assert.Equal(t, data, dst.buf) @@ -241,19 +278,216 @@ func TestParallelGetPinsChunksWithIfMatch(t *testing.T) { assert.Equal(t, expected, c.opens) } -func TestParallelGetNoETagSizeChangedBetweenRequests(t *testing.T) { - c := &changingSizeReader{discovery: make([]byte, 1000), rewritten: []byte("changed")} - var dst bufferAt - err := client.ParallelGet(context.Background(), c, client.NewKey("k"), &dst, 100, 4) +func TestParallelGetEmptyObject(t *testing.T) { + c := &recordingReader{data: nil, etag: `"v1"`} + got, err := collect(c, client.NewKey("k"), 100, 4) assert.NoError(t, err) - assert.Equal(t, c.rewritten, dst.buf) + assert.Equal(t, 0, len(got)) } -func TestParallelGetClient(t *testing.T) { - content := make([]byte, 1000) - for i := range content { - content[i] = byte(i % 251) +func TestParallelGetServerIgnoresRange(t *testing.T) { + data := patternBytes(1000) + c := &ignoreRangeReader{data: data} + got, err := collect(c, client.NewKey("k"), 100, 4) + assert.NoError(t, err) + assert.Equal(t, data, got) +} + +func TestParallelGetOutOfOrderCompletion(t *testing.T) { + data := patternBytes(10_000) + c := &reorderReader{data: data, etag: `"v1"`, chunkSize: 1000} + got, err := collect(c, client.NewKey("k"), 1000, 4) + assert.NoError(t, err) + assert.Equal(t, data, got) +} + +func TestParallelGetPropagatesOpenError(t *testing.T) { + c := &failingChunkReader{data: patternBytes(10_000), etag: `"v1"`, failAtStart: 5000} + _, err := collect(c, client.NewKey("k"), 1000, 4) + assert.Error(t, err) + assert.Contains(t, err.Error(), "boom") +} + +func TestParallelGetRejectsOverlongChunk(t *testing.T) { + c := &fullBodyOnChunkReader{data: patternBytes(10_000), etag: `"v1"`} + _, err := collect(c, client.NewKey("k"), 1000, 4) + assert.Error(t, err) + assert.Contains(t, err.Error(), "more than the expected 1000 bytes") +} + +func TestParallelGetWriterAtReassembly(t *testing.T) { + data := patternBytes(10_000) + c := &recordingReader{data: data, etag: `"v1"`} + dst := &bufferAt{} + err := client.ParallelGet(context.Background(), c, client.NewKey("k"), client.DiskSink{W: dst}, 1000, 4) + assert.NoError(t, err) + assert.Equal(t, data, dst.buf) +} + +func TestParallelGetWriterAtOutOfOrder(t *testing.T) { + data := patternBytes(10_000) + c := &reorderReader{data: data, etag: `"v1"`, chunkSize: 1000} + dst := &bufferAt{} + err := client.ParallelGet(context.Background(), c, client.NewKey("k"), client.DiskSink{W: dst}, 1000, 4) + assert.NoError(t, err) + assert.Equal(t, data, dst.buf) +} + +func TestParallelGetRejectsMidflightRangeIgnore(t *testing.T) { + c := &midflightRangeIgnoreReader{data: patternBytes(10_000), etag: `"v1"`} + _, err := collect(c, client.NewKey("k"), 1000, 4) + assert.Error(t, err) + assert.Contains(t, err.Error(), "server ignored range") +} + +func TestParallelGetWriterAtRejectsOverlongChunk(t *testing.T) { + c := &fullBodyOnChunkReader{data: patternBytes(10_000), etag: `"v1"`} + dst := &bufferAt{} + err := client.ParallelGet(context.Background(), c, client.NewKey("k"), client.DiskSink{W: dst}, 1000, 4) + assert.Error(t, err) + assert.Contains(t, err.Error(), "more than the expected 1000 bytes") +} + +// bufferAt is an in-memory io.WriterAt that extends like a file, zero-filling +// gaps, so tests can exercise DiskSink's concurrent WriteAt path without +// touching disk. +type bufferAt struct { + mu sync.Mutex + buf []byte +} + +func (b *bufferAt) WriteAt(p []byte, off int64) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + if end := int(off) + len(p); end > len(b.buf) { + b.buf = append(b.buf, make([]byte, end-len(b.buf))...) } + copy(b.buf[off:], p) + return len(p), nil +} + +// fullBodyOnChunkReader honours the discovery range (start 0) with a proper 206 +// but ignores the range on any later chunk, returning the entire object with the +// same ETag — modelling a backend that degrades to full responses mid-download. +type fullBodyOnChunkReader struct { + data []byte + etag string +} + +func (r *fullBodyOnChunkReader) Open(_ context.Context, _ client.Key, opts ...client.RequestOption) (io.ReadCloser, http.Header, error) { + size := int64(len(r.data)) + start, length, outcome := client.NewRequestOptions(opts...).ResolveRange(size, r.etag) + headers := http.Header{} + headers.Set(client.ETagKey, r.etag) + if outcome == client.RangePartial && start == 0 { + headers.Set("Content-Length", strconv.FormatInt(length, 10)) + headers.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, start+length-1, size)) + return io.NopCloser(bytes.NewReader(r.data[start : start+length])), headers, nil + } + if outcome == client.RangePartial { + headers.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, start+length-1, size)) + } + headers.Set("Content-Length", strconv.FormatInt(size, 10)) + return io.NopCloser(bytes.NewReader(r.data)), headers, nil +} + +// midflightRangeIgnoreReader ranges the discovery chunk but answers subsequent +// chunk requests with the full body and no Content-Range, modelling a replica +// that stops honouring ranges mid-download. +type midflightRangeIgnoreReader struct { + data []byte + etag string +} + +func (r *midflightRangeIgnoreReader) Open(_ context.Context, _ client.Key, opts ...client.RequestOption) (io.ReadCloser, http.Header, error) { + size := int64(len(r.data)) + start, length, outcome := client.NewRequestOptions(opts...).ResolveRange(size, r.etag) + headers := http.Header{} + headers.Set(client.ETagKey, r.etag) + if outcome == client.RangePartial && start == 0 { + headers.Set("Content-Length", strconv.FormatInt(length, 10)) + headers.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, start+length-1, size)) + return io.NopCloser(bytes.NewReader(r.data[start : start+length])), headers, nil + } + headers.Set("Content-Length", strconv.FormatInt(size, 10)) + return io.NopCloser(bytes.NewReader(r.data)), headers, nil +} + +// ignoreRangeReader returns the whole object with no Content-Range regardless of +// the requested range, modelling a backend that doesn't honour ranges. +type ignoreRangeReader struct{ data []byte } + +func (r *ignoreRangeReader) Open(_ context.Context, _ client.Key, _ ...client.RequestOption) (io.ReadCloser, http.Header, error) { + headers := http.Header{} + headers.Set("Content-Length", strconv.Itoa(len(r.data))) + return io.NopCloser(bytes.NewReader(r.data)), headers, nil +} + +// reorderReader serves correct byte ranges but delays earlier offsets longer +// than later ones, so within the in-flight window chunks complete out of order +// and the writer must buffer and reorder them. +type reorderReader struct { + data []byte + etag string + chunkSize int64 +} + +func (r *reorderReader) Open(_ context.Context, _ client.Key, opts ...client.RequestOption) (io.ReadCloser, http.Header, error) { + size := int64(len(r.data)) + o := client.NewRequestOptions(opts...) + start, length, outcome := o.ResolveRange(size, r.etag) + headers := http.Header{} + if outcome == client.RangeNotSatisfiable { + headers.Set("Content-Range", fmt.Sprintf("bytes */%d", size)) + return nil, headers, client.ErrRangeNotSatisfiable + } + // Earlier chunks within a window sleep longer, so higher offsets finish + // first and the writer is forced to reorder. + if outcome == client.RangePartial { + chunks := (size - start) / r.chunkSize + time.Sleep(time.Duration(chunks) * time.Millisecond) + } + headers.Set(client.ETagKey, r.etag) + headers.Set("Content-Length", strconv.FormatInt(length, 10)) + if outcome == client.RangePartial { + headers.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, start+length-1, size)) + } + return io.NopCloser(bytes.NewReader(r.data[start : start+length])), headers, nil +} + +// failingChunkReader serves ranges normally but errors when the requested range +// starts at failAtStart, modelling a mid-download fetch failure. +type failingChunkReader struct { + data []byte + etag string + failAtStart int64 + + opens atomic.Int64 +} + +func (r *failingChunkReader) Open(_ context.Context, _ client.Key, opts ...client.RequestOption) (io.ReadCloser, http.Header, error) { + r.opens.Add(1) + size := int64(len(r.data)) + o := client.NewRequestOptions(opts...) + start, length, outcome := o.ResolveRange(size, r.etag) + if outcome == client.RangePartial && start == r.failAtStart { + return nil, nil, errors.New("boom") + } + headers := http.Header{} + if outcome == client.RangeNotSatisfiable { + headers.Set("Content-Range", fmt.Sprintf("bytes */%d", size)) + return nil, headers, client.ErrRangeNotSatisfiable + } + headers.Set(client.ETagKey, r.etag) + headers.Set("Content-Length", strconv.FormatInt(length, 10)) + if outcome == client.RangePartial { + headers.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, start+length-1, size)) + } + return io.NopCloser(bytes.NewReader(r.data[start : start+length])), headers, nil +} + +func TestParallelGetClient(t *testing.T) { + content := patternBytes(1000) etag := fakeETag(content) mux := http.NewServeMux() @@ -291,8 +525,8 @@ func TestParallelGetClient(t *testing.T) { c := client.New(srv.URL, nil).Namespace("test") defer c.Close() - var dst bufferAt - err := client.ParallelGet(context.Background(), c, client.NewKey("parallel-client"), &dst, 100, 4) + dst := &bufferAt{} + err := client.ParallelGet(context.Background(), c, client.NewKey("parallel-client"), client.DiskSink{W: dst}, 100, 4) assert.NoError(t, err) assert.Equal(t, content, dst.buf) } diff --git a/client/stream_sink.go b/client/stream_sink.go new file mode 100644 index 00000000..3cfc5576 --- /dev/null +++ b/client/stream_sink.go @@ -0,0 +1,208 @@ +package client + +import ( + "context" + "io" + "sync" + + "github.com/alecthomas/errors" +) + +// StreamSink reorders concurrently-fetched chunks back into the original byte +// stream, exposed via Read. Chunks land in a fixed arena of 2*concurrency +// reusable in-memory slots, so a slow consumer applies backpressure to the +// fetchers and peak memory is capped at 2*concurrency*chunkSize regardless of +// object size. +// +// A StreamSink must be read concurrently while ParallelGet runs — fetchers +// block once they are a window ahead of the reader. After the download the +// caller signals completion with Done; Read then drains the remaining chunks +// and returns io.EOF, or the download error. +type StreamSink struct { + chunkSize int64 + n int // slot count = 2*concurrency + + mu sync.Mutex + cond *sync.Cond // signals Read that a chunk was deposited (or Done) + advance chan struct{} // closed and replaced when readSeq advances, waking blocked Place + bufs [][]byte // n reusable backing buffers, indexed by seq%n (nil until first use) + ready []bool // ready[slot] => bufs[slot] holds the chunk for its current seq + readSeq int64 // sequence number of the chunk Read is emitting next + cur []byte // chunk currently being emitted (aliases bufs[readSeq%n]) + curPos int + + passthru io.ReadCloser // set in single-stream fallback mode (length < 0) + done bool + err error + closed bool +} + +// NewStreamSink returns a StreamSink holding up to 2*concurrency lazily +// allocated chunk buffers. +func NewStreamSink(chunkSize int64, concurrency int) *StreamSink { + n := 2 * max(concurrency, 1) + s := &StreamSink{ + chunkSize: chunkSize, + n: n, + bufs: make([][]byte, n), + ready: make([]bool, n), + advance: make(chan struct{}), + } + s.cond = sync.NewCond(&s.mu) + return s +} + +// Place reads the chunk into its slot for in-order delivery to Read, blocking +// while the chunk is more than a window ahead of the read cursor. A negative +// length switches to pass-through mode: the body is handed to Read directly, +// since a single-stream fallback has unknown size and must not be buffered. +func (s *StreamSink) Place(ctx context.Context, off, length int64, body io.ReadCloser) error { + if length < 0 { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + return errors.Join(errors.New("stream sink closed"), body.Close()) + } + s.passthru = body + s.cond.Broadcast() + s.mu.Unlock() + return nil + } + + seq := off / s.chunkSize + slot := int(seq % int64(s.n)) + + // A chunk may occupy its slot only once the previous occupant (seq-n) has + // been read; this bounds run-ahead and guarantees exclusive slot ownership. + s.mu.Lock() + for seq >= s.readSeq+int64(s.n) { + if s.closed { + s.mu.Unlock() + return errors.Join(errors.New("stream sink closed"), body.Close()) + } + ch := s.advance + s.mu.Unlock() + select { + case <-ch: + case <-ctx.Done(): + return errors.Join(errors.WithStack(ctx.Err()), body.Close()) + } + s.mu.Lock() + } + buf := s.bufs[slot] + s.mu.Unlock() + + if int64(cap(buf)) < length { + buf = make([]byte, s.chunkSize) + } + buf = buf[:length] + if err := readChunk(off, buf, body); err != nil { + return err + } + + s.mu.Lock() + // A Close racing the body read leaves no reader to drain this slot; drop + // the chunk. readChunk already closed body. + if s.closed { + s.mu.Unlock() + return errors.New("stream sink closed") + } + s.bufs[slot] = buf + s.ready[slot] = true + s.cond.Broadcast() + s.mu.Unlock() + return nil +} + +// Read emits the reassembled object in order, blocking until the next chunk is +// available. It returns io.EOF once Done has been called and every chunk has +// been read, or the download error reported to Done. +func (s *StreamSink) Read(p []byte) (int, error) { + s.mu.Lock() + for { + if s.passthru != nil { + body := s.passthru + s.mu.Unlock() + return body.Read(p) //nolint:wrapcheck // must return io.EOF verbatim for io.ReadAll + } + if s.cur != nil { + n := copy(p, s.cur[s.curPos:]) + s.curPos += n + if s.curPos >= len(s.cur) { + // Free the slot and wake any Place blocked on the window. + slot := int(s.readSeq % int64(s.n)) + s.ready[slot] = false + s.readSeq++ + s.cur = nil + s.curPos = 0 + close(s.advance) + s.advance = make(chan struct{}) + } + s.mu.Unlock() + return n, nil + } + slot := int(s.readSeq % int64(s.n)) + if s.ready[slot] { + s.cur = s.bufs[slot] + s.curPos = 0 + continue + } + if s.err != nil { + err := s.err + s.mu.Unlock() + return 0, err + } + if s.done { + s.mu.Unlock() + return 0, io.EOF + } + // Closed mid-download: stop rather than block forever on cond. + if s.closed { + s.mu.Unlock() + return 0, errors.WithStack(io.ErrClosedPipe) + } + s.cond.Wait() + } +} + +// Done signals that no further chunks will be placed; err (nil on success) is +// surfaced to Read after the buffered chunks drain. +func (s *StreamSink) Done(err error) { + s.mu.Lock() + s.done = true + if err != nil && s.err == nil { + s.err = err + } + s.cond.Broadcast() + s.mu.Unlock() +} + +// Close releases the sink, unblocking any in-flight Place and closing the +// pass-through body if one is set. Cancelling the download itself is the +// caller's responsibility. +func (s *StreamSink) Close() error { + s.mu.Lock() + s.closed = true + body := s.passthru + s.passthru = nil + close(s.advance) + s.advance = make(chan struct{}) + s.cond.Broadcast() + s.mu.Unlock() + if body != nil { + return errors.WithStack(body.Close()) + } + return nil +} + +// readChunk reads exactly len(buf) bytes from body then closes it, reporting a +// short or overlong body as an error. +func readChunk(off int64, buf []byte, body io.ReadCloser) error { + if _, err := io.ReadFull(body, buf); err != nil { + return errors.Join(errors.Errorf("read chunk at offset %d: %w", off, err), body.Close()) + } + if overlong(body) { + return errors.Join(errors.Errorf("chunk at offset %d: read more than the expected %d bytes", off, len(buf)), body.Close()) + } + return errors.WithStack(body.Close()) +} diff --git a/cmd/cachew/git.go b/cmd/cachew/git.go index 0774314c..be01caca 100644 --- a/cmd/cachew/git.go +++ b/cmd/cachew/git.go @@ -47,17 +47,14 @@ type GitCmd struct { // ensure those refs/commits are fresh and runs `git pull --ff-only` so the // working tree catches up to upstream. type GitRestoreCmd struct { - RepoURL string `arg:"" help:"Repository URL (e.g. https://github.com/org/repo)."` - Directory string `arg:"" help:"Target directory for the clone." type:"path"` - Ref map[string]string `help:"Required refs to freshen on the server before pulling, in the form 'name=sha' (e.g. 'refs/heads/main=abc123'). An empty SHA means any SHA is acceptable. Setting this (or --commit) runs a final 'git pull' from origin so the working tree is brought up to date."` - Commit []string `help:"Required commit SHAs that must exist on the server, regardless of which ref points at them. May be repeated."` - NoBundle bool `help:"Skip applying delta bundle."` - ZstdThreads int `help:"Threads for zstd decompression (0 = all CPU cores)." default:"0"` - // DownloadConcurrency > 1 fetches the snapshot with that many concurrent - // range requests (requires server range support; falls back to a single - // request otherwise). 1 keeps the streaming single-request download. - DownloadConcurrency int `help:"Concurrent range requests for the snapshot download (1 = single streaming request)." default:"1"` - DownloadChunkSizeMB int `help:"Chunk size in MiB for parallel snapshot downloads." default:"8"` + RepoURL string `arg:"" help:"Repository URL (e.g. https://github.com/org/repo)."` + Directory string `arg:"" help:"Target directory for the clone." type:"path"` + Ref map[string]string `help:"Required refs to freshen on the server before pulling, in the form 'name=sha' (e.g. 'refs/heads/main=abc123'). An empty SHA means any SHA is acceptable. Setting this (or --commit) runs a final 'git pull' from origin so the working tree is brought up to date."` + Commit []string `help:"Required commit SHAs that must exist on the server, regardless of which ref points at them. May be repeated."` + NoBundle bool `help:"Skip applying delta bundle."` + ZstdThreads int `help:"Threads for zstd decompression (0 = all CPU cores)." default:"0"` + DownloadConcurrency int `help:"Concurrent range requests for the snapshot download (1 = single streaming request)." default:"8"` + DownloadChunkSizeMB int `help:"Chunk size in MiB for parallel snapshot downloads." default:"16"` } func (c *GitRestoreCmd) Run(ctx context.Context, api *client.Client) error { @@ -116,27 +113,19 @@ func (c *GitRestoreCmd) Run(ctx context.Context, api *client.Client) error { return nil } -// fetchAndExtractSnapshot downloads the snapshot and extracts it into the target -// directory, returning its freshen metadata (commit and bundle URL). With a -// download concurrency above 1 it downloads in parallel into a temp file, since -// ParallelGet needs a WriterAt; otherwise it streams the single response -// directly into extraction. +// fetchAndExtractSnapshot pipes the snapshot download straight into extraction +// and returns its freshen metadata (commit and bundle URL). func (c *GitRestoreCmd) fetchAndExtractSnapshot(ctx context.Context, api *client.Client) (commit, bundleURL string, err error) { - if c.DownloadConcurrency > 1 { - return c.parallelFetchAndExtract(ctx, api) - } - return c.streamFetchAndExtract(ctx, api) -} - -// streamFetchAndExtract downloads the snapshot in a single request and pipes the -// response body straight into extraction, overlapping download and extraction. -func (c *GitRestoreCmd) streamFetchAndExtract(ctx context.Context, api *client.Client) (string, string, error) { var snap *client.GitSnapshot if err := inSpan(ctx, "cachew.download_snapshot", - []attribute.KeyValue{attribute.String("cachew.repo_url", c.RepoURL)}, + []attribute.KeyValue{ + attribute.String("cachew.repo_url", c.RepoURL), + attribute.Int("cachew.download_concurrency", c.DownloadConcurrency), + attribute.Int("cachew.download_chunk_size_mb", c.DownloadChunkSizeMB), + }, func(ctx context.Context) error { downloadStart := time.Now() - s, err := api.OpenGitSnapshot(ctx, c.RepoURL) + s, err := api.OpenGitSnapshotParallel(ctx, c.RepoURL, int64(c.DownloadChunkSizeMB)<<20, c.DownloadConcurrency) if err != nil { return err //nolint:wrapcheck // wrapped by caller } @@ -158,76 +147,39 @@ func (c *GitRestoreCmd) streamFetchAndExtract(ctx context.Context, api *client.C return snap.Commit, snap.BundleURL, nil } -// parallelFetchAndExtract downloads the snapshot into a temp file using bounded -// concurrent range requests, then extracts from the file. ParallelGet writes via -// WriteAt so it cannot stream into extraction; the temp file is removed on -// return. -func (c *GitRestoreCmd) parallelFetchAndExtract(ctx context.Context, api *client.Client) (string, string, error) { - // Stage the temp snapshot on the same filesystem as the restore target so a - // small or separate /tmp can't fail a restore the target directory has room - // for. The parent of c.Directory shares its filesystem and is created by - // extraction anyway. - tmpDir := filepath.Dir(c.Directory) - if err := os.MkdirAll(tmpDir, 0o750); err != nil { - return "", "", errors.Wrap(err, "create snapshot temp dir") +// extract unpacks the snapshot into a staging directory renamed into place on +// success, so a mid-download failure cannot leave a half-written checkout. +func (c *GitRestoreCmd) extract(ctx context.Context, body io.Reader) error { + parent := filepath.Dir(c.Directory) + if err := os.MkdirAll(parent, 0o750); err != nil { + return errors.Wrap(err, "create restore parent directory") } - tmp, err := os.CreateTemp(tmpDir, ".cachew-snapshot-*.tar.zst") + staging, err := os.MkdirTemp(parent, ".cachew-restore-*") if err != nil { - return "", "", errors.Wrap(err, "create snapshot temp file") + return errors.Wrap(err, "create restore staging directory") } - defer func() { - _ = tmp.Close() - _ = os.Remove(tmp.Name()) //nolint:gosec // name is from os.CreateTemp, not external input - }() + defer os.RemoveAll(staging) //nolint:errcheck // best-effort cleanup; a no-op once renamed into place - var meta client.GitSnapshotMetadata - if err := inSpan(ctx, "cachew.download_snapshot", - []attribute.KeyValue{ - attribute.String("cachew.repo_url", c.RepoURL), - attribute.Int("cachew.download_concurrency", c.DownloadConcurrency), - attribute.Int("cachew.download_chunk_size_mb", c.DownloadChunkSizeMB), - }, - func(ctx context.Context) error { - downloadStart := time.Now() - m, err := api.DownloadGitSnapshot(ctx, c.RepoURL, tmp, int64(c.DownloadChunkSizeMB)<<20, c.DownloadConcurrency) - if err != nil { - return err //nolint:wrapcheck // wrapped by caller - } - meta = m - trace.SpanFromContext(ctx).SetAttributes( - attribute.String("cachew.snapshot_commit", m.Commit), - attribute.String("cachew.bundle_url", m.BundleURL), - attribute.Float64("cachew.elapsed_seconds", time.Since(downloadStart).Seconds()), - ) - return nil - }); err != nil { - return "", "", err - } - - if _, err := tmp.Seek(0, io.SeekStart); err != nil { - return "", "", errors.Wrap(err, "rewind snapshot temp file") - } - if err := c.extract(ctx, tmp); err != nil { - return "", "", err - } - return meta.Commit, meta.BundleURL, nil -} - -// extract decompresses and unpacks the snapshot body into the target directory. -func (c *GitRestoreCmd) extract(ctx context.Context, body io.Reader) error { fmt.Fprintf(os.Stderr, "Extracting to %s...\n", c.Directory) //nolint:forbidigo,gosec // c.Directory is an operator-supplied CLI path - return inSpan(ctx, "cachew.extract", + if err := inSpan(ctx, "cachew.extract", []attribute.KeyValue{attribute.String("cachew.directory", c.Directory)}, func(ctx context.Context) error { extractStart := time.Now() - if err := snapshot.Extract(ctx, body, c.Directory, c.ZstdThreads); err != nil { + if err := snapshot.Extract(ctx, body, staging, c.ZstdThreads); err != nil { return err //nolint:wrapcheck // wrapped by caller } elapsed := time.Since(extractStart) trace.SpanFromContext(ctx).SetAttributes(attribute.Float64("cachew.elapsed_seconds", elapsed.Seconds())) fmt.Fprintf(os.Stderr, "Snapshot extracted in %s\n", elapsed) //nolint:forbidigo return nil - }) + }); err != nil { + return err + } + + if err := os.Rename(staging, c.Directory); err != nil { + return errors.Wrap(err, "move restored snapshot into place") + } + return nil } // satisfyRefs ensures the working tree contains every requested ref and diff --git a/internal/cache/parallel_get.go b/internal/cache/parallel_get.go index 55458d07..0276e94d 100644 --- a/internal/cache/parallel_get.go +++ b/internal/cache/parallel_get.go @@ -2,16 +2,15 @@ package cache import ( "context" - "io" "github.com/alecthomas/errors" "github.com/block/cachew/client" ) -// ParallelGet downloads an object from any Range-capable Cache into dst, -// fetching it in chunkSize-byte chunks concurrently. It delegates to +// ParallelGet downloads an object from any Range-capable Cache, fetching it in +// chunkSize-byte chunks concurrently and handing each to sink. It delegates to // [client.ParallelGet]; see that function for the full semantics. -func ParallelGet(ctx context.Context, c Cache, key Key, dst io.WriterAt, chunkSize int64, concurrency int) error { - return errors.WithStack(client.ParallelGet(ctx, c, key, dst, chunkSize, concurrency)) +func ParallelGet(ctx context.Context, c Cache, key Key, sink client.ChunkSink, chunkSize int64, concurrency int) error { + return errors.WithStack(client.ParallelGet(ctx, c, key, sink, chunkSize, concurrency)) } diff --git a/internal/cache/parallel_get_test.go b/internal/cache/parallel_get_test.go index 1baa3bb8..783f7e71 100644 --- a/internal/cache/parallel_get_test.go +++ b/internal/cache/parallel_get_test.go @@ -5,31 +5,37 @@ import ( "io" "log/slog" "os" - "sync" "testing" "time" "github.com/alecthomas/assert/v2" + "github.com/block/cachew/client" "github.com/block/cachew/internal/cache" "github.com/block/cachew/internal/logging" ) -// bufferAt is an in-memory io.WriterAt that extends like a file, zero-filling -// any gap, so tests can assert reassembly without touching disk. -type bufferAt struct { - mu sync.Mutex - buf []byte -} - -func (b *bufferAt) WriteAt(p []byte, off int64) (int, error) { - b.mu.Lock() - defer b.mu.Unlock() - if end := int(off) + len(p); end > len(b.buf) { - b.buf = append(b.buf, make([]byte, end-len(b.buf))...) +// collect runs cache.ParallelGet into a StreamSink and returns the +// reassembled bytes, reading the sink concurrently as the engine requires. +func collect(ctx context.Context, c cache.Cache, key cache.Key, chunkSize int64, concurrency int) ([]byte, error) { + sink := client.NewStreamSink(chunkSize, concurrency) + defer sink.Close() //nolint:errcheck + type result struct { + data []byte + err error + } + rc := make(chan result, 1) + go func() { + data, err := io.ReadAll(sink) + rc <- result{data: data, err: err} + }() + err := cache.ParallelGet(ctx, c, key, sink, chunkSize, concurrency) + sink.Done(err) + res := <-rc + if err != nil { + return res.data, err } - copy(b.buf[off:], p) - return len(p), nil + return res.data, res.err } func TestParallelGet(t *testing.T) { @@ -61,10 +67,9 @@ func TestParallelGet(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - var dst bufferAt - err := cache.ParallelGet(ctx, c, key, &dst, tt.chunkSize, tt.concurrency) + got, err := collect(ctx, c, key, tt.chunkSize, tt.concurrency) assert.NoError(t, err) - assert.Equal(t, content, dst.buf) + assert.Equal(t, content, got) }) } } @@ -80,12 +85,10 @@ func TestParallelGetEmptyObject(t *testing.T) { assert.NoError(t, err) assert.NoError(t, w.Close()) - // concurrency 4 takes the ranged discovery path (ErrRangeNotSatisfiable), - // concurrency 1 takes the up-front full-read path; both must yield nothing. for _, concurrency := range []int{4, 1} { - var dst bufferAt - assert.NoError(t, cache.ParallelGet(ctx, c, key, &dst, 100, concurrency)) - assert.Equal(t, 0, len(dst.buf)) + got, err := collect(ctx, c, key, 100, concurrency) + assert.NoError(t, err) + assert.Equal(t, 0, len(got)) } } @@ -95,7 +98,6 @@ func TestParallelGetNotFound(t *testing.T) { assert.NoError(t, err) defer c.Close() - var dst bufferAt - err = cache.ParallelGet(ctx, c, cache.NewKey("missing"), &dst, 100, 4) + _, err = collect(ctx, c, cache.NewKey("missing"), 100, 4) assert.IsError(t, err, os.ErrNotExist) }