From eb782189c7def38928f72f8303eee606e01c78db Mon Sep 17 00:00:00 2001
From: Joe Dollard <jdollard@akamai.com>
Date: Mon, 19 Aug 2019 09:12:03 -0400
Subject: [PATCH 1/2] support setting the retry policy to use

---
 api/client.go | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/api/client.go b/api/client.go
index fb50f2073983..0460f7ed6d07 100644
--- a/api/client.go
+++ b/api/client.go
@@ -88,6 +88,9 @@ type Config struct {
 	// The Backoff function to use; a default is used if not provided
 	Backoff retryablehttp.Backoff
 
+	// The CheckRetry function to use; a default is used if not provided
+	CheckRetry retryablehttp.CheckRetry
+
 	// Limiter is the rate limiter used by the client.
 	// If this pointer is nil, then there will be no limit set.
 	// In contrast, if this pointer is set, even to an empty struct,
@@ -171,6 +174,7 @@ func DefaultConfig() *Config {
 	}
 
 	config.Backoff = retryablehttp.LinearJitterBackoff
+	config.CheckRetry = retryablehttp.DefaultRetryPolicy
 	config.MaxRetries = 2
 
 	return config
@@ -488,6 +492,16 @@ func (c *Client) SetMaxRetries(retries int) {
 	c.config.MaxRetries = retries
 }
 
+// SetCheckRetry sets the CheckRetry function to be used for future requests.
+func (c *Client) SetCheckRetry(checkRetry retryablehttp.CheckRetry) {
+	c.modifyLock.RLock()
+	c.config.modifyLock.Lock()
+	defer c.config.modifyLock.Unlock()
+	c.modifyLock.RUnlock()
+
+	c.config.CheckRetry = checkRetry
+}
+
 // SetClientTimeout sets the client request timeout
 func (c *Client) SetClientTimeout(timeout time.Duration) {
 	c.modifyLock.RLock()
@@ -643,6 +657,7 @@ func (c *Client) Clone() (*Client, error) {
 		MaxRetries: config.MaxRetries,
 		Timeout:    config.Timeout,
 		Backoff:    config.Backoff,
+		CheckRetry: config.CheckRetry,
 		Limiter:    config.Limiter,
 	}
 	config.modifyLock.RUnlock()
@@ -740,6 +755,7 @@ func (c *Client) RawRequestWithContext(ctx context.Context, r *Request) (*Respon
 	c.config.modifyLock.RLock()
 	limiter := c.config.Limiter
 	maxRetries := c.config.MaxRetries
+	checkRetry := c.config.CheckRetry
 	backoff := c.config.Backoff
 	httpClient := c.config.HttpClient
 	timeout := c.config.Timeout
@@ -784,13 +800,17 @@ START:
 		backoff = retryablehttp.LinearJitterBackoff
 	}
 
+	if checkRetry == nil {
+		checkRetry = retryablehttp.DefaultRetryPolicy
+	}
+
 	client := &retryablehttp.Client{
 		HTTPClient:   httpClient,
 		RetryWaitMin: 1000 * time.Millisecond,
 		RetryWaitMax: 1500 * time.Millisecond,
 		RetryMax:     maxRetries,
-		CheckRetry:   retryablehttp.DefaultRetryPolicy,
 		Backoff:      backoff,
+		CheckRetry:   checkRetry,
 		ErrorHandler: retryablehttp.PassthroughErrorHandler,
 	}
 

From 191943ae6ac2c67955efa1fe4fba8679b66a69e6 Mon Sep 17 00:00:00 2001
From: Joe Dollard <jdollard@akamai.com>
Date: Mon, 19 Aug 2019 11:09:50 -0400
Subject: [PATCH 2/2] do not set CheckRetry in DefaultConfig func

---
 api/client.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/api/client.go b/api/client.go
index 0460f7ed6d07..088cb0215af8 100644
--- a/api/client.go
+++ b/api/client.go
@@ -174,7 +174,6 @@ func DefaultConfig() *Config {
 	}
 
 	config.Backoff = retryablehttp.LinearJitterBackoff
-	config.CheckRetry = retryablehttp.DefaultRetryPolicy
 	config.MaxRetries = 2
 
 	return config