fix: Correct retry logic

internal/ghclient/body.go

@@ -6,6 +6,29 @@ import (
 	"net/http"
 )
 
+// rewindRequestBody rewinds the request body to the beginning, if possible. This is necessary for retrying requests with a non-nil body.
+func rewindRequestBody(req *http.Request) error {
+	if req == nil {
+		return fmt.Errorf("request is nil")
+	}
+
+	if req.Body == nil || req.Body == http.NoBody {
+		return nil
+	}
+
+	if req.GetBody == nil {
+		return fmt.Errorf("cannot rewind request body: get body is nil")
+	}
+
+	body, err := req.GetBody()
+	if err != nil {
+		return fmt.Errorf("failed to get request body: %w", err)
+	}
+
+	req.Body = body
+	return nil
+}
+
 // drainResponseBody reads and closes a [http.Response] body, up to 4KB, to avoid memory leaks.
 func drainResponseBody(resp *http.Response) error {
 	if resp == nil {

internal/ghclient/body_test.go

@@ -8,6 +8,88 @@ import (
 	"testing"
 )
 
+func Test_rewindRequestBody(t *testing.T) {
+	t.Parallel()
+
+	body := []byte("TEST")
+
+	for _, tt := range []struct {
+		name     string
+		req      *http.Request
+		wantBody bool
+		wantErr  *string
+	}{
+		{
+			name:    "nil_request_errors",
+			req:     nil,
+			wantErr: new("request is nil"),
+		},
+		{
+			name:    "request_with_no_body_succeeds",
+			req:     &http.Request{},
+			wantErr: nil,
+		},
+		{
+			name: "request_with_get_body_succeeds",
+			req: &http.Request{
+				Body:    io.NopCloser(bytes.NewBuffer(body)),
+				GetBody: func() (io.ReadCloser, error) { return io.NopCloser(bytes.NewBuffer(body)), nil },
+			},
+			wantBody: true,
+		},
+		{
+			name: "request_without_get_body_errors",
+			req: &http.Request{
+				Body: io.NopCloser(bytes.NewBuffer(body)),
+			},
+			wantErr: new("cannot rewind request body: get body is nil"),
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			err := rewindRequestBody(tt.req)
+			if err != nil {
+				if tt.wantErr == nil {
+					t.Fatalf("unexpected error: %s", err)
+				}
+				if !regexp.MustCompile(regexp.QuoteMeta(*tt.wantErr)).MatchString(err.Error()) {
+					t.Fatalf("unexpected error: %s", err)
+				}
+				return
+			}
+
+			if tt.wantErr != nil {
+				t.Fatalf("expected error: %s, got: %v", *tt.wantErr, err)
+			}
+
+			if tt.req == nil {
+				t.Fatalf("request is nil after rewind")
+			}
+
+			if !tt.wantBody {
+				return
+			}
+
+			if tt.req.Body == nil {
+				t.Fatalf("request body is nil after rewind")
+			}
+
+			if tt.req.GetBody == nil {
+				t.Fatalf("request GetBody is nil after rewind")
+			}
+
+			got, err := io.ReadAll(tt.req.Body)
+			if err != nil {
+				t.Fatalf("failed to read request body: %s", err)
+			}
+			if !bytes.Equal(got, body) {
+				t.Fatalf("unexpected request body: got %s, want %s", got, body)
+			}
+		})
+	}
+}
+
 func Test_drainResponseBody(t *testing.T) {
 	t.Parallel()
 

internal/ghclient/retry.go

@@ -36,24 +36,37 @@ type retryTransport struct {
 
 // RoundTrip implements the [http.RoundTripper] interface for the [retryTransport]. It retries failed requests using the configured backoff strategy.
 func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	var resp *http.Response
-
 	backoff, err := newBackoff(t.opts)
 	if err != nil {
 		return nil, err
 	}
 
+	var resp *http.Response
+	tries := 0
+
 	err = retry.Do(req.Context(), backoff, func(ctx context.Context) error {
-		// If there was a previous response, drain it and reset the body
+		// If there was a previous response, drain it so the connection can be reused
 		if resp != nil {
 			_ = drainResponseBody(resp)
 			resp = nil
 		}
 
+		// Shallow copy so we never mutate the caller's Request
+		reqCopy := *req
+
+		// If this is a retry, rewind the request body so it can be sent again
+		if tries > 0 {
+			if err := rewindRequestBody(&reqCopy); err != nil {
+				return fmt.Errorf("failed to rewind request body: %w", err)
+			}
+		}
+		tries++
+
 		var respErr error
-		resp, respErr = t.inner.RoundTrip(req)
+		resp, respErr = t.inner.RoundTrip(&reqCopy)
 		if respErr != nil {
 			_ = drainResponseBody(resp)
+			resp = nil
 			if isRetryableNetworkError(respErr) {
 				return retry.RetryableError(respErr)
 			}
@@ -64,13 +77,20 @@ func (t *retryTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 			return retry.RetryableError(errRetryServerError)
 		}
 
-		return respErr
+		return nil
 	})
 
+	// Cleanup error used for retryable 5xx responses so we don't return a server error when the request was actually successful
 	if errors.Is(err, errRetryServerError) {
 		err = nil
 	}
 
+	// Clean up any lingering 5xx response if the context was canceled during backoff sleep
+	if err != nil && resp != nil {
+		_ = drainResponseBody(resp)
+		resp = nil
+	}
+
 	return resp, err
 }
 

internal/ghclient/retry_test.go

@@ -1,8 +1,10 @@
 package ghclient
 
 import (
+	"bytes"
 	"context"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"regexp"
@@ -82,59 +84,102 @@ func Test_retryTransport_RoundTrip(t *testing.T) {
 	for _, tt := range []struct {
 		name           string
 		retryMax       int
+		wait           time.Duration
 		failures       []fakeHTTPResponse
 		cancel         bool
+		cancelDelay    time.Duration
 		wantStatusCode int
+		wantCalls      int
 		wantErr        *string
 	}{
 		{
-			name:    "errors_if_context_canceled",
+			name:    "errors_if_context_canceled_early",
 			cancel:  true,
 			wantErr: new("context canceled"),
 		},
+		{
+			name:        "errors_if_context_canceled_in_flight",
+			retryMax:    1,
+			wait:        time.Millisecond * 2,
+			cancel:      true,
+			cancelDelay: time.Millisecond,
+			failures:    []fakeHTTPResponse{{statusCode: http.StatusInternalServerError}},
+			wantErr:     new("context canceled"),
+		},
 		{
 			name:           "no_failures",
 			retryMax:       3,
+			wait:           time.Millisecond,
 			wantStatusCode: http.StatusOK,
+			wantCalls:      1,
 		},
 		{
 			name:           "no_failures_no_retries",
 			retryMax:       0,
+			wait:           time.Millisecond,
 			wantStatusCode: http.StatusOK,
+			wantCalls:      1,
 		},
 		{
 			name:           "failure_no_error",
 			retryMax:       0,
+			wait:           time.Millisecond,
 			failures:       []fakeHTTPResponse{{statusCode: http.StatusInternalServerError}},
 			wantStatusCode: http.StatusInternalServerError,
+			wantCalls:      1,
 		},
 		{
 			name:           "retries_until_success",
 			retryMax:       3,
+			wait:           time.Millisecond,
 			failures:       []fakeHTTPResponse{{statusCode: http.StatusInternalServerError}, {statusCode: http.StatusInternalServerError}},
 			wantStatusCode: http.StatusOK,
+			wantCalls:      3,
 		},
 		{
 			name:           "retries_until_failure",
 			retryMax:       1,
+			wait:           time.Millisecond,
 			failures:       []fakeHTTPResponse{{statusCode: http.StatusInternalServerError}, {statusCode: http.StatusInternalServerError}},
 			wantStatusCode: http.StatusInternalServerError,
+			wantCalls:      2,
+		},
+		{
+			name:           "retries_closed_connection",
+			retryMax:       2,
+			wait:           time.Millisecond,
+			failures:       []fakeHTTPResponse{{statusCode: http.StatusInternalServerError, headers: map[string]string{"Connection": "close"}}},
+			wantStatusCode: http.StatusOK,
+			wantCalls:      2,
 		},
 		{
 			name:           "does_not_retry_on_4xx",
 			retryMax:       3,
+			wait:           time.Millisecond,
 			failures:       []fakeHTTPResponse{{statusCode: http.StatusBadRequest}},
 			wantStatusCode: http.StatusBadRequest,
+			wantCalls:      1,
 		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
 			t.Parallel()
+
+			reqBody := "REQUEST"
 			wantBody := "PASS"
 
 			called := atomic.Int32{}
 			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 				call := int(called.Add(1))
 
+				by, err := io.ReadAll(r.Body)
+				if err != nil {
+					t.Fatalf("failed to read request body: %v", err)
+				}
+
+				if string(by) != reqBody {
+					t.Fatalf("expected request body %q, got %q", reqBody, string(by))
+				}
+
 				if call <= len(tt.failures) {
 					f := tt.failures[call-1]
 
@@ -154,7 +199,7 @@ func Test_retryTransport_RoundTrip(t *testing.T) {
 			}))
 			defer ts.Close()
 
-			tr, err := newRetryTransport(ts.Client().Transport, RetryOptions{Max: tt.retryMax, WaitMin: time.Millisecond, WaitMax: time.Millisecond, Jitter: time.Microsecond})
+			tr, err := newRetryTransport(ts.Client().Transport, RetryOptions{Max: tt.retryMax, WaitMin: tt.wait, WaitMax: tt.wait, Jitter: time.Microsecond})
 			if err != nil || tr == nil {
 				t.Fatalf("failed to create retry transport: %v", err)
 			}
@@ -170,13 +215,20 @@ func Test_retryTransport_RoundTrip(t *testing.T) {
 				ctx = t.Context()
 			}
 
-			req, err := http.NewRequestWithContext(ctx, "GET", ts.URL, nil)
+			req, err := http.NewRequestWithContext(ctx, "POST", ts.URL, bytes.NewReader([]byte(reqBody)))
 			if err != nil {
 				t.Fatalf("failed to create request: %v", err)
 			}
 
 			if tt.cancel {
-				cancel()
+				if tt.cancelDelay == 0 {
+					cancel()
+				} else {
+					go func() {
+						time.Sleep(tt.cancelDelay)
+						cancel()
+					}()
+				}
 			}
 
 			res, err := client.Do(req)
@@ -201,6 +253,11 @@ func Test_retryTransport_RoundTrip(t *testing.T) {
 			if res.StatusCode != tt.wantStatusCode {
 				t.Fatalf("expected status code %d, got %d", tt.wantStatusCode, res.StatusCode)
 			}
+
+			gotCalls := int(called.Load())
+			if gotCalls != tt.wantCalls {
+				t.Fatalf("expected %d calls, got %d", tt.wantCalls, called.Load())
+			}
 		})
 	}
 }

internal/ghclient/rewind.go

@@ -28,11 +28,14 @@ type rewindableTransport struct {
 
 // RoundTrip implements the [http.RoundTripper] interface for the [rewindableTransport]. It sets [request.GetBody] if it's unset and can either be rewound or fully buffered within the memory limit provided.
 func (t *rewindableTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-	if err := ensureRequestGetBody(req, t.maxBufferBytes); err != nil {
+	// Shallow copy so we never mutate the caller's Request
+	reqCopy := *req
+
+	if err := ensureRequestGetBody(&reqCopy, t.maxBufferBytes); err != nil {
 		return nil, err
 	}
 
-	return t.inner.RoundTrip(req)
+	return t.inner.RoundTrip(&reqCopy)
 }
 
 // ensureRequestGetBody sets [request.GetBody] if it's unset and can either be rewound or fully buffered within the memory limit provided.

internal/ghclient/transport.go

@@ -69,14 +69,6 @@ func newTransport(tokenSource oauth2.TokenSource, opts ClientOptions) (http.Roun
 		}
 	}
 
-	tr = logging.NewLoggingHTTPTransport(tr)
-
-	rtr, err := newRewindableTransport(tr, 0)
-	if err != nil {
-		return nil, err
-	}
-	tr = rtr
-
 	if opts.Retry.Max > 0 {
 		rtr, err := newRetryTransport(tr, opts.Retry)
 		if err != nil {
@@ -87,6 +79,12 @@ func newTransport(tokenSource oauth2.TokenSource, opts ClientOptions) (http.Roun
 
 	tr = newRateLimitTransport(tr)
 
+	rtr, err := newRewindableTransport(tr, 0)
+	if err != nil {
+		return nil, err
+	}
+	tr = rtr
+
 	if opts.Concurrency.Max > 0 {
 		ctr, err := newThrottleTransport(tr, opts.Concurrency)
 		if err != nil {
@@ -95,5 +93,7 @@ func newTransport(tokenSource oauth2.TokenSource, opts ClientOptions) (http.Roun
 		tr = ctr
 	}
 
+	tr = logging.NewLoggingHTTPTransport(tr)
+
 	return tr, nil
 }