diff --git a/gradlecache/gradlecache_test.go b/gradlecache/gradlecache_test.go index ffb66e1..c2265b4 100644 --- a/gradlecache/gradlecache_test.go +++ b/gradlecache/gradlecache_test.go @@ -14,6 +14,7 @@ import ( "os/exec" "path/filepath" "sort" + "strconv" "strings" "sync" "testing" @@ -1741,6 +1742,197 @@ func TestS3BundleStoreRoundTrip(t *testing.T) { } } +func TestS3ParallelDownloadBackpressure(t *testing.T) { + const chunkSize = 257 + const workers = 3 + const bufferedChunks = 2 * workers + payload := make([]byte, 19*chunkSize+37) + for i := range payload { + payload[i] = byte(i*17 + i/chunkSize) + } + fs := newFakeS3() + fs.objects["test-bucket/key"] = payload + requests := make(chan int, 20) + laterCompleted := make(chan struct{}, workers-1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var start, end int + if n, _ := fmt.Sscanf(r.Header.Get("Range"), "bytes=%d-%d", &start, &end); n != 2 { + t.Error("missing Range") + http.Error(w, "missing Range", http.StatusBadRequest) + return + } + if r.Header.Get("If-Match") != `"pinned"` { + t.Error("missing If-Match") + } + requests <- start + // Complete later chunks first to exercise ordered delivery and buffer ownership. + if start == 0 { + for range workers - 1 { + select { + case <-laterCompleted: + case <-r.Context().Done(): + return + } + } + } + fs.ServeHTTP(w, r) + if start > 0 && start < workers*chunkSize { + laterCompleted <- struct{}{} + } + })) + defer srv.Close() + client := newTestS3BundleStore(srv).client + client.chunkSize, client.dlWorkers = chunkSize, workers + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + r, err := client.get(ctx, "test-bucket", "key", s3ObjInfo{Size: int64(len(payload)), ETag: `"pinned"`}) + must(t, err) + defer r.Close() //nolint:errcheck + for range bufferedChunks { + select { + case <-requests: + case <-ctx.Done(): + t.Fatal("initial download window did not start") + } + } + // Without a reader, fetched chunks must consume the entire read-ahead budget. + select { + case start := <-requests: + t.Fatalf("download read ahead beyond %d chunks: range starts at %d", bufferedChunks, start) + case <-time.After(100 * time.Millisecond): + } + got, err := io.ReadAll(r) + must(t, err) + if !bytes.Equal(got, payload) { + t.Fatal("ordered download differs from source") + } +} + +func TestS3ParallelDownloadStreamsRange(t *testing.T) { + const chunkSize = 257 + const prefixSize = chunkSize - 1 + payload := bytes.Repeat([]byte("abc"), chunkSize) + fs := newFakeS3() + fs.objects["bucket/key"] = payload + consumedPrefix := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("Range") != "bytes=0-256" { + fs.ServeHTTP(w, r) + return + } + w.Header().Set("Content-Length", strconv.Itoa(chunkSize)) + w.WriteHeader(http.StatusPartialContent) + _, _ = w.Write(payload[:prefixSize]) + w.(http.Flusher).Flush() + select { + case <-consumedPrefix: + _, _ = w.Write(payload[prefixSize:chunkSize]) + case <-r.Context().Done(): + } + })) + defer srv.Close() + client := newTestS3BundleStore(srv).client + client.chunkSize, client.dlWorkers = chunkSize, 2 + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + r, err := client.get(ctx, "bucket", "key", s3ObjInfo{Size: int64(len(payload))}) + must(t, err) + defer r.Close() //nolint:errcheck + prefix := make([]byte, 1) + _, err = io.ReadFull(r, prefix) + must(t, err) + close(consumedPrefix) + rest, err := io.ReadAll(r) + must(t, err) + if !bytes.Equal(append(prefix, rest...), payload) { + t.Fatal("streamed download differs from source") + } +} + +func TestS3ParallelDownloadInvalidRange(t *testing.T) { + const chunkSize = 257 + for _, tc := range []struct { + name string + status int + size int + chunked bool + }{ + {"empty", http.StatusPartialContent, 0, false}, + {"short", http.StatusPartialContent, chunkSize - 1, false}, + {"long", http.StatusPartialContent, chunkSize + 1, false}, + {"chunked-empty", http.StatusPartialContent, 0, true}, + {"chunked-short", http.StatusPartialContent, chunkSize - 1, true}, + {"chunked-long", http.StatusPartialContent, chunkSize + 1, true}, + {"ignored-range", http.StatusOK, chunkSize, false}, + {"server-error", http.StatusInternalServerError, chunkSize, false}, + } { + t.Run(tc.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(tc.status) + if tc.chunked { + w.(http.Flusher).Flush() + } + _, _ = w.Write(make([]byte, tc.size)) + })) + defer srv.Close() + client := newTestS3BundleStore(srv).client + client.chunkSize, client.dlWorkers = chunkSize, 2 + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + r, err := client.get(ctx, "bucket", "key", s3ObjInfo{Size: 5 * chunkSize}) + must(t, err) + defer r.Close() //nolint:errcheck + if _, err := io.ReadAll(r); err == nil || stderrors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected a range error, got %v", err) + } + }) + } +} + +func TestS3ParallelDownloadCancellation(t *testing.T) { + for _, closeReader := range []bool{false, true} { + t.Run(fmt.Sprintf("close=%v", closeReader), func(t *testing.T) { + started := make(chan struct{}, 2) + stopped := make(chan struct{}, 2) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + started <- struct{}{} + <-r.Context().Done() + stopped <- struct{}{} + })) + defer srv.Close() + client := newTestS3BundleStore(srv).client + client.chunkSize, client.dlWorkers = 257, 2 + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + r, err := client.get(ctx, "bucket", "key", s3ObjInfo{Size: 10 * client.chunkSize}) + must(t, err) + defer r.Close() //nolint:errcheck + for range client.dlWorkers { + select { + case <-started: + case <-ctx.Done(): + t.Fatal("requests did not start") + } + } + if closeReader { + must(t, r.Close()) + } else { + cancel() + } + for range client.dlWorkers { + select { + case <-stopped: + case <-time.After(5 * time.Second): + t.Fatal("request did not stop") + } + } + if _, err := io.ReadAll(r); err == nil { + t.Fatal("cancellation returned success") + } + }) + } +} + // ─── Helpers ───────────────────────────────────────────────────────────────── func must(t *testing.T, err error) { diff --git a/gradlecache/s3.go b/gradlecache/s3.go index a6bef1b..40491e3 100644 --- a/gradlecache/s3.go +++ b/gradlecache/s3.go @@ -17,10 +17,10 @@ import ( "path/filepath" "sort" "strings" + "sync" "time" "github.com/alecthomas/errors" - "golang.org/x/sync/errgroup" ) type awsCreds struct { @@ -41,6 +41,7 @@ type s3Client struct { const ( defaultDownloadChunkSize = 32 << 20 defaultDownloadWorkers = 8 + downloadBlocksPerChunk = 8 ) func newS3Client(region string) (*s3Client, error) { @@ -123,10 +124,12 @@ func (c *s3Client) get(ctx context.Context, bucket, key string, info s3ObjInfo) dlCtx, cancel := context.WithCancel(ctx) pr, pw := io.Pipe() + stopCancel := context.AfterFunc(dlCtx, func() { pw.CloseWithError(dlCtx.Err()) }) go func() { err := c.parallelGet(dlCtx, bucket, key, info, pw) - cancel() + stopCancel() pw.CloseWithError(err) + cancel() }() return &cancelReadCloser{ReadCloser: pr, cancel: cancel}, nil } @@ -142,17 +145,26 @@ func (c *cancelReadCloser) Close() error { } func (c *s3Client) parallelGet(ctx context.Context, bucket, key string, info s3ObjInfo, w io.Writer) error { + ctx, cancel := context.WithCancel(ctx) + var workers sync.WaitGroup + defer func() { + cancel() + workers.Wait() + }() + numChunks := int((info.Size + c.chunkSize - 1) / c.chunkSize) numWorkers := min(c.dlWorkers, numChunks) + blockSize := (c.chunkSize + downloadBlocksPerChunk - 1) / downloadBlocksPerChunk type chunkResult struct { - data []byte - err error + data []byte + recycle chan []byte + err error } results := make([]chan chunkResult, numChunks) for i := range results { - results[i] = make(chan chunkResult, 1) + results[i] = make(chan chunkResult, downloadBlocksPerChunk+1) } work := make(chan int, numChunks) @@ -161,73 +173,101 @@ func (c *s3Client) parallelGet(ctx context.Context, bucket, key string, info s3O } close(work) - eg, egCtx := errgroup.WithContext(ctx) + // Stream blocks before a whole range arrives. Each worker retains at most one + // chunk plus one block, leased until the ordered writer consumes them. Pools + // are worker-local: later chunks cannot take the buffers an earlier one needs. for range numWorkers { - eg.Go(func() error { + workers.Go(func() { + buffers := make(chan []byte, downloadBlocksPerChunk+1) + for range cap(buffers) { + buffers <- make([]byte, blockSize) + } for seq := range work { - if egCtx.Err() != nil { - results[seq] <- chunkResult{err: egCtx.Err()} - continue - } - start := int64(seq) * c.chunkSize - end := min(start+c.chunkSize-1, info.Size-1) - - req, err := http.NewRequestWithContext(egCtx, http.MethodGet, c.objectURL(bucket, key), nil) + err := func() error { + var data []byte + select { + case data = <-buffers: + case <-ctx.Done(): + return ctx.Err() + } + start := int64(seq) * c.chunkSize + end := min(start+c.chunkSize-1, info.Size-1) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.objectURL(bucket, key), nil) + if err != nil { + return err + } + req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", start, end)) + if info.ETag != "" { + req.Header.Set("If-Match", info.ETag) + } + c.sign(req) + resp, err := c.http.Do(req) //nolint:gosec + if err != nil { + return err + } + defer resp.Body.Close() //nolint:errcheck,gosec + if resp.StatusCode != http.StatusPartialContent { + msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + return errors.Errorf("s3 GET range %d-%d: status %d: %s", start, end, resp.StatusCode, msg) + } + if resp.ContentLength >= 0 && resp.ContentLength != end-start+1 { + return errors.Errorf("s3 GET range %d-%d: unexpected length %d", start, end, resp.ContentLength) + } + for remaining := end - start + 1; remaining > 0; { + data = data[:min(blockSize, remaining)] + if _, err := io.ReadFull(resp.Body, data); err != nil { + return errors.Wrap(err, fmt.Sprintf("s3 GET range %d-%d", start, end)) + } + results[seq] <- chunkResult{data: data, recycle: buffers} + remaining -= int64(len(data)) + if remaining > 0 { + select { + case data = <-buffers: + case <-ctx.Done(): + return ctx.Err() + } + } + } + var extra [1]byte + if n, err := io.ReadFull(resp.Body, extra[:]); n != 0 { + return errors.Errorf("s3 GET range %d-%d: body exceeds requested length", start, end) + } else if err != io.EOF { + return errors.Wrap(err, fmt.Sprintf("s3 GET range %d-%d", start, end)) + } + return nil + }() if err != nil { results[seq] <- chunkResult{err: err} - continue + close(results[seq]) + return } - req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", start, end)) - if info.ETag != "" { - req.Header.Set("If-Match", info.ETag) - } - c.sign(req) - - resp, err := c.http.Do(req) //nolint:gosec - if err != nil { - results[seq] <- chunkResult{err: err} - continue - } - if resp.StatusCode != http.StatusPartialContent && resp.StatusCode != http.StatusOK { - msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) - resp.Body.Close() //nolint:errcheck,gosec - results[seq] <- chunkResult{err: errors.Errorf("s3 GET range %d-%d: status %d: %s", start, end, resp.StatusCode, msg)} - continue - } - data, readErr := io.ReadAll(resp.Body) - resp.Body.Close() //nolint:errcheck,gosec - results[seq] <- chunkResult{data: data, err: readErr} + close(results[seq]) } - return nil }) } - var writeErr error - writeDone := make(chan struct{}) - go func() { - defer close(writeDone) - for _, ch := range results { - r := <-ch - if writeErr != nil { - continue + for _, ch := range results { + for { + var r chunkResult + var ok bool + select { + case r, ok = <-ch: + case <-ctx.Done(): + return ctx.Err() + } + if !ok { + break } if r.err != nil { - writeErr = r.err - continue + return r.err } if _, err := w.Write(r.data); err != nil { - writeErr = err + return err } + r.recycle <- r.data[:cap(r.data)] } - }() - - egErr := eg.Wait() - <-writeDone - - if writeErr != nil { - return writeErr } - return egErr + return nil } func (c *s3Client) objectURL(bucket, key string) string {