diff --git a/.golangci.yml b/.golangci.yml index a83a1ea7..5edd5b62 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -480,7 +480,7 @@ linters: # components intentionally forward URLs and construct paths from config. - text: "G703:" linters: [gosec] - path: "(internal/cache/disk\\.go|internal/strategy/git/spool\\.go)" + path: "(internal/cache/disk\\.go|internal/strategy/git/spool\\.go|client/spill\\.go)" - text: "G704:" linters: [gosec] - text: "avoid package names that conflict with Go standard library" diff --git a/client/git.go b/client/git.go index 999509d9..8cb16aa6 100644 --- a/client/git.go +++ b/client/git.go @@ -201,12 +201,36 @@ type GitSnapshotMetadata struct { // 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) { + return c.downloadGitSnapshot(repoURL, func(reader *gitArtifactRangeReader) error { + return errors.WithStack(ParallelGet(ctx, reader, NewKey(repoURL), dst, chunkSize, concurrency)) + }) +} + +// DownloadGitSnapshotStream fetches the working-tree snapshot for repoURL and +// writes it to w as a sequential byte stream, fetching up to concurrency +// chunkSize-byte ranges concurrently and reordering them through a spill file +// in spillDir (see [ParallelGetStream]). When concurrency is 1 the single +// response is streamed directly to w with no spill file; a server without +// range support falls back to a single full download, which still round-trips +// through the spill file. 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) DownloadGitSnapshotStream(ctx context.Context, repoURL string, w io.Writer, chunkSize int64, concurrency int, spillDir string) (GitSnapshotMetadata, error) { + return c.downloadGitSnapshot(repoURL, func(reader *gitArtifactRangeReader) error { + return errors.WithStack(ParallelGetStream(ctx, reader, NewKey(repoURL), w, chunkSize, concurrency, spillDir)) + }) +} + +// downloadGitSnapshot resolves repoURL's snapshot endpoint, runs fetch against +// a metadata-recording RangeReader for it, and returns the recorded freshen +// metadata. +func (c *Client) downloadGitSnapshot(repoURL string, fetch func(*gitArtifactRangeReader) error) (GitSnapshotMetadata, error) { endpoint, err := gitEndpointURL(c.baseURL, repoURL, "snapshot.tar.zst") if err != nil { return GitSnapshotMetadata{}, err } reader := &gitArtifactRangeReader{client: c, endpoint: endpoint} - if err := ParallelGet(ctx, reader, NewKey(repoURL), dst, chunkSize, concurrency); err != nil { + if err := fetch(reader); err != nil { return GitSnapshotMetadata{}, errors.Wrap(err, "download snapshot") } return reader.metadata(), nil diff --git a/client/git_test.go b/client/git_test.go index 0e2ac435..8097ec00 100644 --- a/client/git_test.go +++ b/client/git_test.go @@ -175,14 +175,11 @@ func TestOpenGitBundleNotFound(t *testing.T) { assert.True(t, errors.Is(err, os.ErrNotExist)) } -func TestDownloadGitSnapshotParallel(t *testing.T) { - body := make([]byte, 1000) - for i := range body { - body[i] = byte(i % 251) - } - const etag = `"snap-v1"` - - var requests atomic.Int64 +// newSnapshotRangeServer serves body as a ranged snapshot endpoint, counting +// requests. ServeContent honours Range/If-Range against the ETag, so it +// returns 206 + Content-Range for the chunked requests ParallelGet makes. +func newSnapshotRangeServer(t *testing.T, body []byte, etag string, requests *atomic.Int64) *httptest.Server { + t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { assert.Equal(t, "/git/github.com/org/repo/snapshot.tar.zst", r.URL.Path) requests.Add(1) @@ -190,11 +187,20 @@ 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() + t.Cleanup(srv.Close) + return srv +} + +func TestDownloadGitSnapshotParallel(t *testing.T) { + body := make([]byte, 1000) + for i := range body { + body[i] = byte(i % 251) + } + + var requests atomic.Int64 + srv := newSnapshotRangeServer(t, body, `"snap-v1"`, &requests) api := client.NewWithHTTPClient(srv.URL, srv.Client()) var dst bufferAt @@ -208,6 +214,75 @@ func TestDownloadGitSnapshotParallel(t *testing.T) { assert.True(t, requests.Load() > 1, "expected multiple range requests, got %d", requests.Load()) } +func TestDownloadGitSnapshotStreamParallel(t *testing.T) { + body := make([]byte, 1000) + for i := range body { + body[i] = byte(i % 251) + } + + var requests atomic.Int64 + srv := newSnapshotRangeServer(t, body, `"snap-v1"`, &requests) + + api := client.NewWithHTTPClient(srv.URL, srv.Client()) + var out bytes.Buffer + meta, err := api.DownloadGitSnapshotStream(context.Background(), "https://github.com/org/repo", &out, 128, 4, t.TempDir()) + assert.NoError(t, err) + assert.Equal(t, body, out.Bytes()) + assert.Equal(t, "deadbeef", meta.Commit) + assert.Equal(t, "/git/github.com/org/repo/snapshot.bundle?base=deadbeef", meta.BundleURL) + assert.True(t, requests.Load() > 1, "expected multiple range requests, got %d", requests.Load()) +} + +func TestDownloadGitSnapshotStreamSingleRequest(t *testing.T) { + body := make([]byte, 1000) + for i := range body { + body[i] = byte(i % 251) + } + + var requests atomic.Int64 + srv := newSnapshotRangeServer(t, body, `"snap-v1"`, &requests) + + api := client.NewWithHTTPClient(srv.URL, srv.Client()) + var out bytes.Buffer + meta, err := api.DownloadGitSnapshotStream(context.Background(), "https://github.com/org/repo", &out, 128, 1, t.TempDir()) + assert.NoError(t, err) + assert.Equal(t, body, out.Bytes()) + assert.Equal(t, "deadbeef", meta.Commit) + assert.Equal(t, int64(1), requests.Load()) +} + +func TestDownloadGitSnapshotStreamFallsBackWithoutRange(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. ParallelGetStream must fall back to a + // single read through the spill file. + 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 out bytes.Buffer + meta, err := api.DownloadGitSnapshotStream(context.Background(), "https://github.com/org/repo", &out, 8, 4, t.TempDir()) + assert.NoError(t, err) + assert.Equal(t, body, out.Bytes()) + assert.Equal(t, "cafe", meta.Commit) +} + +func TestDownloadGitSnapshotStreamNotFound(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 out bytes.Buffer + _, err := api.DownloadGitSnapshotStream(context.Background(), "https://github.com/org/repo", &out, 8, 4, t.TempDir()) + assert.True(t, errors.Is(err, os.ErrNotExist)) +} + func TestDownloadGitSnapshotFallsBackWithoutRange(t *testing.T) { body := []byte("full body, server ignores ranges") srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { diff --git a/client/internal/punchholelayout/layout.go b/client/internal/punchholelayout/layout.go new file mode 100644 index 00000000..9f6b83ce --- /dev/null +++ b/client/internal/punchholelayout/layout.go @@ -0,0 +1,18 @@ +// Package punchholelayout captures the memory layout of the macOS SDK's +// fpunchhole_t via cgo, so tests can verify that the hand-defined Go mirror +// in package client stays identical to the C struct. +package punchholelayout + +// Layout describes the size and field layout of a punch-hole argument struct, +// in bytes. +type Layout struct { + Size uintptr + FlagsOffset uintptr + FlagsSize uintptr + ReservedOffset uintptr + ReservedSize uintptr + OffsetOffset uintptr + OffsetSize uintptr + LengthOffset uintptr + LengthSize uintptr +} diff --git a/client/internal/punchholelayout/layout_darwin_cgo.go b/client/internal/punchholelayout/layout_darwin_cgo.go new file mode 100644 index 00000000..490e357f --- /dev/null +++ b/client/internal/punchholelayout/layout_darwin_cgo.go @@ -0,0 +1,27 @@ +//go:build darwin && cgo + +package punchholelayout + +/* +#include +*/ +import "C" + +import "unsafe" + +// SDKLayout returns the layout of struct fpunchhole_t as compiled from the +// SDK's . +func SDKLayout() (Layout, bool) { + var v C.fpunchhole_t + return Layout{ + Size: unsafe.Sizeof(v), + FlagsOffset: unsafe.Offsetof(v.fp_flags), + FlagsSize: unsafe.Sizeof(v.fp_flags), + ReservedOffset: unsafe.Offsetof(v.reserved), + ReservedSize: unsafe.Sizeof(v.reserved), + OffsetOffset: unsafe.Offsetof(v.fp_offset), + OffsetSize: unsafe.Sizeof(v.fp_offset), + LengthOffset: unsafe.Offsetof(v.fp_length), + LengthSize: unsafe.Sizeof(v.fp_length), + }, true +} diff --git a/client/internal/punchholelayout/layout_stub.go b/client/internal/punchholelayout/layout_stub.go new file mode 100644 index 00000000..4617e793 --- /dev/null +++ b/client/internal/punchholelayout/layout_stub.go @@ -0,0 +1,9 @@ +//go:build !darwin || !cgo + +package punchholelayout + +// SDKLayout reports false: without cgo on darwin there is no SDK struct to +// compare against. +func SDKLayout() (Layout, bool) { + return Layout{}, false +} diff --git a/client/parallel_get.go b/client/parallel_get.go index 9db919b7..9523a747 100644 --- a/client/parallel_get.go +++ b/client/parallel_get.go @@ -91,6 +91,39 @@ func ParallelGet(ctx context.Context, c RangeReader, key Key, dst io.WriterAt, c return errors.Wrap(eg.Wait(), "parallel get") } +// ParallelGetStream downloads an object from any Range-capable RangeReader +// into w as a sequential byte stream, while still fetching chunkSize-byte +// chunks with up to concurrency concurrent requests. Out-of-order chunks are +// reordered through an unlinked spill file created in spillDir rather than +// held in RAM; consumed regions of the spill file are hole-punched where the +// platform supports it, so its disk footprint tracks the gap between download +// and consumption. A concurrency of 1 streams the single response directly to +// w and touches no disk. Revision-safety semantics match [ParallelGet]: a +// partially written w must be discarded by the caller on failure. +func ParallelGetStream(ctx context.Context, c RangeReader, key Key, w io.Writer, chunkSize int64, concurrency int, spillDir string) error { + // chunkSize is irrelevant to the single-request path, so it is validated + // by ParallelGet only when actually chunking. + if concurrency <= 1 { + rc, _, err := c.Open(ctx, key) + if err != nil { + return errors.Wrap(err, "parallel get: full read") + } + _, copyErr := io.Copy(w, rc) + return errors.Wrap(errors.Join(copyErr, rc.Close()), "parallel get") + } + + sb, err := newSpillBuffer(spillDir) + if err != nil { + return errors.Wrap(err, "parallel get") + } + eg, egCtx := errgroup.WithContext(ctx) + eg.Go(func() error { + return sb.closeWrite(ParallelGet(egCtx, c, key, sb, chunkSize, concurrency)) + }) + eg.Go(func() error { return sb.streamTo(egCtx, w) }) + return errors.Join(eg.Wait(), sb.Close()) +} + func fullRead(ctx context.Context, c RangeReader, key Key, dst io.WriterAt) error { rc, _, err := c.Open(ctx, key) if err != nil { diff --git a/client/parallel_get_test.go b/client/parallel_get_test.go index 3d74ddfb..bd3aca11 100644 --- a/client/parallel_get_test.go +++ b/client/parallel_get_test.go @@ -14,6 +14,7 @@ import ( "testing" "github.com/alecthomas/assert/v2" + "github.com/alecthomas/errors" "github.com/block/cachew/client" ) @@ -249,6 +250,62 @@ func TestParallelGetNoETagSizeChangedBetweenRequests(t *testing.T) { assert.Equal(t, c.rewritten, dst.buf) } +func TestParallelGetStream(t *testing.T) { + data := make([]byte, 1000) + for i := range data { + data[i] = byte(i % 251) + } + + tests := []struct { + name string + reader client.RangeReader + concurrency int + want []byte + wantErrIs error + wantErrText string + }{ + {name: "MultiChunk", reader: &recordingReader{data: data, etag: `"v1"`}, concurrency: 4, want: data}, + {name: "SingleWorker", reader: &recordingReader{data: data, etag: `"v1"`}, concurrency: 1, want: data}, + {name: "NoETagFallback", reader: &noETagReader{data: data}, concurrency: 4, want: data}, + {name: "Empty", reader: &recordingReader{etag: `"v1"`}, concurrency: 4}, + {name: "ETagMismatch", reader: &rangeFlipReader{data: data, firstETag: `"v1"`, restETag: `"v2"`}, + concurrency: 4, wantErrText: "object changed during read"}, + {name: "PreconditionFailed", reader: &ifMatchFlipReader{data: data, firstETag: `"v1"`, restETag: `"v2"`}, + concurrency: 4, wantErrIs: client.ErrPreconditionFailed}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var out bytes.Buffer + err := client.ParallelGetStream(context.Background(), tt.reader, client.NewKey("k"), &out, 100, tt.concurrency, t.TempDir()) + switch { + case tt.wantErrIs != nil: + assert.IsError(t, err, tt.wantErrIs) + case tt.wantErrText != "": + assert.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErrText) + default: + assert.NoError(t, err) + assert.Equal(t, tt.want, out.Bytes()) + } + }) + } +} + +type failingWriter struct{ err error } + +func (f *failingWriter) Write([]byte) (int, error) { return 0, f.err } + +func TestParallelGetStreamWriterError(t *testing.T) { + data := make([]byte, 1000) + for i := range data { + data[i] = byte(i % 251) + } + c := &recordingReader{data: data, etag: `"v1"`} + sinkErr := errors.New("sink failed") + err := client.ParallelGetStream(context.Background(), c, client.NewKey("k"), &failingWriter{err: sinkErr}, 100, 4, t.TempDir()) + assert.IsError(t, err, sinkErr) +} + func TestParallelGetClient(t *testing.T) { content := make([]byte, 1000) for i := range content { diff --git a/client/punchhole_darwin.go b/client/punchhole_darwin.go new file mode 100644 index 00000000..67d7e54d --- /dev/null +++ b/client/punchhole_darwin.go @@ -0,0 +1,34 @@ +//go:build darwin + +package client + +import ( + "os" + "runtime" + "unsafe" + + "github.com/alecthomas/errors" + "golang.org/x/sys/unix" +) + +// fpunchhole mirrors struct fpunchhole_t from . x/sys/unix +// defines F_PUNCHHOLE for darwin but neither the argument struct nor a +// wrapper, so we pass the pointer through FcntlInt the same way +// unix.FcntlFstore does for F_PREALLOCATE. +type fpunchhole struct { + flags uint32 // fp_flags: no flags are currently defined + reserved uint32 + offset int64 + length int64 +} + +// punchHole deallocates the byte range [off, off+length) of f, leaving a hole +// that reads as zeros without changing the file size. off and length must be +// multiples of the filesystem block size. Requires APFS. +func punchHole(f *os.File, off, length int64) error { + arg := fpunchhole{offset: off, length: length} + //nolint:gosec // pointers fit in int on darwin; the same pattern unix.FcntlFstore uses + _, err := unix.FcntlInt(f.Fd(), unix.F_PUNCHHOLE, int(uintptr(unsafe.Pointer(&arg)))) + runtime.KeepAlive(&arg) + return errors.Wrap(err, "punch hole") +} diff --git a/client/punchhole_internal_test.go b/client/punchhole_internal_test.go new file mode 100644 index 00000000..ad414ce3 --- /dev/null +++ b/client/punchhole_internal_test.go @@ -0,0 +1,123 @@ +//go:build linux || darwin + +package client + +import ( + "bytes" + "context" + "os" + "path/filepath" + "syscall" + "testing" + + "github.com/alecthomas/assert/v2" +) + +func TestPunchHole(t *testing.T) { + f := newPatternFile(t, 8<<20) + before := allocatedBytes(t, f) + assert.True(t, before >= 7<<20, "expected ~8MiB allocated, got %d", before) + + if err := punchHole(f, 0, 4<<20); err != nil { + t.Skipf("hole punching not supported on this filesystem: %v", err) + } + assert.NoError(t, f.Sync()) + after := allocatedBytes(t, f) + assert.True(t, after <= before-3<<20, "expected at least 3MiB deallocated, before=%d after=%d", before, after) + + fi, err := f.Stat() + assert.NoError(t, err) + assert.Equal(t, int64(8<<20), fi.Size()) + buf := make([]byte, 4096) + _, err = f.ReadAt(buf, 4096) + assert.NoError(t, err) + assert.Equal(t, make([]byte, 4096), buf) +} + +func TestSpillBufferPunchesConsumedRegions(t *testing.T) { + requirePunchSupport(t) + + sb, err := newSpillBuffer(t.TempDir()) + assert.NoError(t, err) + defer sb.Close() + + const size = 16 << 20 + const half = size / 2 + const chunk = 1 << 20 + data := patternBytes(size) + for off := 0; off < half; off += chunk { + _, err := sb.WriteAt(data[off:off+chunk], int64(off)) + assert.NoError(t, err) + } + + gw := &gatedWriter{gateAt: 6 << 20, blocked: make(chan struct{}), release: make(chan struct{})} + done := make(chan error, 1) + go func() { done <- sb.streamTo(context.Background(), gw) }() + + <-gw.blocked + for off := half; off < size; off += chunk { + _, err := sb.WriteAt(data[off:off+chunk], int64(off)) + assert.NoError(t, err) + } + assert.NoError(t, sb.f.Sync()) + before := allocatedBytes(t, sb.f) + assert.True(t, before >= size*3/4, "expected ~16MiB allocated, got %d", before) + assert.NoError(t, sb.closeWrite(nil)) + close(gw.release) + + assert.NoError(t, <-done) + assert.Equal(t, data, gw.buf.Bytes()) + assert.NoError(t, sb.f.Sync()) + after := allocatedBytes(t, sb.f) + assert.True(t, after <= size/8, "expected consumed regions deallocated, before=%d after=%d", before, after) +} + +// gatedWriter blocks streamTo mid-drain: crossing gateAt closes blocked and +// waits for release, so the test can land more writes and force the first +// punch to happen before the second half is read back. A punch that zeroed +// not-yet-streamed regions would then corrupt the second half of buf. +type gatedWriter struct { + buf bytes.Buffer + gateAt int + blocked chan struct{} + release chan struct{} + gated bool +} + +func (g *gatedWriter) Write(p []byte) (int, error) { + if !g.gated && g.buf.Len()+len(p) >= g.gateAt { + g.gated = true + close(g.blocked) + <-g.release + } + return g.buf.Write(p) +} + +func requirePunchSupport(t *testing.T) { + t.Helper() + f := newPatternFile(t, spillPunchAlign) + if err := punchHole(f, 0, spillPunchAlign); err != nil { + t.Skipf("hole punching not supported on this filesystem: %v", err) + } +} + +func newPatternFile(t *testing.T, size int64) *os.File { + t.Helper() + f, err := os.Create(filepath.Join(t.TempDir(), "punch")) + assert.NoError(t, err) + t.Cleanup(func() { _ = f.Close() }) + _, err = f.Write(patternBytes(int(size))) + assert.NoError(t, err) + assert.NoError(t, f.Sync()) + return f +} + +func allocatedBytes(t *testing.T, f *os.File) int64 { + t.Helper() + fi, err := f.Stat() + assert.NoError(t, err) + st, ok := fi.Sys().(*syscall.Stat_t) + assert.True(t, ok) + // Blocks counts 512-byte units regardless of the filesystem block size. + return st.Blocks * 512 +} diff --git a/client/punchhole_linux.go b/client/punchhole_linux.go new file mode 100644 index 00000000..6696f746 --- /dev/null +++ b/client/punchhole_linux.go @@ -0,0 +1,19 @@ +//go:build linux + +package client + +import ( + "os" + + "github.com/alecthomas/errors" + "golang.org/x/sys/unix" +) + +// punchHole deallocates the byte range [off, off+length) of f, leaving a hole +// that reads as zeros without changing the file size. off and length must be +// multiples of the filesystem block size. +func punchHole(f *os.File, off, length int64) error { + //nolint:gosec // file descriptors fit in int; the conversion every unix.* call site uses + err := unix.Fallocate(int(f.Fd()), unix.FALLOC_FL_PUNCH_HOLE|unix.FALLOC_FL_KEEP_SIZE, off, length) + return errors.Wrap(err, "punch hole") +} diff --git a/client/punchhole_unsupported.go b/client/punchhole_unsupported.go new file mode 100644 index 00000000..580a1bc4 --- /dev/null +++ b/client/punchhole_unsupported.go @@ -0,0 +1,16 @@ +//go:build !linux && !darwin + +package client + +import ( + "os" + + "github.com/alecthomas/errors" +) + +// punchHole reports that hole punching is unavailable. Callers treat it as +// best-effort: the spilled data stays allocated, costing no more disk than a +// plain temporary file. +func punchHole(_ *os.File, _, _ int64) error { + return errors.New("hole punching unsupported on this platform") +} diff --git a/client/punchholelayout_internal_test.go b/client/punchholelayout_internal_test.go new file mode 100644 index 00000000..93226ea3 --- /dev/null +++ b/client/punchholelayout_internal_test.go @@ -0,0 +1,31 @@ +//go:build darwin + +package client + +import ( + "testing" + "unsafe" + + "github.com/alecthomas/assert/v2" + + "github.com/block/cachew/client/internal/punchholelayout" +) + +func TestFpunchholeMatchesSDKLayout(t *testing.T) { + sdk, ok := punchholelayout.SDKLayout() + if !ok { + t.Skip("cgo unavailable, cannot compare against the SDK's fpunchhole_t") + } + var arg fpunchhole + assert.Equal(t, punchholelayout.Layout{ + Size: unsafe.Sizeof(arg), + FlagsOffset: unsafe.Offsetof(arg.flags), + FlagsSize: unsafe.Sizeof(arg.flags), + ReservedOffset: unsafe.Offsetof(arg.reserved), + ReservedSize: unsafe.Sizeof(arg.reserved), + OffsetOffset: unsafe.Offsetof(arg.offset), + OffsetSize: unsafe.Sizeof(arg.offset), + LengthOffset: unsafe.Offsetof(arg.length), + LengthSize: unsafe.Sizeof(arg.length), + }, sdk) +} diff --git a/client/spill.go b/client/spill.go new file mode 100644 index 00000000..e7d7ea69 --- /dev/null +++ b/client/spill.go @@ -0,0 +1,179 @@ +package client + +import ( + "cmp" + "context" + "io" + "os" + "slices" + "sync" + + "github.com/alecthomas/errors" +) + +const ( + // spillPunchAlign keeps hole boundaries at a multiple of any plausible + // filesystem block size, as required by fallocate and F_PUNCHHOLE. + spillPunchAlign int64 = 1 << 20 + // spillPunchThreshold batches hole punching to one syscall per few + // megabytes consumed rather than one per read. + spillPunchThreshold int64 = 4 << 20 + + spillReadBufferSize = 256 << 10 +) + +// spillBuffer reorders concurrent WriteAt calls into a sequential byte stream +// through an unlinked temporary file, so a parallel ranged download can feed a +// pipe without holding chunks in RAM. streamTo follows the contiguous write +// frontier and punches holes in consumed regions (best effort), keeping the +// file's disk footprint near the gap between download and consumption. +type spillBuffer struct { + f *os.File + name string + + mu sync.Mutex + changed chan struct{} // closed and replaced whenever frontier or closed changes + pending []spillInterval // written ranges not yet contiguous with the frontier + frontier int64 // contiguous bytes written from offset 0 + closed bool + writeErr error +} + +type spillInterval struct{ start, end int64 } + +func newSpillBuffer(dir string) (*spillBuffer, error) { + f, err := os.CreateTemp(dir, ".cachew-spill-*") + if err != nil { + return nil, errors.Wrap(err, "create spill file") + } + // Unlink immediately so the kernel reclaims the space even if the process + // dies mid-download; Close removes it again for platforms where unlinking + // an open file fails. + _ = os.Remove(f.Name()) + return &spillBuffer{f: f, name: f.Name(), changed: make(chan struct{})}, nil +} + +// Close releases the spill file. Callers must not call WriteAt or streamTo +// after Close. +func (s *spillBuffer) Close() error { + err := s.f.Close() + if rmErr := os.Remove(s.name); rmErr != nil && !errors.Is(rmErr, os.ErrNotExist) { + err = errors.Join(err, rmErr) + } + return errors.WithStack(err) +} + +// WriteAt implements io.WriterAt. Writes at non-overlapping offsets may run +// concurrently. +func (s *spillBuffer) WriteAt(p []byte, off int64) (int, error) { + n, err := s.f.WriteAt(p, off) + if n > 0 { + s.extend(off, off+int64(n)) + } + return n, errors.WithStack(err) +} + +// closeWrite marks the write side finished. A nil err with incomplete +// coverage from offset 0 is reported as an error so a gap can never silently +// truncate the stream. It returns the error recorded for the reader. +func (s *spillBuffer) closeWrite(err error) error { + s.mu.Lock() + defer s.mu.Unlock() + if err == nil && len(s.pending) > 0 { + err = errors.Errorf("spill: write coverage gap at offset %d", s.frontier) + } + s.closed = true + s.writeErr = err + s.wake() + return err +} + +// streamTo copies bytes to w in offset order as the contiguous write frontier +// advances, punching holes in consumed regions when the platform supports it. +// It returns nil once the write side is closed and all contiguous bytes have +// been flushed; a write-side error also ends the stream with nil, since the +// writer reports its own failure. +func (s *spillBuffer) streamTo(ctx context.Context, w io.Writer) error { + buf := make([]byte, spillReadBufferSize) + var readOff, punched int64 + punchable := true + for { + frontier, done, err := s.waitFrontier(ctx, readOff) + if err != nil { + return err + } + if done { + return nil + } + for readOff < frontier { + n := int(min(frontier-readOff, int64(len(buf)))) + if _, err := s.f.ReadAt(buf[:n], readOff); err != nil { + return errors.Wrap(err, "read spill file") + } + if _, err := w.Write(buf[:n]); err != nil { + return errors.Wrap(err, "write stream") + } + readOff += int64(n) + } + if aligned := readOff &^ (spillPunchAlign - 1); punchable && aligned-punched >= spillPunchThreshold { + // Best effort: on filesystems without hole support the spill + // costs no more disk than a plain temporary file. + punchable = punchHole(s.f, punched, aligned-punched) == nil + if punchable { + punched = aligned + } + } + } +} + +// waitFrontier blocks until the frontier moves past readOff, the write side +// closes, or ctx is cancelled. done reports that no more bytes will arrive. +func (s *spillBuffer) waitFrontier(ctx context.Context, readOff int64) (frontier int64, done bool, err error) { + s.mu.Lock() + for s.frontier == readOff && !s.closed { + ch := s.changed + s.mu.Unlock() + select { + case <-ch: + case <-ctx.Done(): + return 0, false, errors.WithStack(ctx.Err()) + } + s.mu.Lock() + } + frontier, closed, writeErr := s.frontier, s.closed, s.writeErr + s.mu.Unlock() + return frontier, (frontier == readOff && closed) || writeErr != nil, nil +} + +// extend records [start, end) as written and advances the contiguous frontier +// over any pending intervals it now reaches. +func (s *spillBuffer) extend(start, end int64) { + s.mu.Lock() + defer s.mu.Unlock() + s.pending = append(s.pending, spillInterval{start, end}) + slices.SortFunc(s.pending, func(a, b spillInterval) int { return cmp.Compare(a.start, b.start) }) + merged := s.pending[:1] + for _, iv := range s.pending[1:] { + if last := &merged[len(merged)-1]; iv.start <= last.end { + last.end = max(last.end, iv.end) + } else { + merged = append(merged, iv) + } + } + s.pending = merged + advanced := false + for len(s.pending) > 0 && s.pending[0].start <= s.frontier { + s.frontier = max(s.frontier, s.pending[0].end) + s.pending = s.pending[1:] + advanced = true + } + if advanced { + s.wake() + } +} + +// wake must be called with mu held. +func (s *spillBuffer) wake() { + close(s.changed) + s.changed = make(chan struct{}) +} diff --git a/client/spill_internal_test.go b/client/spill_internal_test.go new file mode 100644 index 00000000..74801936 --- /dev/null +++ b/client/spill_internal_test.go @@ -0,0 +1,129 @@ +package client + +import ( + "bytes" + "context" + "sync" + "testing" + + "github.com/alecthomas/assert/v2" + "github.com/alecthomas/errors" +) + +func TestSpillBufferReordersWrites(t *testing.T) { + data := patternBytes(1 << 20) + sb, err := newSpillBuffer(t.TempDir()) + assert.NoError(t, err) + defer sb.Close() + + var out bytes.Buffer + done := make(chan error, 1) + go func() { done <- sb.streamTo(context.Background(), &out) }() + + const chunk = 64 << 10 + for _, off := range scrambledOffsets(len(data), chunk) { + _, err := sb.WriteAt(data[off:off+chunk], int64(off)) + assert.NoError(t, err) + } + assert.NoError(t, sb.closeWrite(nil)) + assert.NoError(t, <-done) + assert.Equal(t, data, out.Bytes()) +} + +func TestSpillBufferConcurrentWriters(t *testing.T) { + const chunk = 256 << 10 + const workers = 8 + data := patternBytes(8 << 20) + sb, err := newSpillBuffer(t.TempDir()) + assert.NoError(t, err) + defer sb.Close() + + var out bytes.Buffer + done := make(chan error, 1) + go func() { done <- sb.streamTo(context.Background(), &out) }() + + offsets := make(chan int, len(data)/chunk) + for _, off := range scrambledOffsets(len(data), chunk) { + offsets <- off + } + close(offsets) + + writeErrs := make(chan error, workers) + var wg sync.WaitGroup + for range workers { + wg.Go(func() { + for off := range offsets { + if _, err := sb.WriteAt(data[off:off+chunk], int64(off)); err != nil { + writeErrs <- err + return + } + } + }) + } + wg.Wait() + close(writeErrs) + for err := range writeErrs { + assert.NoError(t, err) + } + assert.NoError(t, sb.closeWrite(nil)) + assert.NoError(t, <-done) + assert.Equal(t, data, out.Bytes()) +} + +func TestSpillBufferCoverageGap(t *testing.T) { + sb, err := newSpillBuffer(t.TempDir()) + assert.NoError(t, err) + defer sb.Close() + + _, err = sb.WriteAt([]byte("abc"), 10) + assert.NoError(t, err) + err = sb.closeWrite(nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "coverage gap") +} + +func TestSpillBufferStreamCancellation(t *testing.T) { + sb, err := newSpillBuffer(t.TempDir()) + assert.NoError(t, err) + defer sb.Close() + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- sb.streamTo(ctx, &bytes.Buffer{}) }() + cancel() + assert.IsError(t, <-done, context.Canceled) +} + +func TestSpillBufferWriteErrorEndsStream(t *testing.T) { + sb, err := newSpillBuffer(t.TempDir()) + assert.NoError(t, err) + defer sb.Close() + + done := make(chan error, 1) + go func() { done <- sb.streamTo(context.Background(), &bytes.Buffer{}) }() + _, err = sb.WriteAt([]byte("prefix"), 0) + assert.NoError(t, err) + boom := errors.New("boom") + assert.IsError(t, sb.closeWrite(boom), boom) + assert.NoError(t, <-done) +} + +// scrambledOffsets returns every chunk offset in [0, size) exactly once, out +// of order: a stride of 7 is coprime with any power-of-two chunk count, so +// the walk is a permutation. +func scrambledOffsets(size, chunk int) []int { + numChunks := size / chunk + offsets := make([]int, numChunks) + for i := range numChunks { + offsets[i] = ((i*7 + 3) % numChunks) * chunk + } + return offsets +} + +func patternBytes(n int) []byte { + data := make([]byte, n) + for i := range data { + data[i] = byte(i % 251) + } + return data +} diff --git a/cmd/cachew/git.go b/cmd/cachew/git.go index 0774314c..4e98cbd5 100644 --- a/cmd/cachew/git.go +++ b/cmd/cachew/git.go @@ -8,6 +8,7 @@ import ( "os/exec" "path/filepath" "strings" + "sync" "time" "github.com/alecthomas/errors" @@ -15,6 +16,7 @@ import ( "go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "go.opentelemetry.io/otel/trace" + "golang.org/x/sync/errgroup" "github.com/block/cachew/client" "github.com/block/cachew/internal/snapshot" @@ -116,99 +118,77 @@ 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 streams the snapshot download straight into +// extraction, overlapping the two, and returns its freshen metadata (commit +// and bundle URL). Concurrent range requests are reordered through a spill +// file rather than staging a full temp copy of the snapshot (see +// [client.ParallelGetStream]). Because extraction overlaps the download, a +// failed restore may leave the target directory partially extracted; callers +// must treat it as disposable on error, as with the streaming restores that +// preceded parallel downloads. func (c *GitRestoreCmd) fetchAndExtractSnapshot(ctx context.Context, api *client.Client) (commit, bundleURL string, err error) { - if c.DownloadConcurrency > 1 { - return c.parallelFetchAndExtract(ctx, api) + // Stage the spill file 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. + spillDir := filepath.Dir(c.Directory) + if err := os.MkdirAll(spillDir, 0o750); err != nil { + return "", "", errors.Wrap(err, "create snapshot spill dir") } - 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)}, - func(ctx context.Context) error { - downloadStart := time.Now() - s, err := api.OpenGitSnapshot(ctx, c.RepoURL) - if err != nil { - return err //nolint:wrapcheck // wrapped by caller - } - snap = s - trace.SpanFromContext(ctx).SetAttributes( - attribute.String("cachew.snapshot_commit", s.Commit), - attribute.String("cachew.bundle_url", s.BundleURL), - attribute.Float64("cachew.elapsed_seconds", time.Since(downloadStart).Seconds()), - ) - return nil - }); err != nil { - return "", "", err - } - defer snap.Close() - - if err := c.extract(ctx, snap.Body); err != nil { - return "", "", err - } - 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") - } - tmp, err := os.CreateTemp(tmpDir, ".cachew-snapshot-*.tar.zst") - if err != nil { - return "", "", errors.Wrap(err, "create snapshot temp file") - } - defer func() { - _ = tmp.Close() - _ = os.Remove(tmp.Name()) //nolint:gosec // name is from os.CreateTemp, not external input - }() + pr, pw := io.Pipe() 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") + // Whichever side of the pipe fails first is the root cause: its failure is + // recorded before it triggers the teardown (pipe close, group cancel) that + // makes the other side fail with an echo such as a closed pipe, a cancelled + // request, or a killed tar process. + var rootCause error + var rootOnce sync.Once + record := func(err error) { + if err != nil { + rootOnce.Do(func() { rootCause = err }) + } } - if err := c.extract(ctx, tmp); err != nil { - return "", "", err + eg, egCtx := errgroup.WithContext(ctx) + eg.Go(func() error { + downloadErr := inSpan(egCtx, "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.DownloadGitSnapshotStream(ctx, c.RepoURL, pw, + int64(c.DownloadChunkSizeMB)<<20, c.DownloadConcurrency, spillDir) + 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 + }) + record(downloadErr) + pw.CloseWithError(downloadErr) + return downloadErr + }) + eg.Go(func() error { + extractErr := c.extract(egCtx, pr) + record(extractErr) + _ = pr.Close() + return extractErr + }) + if eg.Wait() != nil { + // External cancellation fails both sides with unhelpful echoes, so + // report it as itself for errors.Is(context.Canceled) callers. + if ctxErr := ctx.Err(); ctxErr != nil { + return "", "", errors.WithStack(ctxErr) + } + return "", "", rootCause } return meta.Commit, meta.BundleURL, nil } diff --git a/cmd/cachew/git_test.go b/cmd/cachew/git_test.go index 4eaeeaaa..2621850c 100644 --- a/cmd/cachew/git_test.go +++ b/cmd/cachew/git_test.go @@ -3,6 +3,7 @@ package main import ( "bytes" "context" + "crypto/rand" "encoding/json" "net/http" "net/http/httptest" @@ -15,6 +16,7 @@ import ( "time" "github.com/alecthomas/assert/v2" + "github.com/alecthomas/errors" "github.com/block/cachew/client" ) @@ -168,13 +170,78 @@ func TestGitRestoreSnapshotParallel(t *testing.T) { content, err = os.ReadFile(filepath.Join(dstDir, "subdir", "nested.txt")) assert.NoError(t, err) assert.Equal(t, "nested content", string(content)) +} - // The temp snapshot is staged on the target filesystem and cleaned up. - entries, err := os.ReadDir(filepath.Dir(dstDir)) +func TestGitRestoreMidDownloadFailureSurfacesDownloadError(t *testing.T) { + srcDir := t.TempDir() + big := make([]byte, 8<<20) + _, err := rand.Read(big) assert.NoError(t, err) - for _, e := range entries { - assert.False(t, strings.HasPrefix(e.Name(), ".cachew-snapshot-"), "temp snapshot left behind: %s", e.Name()) + initGitRepo(t, srcDir, map[string]string{"big.bin": string(big)}) + snapshotData := createTarZst(t, srcDir) + assert.True(t, len(snapshotData) > 3<<20, "snapshot too small to span multiple chunks: %d", len(snapshotData)) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if !strings.HasSuffix(r.URL.Path, "/snapshot.tar.zst") { + http.NotFound(w, r) + return + } + if rng := r.Header.Get("Range"); rng != "" && !strings.HasPrefix(rng, "bytes=0-") { + w.WriteHeader(http.StatusInternalServerError) + return + } + w.Header().Set("Content-Type", "application/zstd") + w.Header().Set("ETag", `"snap-v1"`) + http.ServeContent(w, r, "snapshot.tar.zst", time.Time{}, bytes.NewReader(snapshotData)) + })) + defer srv.Close() + + dstDir := filepath.Join(t.TempDir(), "restored") + cmd := &GitRestoreCmd{ + RepoURL: "https://github.com/test/repo", + Directory: dstDir, + DownloadConcurrency: 4, + DownloadChunkSizeMB: 1, } + api := client.NewWithHTTPClient(srv.URL, srv.Client()) + err = cmd.Run(context.Background(), api) + assert.Error(t, err) + var statusErr *client.HTTPStatusError + assert.True(t, errors.As(err, &statusErr), "expected HTTP status error, got: %v", err) + assert.Equal(t, http.StatusInternalServerError, statusErr.StatusCode) + assert.NotContains(t, err.Error(), "zstd failed") +} + +func TestGitRestoreCorruptSnapshotSurfacesExtractionError(t *testing.T) { + garbage := make([]byte, 8<<20) + for i := range garbage { + garbage[i] = byte(i % 251) + } + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/snapshot.tar.zst") { + w.Header().Set("Content-Type", "application/zstd") + w.Header().Set("ETag", `"snap-v1"`) + http.ServeContent(w, r, "snapshot.tar.zst", time.Time{}, bytes.NewReader(garbage)) + return + } + http.NotFound(w, r) + })) + defer srv.Close() + + dstDir := filepath.Join(t.TempDir(), "restored") + cmd := &GitRestoreCmd{ + RepoURL: "https://github.com/test/repo", + Directory: dstDir, + DownloadConcurrency: 4, + DownloadChunkSizeMB: 1, + } + api := client.NewWithHTTPClient(srv.URL, srv.Client()) + err := cmd.Run(context.Background(), api) + assert.Error(t, err) + assert.Contains(t, err.Error(), "zstd failed") + assert.NotContains(t, err.Error(), "closed pipe") + assert.NotContains(t, err.Error(), "context canceled") } func TestGitRestoreWithBundle(t *testing.T) { diff --git a/go.mod b/go.mod index 52a03f2c..2ada4663 100644 --- a/go.mod +++ b/go.mod @@ -26,6 +26,7 @@ require ( go.opentelemetry.io/otel/trace v1.43.0 golang.org/x/mod v0.35.0 golang.org/x/sync v0.20.0 + golang.org/x/sys v0.42.0 ) require ( @@ -81,7 +82,6 @@ require ( go.yaml.in/yaml/v3 v3.0.4 // indirect golang.org/x/crypto v0.49.0 // indirect golang.org/x/net v0.52.0 // indirect - golang.org/x/sys v0.42.0 // indirect golang.org/x/text v0.35.0 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260401024825-9d38bb4040a9 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260401024825-9d38bb4040a9 // indirect diff --git a/internal/cache/parallel_get.go b/internal/cache/parallel_get.go index 55458d07..ec26cd86 100644 --- a/internal/cache/parallel_get.go +++ b/internal/cache/parallel_get.go @@ -15,3 +15,11 @@ import ( 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)) } + +// ParallelGetStream downloads an object from any Range-capable Cache into w +// as a sequential byte stream, reordering concurrently fetched chunks through +// a spill file in spillDir. It delegates to [client.ParallelGetStream]; see +// that function for the full semantics. +func ParallelGetStream(ctx context.Context, c Cache, key Key, w io.Writer, chunkSize int64, concurrency int, spillDir string) error { + return errors.WithStack(client.ParallelGetStream(ctx, c, key, w, chunkSize, concurrency, spillDir)) +} diff --git a/internal/cache/parallel_get_bench_test.go b/internal/cache/parallel_get_bench_test.go new file mode 100644 index 00000000..cef7c494 --- /dev/null +++ b/internal/cache/parallel_get_bench_test.go @@ -0,0 +1,120 @@ +package cache_test + +import ( + "context" + "io" + "log/slog" + "testing" + "time" + + "github.com/alecthomas/assert/v2" + + "github.com/block/cachew/internal/cache" + "github.com/block/cachew/internal/logging" + "github.com/block/cachew/internal/s3client" + "github.com/block/cachew/internal/s3client/s3clienttest" +) + +const ( + // benchObjectSize matches real git snapshots, which run to several GB. + benchObjectSize = 4 << 30 + benchChunkSize = 8 << 20 // the cachew git restore default + benchConcurrency = 8 +) + +// discardWriterAt is the io.WriterAt analogue of io.Discard, so benchmarks +// measure the get strategies rather than a destination file. +type discardWriterAt struct{} + +func (discardWriterAt) WriteAt(p []byte, _ int64) (int, error) { return len(p), nil } + +func newBenchDiskCache(b *testing.B) (context.Context, cache.Cache, cache.Key) { + b.Helper() + _, ctx := logging.Configure(context.Background(), logging.Config{Level: slog.LevelError}) + c, err := cache.NewDisk(ctx, cache.DiskConfig{ + Root: b.TempDir(), + LimitMB: 2 * benchObjectSize >> 20, + MaxTTL: time.Hour, + }) + assert.NoError(b, err) + b.Cleanup(func() { _ = c.Close() }) + return ctx, c, writeBenchObject(ctx, b, c) +} + +func newBenchS3Cache(b *testing.B) (context.Context, cache.Cache, cache.Key) { + b.Helper() + bucket := s3clienttest.Start(b) + _, ctx := logging.Configure(context.Background(), logging.Config{Level: slog.LevelError}) + clientProvider := s3client.NewClientProvider(ctx, s3client.Config{ + Endpoint: s3clienttest.Addr, + UseSSL: false, + }) + c, err := cache.NewS3(ctx, cache.S3Config{ + Bucket: bucket, + MaxTTL: time.Hour, + UploadPartSizeMB: 16, + UploadConcurrency: 0, // all cores, to speed up the 4GiB setup upload + }, clientProvider) + assert.NoError(b, err) + b.Cleanup(func() { _ = c.Close() }) + return ctx, c, writeBenchObject(ctx, b, c) +} + +// writeBenchObject writes the object in pattern-block increments so setup +// never holds gigabytes in RAM. +func writeBenchObject(ctx context.Context, b *testing.B, c cache.Cache) cache.Key { + b.Helper() + pattern := make([]byte, 1<<20) + for i := range pattern { + pattern[i] = byte(i % 251) + } + key := cache.NewKey("bench-object") + assert.NoError(b, cache.WriteFunc(ctx, c, key, nil, time.Hour, func(w io.Writer) error { + for written := 0; written < benchObjectSize; written += len(pattern) { + if _, err := w.Write(pattern); err != nil { + return err + } + } + return nil + })) + return key +} + +func runGetBenchmarks(ctx context.Context, b *testing.B, c cache.Cache, key cache.Key) { + b.Run("Sequential", func(b *testing.B) { + b.SetBytes(benchObjectSize) + for b.Loop() { + rc, _, err := c.Open(ctx, key) + assert.NoError(b, err) + _, err = io.Copy(io.Discard, rc) + assert.NoError(b, err) + assert.NoError(b, rc.Close()) + } + }) + + b.Run("Parallel", func(b *testing.B) { + b.SetBytes(benchObjectSize) + for b.Loop() { + assert.NoError(b, cache.ParallelGet(ctx, c, key, discardWriterAt{}, benchChunkSize, benchConcurrency)) + } + }) + + b.Run("ParallelStream", func(b *testing.B) { + spillDir := b.TempDir() + b.SetBytes(benchObjectSize) + for b.Loop() { + assert.NoError(b, cache.ParallelGetStream(ctx, c, key, io.Discard, benchChunkSize, benchConcurrency, spillDir)) + } + }) +} + +func BenchmarkGet(b *testing.B) { + b.Run("Disk", func(b *testing.B) { + ctx, c, key := newBenchDiskCache(b) + runGetBenchmarks(ctx, b, c, key) + }) + b.Run("S3", func(b *testing.B) { + ctx, c, key := newBenchS3Cache(b) + runGetBenchmarks(ctx, b, c, key) + }) +} diff --git a/internal/cache/parallel_get_test.go b/internal/cache/parallel_get_test.go index 1baa3bb8..51d48a77 100644 --- a/internal/cache/parallel_get_test.go +++ b/internal/cache/parallel_get_test.go @@ -1,6 +1,7 @@ package cache_test import ( + "bytes" "context" "io" "log/slog" @@ -65,6 +66,11 @@ func TestParallelGet(t *testing.T) { err := cache.ParallelGet(ctx, c, key, &dst, tt.chunkSize, tt.concurrency) assert.NoError(t, err) assert.Equal(t, content, dst.buf) + + var stream bytes.Buffer + err = cache.ParallelGetStream(ctx, c, key, &stream, tt.chunkSize, tt.concurrency, t.TempDir()) + assert.NoError(t, err) + assert.Equal(t, content, stream.Bytes()) }) } } @@ -86,6 +92,10 @@ func TestParallelGetEmptyObject(t *testing.T) { var dst bufferAt assert.NoError(t, cache.ParallelGet(ctx, c, key, &dst, 100, concurrency)) assert.Equal(t, 0, len(dst.buf)) + + var stream bytes.Buffer + assert.NoError(t, cache.ParallelGetStream(ctx, c, key, &stream, 100, concurrency, t.TempDir())) + assert.Equal(t, 0, stream.Len()) } } @@ -98,4 +108,7 @@ func TestParallelGetNotFound(t *testing.T) { var dst bufferAt err = cache.ParallelGet(ctx, c, cache.NewKey("missing"), &dst, 100, 4) assert.IsError(t, err, os.ErrNotExist) + + err = cache.ParallelGetStream(ctx, c, cache.NewKey("missing"), &bytes.Buffer{}, 100, 4, t.TempDir()) + assert.IsError(t, err, os.ErrNotExist) } diff --git a/internal/s3client/s3clienttest/s3clienttest.go b/internal/s3client/s3clienttest/s3clienttest.go index 71ecb693..c1351226 100644 --- a/internal/s3client/s3clienttest/s3clienttest.go +++ b/internal/s3client/s3clienttest/s3clienttest.go @@ -33,7 +33,7 @@ var bucketNameRe = regexp.MustCompile(`[^a-z0-9-]`) // Any stale objects left over from a previous (possibly crashed) run // are removed up front, and t.Cleanup removes objects when the test // finishes normally. -func Start(t *testing.T) string { +func Start(t testing.TB) string { t.Helper() t.Setenv("AWS_ACCESS_KEY_ID", Username) @@ -67,7 +67,7 @@ const modulePrefix = "github.com/block/cachew/" // incorporates the calling package path and the test name. This // ensures uniqueness across packages even though t.Name() alone does // not include the package. -func bucketName(t *testing.T) string { +func bucketName(t testing.TB) string { t.Helper() // Walk up the call stack past this package to find the caller's function. @@ -95,7 +95,7 @@ func bucketName(t *testing.T) string { } // CleanBucket removes all objects from the given bucket. -func CleanBucket(t *testing.T, bucket string) { +func CleanBucket(t testing.TB, bucket string) { t.Helper() client := Client(t) for obj := range client.ListObjects(t.Context(), bucket, minio.ListObjectsOptions{Recursive: true}) { @@ -109,7 +109,7 @@ func CleanBucket(t *testing.T, bucket string) { } // Client returns a minio client connected to the test server. -func Client(t *testing.T) *minio.Client { +func Client(t testing.TB) *minio.Client { t.Helper() client, err := minio.New(Addr, &minio.Options{ Creds: credentials.NewStaticV4(Username, Password, ""), @@ -119,7 +119,7 @@ func Client(t *testing.T) *minio.Client { return client } -func startContainer(t *testing.T) { +func startContainer(t testing.TB) { t.Helper() cmd := exec.CommandContext(t.Context(), "docker", "run", "-d", "--name", containerName, @@ -141,14 +141,14 @@ func startContainer(t *testing.T) { } } -func isHealthy(t *testing.T) bool { +func isHealthy(t testing.TB) bool { t.Helper() client := Client(t) _, err := client.ListBuckets(t.Context()) return err == nil } -func waitForReady(t *testing.T) { +func waitForReady(t testing.TB) { t.Helper() client := Client(t) timeout := time.After(30 * time.Second)