Merge pull request #35 from pritamsarkar007/master

cofigurable retries for HTTP 429

Author: Seth Ammons

pester.go

@@ -48,7 +48,8 @@ type Client struct {
 	wg *sync.WaitGroup
 
 	sync.Mutex
-	ErrLog []ErrEntry
+	ErrLog         []ErrEntry
+	RetryOnHTTP429 bool
 }
 
 // ErrEntry is used to provide the LogString() data and is populated
@@ -93,11 +94,12 @@ func init() {
 // New constructs a new DefaultClient with sensible default values
 func New() *Client {
 	return &Client{
-		Concurrency: DefaultClient.Concurrency,
-		MaxRetries:  DefaultClient.MaxRetries,
-		Backoff:     DefaultClient.Backoff,
-		ErrLog:      DefaultClient.ErrLog,
-		wg:          &sync.WaitGroup{},
+		Concurrency:    DefaultClient.Concurrency,
+		MaxRetries:     DefaultClient.MaxRetries,
+		Backoff:        DefaultClient.Backoff,
+		ErrLog:         DefaultClient.ErrLog,
+		wg:             &sync.WaitGroup{},
+		RetryOnHTTP429: false,
 	}
 }
 
@@ -279,7 +281,7 @@ func (c *Client) pester(p params) (*http.Response, error) {
 				// Early return if we have a valid result
 				// Only retry (ie, continue the loop) on 5xx status codes and 429
 
-				if err == nil && resp.StatusCode < 500 && resp.StatusCode != 429 {
+				if err == nil && resp.StatusCode < 500 && (resp.StatusCode != 429 || (resp.StatusCode == 429 && !c.RetryOnHTTP429)) {
 					multiplexCh <- result{resp: resp, err: err, req: n, retry: i}
 					return
 				}
@@ -422,6 +424,11 @@ func (c *Client) PostForm(url string, data url.Values) (resp *http.Response, err
 	return c.pester(params{method: "PostForm", url: url, data: data, verb: "POST"})
 }
 
+// set RetryOnHTTP429 for clients,
+func (c *Client) SetRetryOnHTTP429(flag bool) {
+	c.RetryOnHTTP429 = flag
+}
+
 ////////////////////////////////////////
 // Provide self-constructing variants //
 ////////////////////////////////////////

pester_test.go

@@ -40,12 +40,76 @@ func TestConcurrentRequests(t *testing.T) {
 	}
 }
 
+func TestConcurrentRequestsWith429DefaultClient(t *testing.T) {
+	t.Parallel()
+
+	c := New()
+	c.Concurrency = 3
+	c.KeepLog = true
+
+	port, err := serverWith429()
+	if err != nil {
+		t.Fatal("unable to start server", err)
+	}
+
+	url := fmt.Sprintf("http://localhost:%d", port)
+
+	response, err := c.Get(url)
+	if err != nil {
+		t.Fatal("unable to GET", err)
+	}
+	c.Wait()
+
+	response.Body.Close()
+	c.Wait()
+
+	// in the event of an error, let's see what the logs were
+	t.Log("\n", c.LogString())
+
+	if got, want := c.LogErrCount(), 0; got != want {
+		t.Errorf("got %d attempts, want %d", got, want)
+	}
+}
+
+func TestConcurrentRequestsWith400(t *testing.T) {
+	t.Parallel()
+
+	c := New()
+	c.Concurrency = 3
+	c.KeepLog = true
+	c.SetRetryOnHTTP429(true)
+
+	port, err := serverWith400()
+	if err != nil {
+		t.Fatal("unable to start server", err)
+	}
+
+	url := fmt.Sprintf("http://localhost:%d", port)
+
+	response, err := c.Get(url)
+	if err != nil {
+		t.Fatal("unable to GET", err)
+	}
+	c.Wait()
+
+	response.Body.Close()
+	c.Wait()
+
+	// in the event of an error, let's see what the logs were
+	t.Log("\n", c.LogString())
+
+	if got, want := c.LogErrCount(), 0; got != want {
+		t.Errorf("got %d attempts, want %d", got, want)
+	}
+}
+
 func TestConcurrentRequestsWith429(t *testing.T) {
 	t.Parallel()
 
 	c := New()
 	c.Concurrency = 3
 	c.KeepLog = true
+	c.SetRetryOnHTTP429(true)
 
 	port, err := serverWith429()
 	if err != nil {
@@ -71,13 +135,79 @@ func TestConcurrentRequestsWith429(t *testing.T) {
 	}
 }
 
+func TestMaxRetriesConcurrentRequestsWith429DefaultClient(t *testing.T) {
+	t.Parallel()
+
+	c := New()
+	c.Concurrency = 3
+	c.KeepLog = true
+	c.MaxRetries = 5
+
+	port, err := serverWith429()
+	if err != nil {
+		t.Fatal("unable to start server", err)
+	}
+
+	url := fmt.Sprintf("http://localhost:%d", port)
+
+	response, err := c.Get(url)
+	if err != nil {
+		t.Fatal("unable to GET", err)
+	}
+	c.Wait()
+
+	response.Body.Close()
+	c.Wait()
+
+	// in the event of an error, let's see what the logs were
+	t.Log("\n", c.LogString())
+
+	if got, want := c.LogErrCount(), 0; got != want {
+		t.Errorf("got %d attempts, want %d", got, want)
+	}
+}
+
+func TestMaxRetriesConcurrentRequestsWith400(t *testing.T) {
+	t.Parallel()
+
+	c := New()
+	c.Concurrency = 3
+	c.KeepLog = true
+	c.MaxRetries = 5
+	c.SetRetryOnHTTP429(true)
+
+	port, err := serverWith400()
+	if err != nil {
+		t.Fatal("unable to start server", err)
+	}
+
+	url := fmt.Sprintf("http://localhost:%d", port)
+
+	response, err := c.Get(url)
+	if err != nil {
+		t.Fatal("unable to GET", err)
+	}
+	c.Wait()
+
+	response.Body.Close()
+	c.Wait()
+
+	// in the event of an error, let's see what the logs were
+	t.Log("\n", c.LogString())
+
+	if got, want := c.LogErrCount(), 0; got != want {
+		t.Errorf("got %d attempts, want %d", got, want)
+	}
+}
+
 func TestMaxRetriesConcurrentRequestsWith429(t *testing.T) {
 	t.Parallel()
 
 	c := New()
 	c.Concurrency = 3
 	c.KeepLog = true
 	c.MaxRetries = 5
+	c.SetRetryOnHTTP429(true)
 
 	port, err := serverWith429()
 	if err != nil {
@@ -127,13 +257,44 @@ func TestConcurrent2Retry0(t *testing.T) {
 	}
 }
 
+func TestConcurrent2Retry0for429DefaultClient(t *testing.T) {
+	t.Parallel()
+
+	c := New()
+	c.Concurrency = 2
+	c.MaxRetries = 0
+	c.KeepLog = true
+
+	port, err := serverWith429()
+	if err != nil {
+		t.Fatal("unable to start server", err)
+	}
+
+	url := fmt.Sprintf("http://localhost:%d", port)
+
+	_, getErr := c.Get(url)
+
+	if getErr != nil {
+		t.Fatal("unable to GET", getErr)
+	}
+	c.Wait()
+
+	// in the event of an error, let's see what the logs were
+	t.Log("\n", c.LogString())
+
+	if got, want := c.LogErrCount(), 0; got != want {
+		t.Errorf("got %d attempts, want %d", got, want)
+	}
+}
+
 func TestConcurrent2Retry0for429(t *testing.T) {
 	t.Parallel()
 
 	c := New()
 	c.Concurrency = 2
 	c.MaxRetries = 0
 	c.KeepLog = true
+	c.SetRetryOnHTTP429(true)
 
 	port, err := serverWith429()
 	if err != nil {
@@ -629,3 +790,33 @@ func serverWith429() (int, error) {
 
 	return port, nil
 }
+
+func serverWith400() (int, error) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+
+		w.WriteHeader(http.StatusBadRequest)
+		w.Write([]byte("400 Bad Request"))
+	})
+	l, err := net.Listen("tcp", ":0")
+	if err != nil {
+		return -1, fmt.Errorf("unable to secure listener %v", err)
+	}
+	go func() {
+		if err := http.Serve(l, mux); err != nil {
+			log.Fatalf("slow-server error %v", err)
+		}
+	}()
+
+	var port int
+	_, sport, err := net.SplitHostPort(l.Addr().String())
+	if err == nil {
+		port, err = strconv.Atoi(sport)
+	}
+
+	if err != nil {
+		return -1, fmt.Errorf("unable to determine port %v", err)
+	}
+
+	return port, nil
+}