From 16f7b9cbf64a7b42af534d6ca1e3f0bcc21b470c Mon Sep 17 00:00:00 2001 From: erezrokah Date: Mon, 1 Apr 2024 19:02:21 +0100 Subject: [PATCH 1/8] feat: Allow blocking until primary rate limit is reset --- README.md | 6 ++++++ github/github.go | 12 ++++++++++++ 2 files changed, 18 insertions(+) diff --git a/README.md b/README.md index 7e429525936..bfe5c9df8ef 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,12 @@ if _, ok := err.(*github.AbuseRateLimitError); ok { } ``` +Alternatively, you can block until the rate limit is reset by using the `context.WithValue` method: + +````go +repos, _, err := client.Repositories.List(context.WithValue(ctx, github.SleepUntilPrimaryRateLimitResetWhenRateLimited, true), "", nil) +``` + You can use [go-github-ratelimit](https://github.com/gofri/go-github-ratelimit) to handle secondary rate limit sleep-and-retry for you. diff --git a/github/github.go b/github/github.go index bf0b5715a89..9f713117851 100644 --- a/github/github.go +++ b/github/github.go @@ -804,6 +804,7 @@ type requestContext uint8 const ( bypassRateLimitCheck requestContext = iota + SleepUntilPrimaryRateLimitResetWhenRateLimited ) // BareDo sends an API request and lets you handle the api response. If an error @@ -889,6 +890,13 @@ func (c *Client) BareDo(ctx context.Context, req *http.Request) (*Response, erro err = aerr } + rateLimitError, ok := err.(*RateLimitError) + if ok && req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { + time.Sleep(time.Until(rateLimitError.Rate.Reset.Time)) + // re-send the request + return c.BareDo(ctx, req) + } + // Update the secondary rate limit if we hit it. rerr, ok := err.(*AbuseRateLimitError) if ok && rerr.RetryAfter != nil { @@ -942,6 +950,10 @@ func (c *Client) checkRateLimitBeforeDo(req *http.Request, rateLimitCategory Rat rate := c.rateLimits[rateLimitCategory] c.rateMu.Unlock() if !rate.Reset.Time.IsZero() && rate.Remaining == 0 && time.Now().Before(rate.Reset.Time) { + if req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { + time.Sleep(time.Until(rate.Reset.Time)) + return nil + } // Create a fake response. resp := &http.Response{ Status: http.StatusText(http.StatusForbidden), From ec3d31c23dbdad0efddd8a09424e606563e50262 Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 2 Apr 2024 16:53:18 +0100 Subject: [PATCH 2/8] test: Add unit tests --- github/github_test.go | 72 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/github/github_test.go b/github/github_test.go index efc42d432fb..b5b2fcfee47 100644 --- a/github/github_test.go +++ b/github/github_test.go @@ -1381,6 +1381,78 @@ func TestDo_rateLimit_ignoredFromCache(t *testing.T) { } } +// Ensure *RateLimitError is returned when API rate limit is exceeded. +func TestDo_rateLimit_sleepUntilResponseResetLimit(t *testing.T) { + client, mux, _, teardown := setup() + defer teardown() + + reset := time.Now().UTC().Add(time.Second) + + var firstRequest = true + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + if firstRequest { + firstRequest = false + w.Header().Set(headerRateLimit, "60") + w.Header().Set(headerRateRemaining, "0") + w.Header().Set(headerRateReset, fmt.Sprint(reset.Unix())) + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(http.StatusForbidden) + fmt.Fprintln(w, `{ + "message": "API rate limit exceeded for xxx.xxx.xxx.xxx. (But here's the good news: Authenticated requests get a higher rate limit. Check out the documentation for more details.)", + "documentation_url": "https://docs.github.com/en/rest/overview/resources-in-the-rest-api#abuse-rate-limits" +}`) + return + } + w.Header().Set(headerRateLimit, "5000") + w.Header().Set(headerRateRemaining, "5000") + w.Header().Set(headerRateReset, fmt.Sprint(reset.Add(time.Hour).Unix())) + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(http.StatusOK) + fmt.Fprintln(w, `{}`) + }) + + req, _ := client.NewRequest("GET", ".", nil) + ctx := context.Background() + resp, err := client.Do(context.WithValue(ctx, SleepUntilPrimaryRateLimitResetWhenRateLimited, true), req, nil) + if err != nil { + t.Errorf("Do returned unexpected error: %v", err) + } + if got, want := resp.StatusCode, http.StatusOK; got != want { + t.Errorf("Response status code = %v, want %v", got, want) + } +} + +// Ensure a network call is not made when it's known that API rate limit is still exceeded. +func TestDo_rateLimit_sleepUntilClientResetLimit(t *testing.T) { + client, mux, _, teardown := setup() + defer teardown() + + reset := time.Now().UTC().Add(time.Second) + client.rateLimits[CoreCategory] = Rate{Limit: 5000, Remaining: 0, Reset: Timestamp{reset}} + requestCount := 0 + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + requestCount++ + w.Header().Set(headerRateLimit, "5000") + w.Header().Set(headerRateRemaining, "5000") + w.Header().Set(headerRateReset, fmt.Sprint(reset.Add(time.Hour).Unix())) + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(http.StatusOK) + fmt.Fprintln(w, `{}`) + }) + req, _ := client.NewRequest("GET", ".", nil) + ctx := context.Background() + resp, err := client.Do(context.WithValue(ctx, SleepUntilPrimaryRateLimitResetWhenRateLimited, true), req, nil) + if err != nil { + t.Errorf("Do returned unexpected error: %v", err) + } + if got, want := resp.StatusCode, http.StatusOK; got != want { + t.Errorf("Response status code = %v, want %v", got, want) + } + if got, want := requestCount, 1; got != want { + t.Errorf("Expected 1 request, got %d", got) + } +} + // Ensure *AbuseRateLimitError is returned when the response indicates that // the client has triggered an abuse detection mechanism. func TestDo_rateLimit_abuseRateLimitError(t *testing.T) { From 05a697ff582b3fb800f11a1addf61cfdd99e0a1a Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 2 Apr 2024 17:07:55 +0100 Subject: [PATCH 3/8] fix: Retry the request only once, add test --- github/github.go | 4 ++-- github/github_test.go | 31 +++++++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/github/github.go b/github/github.go index 4a7cd033ce1..86ac3a95db1 100644 --- a/github/github.go +++ b/github/github.go @@ -893,8 +893,8 @@ func (c *Client) BareDo(ctx context.Context, req *http.Request) (*Response, erro rateLimitError, ok := err.(*RateLimitError) if ok && req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { time.Sleep(time.Until(rateLimitError.Rate.Reset.Time)) - // re-send the request - return c.BareDo(ctx, req) + // retry the request once the rate limit has reset + return c.BareDo(context.WithValue(req.Context(), SleepUntilPrimaryRateLimitResetWhenRateLimited, nil), req) } // Update the secondary rate limit if we hit it. diff --git a/github/github_test.go b/github/github_test.go index b5b2fcfee47..993940191de 100644 --- a/github/github_test.go +++ b/github/github_test.go @@ -1422,6 +1422,37 @@ func TestDo_rateLimit_sleepUntilResponseResetLimit(t *testing.T) { } } +func TestDo_rateLimit_sleepUntilResponseResetLimitRetryOnce(t *testing.T) { + client, mux, _, teardown := setup() + defer teardown() + + reset := time.Now().UTC().Add(time.Second) + + requestCount := 0 + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + requestCount++ + w.Header().Set(headerRateLimit, "60") + w.Header().Set(headerRateRemaining, "0") + w.Header().Set(headerRateReset, fmt.Sprint(reset.Unix())) + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(http.StatusForbidden) + fmt.Fprintln(w, `{ + "message": "API rate limit exceeded for xxx.xxx.xxx.xxx. (But here's the good news: Authenticated requests get a higher rate limit. Check out the documentation for more details.)", + "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(context.WithValue(ctx, SleepUntilPrimaryRateLimitResetWhenRateLimited, true), req, nil) + if err == nil { + t.Error("Expected error to be returned.") + } + if got, want := requestCount, 2; got != want { + t.Errorf("Expected 2 requests, got %d", got) + } +} + // Ensure a network call is not made when it's known that API rate limit is still exceeded. func TestDo_rateLimit_sleepUntilClientResetLimit(t *testing.T) { client, mux, _, teardown := setup() From d9d555fa49aefe330ff3daad3927e38503644bc7 Mon Sep 17 00:00:00 2001 From: erezrokah Date: Tue, 2 Apr 2024 17:12:29 +0100 Subject: [PATCH 4/8] refactor: Extract to method, add buffer --- github/github.go | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/github/github.go b/github/github.go index 86ac3a95db1..67a539a8f6f 100644 --- a/github/github.go +++ b/github/github.go @@ -892,8 +892,8 @@ func (c *Client) BareDo(ctx context.Context, req *http.Request) (*Response, erro rateLimitError, ok := err.(*RateLimitError) if ok && req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { - time.Sleep(time.Until(rateLimitError.Rate.Reset.Time)) - // retry the request once the rate limit has reset + sleepUntilResetWithBuffer(rateLimitError.Rate.Reset.Time) + // retry the request once (and only once) the rate limit has reset return c.BareDo(context.WithValue(req.Context(), SleepUntilPrimaryRateLimitResetWhenRateLimited, nil), req) } @@ -951,7 +951,7 @@ func (c *Client) checkRateLimitBeforeDo(req *http.Request, rateLimitCategory Rat c.rateMu.Unlock() if !rate.Reset.Time.IsZero() && rate.Remaining == 0 && time.Now().Before(rate.Reset.Time) { if req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { - time.Sleep(time.Until(rate.Reset.Time)) + sleepUntilResetWithBuffer(rate.Reset.Time) return nil } // Create a fake response. @@ -1526,6 +1526,11 @@ func formatRateReset(d time.Duration) string { return fmt.Sprintf("[rate reset in %v]", timeString) } +func sleepUntilResetWithBuffer(reset time.Time) { + buffer := time.Second + time.Sleep(time.Until(reset) + buffer) +} + // When using roundTripWithOptionalFollowRedirect, note that it // is the responsibility of the caller to close the response body. func (c *Client) roundTripWithOptionalFollowRedirect(ctx context.Context, u string, maxRedirects int, opts ...RequestOption) (*http.Response, error) { From 1579468b70a94704b81f47dacbd346eac4133b88 Mon Sep 17 00:00:00 2001 From: Erez Rokah Date: Tue, 2 Apr 2024 19:18:01 +0300 Subject: [PATCH 5/8] Update github.go --- github/github.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/github/github.go b/github/github.go index 67a539a8f6f..9264bb585f7 100644 --- a/github/github.go +++ b/github/github.go @@ -893,7 +893,7 @@ func (c *Client) BareDo(ctx context.Context, req *http.Request) (*Response, erro rateLimitError, ok := err.(*RateLimitError) if ok && req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { sleepUntilResetWithBuffer(rateLimitError.Rate.Reset.Time) - // retry the request once (and only once) the rate limit has reset + // retry the request once when the rate limit has reset return c.BareDo(context.WithValue(req.Context(), SleepUntilPrimaryRateLimitResetWhenRateLimited, nil), req) } From 22eba4c83515649fe06e82b4b7dad938a06b5044 Mon Sep 17 00:00:00 2001 From: erezrokah Date: Thu, 18 Apr 2024 17:24:16 +0100 Subject: [PATCH 6/8] feat: Add context cancellation --- github/github.go | 35 +++++++++++++++++----- github/github_test.go | 69 ++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 96 insertions(+), 8 deletions(-) diff --git a/github/github.go b/github/github.go index 9264bb585f7..f9054aa489e 100644 --- a/github/github.go +++ b/github/github.go @@ -892,7 +892,10 @@ func (c *Client) BareDo(ctx context.Context, req *http.Request) (*Response, erro rateLimitError, ok := err.(*RateLimitError) if ok && req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { - sleepUntilResetWithBuffer(rateLimitError.Rate.Reset.Time) + err := sleepUntilResetWithBuffer(req.Context(), rateLimitError.Rate.Reset.Time) + if err != nil { + return response, err + } // retry the request once when the rate limit has reset return c.BareDo(context.WithValue(req.Context(), SleepUntilPrimaryRateLimitResetWhenRateLimited, nil), req) } @@ -950,10 +953,6 @@ func (c *Client) checkRateLimitBeforeDo(req *http.Request, rateLimitCategory Rat rate := c.rateLimits[rateLimitCategory] c.rateMu.Unlock() if !rate.Reset.Time.IsZero() && rate.Remaining == 0 && time.Now().Before(rate.Reset.Time) { - if req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { - sleepUntilResetWithBuffer(rate.Reset.Time) - return nil - } // Create a fake response. resp := &http.Response{ Status: http.StatusText(http.StatusForbidden), @@ -962,6 +961,19 @@ func (c *Client) checkRateLimitBeforeDo(req *http.Request, rateLimitCategory Rat Header: make(http.Header), Body: io.NopCloser(strings.NewReader("")), } + + if req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { + err := sleepUntilResetWithBuffer(req.Context(), rate.Reset.Time) + if err == nil { + return nil + } + return &RateLimitError{ + Rate: rate, + Response: resp, + Message: fmt.Sprintf("Context cancelled while waiting for rate limit to reset until %v, not making remote request.", rate.Reset.Time), + } + } + return &RateLimitError{ Rate: rate, Response: resp, @@ -1526,9 +1538,18 @@ func formatRateReset(d time.Duration) string { return fmt.Sprintf("[rate reset in %v]", timeString) } -func sleepUntilResetWithBuffer(reset time.Time) { +func sleepUntilResetWithBuffer(ctx context.Context, reset time.Time) error { buffer := time.Second - time.Sleep(time.Until(reset) + buffer) + timer := time.NewTimer(time.Until(reset) + buffer) + select { + case <-ctx.Done(): + if !timer.Stop() { + <-timer.C + } + return ctx.Err() + case <-timer.C: + } + return nil } // When using roundTripWithOptionalFollowRedirect, note that it diff --git a/github/github_test.go b/github/github_test.go index 993940191de..8c6c38a84f0 100644 --- a/github/github_test.go +++ b/github/github_test.go @@ -1381,7 +1381,7 @@ func TestDo_rateLimit_ignoredFromCache(t *testing.T) { } } -// Ensure *RateLimitError is returned when API rate limit is exceeded. +// Ensure sleeps until the rate limit is reset when the client is rate limited. func TestDo_rateLimit_sleepUntilResponseResetLimit(t *testing.T) { client, mux, _, teardown := setup() defer teardown() @@ -1422,6 +1422,7 @@ func TestDo_rateLimit_sleepUntilResponseResetLimit(t *testing.T) { } } +// Ensure tries to sleep until the rate limit is reset when the client is rate limited, but only once. func TestDo_rateLimit_sleepUntilResponseResetLimitRetryOnce(t *testing.T) { client, mux, _, teardown := setup() defer teardown() @@ -1484,6 +1485,72 @@ func TestDo_rateLimit_sleepUntilClientResetLimit(t *testing.T) { } } +// Ensure sleep is aborted when the context is cancelled. +func TestDo_rateLimit_abortSleepContextCancelled(t *testing.T) { + client, mux, _, teardown := setup() + defer teardown() + + // We use a 1 minute reset time to ensure the sleep is not completed. + reset := time.Now().UTC().Add(time.Minute) + requestCount := 0 + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + requestCount++ + w.Header().Set(headerRateLimit, "60") + w.Header().Set(headerRateRemaining, "0") + w.Header().Set(headerRateReset, fmt.Sprint(reset.Unix())) + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(http.StatusForbidden) + fmt.Fprintln(w, `{ + "message": "API rate limit exceeded for xxx.xxx.xxx.xxx. (But here's the good news: Authenticated requests get a higher rate limit. Check out the documentation for more details.)", + "documentation_url": "https://docs.github.com/en/rest/overview/resources-in-the-rest-api#abuse-rate-limits" +}`) + }) + + req, _ := client.NewRequest("GET", ".", nil) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + _, err := client.Do(context.WithValue(ctx, SleepUntilPrimaryRateLimitResetWhenRateLimited, true), req, nil) + if !errors.Is(err, context.DeadlineExceeded) { + t.Error("Expected context deadline exceeded error.") + } + if got, want := requestCount, 1; got != want { + t.Errorf("Expected 1 requests, got %d", got) + } +} + +// Ensure sleep is aborted when the context is cancelled on initial request. +func TestDo_rateLimit_abortSleepContextCancelledClientLimit(t *testing.T) { + client, mux, _, teardown := setup() + defer teardown() + + reset := time.Now().UTC().Add(time.Minute) + client.rateLimits[CoreCategory] = Rate{Limit: 5000, Remaining: 0, Reset: Timestamp{reset}} + requestCount := 0 + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + requestCount++ + w.Header().Set(headerRateLimit, "5000") + w.Header().Set(headerRateRemaining, "5000") + w.Header().Set(headerRateReset, fmt.Sprint(reset.Add(time.Hour).Unix())) + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(http.StatusOK) + fmt.Fprintln(w, `{}`) + }) + req, _ := client.NewRequest("GET", ".", nil) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond) + defer cancel() + _, err := client.Do(context.WithValue(ctx, SleepUntilPrimaryRateLimitResetWhenRateLimited, true), req, nil) + rateLimitError, ok := err.(*RateLimitError) + if !ok { + t.Fatalf("Expected a *rateLimitError error; got %#v.", err) + } + if got, wantSuffix := rateLimitError.Message, "Context cancelled while waiting for rate limit to reset until"; !strings.HasPrefix(got, wantSuffix) { + t.Errorf("Expected request to be prevented because context cancellation, got: %v.", got) + } + if got, want := requestCount, 0; got != want { + t.Errorf("Expected 1 requests, got %d", got) + } +} + // Ensure *AbuseRateLimitError is returned when the response indicates that // the client has triggered an abuse detection mechanism. func TestDo_rateLimit_abuseRateLimitError(t *testing.T) { From aa0d4b92f7682e28cad122dc41c1b47d841a58ba Mon Sep 17 00:00:00 2001 From: Erez Rokah Date: Thu, 18 Apr 2024 20:57:14 +0300 Subject: [PATCH 7/8] Update github/github.go Co-authored-by: Glenn Lewis <6598971+gmlewis@users.noreply.github.com> --- github/github.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/github/github.go b/github/github.go index f9054aa489e..e64d89c476d 100644 --- a/github/github.go +++ b/github/github.go @@ -963,8 +963,7 @@ func (c *Client) checkRateLimitBeforeDo(req *http.Request, rateLimitCategory Rat } if req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { - err := sleepUntilResetWithBuffer(req.Context(), rate.Reset.Time) - if err == nil { + if err := sleepUntilResetWithBuffer(req.Context(), rate.Reset.Time); err == nil { return nil } return &RateLimitError{ From 960ce81d7ecf2f12953281bf0222a1e119c35f90 Mon Sep 17 00:00:00 2001 From: Erez Rokah Date: Thu, 18 Apr 2024 20:57:20 +0300 Subject: [PATCH 8/8] Update github/github.go Co-authored-by: Glenn Lewis <6598971+gmlewis@users.noreply.github.com> --- github/github.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/github/github.go b/github/github.go index e64d89c476d..ad4269595bf 100644 --- a/github/github.go +++ b/github/github.go @@ -892,8 +892,7 @@ func (c *Client) BareDo(ctx context.Context, req *http.Request) (*Response, erro rateLimitError, ok := err.(*RateLimitError) if ok && req.Context().Value(SleepUntilPrimaryRateLimitResetWhenRateLimited) != nil { - err := sleepUntilResetWithBuffer(req.Context(), rateLimitError.Rate.Reset.Time) - if err != nil { + if err := sleepUntilResetWithBuffer(req.Context(), rateLimitError.Rate.Reset.Time); err != nil { return response, err } // retry the request once when the rate limit has reset