Fix request timeouts: retries, body reads, and the flaky timeout test

cli/request.go

@@ -368,18 +368,23 @@ func isRetryable(code int) bool {
 
 // doRequestWithRetry logs and makes a request, retrying as needed (if
 // configured) and returning the last response.
+//
+// Each attempt gets its own rsh-timeout, derived from the request's original
+// context so a retry isn't born already expired. The timeout covers reading
+// the body too: it is released when the returned body is closed, not when
+// this function returns.
 func doRequestWithRetry(log bool, client *http.Client, req *http.Request) (*http.Response, error) {
 	retries := viper.GetInt("rsh-retry")
+	timeout := viper.GetDuration("rsh-timeout")
 
-	if retries == 0 {
-		return client.Do(req)
-	}
-
+	// Only buffer the body when it may have to be sent more than once.
 	var bodyContents []byte
-	if req.Body != nil {
+	if retries > 0 && req.Body != nil {
 		bodyContents, _ = io.ReadAll(req.Body)
 	}
 
+	parent := req.Context()
+
 	var resp *http.Response
 	var err error
 	triesLeft := 1 + retries
@@ -395,24 +400,26 @@ func doRequestWithRetry(log bool, client *http.Client, req *http.Request) (*http
 			LogDebugRequest(req)
 		}
 
-		if timeout := viper.GetDuration("rsh-timeout"); timeout > 0 {
-			ctx, cancel := context.WithTimeout(req.Context(), timeout)
-			defer cancel()
-			req = req.WithContext(ctx)
+		attempt, cancel := req, context.CancelFunc(func() {})
+		if timeout > 0 {
+			var ctx context.Context
+			ctx, cancel = context.WithTimeout(parent, timeout)
+			attempt = req.WithContext(ctx)
 		}
 
 		start := time.Now()
-		resp, err = client.Do(req)
+		resp, err = client.Do(attempt)
 		if err != nil {
+			cancel()
 			if errors.Is(err, context.DeadlineExceeded) {
 				if triesLeft > 0 {
 					// Try again after letting the user know.
-					LogWarning("Got request timeout after %s, retrying", viper.GetDuration("rsh-timeout").Truncate(time.Millisecond))
+					LogWarning("Got request timeout after %s, retrying", timeout.Truncate(time.Millisecond))
 					continue
 				} else {
 					// Add a human-friendly error before the original (context deadline
 					// exceeded).
-					err = fmt.Errorf("Request timed out after %s: %w", viper.GetDuration("rsh-timeout"), err)
+					err = fmt.Errorf("Request timed out after %s: %w", timeout, err)
 				}
 			}
 			return resp, err
@@ -443,17 +450,36 @@ func doRequestWithRetry(log bool, client *http.Client, req *http.Request) (*http
 				}
 			}
 
+			// This response is being discarded, so release it and its timeout.
+			resp.Body.Close()
+			cancel()
+
 			LogWarning("Got %s, retrying in %s", resp.Status, retryAfter.Truncate(time.Millisecond))
 			time.Sleep(retryAfter)
 
 			continue
 		}
+
+		resp.Body = &cancelOnClose{ReadCloser: resp.Body, cancel: cancel}
 		break
 	}
 
 	return resp, err
 }
 
+// cancelOnClose releases a request's timeout context once its response body
+// has been closed, so the body can still be read after the request returns.
+type cancelOnClose struct {
+	io.ReadCloser
+	cancel context.CancelFunc
+}
+
+func (c *cancelOnClose) Close() error {
+	err := c.ReadCloser.Close()
+	c.cancel()
+	return err
+}
+
 // Response describes a parsed HTTP response which can be marshalled to enable
 // printing and filtering/projection.
 type Response struct {

cli/request_test.go

@@ -3,12 +3,16 @@ package cli
 import (
 	"bytes"
 	"errors"
+	"io"
 	"net/http"
+	"net/http/httptest"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"gopkg.in/h2non/gock.v1"
 )
 
@@ -233,20 +237,70 @@ func TestRequestRetryTimeout(t *testing.T) {
 
 	reset(false)
 	viper.Set("rsh-retry", 1)
-	viper.Set("rsh-timeout", 1*time.Millisecond)
-
-	// Duration string value (with units)
+	viper.Set("rsh-timeout", 10*time.Millisecond)
+
+	// The delay must be far longer than the timeout. The context is only
+	// marked expired when its own timer callback runs, so with the two timers
+	// close together (it used to be 1ms against 2ms) a busy machine can fire
+	// gock's first and get a 200 back. gock stops waiting as soon as the
+	// context is done, so the long delay doesn't slow the test down. Both
+	// attempts time out, so the mock has to reply twice.
 	gock.New("http://example.com").
 		Get("/").
 		Times(2).
 		Reply(http.StatusOK).
-		Delay(2 * time.Millisecond)
-		// Note: delay seems to have a bug where subsequent requests without the
-		// delay are still delayed... For now just have it reply twice.
+		Delay(10 * time.Second)
 
 	req, _ := http.NewRequest(http.MethodGet, "http://example.com/", nil)
+	start := time.Now()
 	_, err := MakeRequest(req)
 
 	assert.Error(t, err)
 	assert.ErrorContains(t, err, "timed out")
+
+	// The retry gets a full timeout of its own rather than inheriting the
+	// first attempt's expired one, so it reaches the server too.
+	assert.True(t, gock.IsDone(), "the retry should have reached the server")
+	assert.GreaterOrEqual(t, time.Since(start), 20*time.Millisecond)
+}
+
+func TestRequestTimeoutWithoutRetries(t *testing.T) {
+	defer gock.Off()
+
+	reset(false)
+	viper.Set("rsh-retry", 0)
+	viper.Set("rsh-timeout", 10*time.Millisecond)
+
+	gock.New("http://example.com").
+		Get("/").
+		Reply(http.StatusOK).
+		Delay(10 * time.Second)
+
+	req, _ := http.NewRequest(http.MethodGet, "http://example.com/", nil)
+	_, err := MakeRequest(req)
+
+	assert.ErrorContains(t, err, "timed out")
+}
+
+func TestRequestTimeoutAllowsReadingBody(t *testing.T) {
+	reset(false)
+	viper.Set("rsh-retry", 1)
+	viper.Set("rsh-timeout", 30*time.Second)
+
+	// Larger than the transport buffers, so most of it is read from the
+	// connection after MakeRequest has returned.
+	body := strings.Repeat("x", 4<<20)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Write([]byte(body))
+	}))
+	defer server.Close()
+
+	req, _ := http.NewRequest(http.MethodGet, server.URL, nil)
+	resp, err := MakeRequest(req, WithClient(server.Client()))
+	require.NoError(t, err)
+	defer resp.Body.Close()
+
+	read, err := io.ReadAll(resp.Body)
+	require.NoError(t, err)
+	assert.Len(t, read, len(body))
 }