Repository: PagerDuty/go-pagerduty
Pull request #522 (2024-03-20 16:06:03)

An initial implementation of client retries.
# Overview

Adds the ability to configure the client to retry certain types of failures.

## Changes
When configured for retries: 
* The client will retry any request that fails with a network error, 5xx, or 429 status code.
* The client will back off before retrying a request:
  * Server-provided indicators of how long to wait will be used when available.
  * Otherwise we use a truncated binary exponential backoff with jitter.

## Other Notes
Note that this functionality is not supported in the following API functions at this time:
* `pagerduty.CreateEvent`
* `pagerduty.CreateEventWithHTTPClient`
* `client.Do`

This likely be optional functionality for the current 1.9 release. Please give us your feedback.




Patch:
diff --git a/client.go b/client.go
index 741a73f..03de90d 100644
--- a/client.go
+++ b/client.go
@@ -14,10 +14,12 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"math"
 	"net"
 	"net/http"
 	"path"
 	"runtime"
+	"strconv"
 	"strings"
 	"sync/atomic"
 	"time"
@@ -273,6 +275,16 @@ type HTTPClient interface {
 // Keep this unexported so consumers of the package can't make changes to it.
 var defaultHTTPClient HTTPClient = newDefaultHTTPClient()
 
+type retryPolicy struct {
+	MaxDelay   time.Duration
+	MaxRetries int
+}
+
+var defaultRetryPolicy = retryPolicy{
+	MaxDelay:   20 * time.Second,
+	MaxRetries: 0,
+}
+
 // Client wraps http client
 type Client struct {
 	debugFlag    *uint64
@@ -291,7 +303,8 @@ type Client struct {
 	// HTTPClient is the HTTP client used for making requests against the
 	// PagerDuty API. You can use either *http.Client here, or your own
 	// implementation.
-	HTTPClient HTTPClient
+	HTTPClient  HTTPClient
+	retryPolicy retryPolicy
 
 	userAgent string
 }
@@ -307,6 +320,7 @@ func NewClient(authToken string, options ...ClientOptions) *Client {
 		v2EventsAPIEndpoint: v2EventsAPIEndpoint,
 		authType:            apiToken,
 		HTTPClient:          defaultHTTPClient,
+		retryPolicy:         defaultRetryPolicy,
 	}
 
 	for _, opt := range options {
@@ -331,6 +345,17 @@ func WithAPIEndpoint(endpoint string) ClientOptions {
 	}
 }
 
+// WithRetryPolicy configures the client with a retry policy. Configuring a
+// retry policy on the client is currently experimental and should be used with care.
+func WithRetryPolicy(maxRetryAttempts int, maxDelaySeconds int) ClientOptions {
+	return func(c *Client) {
+		c.retryPolicy = retryPolicy{
+			MaxDelay:   time.Duration(maxDelaySeconds) * time.Second,
+			MaxRetries: maxRetryAttempts,
+		}
+	}
+}
+
 // WithTerraformProvider configures the client to be used as the PagerDuty
 // Terraform provider
 func WithTerraformProvider(version string) ClientOptions {
@@ -580,9 +605,61 @@ func dupeRequest(r *http.Request) (*http.Request, error) {
 }
 
 // needed where pagerduty use a different endpoint for certain actions (eg: v2 events)
-func (c *Client) doWithEndpoint(ctx context.Context, endpoint, method, path string, authRequired bool, body io.Reader, headers map[string]string) (*http.Response, error) {
-	var dreq *http.Request
+func (c *Client) doWithEndpoint(
+	ctx context.Context,
+	endpoint,
+	method,
+	path string,
+	authRequired bool,
+	body io.Reader,
+	headers map[string]string,
+) (*http.Response, error) {
 	var resp *http.Response
+	var respErr error
+
+	// Attempt with optional retries
+	for attempt := 0; ; attempt++ {
+		// Build a new request for each atempt.
+		req, err := http.NewRequestWithContext(ctx, method, endpoint+path, body)
+		if err != nil {
+			return nil, fmt.Errorf("failed to build request: %w", err)
+		}
+
+		resp, respErr = c.doSingleRequest(req, authRequired, headers)
+
+		// Handle retry if applicable
+		shouldRetry, delay := c.shouldRetry(resp, respErr, attempt)
+		if !shouldRetry {
+			break
+		}
+
+		select {
+		case <-time.After(delay):
+			continue
+		case <-ctx.Done():
+			err = fmt.Errorf("context completed during retry: %w", ctx.Err())
+			return nil, err
+		}
+	}
+
+	// Handle final results
+	if respErr != nil {
+		return nil, respErr
+	}
+
+	if resp.StatusCode < 200 || resp.StatusCode > 299 {
+		return resp, c.getErrorFromResponse(resp)
+	}
+
+	return resp, nil
+}
+
+func (c *Client) doSingleRequest(
+	req *http.Request,
+	authRequired bool,
+	headers map[string]string,
+) (resp *http.Response, err error) {
+	var dreq *http.Request
 
 	// so that the last request and response can be nil if there was an error
 	// before the request could be fully processed by the origin, we defer these
@@ -599,11 +676,6 @@ func (c *Client) doWithEndpoint(ctx context.Context, endpoint, method, path stri
 		}()
 	}
 
-	req, err := http.NewRequestWithContext(ctx, method, endpoint+path, body)
-	if err != nil {
-		return nil, fmt.Errorf("failed to build request: %w", err)
-	}
-
 	if err := c.prepRequest(req, authRequired, headers); err != nil {
 		return nil, err
 	}
@@ -616,8 +688,44 @@ func (c *Client) doWithEndpoint(ctx context.Context, endpoint, method, path stri
 	}
 
 	resp, err = c.HTTPClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("error calling the API endpoint: %v", err)
+	}
 
-	return c.checkResponse(resp, err)
+	return resp, nil
+}
+
+// shouldRetry reports whether a request should be retried, and if so, how long to wait before retrying.
+func (c *Client) shouldRetry(resp *http.Response, err error, attempt int) (shouldRetry bool, delay time.Duration) {
+	if attempt >= c.retryPolicy.MaxRetries {
+		return false, 0
+	}
+
+	// For now we only retry on a few known error conditions such as the network errors,
+	// 5xx responses, and rate limiting.
+	if err != nil || resp.StatusCode >= 500 {
+		return true, calculateRetryDelay(attempt, c.retryPolicy)
+	} else if resp.StatusCode == http.StatusTooManyRequests {
+		// The REST API rate limits usually return an indication of how long to wait before retrying
+		resetStr := resp.Header.Get("ratelimit-reset")
+		if delaySeconds, err := strconv.Atoi(resetStr); err == nil {
+			return true, time.Duration(delaySeconds) * time.Second
+		}
+		// otherwise use the default retry delay
+		return true, calculateRetryDelay(attempt, c.retryPolicy)
+	} else {
+		return false, 0
+	}
+}
+
+// calculateRetryDelay uses a binary exponential backoff with jitter algorithm
+func calculateRetryDelay(attempt int, retryPolicy retryPolicy) time.Duration {
+	delay := time.Duration(math.Exp2(float64(attempt))) * time.Second
+	if delay > retryPolicy.MaxDelay {
+		delay = retryPolicy.MaxDelay
+	}
+
+	return delay
 }
 
 func (c *Client) do(ctx context.Context, method, path string, body io.Reader, headers map[string]string) (*http.Response, error) {
@@ -642,18 +750,6 @@ func (c *Client) decodeJSON(resp *http.Response, payload interface{}) error {
 	return json.Unmarshal(body, payload)
 }
 
-func (c *Client) checkResponse(resp *http.Response, err error) (*http.Response, error) {
-	if err != nil {
-		return resp, fmt.Errorf("error calling the API endpoint: %v", err)
-	}
-
-	if resp.StatusCode < 200 || resp.StatusCode > 299 {
-		return resp, c.getErrorFromResponse(resp)
-	}
-
-	return resp, nil
-}
-
 func (c *Client) getErrorFromResponse(resp *http.Response) APIError {
 	// check whether the error response is declared as JSON
 	if !strings.HasPrefix(resp.Header.Get("Content-Type"), "application/json") {


Test patch:
diff --git a/client_test.go b/client_test.go
index e7e3aff..e795a99 100644
--- a/client_test.go
+++ b/client_test.go
@@ -48,6 +48,7 @@ func defaultTestClient(serverURL, authToken string) *Client {
 		debugFlag:           new(uint64),
 		lastRequest:         &atomic.Value{},
 		lastResponse:        &atomic.Value{},
+		retryPolicy:         defaultRetryPolicy,
 	}
 }
 
@@ -688,6 +689,109 @@ func TestClient_Do(t *testing.T) {
 	}
 }
 
+func TestClient_RetriesHandleRequestBodies(t *testing.T) {
+	setup()
+	defer teardown()
+
+	attempt := 1
+
+	mux.HandleFunc("/test", func(w http.ResponseWriter, r *http.Request) {
+		defer r.Body.Close()
+		_, err := io.Copy(io.Discard, r.Body)
+
+		if err != nil {
+			t.Fatalf("expected no error, got %v", err)
+		}
+
+		if attempt == 1 {
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusTooManyRequests)
+			w.Write([]byte(`{"empty":"object"}`))
+			attempt++
+		} else {
+			w.Header().Set("Content-Type", "application/json")
+			w.WriteHeader(http.StatusCreated)
+			w.Write([]byte(`{"empty":"object"}`))
+		}
+	})
+
+	client := NewClient("foo",
+		WithAPIEndpoint(server.URL),
+		WithRetryPolicy(2, 1),
+	)
+
+	_, err := client.do(context.Background(), "POST", "/test", strings.NewReader("some data\n"), nil)
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	if attempt < 2 {
+		t.Fatalf("expected 2 attempts, got %d", attempt)
+	}
+}
+
+func TestClient_RetriesOnVariousConditions(t *testing.T) {
+	tests := []struct {
+		name             string
+		status           int
+		expectedAttempts int
+	}{
+		{
+			name:             "rate_limited",
+			status:           http.StatusTooManyRequests,
+			expectedAttempts: 3,
+		},
+		{
+			name:             "server_error",
+			status:           http.StatusInternalServerError,
+			expectedAttempts: 3,
+		},
+		{
+			name:             "bad_request",
+			status:           http.StatusBadRequest,
+			expectedAttempts: 1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			setup()
+			defer teardown()
+
+			c := NewClient("foo",
+				WithAPIEndpoint(server.URL),
+				WithRetryPolicy(2, 1),
+			)
+
+			attempt := 0
+
+			mux.HandleFunc("/test", func(w http.ResponseWriter, r *http.Request) {
+				attempt++
+				defer r.Body.Close()
+				_, err := io.Copy(io.Discard, r.Body)
+
+				if err != nil {
+					t.Fatalf("expected no error, got %v", err)
+				}
+
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(tt.status)
+				w.Write([]byte(`{"empty":"object"}`))
+			})
+
+			resp, err := c.do(context.Background(), "GET", "/test", nil, nil)
+			testErrCheck(t, "client.do()", "response failed with status code", err)
+
+			defer resp.Body.Close()
+			_, _ = io.Copy(io.Discard, resp.Body)
+
+			if attempt != tt.expectedAttempts {
+				t.Fatalf("expected %d attempts, got %d", tt.expectedAttempts, attempt)
+			}
+		})
+	}
+}
+
 func TestClient_UserAgentDefault(t *testing.T) {
 	setup()
 	defer teardown()