Retry extra handling

README.md

@@ -45,6 +45,21 @@ The returned response object is an `*http.Response`, the same thing you would
 usually get from `net/http`. Had the request failed one or more times, the above
 call would block and retry with exponential backoff.
 
+## Retrying cases that fail after a seeming success
+
+It's possible for a request to succeed, but then for later processing to fail, and sometimes this requires retrying the full request. E.g. a GET that returns a lot of data may "succeed" but the connection may drop while reading all of the data.
+
+In such a case, you can use `DoWithResponseHandler` (or `GetWithResponseHandler`) rather than `Do` (or `Get`). A toy example (which will retry the full request and succeed on the second attempt) is shown below:
+
+```go
+c := retryablehttp.NewClient()
+handlerShouldRetry := false
+c.GetWithResponseHandler("/foo", func(*http.Response) bool {
+    handlerShouldRetry = !handlerShouldRetry
+    return handlerShouldRetry
+})
+```
+
 ## Getting a stdlib `*http.Client` with retries
 
 It's possible to convert a `*retryablehttp.Client` directly to a `*http.Client`.

client.go

@@ -69,6 +69,11 @@ var (
 	// scheme specified in the URL is invalid. This error isn't typed
 	// specifically so we resort to matching on the error string.
 	schemeErrorRe = regexp.MustCompile(`unsupported protocol scheme`)
+
+	// A regular expression to match the error returned by net/http when the
+	// TLS certificate is not trusted. This error isn't typed
+	// specifically so we resort to matching on the error string.
+	notTrustedErrorRe = regexp.MustCompile(`certificate is not trusted`)
 )
 
 // ReaderFunc is the type of function that can be given natively to NewRequest
@@ -445,6 +450,9 @@ func baseRetryPolicy(resp *http.Response, err error) (bool, error) {
 			}
 
 			// Don't retry if the error was due to TLS cert verification failure.
+			if notTrustedErrorRe.MatchString(v.Error()) {
+				return false, v
+			}
 			if _, ok := v.Err.(x509.UnknownAuthorityError); ok {
 				return false, v
 			}
