|
| 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