Skip to content

Commit 8097115

Browse files
worstellampagent
andcommitted
feat(client): parallel git snapshot download APIs
Add ParallelGet — a concurrent chunked range download that fetches an object in chunkSize-byte chunks (up to concurrency requests in flight) and hands each chunk to a ChunkSink, which owns where the bytes land: - StreamSink reassembles the in-order byte stream for a streaming consumer, buffering out-of-order chunks in a fixed arena of 2*concurrency reusable slots. A slow reader applies backpressure to the fetchers, so peak memory is bounded at O(concurrency*chunkSize) regardless of object size — letting a decompress/extract pipeline overlap the download over a non-seekable sink without staging to disk or RAM. - DiskSink scatters each chunk straight to its offset in an io.WriterAt (e.g. *os.File), concurrently and unordered — the faster path for seekable sinks such as cache-to-cache backfill. ETag-pinned chunks reject mid-download rewrites; over/under-length chunks are rejected; a missing ETag, a range-ignoring backend, an object that fits in the first chunk, or concurrency 1 all fall back to a single full read. Add OpenGitSnapshotParallel: streaming git-snapshot helper returning a GitSnapshot whose Commit/BundleURL are available immediately while bytes stream in the background via a StreamSink; closing it cancels the download. The cachew CLI now extracts directly from the snapshot body (no temp file). Drops the redundant DownloadGitSnapshot. Co-authored-by: Amp <amp@ampcode.com> Amp-Thread-ID: https://ampcode.com/threads/T-019ef6a9-a407-7389-bc43-001405e3ae9e
1 parent 3e830a1 commit 8097115

8 files changed

Lines changed: 836 additions & 278 deletions

File tree

client/chunk_sink.go

