Use a max retry after duration for secondary rate limit if specified (#3438)

Author: Letiste

github/github.go

@@ -176,6 +176,10 @@ type Client struct {
 	rateLimits              [Categories]Rate // Rate limits for the client as determined by the most recent API calls.
 	secondaryRateLimitReset time.Time        // Secondary rate limit reset for the client as determined by the most recent API calls.
 
+	// If specified, Client will block requests for at most this duration in case of reaching a secondary
+	// rate limit
+	MaxSecondaryRateLimitRetryAfterDuration time.Duration
+
 	// Whether to respect rate limit headers on endpoints that return 302 redirections to artifacts
 	RateLimitRedirectionalEndpoints bool
 
@@ -928,6 +932,10 @@ func (c *Client) bareDo(ctx context.Context, caller *http.Client, req *http.Requ
 		// Update the secondary rate limit if we hit it.
 		rerr, ok := err.(*AbuseRateLimitError)
 		if ok && rerr.RetryAfter != nil {
+			// if a max duration is specified, make sure that we are waiting at most this duration
+			if c.MaxSecondaryRateLimitRetryAfterDuration > 0 && rerr.GetRetryAfter() > c.MaxSecondaryRateLimitRetryAfterDuration {
+				rerr.RetryAfter = &c.MaxSecondaryRateLimitRetryAfterDuration
+			}
 			c.rateMu.Lock()
 			c.secondaryRateLimitReset = time.Now().Add(*rerr.RetryAfter)
 			c.rateMu.Unlock()

github/github_test.go

@@ -1852,6 +1852,47 @@ func TestDo_rateLimit_abuseRateLimitError_xRateLimitReset(t *testing.T) {
 	}
 }
 
+// Ensure *AbuseRateLimitError.RetryAfter respect a max duration if specified.
+func TestDo_rateLimit_abuseRateLimitError_maxDuration(t *testing.T) {
+	t.Parallel()
+	client, mux, _ := setup(t)
+	// specify a max retry after duration of 1 min
+	client.MaxSecondaryRateLimitRetryAfterDuration = 60 * time.Second
+
+	// x-ratelimit-reset value of 1h into the future, to make sure we are way over the max wait time duration.
+	blockUntil := time.Now().Add(1 * time.Hour).Unix()
+
+	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json; charset=utf-8")
+		w.Header().Set(headerRateReset, strconv.Itoa(int(blockUntil)))
+		w.Header().Set(headerRateRemaining, "1") // set remaining to a value > 0 to distinct from a primary rate limit
+		w.WriteHeader(http.StatusForbidden)
+		fmt.Fprintln(w, `{
+   "message": "You have triggered an abuse detection mechanism ...",
+   "documentation_url": "https://docs.github.com/en/rest/overview/resources-in-the-rest-api#abuse-rate-limits"
+}`)
+	})
+
+	req, _ := client.NewRequest("GET", ".", nil)
+	ctx := context.Background()
+	_, err := client.Do(ctx, req, nil)
+
+	if err == nil {
+		t.Error("Expected error to be returned.")
+	}
+	abuseRateLimitErr, ok := err.(*AbuseRateLimitError)
+	if !ok {
+		t.Fatalf("Expected a *AbuseRateLimitError error; got %#v.", err)
+	}
+	if abuseRateLimitErr.RetryAfter == nil {
+		t.Fatalf("abuseRateLimitErr RetryAfter is nil, expected not-nil")
+	}
+	// check that the retry after is set to be the max allowed duration
+	if got, want := *abuseRateLimitErr.RetryAfter, client.MaxSecondaryRateLimitRetryAfterDuration; got != want {
+		t.Errorf("abuseRateLimitErr RetryAfter = %v, want %v", got, want)
+	}
+}
+
 func TestDo_noContent(t *testing.T) {
 	t.Parallel()
 	client, mux, _ := setup(t)