From ea25c4172f2aa7779d051602e522903fcabb8674 Mon Sep 17 00:00:00 2001
From: panlq01 <panlq01@mingyuanyun.com>
Date: Wed, 6 Apr 2022 15:22:04 +0800
Subject: [PATCH] fix issue #89/#94: retrier called even with 0 retry count and
 tiem sleep will be called even when the retries are exhausted

---
 httpclient/client.go       | 49 ++++++++++++++++++++++++++++----------
 httpclient/options_test.go |  1 -
 2 files changed, 37 insertions(+), 13 deletions(-)

diff --git a/httpclient/client.go b/httpclient/client.go
index 6880776..a49bdcb 100644
--- a/httpclient/client.go
+++ b/httpclient/client.go
@@ -2,6 +2,7 @@ package httpclient
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"io/ioutil"
 	"net/http"
@@ -136,15 +137,17 @@ func (c *Client) Do(request *http.Request) (*http.Response, error) {
 	}
 
 	multiErr := &valkyrie.MultiError{}
+	var err error
+	var shouldRetry bool
 	var response *http.Response
 
-	for i := 0; i <= c.retryCount; i++ {
+	for i := 0; ; i++ {
 		if response != nil {
 			response.Body.Close()
 		}
 
 		c.reportRequestStart(request)
-		var err error
+
 		response, err = c.client.Do(request)
 		if bodyReader != nil {
 			// Reset the body reader after the request since at this point it's already read
@@ -152,28 +155,50 @@ func (c *Client) Do(request *http.Request) (*http.Response, error) {
 			_, _ = bodyReader.Seek(0, 0)
 		}
 
+		shouldRetry, err = c.checkRetry(request.Context(), response, err)
+
 		if err != nil {
 			multiErr.Push(err.Error())
 			c.reportError(request, err)
-			backoffTime := c.retrier.NextInterval(i)
-			time.Sleep(backoffTime)
-			continue
+		} else {
+			c.reportRequestEnd(request, response)
+		}
+
+		if !shouldRetry {
+			break
 		}
-		c.reportRequestEnd(request, response)
 
-		if response.StatusCode >= http.StatusInternalServerError {
-			backoffTime := c.retrier.NextInterval(i)
-			time.Sleep(backoffTime)
-			continue
+		if c.retryCount-i <= 0 {
+			break
 		}
 
-		multiErr = &valkyrie.MultiError{} // Clear errors if any iteration succeeds
-		break
+		backoffTime := c.retrier.NextInterval(i)
+		time.Sleep(backoffTime)
+	}
+
+	if !shouldRetry && err == nil {
+		// Clear errors if any iteration succeeds
+		multiErr = &valkyrie.MultiError{}
 	}
 
 	return response, multiErr.HasError()
 }
 
+func (c *Client) checkRetry(ctx context.Context, resp *http.Response, err error) (bool, error) {
+	if err != nil {
+		return true, err
+	}
+
+	// 429 Too Many Requests is recoverable. Sometimes the server puts
+	// a Retry-After response header to indicate when the server is
+	// available to start processing request from client.
+	if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= http.StatusInternalServerError {
+		return true, nil
+	}
+
+	return false, nil
+}
+
 func (c *Client) reportRequestStart(request *http.Request) {
 	for _, plugin := range c.plugins {
 		plugin.OnRequestStart(request)
diff --git a/httpclient/options_test.go b/httpclient/options_test.go
index a9eb5ba..d1bef15 100644
--- a/httpclient/options_test.go
+++ b/httpclient/options_test.go
@@ -128,6 +128,5 @@ func ExampleWithRetrier() {
 	// Output: retry attempt 0
 	// retry attempt 1
 	// retry attempt 2
-	// retry attempt 3
 	// error
 }