Lines changed: 275 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,275 @@
1+
package client
2+
3+
import (
4+
"context"
5+
"io"
6+
"sync"
7+
8+
"github.com/alecthomas/errors"
9+
)
10+
11+
// ChunkSink is the destination ParallelGet places fetched chunks into. The
12+
// engine calls Place once per chunk, concurrently from up to `concurrency`
13+
// goroutines, with the chunk's absolute byte offset and an open body holding
14+
// exactly length bytes (length < 0 means "read the whole body", used for the
15+
// single-stream fallback when the object cannot be chunked). Place must read the
16+
// chunk from body and close body. Implementations own where the bytes land and
17+
// may block in Place to bound memory; a blocked Place must abort when ctx is
18+
// cancelled.
19+
//
20+
// Two implementations cover the cases in this package: StreamSink reassembles
21+
// the in-order byte stream for a streaming consumer, and DiskSink scatters
22+
// chunks to their offsets in a file.
23+
type ChunkSink interface {
24+
Place(ctx context.Context, off, length int64, body io.ReadCloser) error
25+
}
26+
27+
// StreamSink is a ChunkSink that reorders concurrently-fetched chunks back into
28+
// the original byte stream, exposed via Read. Chunks land in a fixed arena of
29+
// 2*concurrency reusable slots indexed by chunk number, so a slow consumer
30+
// applies backpressure to the fetchers (capping memory) instead of letting
31+
// fetched-but-unread chunks pile up. The doubled slot count lets the fetchers
32+
// run a full window ahead of the consumer rather than stalling on it.
33+
//
34+
// A StreamSink must be read concurrently while ParallelGet runs — the fetchers
35+
// block once they get a window ahead of the reader, so a caller that does not
36+
// read will deadlock. After the download finishes the caller signals completion
37+
// with Done; Read then drains the remaining buffered chunks and returns io.EOF,
38+
// or the download error.
39+
type StreamSink struct {
40+
chunkSize int64
41+
n int // slot count = 2*concurrency
42+
43+
mu sync.Mutex
44+
cond *sync.Cond // signals Read that a chunk was deposited (or Done)
45+
advance chan struct{} // closed and replaced when readSeq advances, waking blocked Place
46+
bufs [][]byte // n reusable backing buffers, indexed by seq%n (nil until first use)
47+
ready []bool // ready[slot] => bufs[slot] holds the chunk for its current seq
48+
readSeq int64 // sequence number of the chunk Read is emitting next
49+
cur []byte // chunk currently being emitted (aliases bufs[readSeq%n])
50+
curPos int
51+
52+
passthru io.ReadCloser // set in single-stream fallback mode (length < 0)
53+
done bool
54+
err error
55+
closed bool
56+
}
57+
58+
// NewStreamSink returns a StreamSink sized for the given chunk size and download
59+
// concurrency. It holds up to 2*concurrency chunk buffers, giving the fetchers a
60+
// full window of run-ahead over the consumer while capping peak memory at
61+
// 2*concurrency*chunkSize. Buffers are allocated lazily, so a small object never
62+
// reserves the full window.
63+
func NewStreamSink(chunkSize int64, concurrency int) *StreamSink {
64+
n := 2 * max(concurrency, 1)
65+
s := &StreamSink{
66+
chunkSize: chunkSize,
67+
n: n,
68+
bufs: make([][]byte, n),
69+
ready: make([]bool, n),
70+
advance: make(chan struct{}),
71+
}
72+
s.cond = sync.NewCond(&s.mu)
73+
return s
74+
}
75+
76+
// Place reads the chunk into its slot and queues it for in-order delivery to
77+
// Read. It blocks until the chunk is within one window of the read cursor
78+
// (backpressure from a slow consumer) and aborts if ctx is cancelled. A negative
79+
// length switches to pass-through mode: the whole body is handed to Read
80+
// directly, since a single-stream fallback has unknown size and must not be
81+
// buffered.
82+
func (s *StreamSink) Place(ctx context.Context, off, length int64, body io.ReadCloser) error {
83+
if length < 0 {
84+
s.mu.Lock()
85+
if s.closed {
86+
s.mu.Unlock()
87+
return errors.Join(errors.New("stream sink closed"), body.Close())
88+
}
89+
s.passthru = body
90+
s.cond.Broadcast()
91+
s.mu.Unlock()
92+
return nil
93+
}
94+
95+
seq := off / s.chunkSize
96+
slot := int(seq % int64(s.n))
97+
98+
// Admission: a chunk may only occupy its slot once the previous occupant
99+
// (seq-n) has been read, i.e. once seq is within n of the read cursor. This
100+
// bounds run-ahead and guarantees no other in-flight chunk maps to this slot,
101+
// so the in-order chunk's slot is always reserved for it.
102+
s.mu.Lock()
103+
for seq >= s.readSeq+int64(s.n) {
104+
if s.closed {
105+
s.mu.Unlock()
106+
return errors.Join(errors.New("stream sink closed"), body.Close())
107+
}
108+
ch := s.advance
109+
s.mu.Unlock()
110+
select {
111+
case <-ch:
112+
case <-ctx.Done():
113+
return errors.Join(errors.WithStack(ctx.Err()), body.Close())
114+
}
115+
s.mu.Lock()
116+
}
117+
buf := s.bufs[slot]
118+
s.mu.Unlock()
119+
120+
if int64(cap(buf)) < length {
121+
buf = make([]byte, s.chunkSize)
122+
}
123+
buf = buf[:length]
124+
if err := readChunk(off, buf, body); err != nil {
125+
return err
126+
}
127+
128+
s.mu.Lock()
129+
// A Close racing the body read above leaves no reader to drain this slot;
130+
// drop the chunk rather than mark it ready. readChunk already closed body.
131+
if s.closed {
132+
s.mu.Unlock()
133+
return errors.New("stream sink closed")
134+
}
135+
s.bufs[slot] = buf
136+
s.ready[slot] = true
137+
s.cond.Broadcast()
138+
s.mu.Unlock()
139+
return nil
140+
}
141+
142+
// Read emits the reassembled object in order. It blocks until the next chunk is
143+
// available, returning io.EOF once every chunk has been read and Done has been
144+
// called, or the download error reported to Done.
145+
func (s *StreamSink) Read(p []byte) (int, error) {
146+
s.mu.Lock()
147+
for {
148+
if s.passthru != nil {
149+
body := s.passthru
150+
s.mu.Unlock()
151+
return body.Read(p) //nolint:wrapcheck // must return io.EOF verbatim for io.ReadAll
152+
}
153+
if s.cur != nil {
154+
n := copy(p, s.cur[s.curPos:])
155+
s.curPos += n
156+
if s.curPos >= len(s.cur) {
157+
// Chunk fully emitted: free its slot and advance, waking any Place
158+
// blocked waiting for this slot's window to open.
159+
slot := int(s.readSeq % int64(s.n))
160+
s.ready[slot] = false
161+
s.readSeq++
162+
s.cur = nil
163+
s.curPos = 0
164+
close(s.advance)
165+
s.advance = make(chan struct{})
166+
}
167+
s.mu.Unlock()
168+
return n, nil
169+
}
170+
slot := int(s.readSeq % int64(s.n))
171+
if s.ready[slot] {
172+
s.cur = s.bufs[slot]
173+
s.curPos = 0
174+
continue
175+
}
176+
if s.err != nil {
177+
err := s.err
178+
s.mu.Unlock()
179+
return 0, err
180+
}
181+
if s.done {
182+
s.mu.Unlock()
183+
return 0, io.EOF
184+
}
185+
// Closed mid-download with no terminal status: stop rather than block
186+
// forever on cond, since the fetchers are being torn down.
187+
if s.closed {
188+
s.mu.Unlock()
189+
return 0, errors.WithStack(io.ErrClosedPipe)
190+
}
191+
s.cond.Wait()
192+
}
193+
}
194+
195+
// Done signals that no further chunks will be placed. err is the download
196+
// outcome (nil on success); it is surfaced to Read after the buffered chunks
197+
// drain.
198+
func (s *StreamSink) Done(err error) {
199+
s.mu.Lock()
200+
s.done = true
201+
if err != nil && s.err == nil {
202+
s.err = err
203+
}
204+
s.cond.Broadcast()
205+
s.mu.Unlock()
206+
}
207+
208+
// Close releases the sink, unblocking any in-flight Place and closing the
209+
// pass-through body if one is set. The arena buffers are released to the garbage
210+
// collector. Cancelling the download itself is the caller's responsibility (see
211+
// OpenGitSnapshotParallel).
212+
func (s *StreamSink) Close() error {
213+
s.mu.Lock()
214+
s.closed = true
215+
body := s.passthru
216+
s.passthru = nil
217+
close(s.advance)
218+
s.advance = make(chan struct{})
219+
s.cond.Broadcast()
220+
s.mu.Unlock()
221+
if body != nil {
222+
return errors.WithStack(body.Close())
223+
}
224+
return nil
225+
}
226+
227+
// DiskSink is a ChunkSink that writes each chunk straight to its offset in an
228+
// io.WriterAt such as an *os.File. io.WriterAt permits concurrent
229+
// non-overlapping writes, so chunks are scattered to disk as they arrive with no
230+
// reordering and negligible memory — the right sink for seekable destinations
231+
// such as cache-to-cache backfill. Unlike StreamSink it needs no concurrent
232+
// reader, so ParallelGet may run to completion synchronously. On error the
233+
// destination is left partially written and must be discarded by the caller.
234+
type DiskSink struct{ W io.WriterAt }
235+
236+
// Place streams the chunk straight to its offset in the underlying WriterAt.
237+
func (d DiskSink) Place(_ context.Context, off, length int64, body io.ReadCloser) error {
238+
dst := io.NewOffsetWriter(d.W, off)
239+
if length < 0 {
240+
_, err := io.Copy(dst, body)
241+
return errors.Join(errors.Wrap(err, "write chunk"), body.Close())
242+
}
243+
n, err := io.Copy(dst, io.LimitReader(body, length))
244+
if err != nil {
245+
return errors.Join(errors.Errorf("write chunk at offset %d: %w", off, err), body.Close())
246+
}
247+
if n != length {
248+
return errors.Join(errors.Errorf("chunk at offset %d: wrote %d of %d bytes", off, n, length), body.Close())
249+
}
250+
if overlong(body) {
251+
return errors.Join(errors.Errorf("chunk at offset %d: read more than the expected %d bytes", off, length), body.Close())
252+
}
253+
return errors.WithStack(body.Close())
254+
}
255+
256+
// readChunk fills buf from body (reading exactly len(buf) bytes) and closes
257+
// body. A body shorter than buf (a truncated chunk) or longer than buf (a
258+
// backend that ignored the range) is reported as an error.
259+
func readChunk(off int64, buf []byte, body io.ReadCloser) error {
260+
if _, err := io.ReadFull(body, buf); err != nil {
261+
return errors.Join(errors.Errorf("read chunk at offset %d: %w", off, err), body.Close())
262+
}
263+
if overlong(body) {
264+
return errors.Join(errors.Errorf("chunk at offset %d: read more than the expected %d bytes", off, len(buf)), body.Close())
265+
}
266+
return errors.WithStack(body.Close())
267+
}
268+
269+
// overlong reports whether r has any bytes left, used to detect a body longer
270+
// than the requested chunk without buffering the excess.
271+
func overlong(r io.Reader) bool {
272+
var probe [1]byte
273+
n, _ := io.ReadFull(r, probe[:]) //nolint:errcheck // any byte past the chunk is overlong, regardless of the error
274+
return n > 0
275+
}

0 commit comments

Comments
 (0)