Merge pull request #35 from keithballdotnet/retry2

Added backoff retry mechanism.

client.go

@@ -82,6 +82,7 @@ type Client struct {
 	Debug       bool
 	DisableWarn bool
 	Log         *log.Logger
+	RetryCount  int
 
 	httpClient       *http.Client
 	transport        *http.Transport
@@ -339,6 +340,12 @@ func (c *Client) SetDebug(d bool) *Client {
 	return c
 }
 
+// SetRetryCount method enables retry on `go-resty` client. Uses a Backoff mechanism.
+func (c *Client) SetRetryCount(count int) *Client {
+	c.RetryCount = count
+	return c
+}
+
 // SetDisableWarn method disables the warning message on `go-resty` client.
 // For example: go-resty warns the user when BasicAuth used on HTTP mode.
 //		resty.SetDisableWarn(true)

default.go

@@ -37,6 +37,7 @@ func New() *Client {
 		httpClient: &http.Client{Jar: cookieJar},
 		transport:  &http.Transport{},
 		mutex:      &sync.Mutex{},
+		RetryCount: 0,
 	}
 
 	// Default redirect policy
@@ -133,6 +134,11 @@ func SetDebug(d bool) *Client {
 	return DefaultClient.SetDebug(d)
 }
 
+// SetRetryCount method set the retry count. See `Client.SetRetryCount` for more information.
+func SetRetryCount(count int) *Client {
+	return DefaultClient.SetRetryCount(count)
+}
+
 // SetDisableWarn method disables warning comes from `go-resty` client. See `Client.SetDisableWarn` for more information.
 func SetDisableWarn(d bool) *Client {
 	return DefaultClient.SetDisableWarn(d)

request.go

@@ -428,7 +428,24 @@ func (r *Request) Execute(method, url string) (*Response, error) {
 	r.Method = method
 	r.URL = url
 
-	return r.client.execute(r)
+	if r.client.RetryCount == 0 {
+		return r.client.execute(r)
+	}
+
+	var resp *Response
+	var err error
+	attempt := 0
+	_ = Backoff(func() error {
+		attempt++
+		resp, err = r.client.execute(r)
+		if err != nil {
+			r.client.Log.Printf("ERROR [%v] Attempt [%v]", err, attempt)
+		}
+
+		return err
+	}, Retries(r.client.RetryCount))
+
+	return resp, err
 }
 
 func (r *Request) fmtBodyString() (body string) {

resty_test.go

@@ -879,6 +879,20 @@ func TestClientTimeout(t *testing.T) {
 	assertEqual(t, true, strings.Contains(err.Error(), "i/o timeout"))
 }
 
+func TestClientRetry(t *testing.T) {
+	ts := createGetServer(t)
+	defer ts.Close()
+
+	c := dc()
+	c.SetHTTPMode().
+		SetTimeout(time.Duration(time.Second * 3)).
+		SetRetryCount(3)
+
+	_, err := c.R().Get(ts.URL + "/set-retrycount-test")
+
+	assertError(t, err)
+}
+
 func TestClientTimeoutInternalError(t *testing.T) {
 	c := dc()
 	c.SetHTTPMode()
@@ -1342,6 +1356,9 @@ func TestClientOptions(t *testing.T) {
 	SetDisableWarn(true)
 	assertEqual(t, DefaultClient.DisableWarn, true)
 
+	SetRetryCount(3)
+	assertEqual(t, 3, DefaultClient.RetryCount)
+
 	err := &AuthError{}
 	SetError(err)
 	if reflect.TypeOf(err) == DefaultClient.Error {
@@ -1383,20 +1400,31 @@ func getTestDataPath() string {
 	return pwd + "/test-data"
 }
 
+// Used for retry testing...
+var attempt int
+
 func createGetServer(t *testing.T) *httptest.Server {
 	ts := createTestServer(func(w http.ResponseWriter, r *http.Request) {
 		t.Logf("Method: %v", r.Method)
 		t.Logf("Path: %v", r.URL.Path)
+
 		if r.Method == GET {
 			if r.URL.Path == "/" {
 				w.Write([]byte("TestGet: text response"))
 			} else if r.URL.Path == "/mypage" {
 				w.WriteHeader(http.StatusBadRequest)
 			} else if r.URL.Path == "/mypage2" {
 				w.Write([]byte("TestGet: text response from mypage2"))
+			} else if r.URL.Path == "/set-retrycount-test" {
+				attempt++
+				if attempt != 3 {
+					time.Sleep(time.Second * 6)
+				}
+				w.Write([]byte("TestClientRetry page"))
 			} else if r.URL.Path == "/set-timeout-test" {
 				time.Sleep(time.Second * 6)
 				w.Write([]byte("TestClientTimeout page"))
+
 			} else if r.URL.Path == "/my-image.png" {
 				fileBytes, _ := ioutil.ReadFile(getTestDataPath() + "/test-img.png")
 				w.Header().Set("Content-Type", "image/png")

retry.go

@@ -0,0 +1,69 @@
+package resty
+
+import (
+	"math"
+	"math/rand"
+	"time"
+)
+
+type function func() error
+
+// Option ...
+type Option func(*Options)
+
+// Options ...
+type Options struct {
+	maxRetries  int
+	waitTime    int
+	maxWaitTime int
+}
+
+// Retries sets the max number of retries
+func Retries(value int) Option {
+	return func(o *Options) {
+		o.maxRetries = value
+	}
+}
+
+// WaitTime sets the default wait time to sleep between requests
+func WaitTime(value int) Option {
+	return func(o *Options) {
+		o.waitTime = value
+	}
+}
+
+// MaxWaitTime sets the max wait time to sleep between requests
+func MaxWaitTime(value int) Option {
+	return func(o *Options) {
+		o.maxWaitTime = value
+	}
+}
+
+//Backoff retries with increasing timeout duration up until X amount of retries (Default is 3 attempts, Override with option Retries(n))
+func Backoff(operation function, options ...Option) error {
+	// Defaults
+	opts := Options{maxRetries: 3, waitTime: 100, maxWaitTime: 2000}
+	for _, o := range options {
+		o(&opts)
+	}
+
+	var err error
+	base := float64(opts.waitTime)        // Time to wait between each attempt
+	capLevel := float64(opts.maxWaitTime) // Maximum amount of wait time for the retry
+	for attempt := 0; attempt < opts.maxRetries; attempt++ {
+		err = operation()
+		if err == nil {
+			return nil
+		}
+		// Adding capped exponential backup with jitter
+		// See the following article...
+		// http://www.awsarchitectureblog.com/2015/03/backoff.html
+		temp := math.Min(capLevel, base*math.Exp2(float64(attempt)))
+		sleepTime := int(temp/2) + rand.Intn(int(temp/2))
+
+		sleepDuration := time.Duration(sleepTime) * time.Millisecond
+		time.Sleep(sleepDuration)
+	}
+
+	return err
+}

retry_test.go

@@ -0,0 +1,36 @@
+package resty
+
+import (
+	"errors"
+	"testing"
+)
+
+func TestBackoffSuccess(t *testing.T) {
+	attempts := 3
+	externalCounter := 0
+	retryErr := Backoff(func() error {
+		externalCounter++
+		if externalCounter < attempts {
+			return errors.New("Not yet got the number we're after...")
+		}
+		return nil
+	})
+
+	assertError(t, retryErr)
+	assertEqual(t, externalCounter, attempts)
+}
+
+func TestBackoffTenAttemptsSuccess(t *testing.T) {
+	attempts := 10
+	externalCounter := 0
+	retryErr := Backoff(func() error {
+		externalCounter++
+		if externalCounter < attempts {
+			return errors.New("Not yet got the number we're after...")
+		}
+		return nil
+	}, Retries(attempts), WaitTime(5), MaxWaitTime(500))
+
+	assertError(t, retryErr)
+	assertEqual(t, externalCounter, attempts)
+}