66 "net/http"
77 "strconv"
88 "strings"
9+ "sync/atomic"
910
1011 "github.com/alecthomas/errors"
1112 "golang.org/x/sync/errgroup"
@@ -19,27 +20,36 @@ type RangeReader interface {
1920 Open (ctx context.Context , key Key , opts ... RequestOption ) (io.ReadCloser , http.Header , error )
2021}
2122
22- // ParallelGet downloads an object from any Range-capable RangeReader into dst,
23- // fetching it in chunkSize-byte chunks concurrently (up to concurrency requests
24- // in flight) and writing each chunk at its offset via dst.WriteAt. Latency-bound
25- // backends such as a remote cache can saturate bandwidth with overlapping reads.
23+ // ParallelGet downloads an object from any Range-capable RangeReader and writes
24+ // it, in order, to dst. It fetches the object in chunkSize-byte chunks
25+ // concurrently (up to concurrency requests in flight) but emits a single
26+ // sequential byte stream, so a streaming consumer (e.g. a decompress/extract
27+ // pipeline) can run concurrently with the download. Latency-bound backends such
28+ // as a remote cache can saturate bandwidth with overlapping reads.
29+ //
30+ // Chunks complete out of order, so a bounded reorder buffer holds fetched chunks
31+ // until their turn to be written. A window caps the number of
32+ // fetched-but-unwritten chunks (and thus in-flight fetches) at concurrency, so
33+ // peak memory is O(concurrency * chunkSize) regardless of object size or
34+ // consumer speed.
2635//
2736// The first chunk is fetched with a ranged Open, whose response yields both the
2837// total size (from Content-Range) and the object's ETag; every remaining chunk
2938// is then requested with IfRange pinned to that ETag. If the object changes
3039// mid-download, a chunk's ETag will differ and ParallelGet returns an error
31- // rather than splicing bytes from two revisions. A missing or truncated chunk
32- // is likewise reported as an error, so a partially written dst must be discarded
33- // by the caller on failure. An object with no ETag to pin to (e.g. one stored
34- // before ETags were recorded) cannot be kept revision-safe across chunks, so it
35- // falls back to a single full read instead of parallelising. A concurrency of
36- // 1 likewise reads the whole object in one request, since chunking a single
37- // worker would only serialise ranged GETs for no benefit.
40+ // rather than splicing bytes from two revisions. A chunk whose length differs
41+ // from the requested range (a short or overlong response, e.g. a backend that
42+ // ignored the range) is likewise reported as an error, so a partially written
43+ // dst must be discarded by the caller on failure. An object with no ETag to pin
44+ // to (e.g. one stored before ETags were recorded) cannot be kept revision-safe
45+ // across chunks, so it falls back to a single full read instead of
46+ // parallelising. A concurrency of 1 likewise reads the whole object in one
47+ // request, since chunking a single worker would only serialise ranged GETs for
48+ // no benefit.
3849//
39- // dst is written via concurrent WriteAt calls at non-overlapping offsets; the
40- // caller owns dst's lifecycle (open, close, cleanup) and need not pre-size it,
41- // as WriteAt extends the destination.
42- func ParallelGet (ctx context.Context , c RangeReader , key Key , dst io.WriterAt , chunkSize int64 , concurrency int ) error {
50+ // dst is written sequentially from a single goroutine, so it need not be safe
51+ // for concurrent writes.
52+ func ParallelGet (ctx context.Context , c RangeReader , key Key , dst io.Writer , chunkSize int64 , concurrency int ) error {
4353 if chunkSize <= 0 {
4454 return errors .Errorf ("parallel get: chunk size must be positive, got %d" , chunkSize )
4555 }
@@ -52,9 +62,16 @@ func ParallelGet(ctx context.Context, c RangeReader, key Key, dst io.WriterAt, c
5262 return fullRead (ctx , c , key , dst )
5363 }
5464
65+ // Build the group before discovery so chunk zero's request shares egCtx and
66+ // a sibling chunk's failure cancels it too. defer cancel covers the early
67+ // returns below, which exit without calling eg.Wait.
68+ ctx , cancel := context .WithCancel (ctx )
69+ defer cancel ()
70+ eg , egCtx := errgroup .WithContext (ctx )
71+
5572 // Discovery: the first ranged Open delivers chunk zero and reveals the total
5673 // size and ETag used to pin the rest.
57- rc , headers , err := c .Open (ctx , key , Range (0 , chunkSize ))
74+ rc , headers , err := c .Open (egCtx , key , Range (0 , chunkSize ))
5875 if errors .Is (err , ErrRangeNotSatisfiable ) {
5976 return nil // Empty object: nothing to write.
6077 }
@@ -74,7 +91,7 @@ func ParallelGet(ctx context.Context, c RangeReader, key Key, dst io.WriterAt, c
7491 firstLen = - 1
7592 }
7693 if ! hasRange || total <= chunkSize {
77- return errors .Wrap (writeChunkAt (dst , 0 , firstLen , rc ), "parallel get" )
94+ return errors .Wrap (copyChunk (dst , firstLen , rc ), "parallel get" )
7895 }
7996
8097 // Subsequent chunks are pinned to the discovery ETag via IfRange. Without a
@@ -86,67 +103,146 @@ func ParallelGet(ctx context.Context, c RangeReader, key Key, dst io.WriterAt, c
86103 if err := rc .Close (); err != nil {
87104 return errors .Wrap (err , "parallel get: close discovery reader" )
88105 }
89- return fullRead (ctx , c , key , dst )
106+ return fullRead (egCtx , c , key , dst )
90107 }
91108
92- // Multiple chunks: copy the already-open first chunk concurrently with the
93- // rest rather than blocking on it here. The first goroutine is scheduled
94- // before the limit can be reached, so it never stalls holding an open body.
95109 numChunks := int ((total + chunkSize - 1 ) / chunkSize )
96- eg , egCtx := errgroup .WithContext (ctx )
97- eg .SetLimit (concurrency )
98- eg .Go (func () error { return writeChunkAt (dst , 0 , firstLen , rc ) })
99- for seq := 1 ; seq < numChunks ; seq ++ {
100- // Stop scheduling once a chunk has failed and cancelled the group.
101- if egCtx .Err () != nil {
102- break
103- }
104- start := int64 (seq ) * chunkSize
105- end := min (start + chunkSize , total )
106- eg .Go (func () error { return fetchChunk (egCtx , c , key , dst , start , end , etag ) })
110+
111+ // slots is a ring of concurrency reorder buffers carrying chunk bytes from a
112+ // fetching worker to the writer. The window bounds fetched-but-unwritten
113+ // chunks to concurrency, so the outstanding sequence numbers always span a
114+ // window of at most concurrency consecutive values and seq%concurrency is
115+ // unique among them — no two outstanding chunks share a slot. Chunk zero is
116+ // streamed directly from the discovery reader and uses no slot.
117+ slots := make ([]chan []byte , concurrency )
118+ for i := range slots {
119+ slots [i ] = make (chan []byte , 1 )
107120 }
121+
122+ // window bounds fetched-but-unwritten chunks (and in-flight fetches) to
123+ // concurrency. A worker takes a token before fetching; the writer returns
124+ // one after writing a chunk, admitting the next fetch. The channel never
125+ // exceeds capacity because a token is always "held" by the chunk in flight
126+ // or being written at the moment it is returned, so sends never block.
127+ window := make (chan struct {}, concurrency )
128+ for range concurrency {
129+ window <- struct {}{}
130+ }
131+
132+ var nextSeq atomic.Int64
133+ nextSeq .Store (1 )
134+ for range concurrency {
135+ eg .Go (func () error {
136+ for {
137+ select {
138+ case <- egCtx .Done ():
139+ return egCtx .Err ()
140+ case <- window :
141+ }
142+ seq := int (nextSeq .Add (1 ) - 1 )
143+ if seq >= numChunks {
144+ window <- struct {}{} // No work for this token; return it.
145+ return nil
146+ }
147+ start := int64 (seq ) * chunkSize
148+ end := min (start + chunkSize , total )
149+ data , err := fetchChunk (egCtx , c , key , start , end , etag )
150+ if err != nil {
151+ return err
152+ }
153+ slots [seq % concurrency ] <- data // Buffered, single producer: never blocks.
154+ }
155+ })
156+ }
157+
158+ // Writer: stream chunk zero from the discovery reader, then emit each
159+ // subsequent chunk in order, returning a window token after each so workers
160+ // can advance.
161+ eg .Go (func () error {
162+ if err := copyChunk (dst , firstLen , rc ); err != nil {
163+ return err
164+ }
165+ for seq := 1 ; seq < numChunks ; seq ++ {
166+ select {
167+ case <- egCtx .Done ():
168+ return egCtx .Err ()
169+ case data := <- slots [seq % concurrency ]:
170+ if err := writeAll (dst , data , int64 (seq )* chunkSize ); err != nil {
171+ return err
172+ }
173+ window <- struct {}{}
174+ }
175+ }
176+ return nil
177+ })
178+
108179 return errors .Wrap (eg .Wait (), "parallel get" )
109180}
110181
111- // fullRead downloads the entire object in a single request and writes it at
112- // offset zero. It is used when chunking would add no value (a single worker) or
113- // cannot be made revision-safe (no ETag to pin). The body is a single
114- // consistent revision, but its length is unknown up front, so writeChunkAt's
115- // length check is skipped (-1).
116- func fullRead (ctx context.Context , c RangeReader , key Key , dst io.WriterAt ) error {
182+ // fullRead downloads the entire object in a single request and copies it to dst.
183+ // It is used when chunking would add no value (a single worker) or cannot be
184+ // made revision-safe (no ETag to pin). The body is a single consistent revision,
185+ // but its length is unknown up front, so copyChunk's length check is skipped (-1).
186+ func fullRead (ctx context.Context , c RangeReader , key Key , dst io.Writer ) error {
117187 rc , _ , err := c .Open (ctx , key )
118188 if err != nil {
119189 return errors .Wrap (err , "parallel get: full read" )
120190 }
121- return errors .Wrap (writeChunkAt (dst , 0 , - 1 , rc ), "parallel get" )
191+ return errors .Wrap (copyChunk (dst , - 1 , rc ), "parallel get" )
122192}
123193
124- // fetchChunk opens the [start, end) range pinned to etag and writes it at start.
125- // An ETag change (the object was rewritten mid-download) or a short read is
126- // reported as an error.
127- func fetchChunk (ctx context.Context , c RangeReader , key Key , dst io.WriterAt , start , end int64 , etag string ) error {
194+ // copyChunk copies src into dst and closes src. It fails if fewer than want
195+ // bytes arrive; a negative want skips that check (total size unknown).
196+ func copyChunk (dst io.Writer , want int64 , src io.ReadCloser ) error {
197+ n , copyErr := io .Copy (dst , src )
198+ if err := errors .Join (copyErr , src .Close ()); err != nil {
199+ return errors .Errorf ("copy chunk: %w" , err )
200+ }
201+ if want >= 0 && n != want {
202+ return errors .Errorf ("short chunk: copied %d of %d bytes" , n , want )
203+ }
204+ return nil
205+ }
206+
207+ // fetchChunk opens the [start, end) range pinned to etag and returns its bytes.
208+ // An ETag change (the object was rewritten mid-download) or a body whose length
209+ // differs from the requested range (a short read, or an overlong response from a
210+ // backend that ignored the range) is reported as an error.
211+ func fetchChunk (ctx context.Context , c RangeReader , key Key , start , end int64 , etag string ) ([]byte , error ) {
128212 rc , headers , err := c .Open (ctx , key , Range (start , end ), IfRange (etag ))
129213 if err != nil {
130- return errors .Errorf ("open range %d-%d: %w" , start , end , err )
214+ return nil , errors .Errorf ("open range %d-%d: %w" , start , end , err )
131215 }
132216 if got := headers .Get (ETagKey ); got != etag {
133- return errors .Join (
217+ return nil , errors .Join (
134218 errors .Errorf ("object changed during read at offset %d: etag %q != %q" , start , got , etag ),
135219 rc .Close (),
136220 )
137221 }
138- return writeChunkAt (dst , start , end - start , rc )
222+ // Read one byte past the expected length so an overlong body (e.g. a backend
223+ // that ignored the range and returned the whole object) is detected rather
224+ // than silently truncated to the wrong bytes.
225+ want := end - start
226+ buf , readErr := io .ReadAll (io .LimitReader (rc , want + 1 ))
227+ if err := errors .Join (readErr , rc .Close ()); err != nil {
228+ return nil , errors .Errorf ("read chunk at offset %d: %w" , start , err )
229+ }
230+ if int64 (len (buf )) != want {
231+ return nil , errors .Errorf ("chunk at offset %d: read %d of %d bytes" , start , len (buf ), want )
232+ }
233+ return buf , nil
139234}
140235
141- // writeChunkAt streams src into dst at off and closes src. It fails if fewer
142- // than want bytes arrive; a negative want skips that check (total size unknown).
143- func writeChunkAt (dst io.WriterAt , off , want int64 , src io.ReadCloser ) error {
144- n , copyErr := io .Copy (io .NewOffsetWriter (dst , off ), src )
145- if err := errors .Join (copyErr , src .Close ()); err != nil {
146- return errors .Errorf ("write chunk at offset %d: %w" , off , err )
236+ // writeAll writes all of data to dst, treating a short write as an error. A
237+ // compliant io.Writer reports short writes via a non-nil error, but the check
238+ // guards against ones that don't.
239+ func writeAll (dst io.Writer , data []byte , offset int64 ) error {
240+ n , err := dst .Write (data )
241+ if err != nil {
242+ return errors .Errorf ("write chunk at offset %d: %w" , offset , err )
147243 }
148- if want >= 0 && n != want {
149- return errors .Errorf ("short chunk at offset %d: wrote %d of %d bytes" , off , n , want )
244+ if n != len ( data ) {
245+ return errors .Errorf ("write chunk at offset %d: short write %d of %d bytes" , offset , n , len ( data ) )
150246 }
151247 return nil
152248}
0 commit comments