Add a new RateLimitLinearJitterBackoff policy

miagilepner · web-flow · commit e1f5485fe847 · 2025-06-18T16:25:10.000+02:00
diff --git a/client.go b/client.go
@@ -638,6 +638,23 @@ func LinearJitterBackoff(min, max time.Duration, attemptNum int, resp *http.Resp
 	return time.Duration(jitterMin * int64(attemptNum))
 }
 
+// RateLimitLinearJitterBackoff wraps the retryablehttp.LinearJitterBackoff.
+// It first checks if the response status code is http.StatusTooManyRequests
+// (HTTP Code 429) or http.StatusServiceUnavailable (HTTP Code 503). If it is
+// and the response contains a Retry-After response header, it will wait the
+// amount of time specified by the header. Otherwise, this calls
+// LinearJitterBackoff.
+func RateLimitLinearJitterBackoff(min, max time.Duration, attemptNum int, resp *http.Response) time.Duration {
+	if resp != nil {
+		if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusServiceUnavailable {
+			if sleep, ok := parseRetryAfterHeader(resp.Header["Retry-After"]); ok {
+				return sleep
+			}
+		}
+	}
+	return LinearJitterBackoff(min, max, attemptNum, resp)
+}
+
 // PassthroughErrorHandler is an ErrorHandler that directly passes through the
 // values from the net/http library for the final request. The body is not
 // closed.
diff --git a/client_test.go b/client_test.go
@@ -1143,6 +1143,113 @@ func TestClient_PostForm(t *testing.T) {
 	resp.Body.Close()
 }
 
+func TestBackoff_RateLimitLinearJitterBackoff(t *testing.T) {
+	testCases := []struct {
+		name         string
+		min          time.Duration
+		max          time.Duration
+		headers      http.Header
+		responseCode int
+		expect       time.Duration
+	}{
+		{
+			name:         "429 no retry header",
+			min:          time.Second,
+			max:          time.Second,
+			headers:      http.Header{},
+			responseCode: http.StatusTooManyRequests,
+			expect:       time.Second,
+		},
+		{
+			name:         "503 no retry header",
+			min:          time.Second,
+			max:          time.Second,
+			headers:      http.Header{},
+			responseCode: http.StatusServiceUnavailable,
+			expect:       time.Second,
+		},
+		{
+			name: "429 retry header",
+			min:  time.Second,
+			max:  time.Second,
+			headers: http.Header{
+				"Retry-After": []string{"2"},
+			},
+			responseCode: http.StatusTooManyRequests,
+			expect:       2 * time.Second,
+		},
+		{
+			name: "503 retry header",
+			min:  time.Second,
+			max:  time.Second,
+			headers: http.Header{
+				"Retry-After": []string{"2"},
+			},
+			responseCode: http.StatusServiceUnavailable,
+			expect:       2 * time.Second,
+		},
+		{
+			name: "502 ignore retry header",
+			min:  time.Second,
+			max:  time.Second,
+			headers: http.Header{
+				"Retry-After": []string{"2"},
+			},
+			responseCode: http.StatusBadGateway,
+			expect:       time.Second,
+		},
+		{
+			name:         "502 no retry header",
+			min:          time.Second,
+			max:          time.Second,
+			headers:      http.Header{},
+			responseCode: http.StatusBadGateway,
+			expect:       time.Second,
+		},
+		{
+			name: "429 retry header with jitter",
+			min:  time.Second,
+			max:  5 * time.Second,
+			headers: http.Header{
+				"Retry-After": []string{"2"},
+			},
+			responseCode: http.StatusTooManyRequests,
+			expect:       2 * time.Second,
+		},
+		{
+			name: "429 retry header less than min",
+			min:  5 * time.Second,
+			max:  10 * time.Second,
+			headers: http.Header{
+				"Retry-After": []string{"2"},
+			},
+			responseCode: http.StatusTooManyRequests,
+			expect:       2 * time.Second,
+		},
+		{
+			name: "429 retry header in range",
+			min:  time.Second,
+			max:  10 * time.Second,
+			headers: http.Header{
+				"Retry-After": []string{"2"},
+			},
+			responseCode: http.StatusTooManyRequests,
+			expect:       2 * time.Second,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			got := RateLimitLinearJitterBackoff(tc.min, tc.max, 0, &http.Response{
+				StatusCode: tc.responseCode,
+				Header:     tc.headers,
+			})
+			if got != tc.expect {
+				t.Fatalf("expected %s, got %s", tc.expect, got)
+			}
+		})
+	}
+}
+
 func TestBackoff(t *testing.T) {
 	type tcase struct {
 		min    time.Duration
