fix: close retried responses in retry middleware

middleware/retry.go

@@ -127,17 +127,17 @@ func (r *RetryMiddleware) RoundTrip(req *http.Request) (*http.Response, error) {
 		}
 
 		resp, err := r.next.RoundTrip(req)
-		if err != nil {
-			lastError = err
-			lastResponse = resp
-			continue
-		}
-
-		if !r.shouldRetry(resp) {
+		if err == nil && !r.shouldRetry(resp) {
 			return resp, nil
 		}
 
+		if err != nil {
+			lastError = err
+		}
 		lastResponse = resp
+		if attempt < attempts-1 {
+			r.discard(resp)
+		}
 	}
 
 	if lastError != nil {
@@ -146,6 +146,15 @@ func (r *RetryMiddleware) RoundTrip(req *http.Request) (*http.Response, error) {
 	return lastResponse, nil
 }
 
+// discard closes the body of a response dropped in favor of another attempt, releasing its connection.
+// the body is not drained first: a read has no time bound and blocks the retry on a stalled upstream.
+func (r *RetryMiddleware) discard(resp *http.Response) {
+	if resp == nil || resp.Body == nil {
+		return
+	}
+	_ = resp.Body.Close()
+}
+
 // bufferRequestBody attempts to buffer the request body for retries
 // this consumes the original body - returns error if body is too large
 func (r *RetryMiddleware) bufferRequestBody(req *http.Request) error {

middleware/retry_test.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"errors"
 	"io"
+	"net"
 	"net/http"
+	"net/http/httptest"
 	"strings"
 	"sync/atomic"
 	"testing"
@@ -530,3 +532,153 @@ func TestRetry_RetryConditions(t *testing.T) {
 		})
 	})
 }
