diff --git a/cli/request.go b/cli/request.go index feed371c..59a221ac 100644 --- a/cli/request.go +++ b/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 { diff --git a/cli/request_test.go b/cli/request_test.go index 6565a351..61c99089 100644 --- a/cli/request_test.go +++ b/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)) }