@@ -545,6 +553,11 @@ func PassthroughErrorHandler(resp *http.Response, err error, _ int) (*http.Respo
 
 // Do wraps calling an HTTP method with retries.
 func (c *Client) Do(req *Request) (*http.Response, error) {
+	return c.DoWithResponseHandler(req, nil)
+}
+
+// DoWithResponseHandler wraps calling an HTTP method plus a response handler with retries.
+func (c *Client) DoWithResponseHandler(req *Request, handler func(*http.Response) (shouldRetry bool)) (*http.Response, error) {
 	c.clientInit.Do(func() {
 		if c.HTTPClient == nil {
 			c.HTTPClient = cleanhttp.DefaultPooledClient()
@@ -598,9 +611,6 @@ func (c *Client) Do(req *Request) (*http.Response, error) {
 		// Attempt the request
 		resp, doErr = c.HTTPClient.Do(req.Request)
 
-		// Check if we should continue with retries.
-		shouldRetry, checkErr = c.CheckRetry(req.Context(), resp, doErr)
-
 		if doErr != nil {
 			switch v := logger.(type) {
 			case LeveledLogger:
@@ -624,6 +634,13 @@ func (c *Client) Do(req *Request) (*http.Response, error) {
 			}
 		}
 
+		// Check if we should continue with retries.
+		shouldRetry, checkErr = c.CheckRetry(req.Context(), resp, doErr)
+
+		successSoFar := !shouldRetry && doErr == nil && checkErr == nil
+		if successSoFar && handler != nil {
+			shouldRetry = handler(resp)
+		}
 		if !shouldRetry {
 			break
 		}
@@ -731,6 +748,16 @@ func (c *Client) Get(url string) (*http.Response, error) {
 	return c.Do(req)
 }
 
+// GetWithResponseHandler is a helper for doing a GET request followed by a function on the response.
+// The intention is for this to be used when errors in the response handling should also be retried.
+func (c *Client) GetWithResponseHandler(url string, handler func(*http.Response) (shouldRetry bool)) (*http.Response, error) {
+	req, err := NewRequest("GET", url, nil)
+	if err != nil {
+		return nil, err
+	}
+	return c.DoWithResponseHandler(req, handler)
+}
+
 // Head is a shortcut for doing a HEAD request without making a new client.
 func Head(url string) (*http.Response, error) {
 	return defaultClient.Head(url)

client_test.go

@@ -167,13 +167,13 @@ func testClientDo(t *testing.T, body interface{}) {
 	// Send the request
 	var resp *http.Response
 	doneCh := make(chan struct{})
+	errCh := make(chan error, 1)
 	go func() {
 		defer close(doneCh)
+		defer close(errCh)
 		var err error
 		resp, err = client.Do(req)
-		if err != nil {
-			t.Fatalf("err: %v", err)
-		}
+		errCh <- err
 	}()
 
 	select {
@@ -247,6 +247,88 @@ func testClientDo(t *testing.T, body interface{}) {
 	if retryCount < 0 {
 		t.Fatal("request log hook was not called")
 	}
+
+	err = <-errCh
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+}
+
+func TestClient_DoWithHandler(t *testing.T) {
+	// Create the client. Use short retry windows so we fail faster.
+	client := NewClient()
+	client.RetryWaitMin = 10 * time.Millisecond
+	client.RetryWaitMax = 10 * time.Millisecond
+	client.RetryMax = 2
+
+	var attempts int
+	client.CheckRetry = func(_ context.Context, resp *http.Response, err error) (bool, error) {
+		attempts++
+		return DefaultRetryPolicy(context.TODO(), resp, err)
+	}
+
+	// Mock server which always responds 200.
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+	}))
+	defer ts.Close()
+
+	alternatingBool := false
+	tests := []struct {
+		name             string
+		handler          func(*http.Response) bool
+		expectedAttempts int
+		err              string
+	}{
+		{
+			name:             "nil handler",
+			handler:          nil,
+			expectedAttempts: 1,
+		},
+		{
+			name:             "handler never should retry",
+			handler:          func(*http.Response) bool { return false },
+			expectedAttempts: 1,
+		},
+		{
+			name: "handler alternates should retry",
+			handler: func(*http.Response) bool {
+				alternatingBool = !alternatingBool
+				return alternatingBool
+			},
+			expectedAttempts: 2,
+		},
+		{
+			name:             "handler always should retry",
+			handler:          func(*http.Response) bool { return true },
+			expectedAttempts: 3,
+			err:              "giving up after 3 attempt(s)",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			attempts = 0
+			// Create the request
+			req, err := NewRequest("GET", ts.URL, nil)
+			if err != nil {
+				t.Fatalf("err: %v", err)
+			}
+
+			// Send the request.
+			_, err = client.DoWithResponseHandler(req, tt.handler)
+			if err != nil && !strings.Contains(err.Error(), tt.err) {
+				t.Fatalf("error does not match expectation, expected: %s, got: %s", tt.err, err.Error())
+			}
+			if err == nil && tt.err != "" {
+				t.Fatalf("no error, expected: %s", tt.err)
+			}
+
+			if attempts != tt.expectedAttempts {
+				t.Fatalf("expected %d attempts, got %d attempts", tt.expectedAttempts, attempts)
+			}
+		})
+	}
 }
 
 func TestClient_Do_fails(t *testing.T) {
@@ -598,7 +680,7 @@ func TestClient_DefaultRetryPolicy_TLS(t *testing.T) {
 
 func TestClient_DefaultRetryPolicy_redirects(t *testing.T) {
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		http.Redirect(w, r, "/", 302)
+		http.Redirect(w, r, "/", http.StatusFound)
 	}))
 	defer ts.Close()
 

roundtripper_test.go

@@ -107,12 +107,12 @@ func TestRoundTripper_TransportFailureErrorHandling(t *testing.T) {
 
 	expectedError := &url.Error{
 		Op:  "Get",
-		URL: "http://this-url-does-not-exist-ed2fb.com/",
+		URL: "http://asdfsa.com/",
 		Err: &net.OpError{
 			Op:  "dial",
 			Net: "tcp",
 			Err: &net.DNSError{
-				Name:       "this-url-does-not-exist-ed2fb.com",
+				Name:       "asdfsa.com",
 				Err:        "no such host",
 				IsNotFound: true,
 			},
@@ -121,10 +121,10 @@ func TestRoundTripper_TransportFailureErrorHandling(t *testing.T) {
 
 	// Get the standard client and execute the request.
 	client := retryClient.StandardClient()
-	_, err := client.Get("http://this-url-does-not-exist-ed2fb.com/")
+	_, err := client.Get("http://asdfsa.com/")
 
 	// assert expectations
-	if !reflect.DeepEqual(normalizeError(err), expectedError) {
+	if !reflect.DeepEqual(expectedError, normalizeError(err)) {
 		t.Fatalf("expected %q, got %q", expectedError, err)
 	}
 }