From b71092d8d8421e8cd61e32d8421e19b2357f97ed Mon Sep 17 00:00:00 2001 From: Josh Friend Date: Tue, 20 Oct 2026 17:35:38 -0400 Subject: [PATCH] Retry transient S3 errors with bounded backoff Signed-off-by: Josh Friend --- gradlecache/s3.go | 78 ++++++++--- gradlecache/s3_test.go | 312 +++++++++++++++++++++++++++++++++++++++++ gradlecache/store.go | 137 +++++++++++------- 3 files changed, 454 insertions(+), 73 deletions(-) create mode 100644 gradlecache/s3_test.go diff --git a/gradlecache/s3.go b/gradlecache/s3.go index 40491e3..bfa19f0 100644 --- a/gradlecache/s3.go +++ b/gradlecache/s3.go @@ -10,6 +10,7 @@ import ( "fmt" "io" "log/slog" + "math/rand/v2" "net/http" "net/url" "os" @@ -44,6 +45,44 @@ const ( downloadBlocksPerChunk = 8 ) +// doRequest sends a freshly built and signed request, retrying transient S3 +// server errors before any of the response body is consumed. The final +// response is returned unchanged so callers can report its status and body. +func (c *s3Client) doRequest(ctx context.Context, newReq func() (*http.Request, error)) (*http.Response, error) { + for attempt, ceiling := 1, 250*time.Millisecond; ; attempt, ceiling = attempt+1, min(2*ceiling, time.Second) { + req, err := newReq() + if err != nil { + return nil, err + } + c.sign(req) + resp, err := c.http.Do(req) //nolint:gosec + if req.Body != nil { + // An early server response can arrive while the transport is still + // reading the upload. Stop that reader before the next attempt rewinds. + req.Body.Close() //nolint:errcheck,gosec + } + if err != nil || attempt == 3 { + return resp, err + } + switch resp.StatusCode { + case http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout: + default: + return resp, nil + } + // Do not wait for an error body to finish before retrying. + resp.Body.Close() //nolint:errcheck,gosec + // Equal jitter keeps a minimum pause while spreading concurrent retries. + delay := ceiling/2 + time.Duration(rand.Int64N(int64(ceiling/2)+1)) //nolint:gosec // not security-sensitive + timer := time.NewTimer(delay) + select { + case <-timer.C: + case <-ctx.Done(): + timer.Stop() + return nil, errors.Errorf("retry after status %d: %w", resp.StatusCode, ctx.Err()) + } + } +} + func newS3Client(region string) (*s3Client, error) { creds, err := resolveAWSCredentials(region) if err != nil { @@ -70,12 +109,9 @@ type s3ObjInfo struct { } func (c *s3Client) stat(ctx context.Context, bucket, key string) (s3ObjInfo, error) { - req, err := http.NewRequestWithContext(ctx, http.MethodHead, c.objectURL(bucket, key), nil) - if err != nil { - return s3ObjInfo{}, err - } - c.sign(req) - resp, err := c.http.Do(req) + resp, err := c.doRequest(ctx, func() (*http.Request, error) { + return http.NewRequestWithContext(ctx, http.MethodHead, c.objectURL(bucket, key), nil) + }) if err != nil { return s3ObjInfo{}, err } @@ -105,12 +141,9 @@ func (c *s3Client) stat(ctx context.Context, bucket, key string) (s3ObjInfo, err func (c *s3Client) get(ctx context.Context, bucket, key string, info s3ObjInfo) (io.ReadCloser, error) { if info.Size <= c.chunkSize { - req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.objectURL(bucket, key), nil) - if err != nil { - return nil, err - } - c.sign(req) - resp, err := c.http.Do(req) //nolint:gosec + resp, err := c.doRequest(ctx, func() (*http.Request, error) { + return http.NewRequestWithContext(ctx, http.MethodGet, c.objectURL(bucket, key), nil) + }) if err != nil { return nil, err } @@ -192,16 +225,17 @@ func (c *s3Client) parallelGet(ctx context.Context, bucket, key string, info s3O } 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 + resp, err := c.doRequest(ctx, func() (*http.Request, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.objectURL(bucket, key), nil) + if err != nil { + return nil, err + } + req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", start, end)) + if info.ETag != "" { + req.Header.Set("If-Match", info.ETag) + } + return req, nil + }) if err != nil { return err } diff --git a/gradlecache/s3_test.go b/gradlecache/s3_test.go new file mode 100644 index 0000000..aa1c0c1 --- /dev/null +++ b/gradlecache/s3_test.go @@ -0,0 +1,312 @@ +package gradlecache + +import ( + "bytes" + "context" + "encoding/xml" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "os" + "strings" + "sync/atomic" + "testing" + "testing/iotest" + "testing/synctest" + "time" +) + +type retryRoundTripFunc func(*http.Request) (*http.Response, error) + +func (f retryRoundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +type retryBody struct { + io.Reader + closed bool +} + +func (b *retryBody) Close() error { b.closed = true; return nil } + +func TestS3DownloadRetries(t *testing.T) { + for _, operation := range []string{"head", "get", "range"} { + for _, tc := range []struct { + name string + status int + failures int + attempts int + wantErr bool + }{ + {"500", 500, 2, 3, false}, + {"502", 502, 2, 3, false}, + {"503", 503, 2, 3, false}, + {"504", 504, 2, 3, false}, + {"exhausted", 503, 3, 3, true}, + {"forbidden", 403, 1, 1, true}, + {"missing", 404, 1, 1, true}, + {"not-implemented", 501, 1, 1, true}, + {"network", 0, 1, 1, true}, + {"truncated", 200, 0, 1, true}, + } { + if operation == "head" && tc.name == "truncated" { + continue + } + t.Run(operation+"/"+tc.name, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + payload := bytes.Repeat([]byte("abc"), 43) + var times []time.Time + var bodies []*retryBody + client := &s3Client{region: "us-east-1", creds: awsCreds{AccessKeyID: "test-id", SecretAccessKey: "test-secret"}, chunkSize: 1024, dlWorkers: 1} + if operation == "range" { + client.chunkSize = 64 + } + client.http = &http.Client{Transport: retryRoundTripFunc(func(r *http.Request) (*http.Response, error) { + if !strings.Contains(r.Header.Get("Authorization"), "Credential=test-id/") || strings.Contains(r.Header.Get("Authorization"), "SignedHeaders=authorization") { + t.Error("retry was not freshly signed") + } + data, status := payload, http.StatusOK + first := r.Header.Get("Range") == "" || r.Header.Get("Range") == "bytes=0-63" + if first { + times = append(times, time.Now()) + if tc.status == 0 { + return nil, errors.New("network failure") + } + } + if r.Header.Get("Range") != "" { + if r.Header.Get("If-Match") != `"etag"` { + t.Error("range lost its ETag precondition") + } + var start, end int + _, _ = fmt.Sscanf(r.Header.Get("Range"), "bytes=%d-%d", &start, &end) + data, status = payload[start:end+1], http.StatusPartialContent + } + if first && len(times) <= tc.failures { + data, status = []byte("final server failure"), tc.status + } + body := &retryBody{Reader: bytes.NewReader(data)} + if tc.name == "truncated" { + body.Reader = io.MultiReader(bytes.NewReader(data[:len(data)/2]), iotest.ErrReader(io.ErrUnexpectedEOF)) + } + bodies = append(bodies, body) + return &http.Response{StatusCode: status, ContentLength: int64(len(data)), Body: body, Header: http.Header{"Etag": {`"etag"`}, "X-Amz-Meta-Kind": {"bundle"}}}, nil + })} + var err error + if operation == "head" { + var info s3ObjInfo + info, err = client.stat(context.Background(), "bucket", "key") + if err == nil && (info.Size != int64(len(payload)) || info.ETag != `"etag"` || info.Metadata["kind"] != "bundle") { + t.Fatalf("incorrect metadata: %+v", info) + } + } else { + var r io.ReadCloser + r, err = client.get(context.Background(), "bucket", "key", s3ObjInfo{Size: int64(len(payload)), ETag: `"etag"`}) + if err == nil { + var got []byte + got, err = io.ReadAll(r) + _ = r.Close() + if err == nil && !bytes.Equal(got, payload) { + t.Fatal("incorrect downloaded bytes") + } + } + } + if (err != nil) != tc.wantErr || len(times) != tc.attempts { + t.Fatalf("error=%v, attempts=%d, want error=%v and attempts=%d", err, len(times), tc.wantErr, tc.attempts) + } + if tc.name == "exhausted" && operation != "head" && !strings.Contains(err.Error(), "final server failure") { + t.Fatalf("lost final response: %v", err) + } + for i := 1; i < len(times); i++ { + elapsed := times[i].Sub(times[i-1]) + lower, upper := 125*time.Millisecond, 250*time.Millisecond + if i == 2 { + lower, upper = 250*time.Millisecond, 500*time.Millisecond + } + if elapsed < lower || elapsed > upper { + t.Errorf("retry %d waited %v, want [%v, %v]", i, elapsed, lower, upper) + } + } + for _, body := range bodies { + if !body.closed { + t.Error("response body was not closed") + } + } + }) + }) + } + } +} + +func TestS3DownloadRetryCancellation(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + attempts := 0 + client := &s3Client{http: &http.Client{Transport: retryRoundTripFunc(func(*http.Request) (*http.Response, error) { + attempts++ + return &http.Response{StatusCode: 503, Body: io.NopCloser(strings.NewReader("busy"))}, nil + })}} + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + go func() { _, err := client.stat(ctx, "bucket", "key"); done <- err }() + synctest.Wait() // The request finished and the retry is sleeping. + cancel() + synctest.Wait() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) || attempts != 1 { + t.Fatalf("error=%v attempts=%d, want cancellation after one attempt", err, attempts) + } + default: + t.Fatal("retry did not stop on cancellation") + } + }) +} + +func TestS3UploadRetries(t *testing.T) { + for _, operation := range []string{"put", "initiate", "part", "complete", "abort"} { + for _, tc := range []struct { + name string + status int + failures int + attempts int + wantErr bool + }{ + {"recover", 503, 2, 3, false}, + {"exhausted", 500, 3, 3, true}, + {"forbidden", 403, 1, 1, true}, + {"network", 0, 1, 1, true}, + } { + t.Run(operation+"/"+tc.name, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + payload := []byte("an asymmetric upload payload\nwith a second line") + var attempts int + var previous io.ReadCloser + var checksum string + client := &s3Client{region: "us-east-1", creds: awsCreds{AccessKeyID: "test-id", SecretAccessKey: "test-secret"}} + client.http = &http.Client{Transport: retryRoundTripFunc(func(r *http.Request) (*http.Response, error) { + attempts++ + if !strings.Contains(r.Header.Get("Authorization"), "Credential=test-id/") || strings.Contains(r.Header.Get("Authorization"), "SignedHeaders=authorization") { + t.Error("upload retry was not freshly signed") + } + if previous != nil { + // A transport may still hold the prior request body after an + // early response. It must not consume the rewound seeker. + if n, err := previous.Read(make([]byte, 1)); n != 0 || err == nil { + t.Error("old upload body can still consume the next attempt") + } + } + if operation == "put" || operation == "part" { + previous = r.Body + if attempts == 1 { + checksum = r.Header.Get("X-Amz-Checksum-Crc64nvme") + } + if checksum == "" || r.Header.Get("X-Amz-Checksum-Crc64nvme") != checksum || r.ContentLength != int64(len(payload)) { + t.Error("upload retry lost checksum or content length") + } + // Fail after consuming a prefix, not only after the full body. + var got []byte + var err error + if attempts <= tc.failures { + got, err = io.ReadAll(io.LimitReader(r.Body, 7)) + } else { + got, err = io.ReadAll(r.Body) + } + must(t, err) + want := payload + if attempts <= tc.failures { + want = payload[:7] + } + if !bytes.Equal(got, want) { + t.Fatalf("upload body was not replayed: got %q, want %q", got, want) + } + } + if operation == "put" || operation == "initiate" { + if r.Header.Get("Content-Type") != "application/zstd" || r.Header.Get("X-Amz-Meta-Base-Commit") != "abc123" { + t.Error("upload retry lost metadata or content type") + } + } + if operation == "part" || operation == "complete" || operation == "abort" { + if r.URL.Query().Get("uploadId") != "id+with/slash" { + t.Error("upload retry changed the upload ID") + } + } + if operation == "complete" { + var request struct { + Parts []struct { + PartNumber int + ETag string + } `xml:"Part"` + } + must(t, xml.NewDecoder(r.Body).Decode(&request)) + if len(request.Parts) != 1 || request.Parts[0].PartNumber != 7 || request.Parts[0].ETag != "part-etag" { + t.Fatalf("completion body changed: %+v", request) + } + } + if tc.status == 0 { + return nil, errors.New("network failure") + } + status, body := 200, "" + if attempts <= tc.failures { + status, body = tc.status, "server failure" + } else if operation == "initiate" { + body = "id+with/slash" + } + return &http.Response{StatusCode: status, Body: io.NopCloser(strings.NewReader(body)), Header: http.Header{"Etag": {"part-etag"}}}, nil + })} + ctx := context.Background() + var err error + switch operation { + case "put": + err = client.put(ctx, "bucket", "key", bytes.NewReader(payload), int64(len(payload)), "application/zstd", map[string]string{"base-commit": "abc123"}) + case "initiate": + _, err = client.createMultipartUpload(ctx, "bucket", "key", "application/zstd", map[string]string{"base-commit": "abc123"}) + case "part": + _, err = client.uploadPart(ctx, "bucket", "key", "id+with/slash", 7, bytes.NewReader(payload), int64(len(payload))) + case "complete": + err = client.completeMultipartUpload(ctx, "bucket", "key", "id+with/slash", []struct { + PartNumber int + ETag string + }{{7, "part-etag"}}) + case "abort": + err = client.abortMultipartUpload(ctx, "bucket", "key", "id+with/slash") + } + if (err != nil) != tc.wantErr || attempts != tc.attempts { + t.Fatalf("error=%v attempts=%d, want error=%v attempts=%d", err, attempts, tc.wantErr, tc.attempts) + } + }) + }) + } + } +} + +func TestS3UploadRetryKeepsFileOpen(t *testing.T) { + payload := bytes.Repeat([]byte("file upload\n"), 1000) + file, err := os.CreateTemp(t.TempDir(), "upload") + must(t, err) + defer file.Close() //nolint:errcheck + _, err = file.Write(payload) + must(t, err) + var attempts atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + attempt := attempts.Add(1) + got, err := io.ReadAll(r.Body) + if err != nil || !bytes.Equal(got, payload) { + t.Errorf("attempt %d: bytes=%d error=%v", attempt, len(got), err) + } + if attempt == 1 { + w.WriteHeader(http.StatusServiceUnavailable) + } + })) + defer server.Close() + _, err = file.Seek(0, io.SeekStart) + must(t, err) + client := newTestS3BundleStore(server).client + must(t, client.put(context.Background(), "bucket", "key", file, int64(len(payload)), "", nil)) + _, err = file.Seek(0, io.SeekStart) + must(t, err) + got, err := io.ReadAll(file) + must(t, err) + if attempts.Load() != 2 || !bytes.Equal(got, payload) { + t.Fatal("file was closed or changed by upload retries") + } +} diff --git a/gradlecache/store.go b/gradlecache/store.go index b969318..d954412 100644 --- a/gradlecache/store.go +++ b/gradlecache/store.go @@ -234,26 +234,55 @@ func crc64Base64(h hash.Hash64) string { return base64.StdEncoding.EncodeToString(buf[:]) } +// uploadBody leaves the caller's seeker open. Close also waits for an active +// transport read and rejects future reads, so the seeker can safely be rewound +// after an early error response even if the transport closes asynchronously. +type uploadBody struct { + mu sync.Mutex + r io.Reader + closed bool +} + +func (b *uploadBody) Read(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return 0, io.ErrClosedPipe + } + return b.r.Read(p) +} + +func (b *uploadBody) Close() error { + b.mu.Lock() + defer b.mu.Unlock() + b.closed = true + return nil +} + func (c *s3Client) putSingle(ctx context.Context, bucket, key string, r io.ReadSeeker, size int64, contentType string, meta map[string]string) error { checksum, err := crc64Of(r) if err != nil { return errors.Wrap(err, "compute CRC64-NVME checksum") } - req, err := http.NewRequestWithContext(ctx, http.MethodPut, c.objectURL(bucket, key), r) - if err != nil { - return err - } - req.ContentLength = size - if contentType != "" { - req.Header.Set("Content-Type", contentType) - } - req.Header.Set("X-Amz-Checksum-Crc64nvme", checksum) - req.Header.Set("X-Amz-Checksum-Algorithm", "CRC64NVME") - for k, v := range meta { - req.Header.Set("X-Amz-Meta-"+k, v) - } - c.sign(req) - resp, err := c.http.Do(req) //nolint:gosec + resp, err := c.doRequest(ctx, func() (*http.Request, error) { + if _, err := r.Seek(0, io.SeekStart); err != nil { + return nil, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPut, c.objectURL(bucket, key), &uploadBody{r: r}) + if err != nil { + return nil, err + } + req.ContentLength = size + if contentType != "" { + req.Header.Set("Content-Type", contentType) + } + req.Header.Set("X-Amz-Checksum-Crc64nvme", checksum) + req.Header.Set("X-Amz-Checksum-Algorithm", "CRC64NVME") + for k, v := range meta { + req.Header.Set("X-Amz-Meta-"+k, v) + } + return req, nil + }) if err != nil { return err } @@ -335,19 +364,20 @@ func (c *s3Client) putMultipart(ctx context.Context, bucket, key string, r io.Re func (c *s3Client) createMultipartUpload(ctx context.Context, bucket, key, contentType string, meta map[string]string) (string, error) { u := c.objectURL(bucket, key) + "?uploads" - req, err := http.NewRequestWithContext(ctx, http.MethodPost, u, nil) - if err != nil { - return "", err - } - if contentType != "" { - req.Header.Set("Content-Type", contentType) - } - req.Header.Set("X-Amz-Checksum-Algorithm", "CRC64NVME") - for k, v := range meta { - req.Header.Set("X-Amz-Meta-"+k, v) - } - c.sign(req) - resp, err := c.http.Do(req) //nolint:gosec + resp, err := c.doRequest(ctx, func() (*http.Request, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, u, nil) + if err != nil { + return nil, err + } + if contentType != "" { + req.Header.Set("Content-Type", contentType) + } + req.Header.Set("X-Amz-Checksum-Algorithm", "CRC64NVME") + for k, v := range meta { + req.Header.Set("X-Amz-Meta-"+k, v) + } + return req, nil + }) if err != nil { return "", err } @@ -382,14 +412,18 @@ func (c *s3Client) uploadPart(ctx context.Context, bucket, key, uploadID string, checksum := crc64Base64(h) u := fmt.Sprintf("%s?partNumber=%d&uploadId=%s", c.objectURL(bucket, key), partNum, url.QueryEscape(uploadID)) - req, err := http.NewRequestWithContext(ctx, http.MethodPut, u, r) - if err != nil { - return uploadPartResult{}, err - } - req.ContentLength = size - req.Header.Set("X-Amz-Checksum-Crc64nvme", checksum) - c.sign(req) - resp, err := c.http.Do(req) //nolint:gosec + resp, err := c.doRequest(ctx, func() (*http.Request, error) { + if _, err := r.Seek(0, io.SeekStart); err != nil { + return nil, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPut, u, &uploadBody{r: r}) + if err != nil { + return nil, err + } + req.ContentLength = size + req.Header.Set("X-Amz-Checksum-Crc64nvme", checksum) + return req, nil + }) if err != nil { return uploadPartResult{}, err } @@ -411,14 +445,15 @@ func (c *s3Client) completeMultipartUpload(ctx context.Context, bucket, key, upl return err } u := fmt.Sprintf("%s?uploadId=%s", c.objectURL(bucket, key), url.QueryEscape(uploadID)) - req, err := http.NewRequestWithContext(ctx, http.MethodPost, u, strings.NewReader(string(xmlBody))) - if err != nil { - return err - } - req.ContentLength = int64(len(xmlBody)) - req.Header.Set("Content-Type", "application/xml") - c.sign(req) - resp, err := c.http.Do(req) //nolint:gosec + resp, err := c.doRequest(ctx, func() (*http.Request, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodPost, u, bytes.NewReader(xmlBody)) + if err != nil { + return nil, err + } + req.ContentLength = int64(len(xmlBody)) + req.Header.Set("Content-Type", "application/xml") + return req, nil + }) if err != nil { return err } @@ -432,17 +467,17 @@ func (c *s3Client) completeMultipartUpload(ctx context.Context, bucket, key, upl func (c *s3Client) abortMultipartUpload(ctx context.Context, bucket, key, uploadID string) error { u := fmt.Sprintf("%s?uploadId=%s", c.objectURL(bucket, key), url.QueryEscape(uploadID)) - req, err := http.NewRequestWithContext(ctx, http.MethodDelete, u, nil) + resp, err := c.doRequest(ctx, func() (*http.Request, error) { + return http.NewRequestWithContext(ctx, http.MethodDelete, u, nil) + }) if err != nil { return err } - c.sign(req) - resp, err := c.http.Do(req) //nolint:gosec - if err != nil { - return err + body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + resp.Body.Close() //nolint:errcheck,gosec + if resp.StatusCode >= http.StatusMultipleChoices { + return errors.Errorf("s3 AbortMultipartUpload %s/%s: status %d: %s", bucket, key, resp.StatusCode, body) } - io.Copy(io.Discard, resp.Body) //nolint:errcheck,gosec - resp.Body.Close() //nolint:errcheck,gosec return nil }