+
+type closeRecorder struct {
+	io.Reader
+	closed atomic.Bool
+}
+
+func (c *closeRecorder) Close() error {
+	c.closed.Store(true)
+	return nil
+}
+
+// retried responses were dropped without closing the body, leaking a connection each
+func TestRetry_DiscardedResponses(t *testing.T) {
+	newBodies := func(n int) []*closeRecorder {
+		res := make([]*closeRecorder, n)
+		for i := range res {
+			res[i] = &closeRecorder{Reader: strings.NewReader("body")}
+		}
+		return res
+	}
+
+	t.Run("closes retried status responses and leaves the returned one open", func(t *testing.T) {
+		bodies := newBodies(3)
+		var attemptCount int32
+		rmock := &mocks.RoundTripper{RoundTripFunc: func(r *http.Request) (*http.Response, error) {
+			count := atomic.AddInt32(&attemptCount, 1)
+			if count < 3 {
+				return &http.Response{StatusCode: 502, Body: bodies[count-1]}, nil
+			}
+			return &http.Response{StatusCode: 200, Body: bodies[count-1]}, nil
+		}}
+
+		h := Retry(3, time.Millisecond)(rmock)
+		req, err := http.NewRequest("GET", "http://example.com/", http.NoBody)
+		require.NoError(t, err)
+
+		resp, err := h.RoundTrip(req)
+		require.NoError(t, err)
+		assert.Equal(t, 200, resp.StatusCode)
+		assert.True(t, bodies[0].closed.Load(), "first retried response should be closed")
+		assert.True(t, bodies[1].closed.Load(), "second retried response should be closed")
+		assert.False(t, bodies[2].closed.Load(), "returned response should stay open")
+	})
+
+	t.Run("closes response returned together with a transport error", func(t *testing.T) {
+		bodies := newBodies(2)
+		var attemptCount int32
+		rmock := &mocks.RoundTripper{RoundTripFunc: func(r *http.Request) (*http.Response, error) {
+			count := atomic.AddInt32(&attemptCount, 1)
+			if count < 2 {
+				return &http.Response{StatusCode: 502, Body: bodies[count-1]}, errors.New("network error")
+			}
+			return &http.Response{StatusCode: 200, Body: bodies[count-1]}, nil
+		}}
+
+		h := Retry(3, time.Millisecond)(rmock)
+		req, err := http.NewRequest("GET", "http://example.com/", http.NoBody)
+		require.NoError(t, err)
+
+		resp, err := h.RoundTrip(req)
+		require.NoError(t, err)
+		assert.Equal(t, 200, resp.StatusCode)
+		assert.True(t, bodies[0].closed.Load(), "response dropped after transport error should be closed")
+		assert.False(t, bodies[1].closed.Load(), "returned response should stay open")
+	})
+
+	t.Run("leaves the last response open when attempts are exhausted", func(t *testing.T) {
+		bodies := newBodies(3)
+		var attemptCount int32
+		rmock := &mocks.RoundTripper{RoundTripFunc: func(r *http.Request) (*http.Response, error) {
+			count := atomic.AddInt32(&attemptCount, 1)
+			return &http.Response{StatusCode: 503, Body: bodies[count-1]}, nil
+		}}
+
+		h := Retry(3, time.Millisecond)(rmock)
+		req, err := http.NewRequest("GET", "http://example.com/", http.NoBody)
+		require.NoError(t, err)
+
+		resp, err := h.RoundTrip(req)
+		require.NoError(t, err)
+		assert.Equal(t, 503, resp.StatusCode)
+		assert.True(t, bodies[0].closed.Load(), "first retried response should be closed")
+		assert.True(t, bodies[1].closed.Load(), "second retried response should be closed")
+		assert.False(t, bodies[2].closed.Load(), "last response is returned to the caller and should stay open")
+
+		data, err := io.ReadAll(resp.Body)
+		require.NoError(t, err)
+		assert.Equal(t, "body", string(data))
+	})
+
+	t.Run("closes retried response when context is cancelled during delay", func(t *testing.T) {
+		bodies := newBodies(1)
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+		rmock := &mocks.RoundTripper{RoundTripFunc: func(r *http.Request) (*http.Response, error) {
+			time.AfterFunc(20*time.Millisecond, cancel)
+			return &http.Response{StatusCode: 503, Body: bodies[0]}, nil
+		}}
+
+		req, err := http.NewRequestWithContext(ctx, "GET", "http://example.com/", http.NoBody)
+		require.NoError(t, err)
+
+		h := Retry(5, 200*time.Millisecond, RetryWithJitter(0))(rmock)
+
+		_, err = h.RoundTrip(req)
+		require.ErrorIs(t, err, context.Canceled)
+		assert.True(t, bodies[0].closed.Load(), "retried response should be closed before the delay")
+	})
+
+	t.Run("releases connections of retried responses", func(t *testing.T) {
+		var calls, opened, closed atomic.Int32
+		srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+			if calls.Add(1)%3 != 0 {
+				w.WriteHeader(http.StatusBadGateway)
+				_, _ = io.WriteString(w, "upstream down")
+				return
+			}
+			_, _ = io.WriteString(w, "ok")
+		}))
+		srv.Config.ConnState = func(_ net.Conn, s http.ConnState) {
+			if s == http.StateNew {
+				opened.Add(1)
+			}
+			if s == http.StateClosed {
+				closed.Add(1)
+			}
+		}
+		srv.Start()
+		defer srv.Close()
+
+		transport := &http.Transport{}
+		h := Retry(3, time.Millisecond)(transport)
+
+		for range 5 {
+			req, err := http.NewRequest("GET", srv.URL, http.NoBody)
+			require.NoError(t, err)
+			resp, err := h.RoundTrip(req)
+			require.NoError(t, err)
+			assert.Equal(t, 200, resp.StatusCode)
+			_, err = io.Copy(io.Discard, resp.Body)
+			require.NoError(t, err)
+			require.NoError(t, resp.Body.Close())
+		}
+
+		assert.Equal(t, int32(15), calls.Load())
+		transport.CloseIdleConnections()
+		assert.Eventually(t, func() bool { return opened.Load() == closed.Load() }, time.Second, 10*time.Millisecond,
+			"every connection should be idle and closed, none held by a retried response")
+	})
+}