add jitter correction for ratelimit handling

heimweh · Dec 12, 2023 · 5415c3a · 5415c3a
1 parent b3dce1e
commit 5415c3a
Show file tree

Hide file tree

Showing 2 changed files with 148 additions and 17 deletions.
diff --git a/pagerduty/pagerduty.go b/pagerduty/pagerduty.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"math/rand"
 	"net/http"
 	"net/url"
 	"strconv"
@@ -23,6 +24,7 @@ const (
 	defaultAppOauthTokenGenerationURL = "https://identity.pagerduty.com/oauth/token"
 	defaultUserAgent                  = "heimweh/go-pagerduty(terraform)"
 	defaultRegion                     = "us"
+	jitterPercent                     = 0.3
 )
 
 // AuthTokenType is an enum of available tokens types
@@ -570,21 +572,23 @@ func (c *Client) decodeErrorResponse(res *Response) error {
 	v := &errorResponse{Error: &Error{ErrorResponse: res}}
 	err := c.DecodeJSON(res, v)
 
-	// Delaying retry based on ratelimit-reset recommended by PagerDuty
-	// https://developer.pagerduty.com/docs/72d3b724589e3-rest-api-rate-limits#reaching-the-limit
-	ratelimitReset := res.Response.Header.Get("ratelimit-reset")
-	if res.Response.StatusCode == http.StatusTooManyRequests && ratelimitReset != "" {
-		waitFor, err := strconv.ParseInt(ratelimitReset, 10, 0)
-		if err == nil {
-			reqMethod := res.Response.Request.Method
-			reqEndpoint := res.Response.Request.URL
-			log.Printf("[INFO] Rate limit hit, throttling by %d seconds until next retry to %s: %s", waitFor, strings.ToUpper(reqMethod), reqEndpoint)
-			time.Sleep(time.Duration(waitFor) * time.Second)
-			v.Error.needToRetry = true
-			return v.Error
-		}
+	if handledError := handleRatelimitError(res, v); handledError != nil {
+		return handledError
+	}
+
+	if handledError := c.handleScopedOAuthError(res, v); handledError != nil {
+		return handledError
+	}
+
+	if err != nil {
+		return fmt.Errorf("%s API call to %s failed: %v", res.Response.Request.Method, res.Response.Request.URL.String(), res.Response.Status)
 	}
+	log.Printf("[INFO] v.Error %+v", v.Error)
+
+	return v.Error
+}
 
+func (c *Client) handleScopedOAuthError(res *Response, v *errorResponse) error {
 	isUsingScopedAPITokenFromCredentials := *c.Config.APIAuthTokenType == AuthTokenTypeUseAppCredentials
 	isOauthScopeMissing := isUsingScopedAPITokenFromCredentials && res.Response.StatusCode == http.StatusForbidden
 	needNewOauthScopedAccessToken := isUsingScopedAPITokenFromCredentials && res.Response.StatusCode == http.StatusUnauthorized
@@ -600,12 +604,51 @@ func (c *Client) decodeErrorResponse(res *Response) error {
 		return v.Error
 	}
 
-	if err != nil {
-		return fmt.Errorf("%s API call to %s failed: %v", res.Response.Request.Method, res.Response.Request.URL.String(), res.Response.Status)
+	return nil
+}
+
+// handleRatelimitError will handle rate limit errors from responses with http
+// code 429. Delaying retry based on ratelimit-reset recommended by PagerDuty
+// https://developer.pagerduty.com/docs/72d3b724589e3-rest-api-rate-limits#reaching-the-limit
+func handleRatelimitError(res *Response, v *errorResponse) error {
+	var markErrorAsRetryable = func(waitFor time.Duration) error {
+		reqMethod := res.Response.Request.Method
+		reqEndpoint := res.Response.Request.URL
+		log.Printf(
+			"[INFO] Rate limit hit, throttling by %v seconds until next retry to %s: %s",
+			strconv.FormatFloat(waitFor.Seconds(), 'f', 1, 64),
+			strings.ToUpper(reqMethod),
+			reqEndpoint)
+		time.Sleep(waitFor)
+		v.Error.needToRetry = true
+		return v.Error
 	}
-	log.Printf("[INFO] v.Error %+v", v.Error)
 
-	return v.Error
+	ratelimitReset := res.Response.Header.Get("ratelimit-reset")
+
+	if res.Response.StatusCode == http.StatusTooManyRequests && ratelimitReset == "" {
+		baseDelay := 5 * time.Second
+		jitter := 1 + (jitterPercent * rand.Float64())
+		waitFor := time.Duration(float64(baseDelay) * jitter)
+
+		return markErrorAsRetryable(waitFor)
+	}
+
+	if res.Response.StatusCode == http.StatusTooManyRequests && ratelimitReset != "" {
+		headerWaitSeconds, err := strconv.ParseInt(ratelimitReset, 10, 0)
+		if err == nil {
+			baseDelay := 500 * time.Millisecond
+			headerWait := time.Duration(headerWaitSeconds) * time.Second
+			jitter := 1 + (jitterPercent * rand.Float64())
+			extraWait := time.Duration(float64(baseDelay) * jitter)
+
+			waitFor := headerWait + extraWait
+
+			return markErrorAsRetryable(waitFor)
+		}
+	}
+
+	return nil
 }
 
 func availableOauthScopes() []string {

diff --git a/pagerduty/pagerduty_test.go b/pagerduty/pagerduty_test.go
@@ -5,6 +5,7 @@ import (
 	"context"
 	"net/http"
 	"net/http/httptest"
+	"reflect"
 	"strings"
 	"testing"
 )
@@ -144,3 +145,90 @@ func TestRetryURL(t *testing.T) {
 		t.Fatalf(err.Error())
 	}
 }
+
+func TestHandleRatelimitErrorWithRatelimitHeaders(t *testing.T) {
+	setup()
+	defer teardown()
+
+	count := 0
+	mux.HandleFunc("/teams", func(w http.ResponseWriter, r *http.Request) {
+		testMethod(t, r, "GET")
+
+		if count > 0 {
+			w.Write([]byte(`{"teams": [{"id": "1"}]}`))
+			return
+		}
+
+		// Expected response ref. https://developer.pagerduty.com/docs/72d3b724589e3-rest-api-rate-limits#reaching-the-limit
+		w.Header().Add("Content-Type", "application/json")
+		w.Header().Add("ratelimit-limit", "960")
+		w.Header().Add("ratelimit-remaining", "0")
+		w.Header().Add("ratelimit-reset", "1")
+		w.WriteHeader(http.StatusTooManyRequests)
+		w.Write([]byte(`{"error":{"message":"Rate Limit Exceeded","code":2020}}`))
+		count++
+	})
+
+	resp, _, err := client.Teams.List(&ListTeamsOptions{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	want := &ListTeamsResponse{
+		Teams: []*Team{
+			{
+				ID: "1",
+			},
+		},
+	}
+
+	if !reflect.DeepEqual(resp, want) {
+		t.Errorf("returned \n\n%#v want \n\n%#v", resp, want)
+	}
+
+}
+
+func TestHandleRatelimitErrorNoRatelimitHeaders(t *testing.T) {
+	setup()
+	defer teardown()
+
+	count := 0
+	mux.HandleFunc("/teams", func(w http.ResponseWriter, r *http.Request) {
+		testMethod(t, r, "GET")
+
+		if count > 0 {
+			w.Write([]byte(`{"teams": [{"id": "1"}]}`))
+			return
+		}
+
+		w.Header().Add("Content-Type", "text/html")
+		w.WriteHeader(http.StatusTooManyRequests)
+		w.Write([]byte(`
+      <html>
+      <head><title>429 Too Many Requests</title></head>
+      <body>
+      <center><h1>429 Too Many Requests</h1></center>
+      <hr><center>nginx</center>
+      </body>
+      </html>`))
+		count++
+	})
+
+	resp, _, err := client.Teams.List(&ListTeamsOptions{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	want := &ListTeamsResponse{
+		Teams: []*Team{
+			{
+				ID: "1",
+			},
+		},
+	}
+
+	if !reflect.DeepEqual(resp, want) {
+		t.Errorf("returned \n\n%#v want \n\n%#v", resp, want)
+	}
+
+}