From 4f7508a01c586b5fe126ec2e0c2b5e8bdcf4b0bb Mon Sep 17 00:00:00 2001 From: Gal Ofri Date: Sat, 13 Jan 2024 13:14:40 +0200 Subject: [PATCH 1/7] readme --- README.md | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 30681e0..a9d264c 100644 --- a/README.md +++ b/README.md @@ -18,23 +18,23 @@ import "golang.org/x/oauth2" import "github.com/gofri/go-github-ratelimit/github_ratelimit" func main() { - ctx := context.Background() - ts := oauth2.StaticTokenSource( - &oauth2.Token{AccessToken: "Your Personal Access Token"}, - ) - tc := oauth2.NewClient(ctx, ts) - rateLimiter, err := github_ratelimit.NewRateLimitWaiterClient(tc.Transport) - if err != nil { - panic(err) - } - client := github.NewClient(rateLimiter) - - // now use the client as you please + ctx := context.Background() + ts := oauth2.StaticTokenSource( + oauth2.Token{AccessToken: "Your Personal Access Token"}, + ) + tc := oauth2.NewClient(ctx, ts) + rateLimiter, err := github_ratelimit.NewRateLimitWaiterClient(tc.Transport) + if err != nil { + panic(err) + } + client := github.NewClient(rateLimiter) + + // now use the client as you please } ``` ## Options -The RoundTripper accepts a set of optional options: +The RoundTripper accepts a set of options: - User Context: provide a context.Context to pass to callbacks. - Single Sleep Limit: limit the sleep time for a single rate limit. - Total Sleep Limit: limit the accumulated sleep time for all rate limits. From b0a44fe68562fbc21af4c23fdae9effd9fed8819 Mon Sep 17 00:00:00 2001 From: Gal Ofri Date: Sat, 13 Jan 2024 13:23:20 +0200 Subject: [PATCH 2/7] readme --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index a9d264c..37987b5 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,8 @@ The RoundTripper accepts a set of optional callbacks: - On Single Limit Exceeded: callback for when a rate limit that exceeds the single sleep limit is detected. - On Total Limit Exceeded: callback for when a rate limit that exceeds the total sleep limit is detected. +note: to detect secondary rate limits and take a custom action without sleeping, use SingleSleepLimit=0 and OnSingleLimitExceeded(). + ## License This package is distributed under the MIT license found in the LICENSE file. Contribution and feedback is welcome. From 771c9406febbc452d601d96c4a9eabd77f273829 Mon Sep 17 00:00:00 2001 From: Gal Ofri Date: Sat, 13 Jan 2024 13:23:55 +0200 Subject: [PATCH 3/7] remove user context (breaking) --- github_ratelimit/callback.go | 2 -- github_ratelimit/github_ratelimit_test/ratelimit_test.go | 9 --------- github_ratelimit/options.go | 8 -------- github_ratelimit/ratelimit.go | 3 --- 4 files changed, 22 deletions(-) diff --git a/github_ratelimit/callback.go b/github_ratelimit/callback.go index e7b07c7..6767077 100644 --- a/github_ratelimit/callback.go +++ b/github_ratelimit/callback.go @@ -1,7 +1,6 @@ package github_ratelimit import ( - "context" "net/http" "time" ) @@ -9,7 +8,6 @@ import ( // CallbackContext is passed to all callbacks. // Fields might be nillable, depending on the specific callback and field. type CallbackContext struct { - UserContext *context.Context RoundTripper *SecondaryRateLimitWaiter SleepUntil *time.Time TotalSleepTime *time.Duration diff --git a/github_ratelimit/github_ratelimit_test/ratelimit_test.go b/github_ratelimit/github_ratelimit_test/ratelimit_test.go index 8eef4b6..8856a42 100644 --- a/github_ratelimit/github_ratelimit_test/ratelimit_test.go +++ b/github_ratelimit/github_ratelimit_test/ratelimit_test.go @@ -1,7 +1,6 @@ package github_ratelimit_test import ( - "context" "fmt" "io" "log" @@ -403,19 +402,12 @@ func TestCallbackContext(t *testing.T) { const sleep = 1 * time.Second i := setupSecondaryLimitInjecter(t, every, sleep) - ctxKey := struct{}{} - ctxVal := 10 - userContext := context.WithValue(context.Background(), ctxKey, ctxVal) var roundTripper *github_ratelimit.SecondaryRateLimitWaiter = nil var requestNum atomic.Int64 requestNum.Add(1) requestsCycle := 1 callback := func(ctx *github_ratelimit.CallbackContext) { - val := (*ctx.UserContext).Value(ctxKey).(int) - if val != ctxVal { - t.Fatalf("user ctx mismatch: %v != %v", val, ctxVal) - } if got, want := ctx.RoundTripper, roundTripper; got != want { t.Fatalf("roundtripper mismatch: %v != %v", got, want) } @@ -432,7 +424,6 @@ func TestCallbackContext(t *testing.T) { } r, err := github_ratelimit.NewRateLimitWaiter(i, - github_ratelimit.WithUserContext(userContext), github_ratelimit.WithLimitDetectedCallback(callback), ) if err != nil { diff --git a/github_ratelimit/options.go b/github_ratelimit/options.go index 83b9579..6f70664 100644 --- a/github_ratelimit/options.go +++ b/github_ratelimit/options.go @@ -1,7 +1,6 @@ package github_ratelimit import ( - "context" "time" ) @@ -32,13 +31,6 @@ func WithTotalSleepLimit(limit time.Duration, callback OnTotalLimitExceeded) Opt } } -// WithUserContext sets the user context to be passed to callbacks. -func WithUserContext(ctx context.Context) Option { - return func(t *SecondaryRateLimitWaiter) { - t.userContext = &ctx - } -} - func applyOptions(w *SecondaryRateLimitWaiter, opts ...Option) { for _, o := range opts { if o == nil { diff --git a/github_ratelimit/ratelimit.go b/github_ratelimit/ratelimit.go index 8e3f97e..9740a06 100644 --- a/github_ratelimit/ratelimit.go +++ b/github_ratelimit/ratelimit.go @@ -1,7 +1,6 @@ package github_ratelimit import ( - "context" "net/http" "strconv" "sync" @@ -19,7 +18,6 @@ type SecondaryRateLimitWaiter struct { totalSleepLimit *time.Duration // callbacks - userContext *context.Context onLimitDetected OnLimitDetected onSingleLimitExceeded OnSingleLimitExceeded onTotalLimitExceeded OnTotalLimitExceeded @@ -149,7 +147,6 @@ func (t *SecondaryRateLimitWaiter) triggerCallback(callback func(*CallbackContex } callbackContext.RoundTripper = t - callbackContext.UserContext = t.userContext callbackContext.SleepUntil = &newSleepUntil callbackContext.TotalSleepTime = &t.totalSleepTime From f127de091eadf37b1601491f0344c760243dd0fb Mon Sep 17 00:00:00 2001 From: Gal Ofri Date: Sat, 13 Jan 2024 19:19:38 +0200 Subject: [PATCH 4/7] config overrides --- github_ratelimit/config.go | 64 +++++++++++++++++++++++++++++++++++ github_ratelimit/options.go | 27 +++++---------- github_ratelimit/ratelimit.go | 48 ++++++++++++++------------ 3 files changed, 100 insertions(+), 39 deletions(-) create mode 100644 github_ratelimit/config.go diff --git a/github_ratelimit/config.go b/github_ratelimit/config.go new file mode 100644 index 0000000..d013077 --- /dev/null +++ b/github_ratelimit/config.go @@ -0,0 +1,64 @@ +package github_ratelimit + +import ( + "context" + "time" +) + +// SecondaryRateLimitConfig is the config for the secondary rate limit waiter. +// Use the options to set the config. +type SecondaryRateLimitConfig struct { + // limits + singleSleepLimit *time.Duration + totalSleepLimit *time.Duration + + // callbacks + onLimitDetected OnLimitDetected + onSingleLimitExceeded OnSingleLimitExceeded + onTotalLimitExceeded OnTotalLimitExceeded +} + +// newConfig creates a new config with the given options. +func newConfig(opts ...Option) *SecondaryRateLimitConfig { + var config SecondaryRateLimitConfig + config.ApplyOptions(opts...) + return &config +} + +// ApplyOptions applies the options to the config. +func (c *SecondaryRateLimitConfig) ApplyOptions(opts ...Option) { + for _, o := range opts { + if o == nil { + continue + } + o(c) + } +} + +// IsAboveSingleSleepLimit returns true if the single sleep time is above the limit. +func (c *SecondaryRateLimitConfig) IsAboveSingleSleepLimit(sleepTime time.Duration) bool { + return c.singleSleepLimit != nil && sleepTime > *c.singleSleepLimit +} + +// IsAboveTotalSleepLimit returns true if the total sleep time is above the limit. +func (c *SecondaryRateLimitConfig) IsAboveTotalSleepLimit(sleepTime time.Duration, totalSleepTime time.Duration) bool { + return c.totalSleepLimit != nil && totalSleepTime+sleepTime > *c.totalSleepLimit +} + +type secondaryRateLimitConfigOverridesKey struct{} + +// WithOverrideConfig adds config overrides to the context. +// The overrides are applied on top of the existing config. +// Allows for request-specific overrides. +func WithOverrideConfig(ctx context.Context, opts ...Option) context.Context { + return context.WithValue(ctx, secondaryRateLimitConfigOverridesKey{}, opts) +} + +// GetConfigOverrides returns the config overrides from the context, if any. +func GetConfigOverrides(ctx context.Context) []Option { + cfg := ctx.Value(secondaryRateLimitConfigOverridesKey{}) + if cfg == nil { + return nil + } + return cfg.([]Option) +} diff --git a/github_ratelimit/options.go b/github_ratelimit/options.go index 6f70664..29ed190 100644 --- a/github_ratelimit/options.go +++ b/github_ratelimit/options.go @@ -4,38 +4,29 @@ import ( "time" ) -type Option func(*SecondaryRateLimitWaiter) +type Option func(*SecondaryRateLimitConfig) // WithLimitDetectedCallback adds a callback to be called when a new active rate limit is detected. func WithLimitDetectedCallback(callback OnLimitDetected) Option { - return func(t *SecondaryRateLimitWaiter) { - t.onLimitDetected = callback + return func(c *SecondaryRateLimitConfig) { + c.onLimitDetected = callback } } // WithSingleSleepLimit adds a limit to the duration allowed to wait for a single sleep (rate limit). // The callback parameter is nillable. func WithSingleSleepLimit(limit time.Duration, callback OnSingleLimitExceeded) Option { - return func(t *SecondaryRateLimitWaiter) { - t.singleSleepLimit = &limit - t.onSingleLimitExceeded = callback + return func(c *SecondaryRateLimitConfig) { + c.singleSleepLimit = &limit + c.onSingleLimitExceeded = callback } } // WithTotalSleepLimit adds a limit to the accumulated duration allowed to wait for all sleeps (one or more rate limits). // The callback parameter is nillable. func WithTotalSleepLimit(limit time.Duration, callback OnTotalLimitExceeded) Option { - return func(t *SecondaryRateLimitWaiter) { - t.totalSleepLimit = &limit - t.onTotalLimitExceeded = callback - } -} - -func applyOptions(w *SecondaryRateLimitWaiter, opts ...Option) { - for _, o := range opts { - if o == nil { - continue - } - o(w) + return func(c *SecondaryRateLimitConfig) { + c.totalSleepLimit = &limit + c.onTotalLimitExceeded = callback } } diff --git a/github_ratelimit/ratelimit.go b/github_ratelimit/ratelimit.go index 9740a06..a1ec562 100644 --- a/github_ratelimit/ratelimit.go +++ b/github_ratelimit/ratelimit.go @@ -8,19 +8,11 @@ import ( ) type SecondaryRateLimitWaiter struct { - Base http.RoundTripper - sleepUntil *time.Time - lock sync.RWMutex - - // limits - totalSleepTime time.Duration - singleSleepLimit *time.Duration - totalSleepLimit *time.Duration - - // callbacks - onLimitDetected OnLimitDetected - onSingleLimitExceeded OnSingleLimitExceeded - onTotalLimitExceeded OnTotalLimitExceeded + Base http.RoundTripper + sleepUntil *time.Time + lock sync.RWMutex + totalSleepTime time.Duration + config *SecondaryRateLimitConfig } func NewRateLimitWaiter(base http.RoundTripper, opts ...Option) (*SecondaryRateLimitWaiter, error) { @@ -29,9 +21,9 @@ func NewRateLimitWaiter(base http.RoundTripper, opts ...Option) (*SecondaryRateL } waiter := SecondaryRateLimitWaiter{ - Base: base, + Base: base, + config: newConfig(opts...), } - applyOptions(&waiter, opts...) return &waiter, nil } @@ -81,6 +73,17 @@ func (t *SecondaryRateLimitWaiter) RoundTrip(request *http.Request) (*http.Respo return t.RoundTrip(request) } +func (t *SecondaryRateLimitWaiter) getRequestConfig(request *http.Request) *SecondaryRateLimitConfig { + overrides := GetConfigOverrides(request.Context()) + if overrides == nil { + // no config override - use the default config (zero-copy) + return t.config + } + reqConfig := *t.config + reqConfig.ApplyOptions(overrides...) + return &reqConfig +} + // waitForRateLimit waits for the cooldown time to finish if a secondary rate limit is active. func (t *SecondaryRateLimitWaiter) waitForRateLimit() { t.lock.RLock() @@ -94,7 +97,7 @@ func (t *SecondaryRateLimitWaiter) waitForRateLimit() { // the rate limit is not updated if there is already an active rate limit. // it never waits because the retry handles sleeping anyway. // returns whether or not to retry the request. -func (t *SecondaryRateLimitWaiter) updateRateLimit(secondaryLimit time.Time, callbackContext *CallbackContext) bool { +func (t *SecondaryRateLimitWaiter) updateRateLimit(secondaryLimit time.Time, callbackContext *CallbackContext) (needRetry bool) { // quick check without the lock: maybe the secondary limit just passed if time.Now().After(secondaryLimit) { return true @@ -114,22 +117,24 @@ func (t *SecondaryRateLimitWaiter) updateRateLimit(secondaryLimit time.Time, cal return true } + config := t.getRequestConfig(callbackContext.Request) + // do not sleep in case it is above the single sleep limit - if t.singleSleepLimit != nil && sleepTime > *t.singleSleepLimit { - t.triggerCallback(t.onSingleLimitExceeded, callbackContext, secondaryLimit) + if config.IsAboveSingleSleepLimit(sleepTime) { + t.triggerCallback(config.onSingleLimitExceeded, callbackContext, secondaryLimit) return false } // do not sleep in case it is above the total sleep limit - if t.totalSleepLimit != nil && t.totalSleepTime+sleepTime > *t.totalSleepLimit { - t.triggerCallback(t.onTotalLimitExceeded, callbackContext, secondaryLimit) + if config.IsAboveTotalSleepLimit(sleepTime, t.totalSleepTime) { + t.triggerCallback(config.onTotalLimitExceeded, callbackContext, secondaryLimit) return false } // a legitimate new limit t.sleepUntil = &secondaryLimit t.totalSleepTime += smoothSleepTime(sleepTime) - t.triggerCallback(t.onLimitDetected, callbackContext, secondaryLimit) + t.triggerCallback(config.onLimitDetected, callbackContext, secondaryLimit) return true } @@ -204,6 +209,7 @@ func parseXRateLimitReset(resp *http.Response) *time.Time { return &sleepUntil } +// httpHeaderIntValue parses an integer value from the given HTTP header. func httpHeaderIntValue(header http.Header, key string) (int64, bool) { val := header.Get(key) if val == "" { From e64f77edaa9a32fc3491aaac8ffa2f1d84c11a84 Mon Sep 17 00:00:00 2001 From: Gal Ofri Date: Sat, 13 Jan 2024 19:23:27 +0200 Subject: [PATCH 5/7] names --- github_ratelimit/ratelimit.go | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/github_ratelimit/ratelimit.go b/github_ratelimit/ratelimit.go index a1ec562..8c8dd0a 100644 --- a/github_ratelimit/ratelimit.go +++ b/github_ratelimit/ratelimit.go @@ -46,7 +46,7 @@ func NewRateLimitWaiterClient(base http.RoundTripper, opts ...Option) (*http.Cli // Nonetheless, there is no way to prevent subtle race conditions without completely serializing the requests, // so we prefer to let some slip in case of a race condition, i.e., // after a retry-after response is received and before it is processed, -// a few other (parallel) requests may be issued. +// a few other (concurrent) requests may be issued. func (t *SecondaryRateLimitWaiter) RoundTrip(request *http.Request) (*http.Response, error) { t.waitForRateLimit() @@ -87,10 +87,10 @@ func (t *SecondaryRateLimitWaiter) getRequestConfig(request *http.Request) *Seco // waitForRateLimit waits for the cooldown time to finish if a secondary rate limit is active. func (t *SecondaryRateLimitWaiter) waitForRateLimit() { t.lock.RLock() - sleepTime := t.currentSleepTimeUnlocked() + sleepDuration := t.currentSleepDurationUnlocked() t.lock.RUnlock() - time.Sleep(sleepTime) + time.Sleep(sleepDuration) } // updateRateLimit updates the active rate limit and triggers user callbacks if needed. @@ -107,39 +107,39 @@ func (t *SecondaryRateLimitWaiter) updateRateLimit(secondaryLimit time.Time, cal defer t.lock.Unlock() // check before update if there is already an active rate limit - if t.currentSleepTimeUnlocked() > 0 { + if t.currentSleepDurationUnlocked() > 0 { return true } // check if the secondary rate limit happened to have passed while we waited for the lock - sleepTime := time.Until(secondaryLimit) - if sleepTime <= 0 { + sleepDuration := time.Until(secondaryLimit) + if sleepDuration <= 0 { return true } config := t.getRequestConfig(callbackContext.Request) // do not sleep in case it is above the single sleep limit - if config.IsAboveSingleSleepLimit(sleepTime) { + if config.IsAboveSingleSleepLimit(sleepDuration) { t.triggerCallback(config.onSingleLimitExceeded, callbackContext, secondaryLimit) return false } // do not sleep in case it is above the total sleep limit - if config.IsAboveTotalSleepLimit(sleepTime, t.totalSleepTime) { + if config.IsAboveTotalSleepLimit(sleepDuration, t.totalSleepTime) { t.triggerCallback(config.onTotalLimitExceeded, callbackContext, secondaryLimit) return false } // a legitimate new limit t.sleepUntil = &secondaryLimit - t.totalSleepTime += smoothSleepTime(sleepTime) + t.totalSleepTime += smoothSleepTime(sleepDuration) t.triggerCallback(config.onLimitDetected, callbackContext, secondaryLimit) return true } -func (t *SecondaryRateLimitWaiter) currentSleepTimeUnlocked() time.Duration { +func (t *SecondaryRateLimitWaiter) currentSleepDurationUnlocked() time.Duration { if t.sleepUntil == nil { return 0 } From 171438878a19eda2067953823a7835e40923ee96 Mon Sep 17 00:00:00 2001 From: Gal Ofri Date: Sat, 13 Jan 2024 19:49:53 +0200 Subject: [PATCH 6/7] testing --- .../github_ratelimit_test/ratelimit_test.go | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/github_ratelimit/github_ratelimit_test/ratelimit_test.go b/github_ratelimit/github_ratelimit_test/ratelimit_test.go index 8856a42..5d9fae6 100644 --- a/github_ratelimit/github_ratelimit_test/ratelimit_test.go +++ b/github_ratelimit/github_ratelimit_test/ratelimit_test.go @@ -1,6 +1,7 @@ package github_ratelimit_test import ( + "context" "fmt" "io" "log" @@ -461,3 +462,79 @@ func TestCallbackContext(t *testing.T) { } close(errChan) } + +func TestRequestConfigOverride(t *testing.T) { + t.Parallel() + const every = 1 * time.Second + const sleep = 1 * time.Second + + slept := false + callback := func(*github_ratelimit.CallbackContext) { + slept = true + } + exceeded := false + onLimitExceeded := func(*github_ratelimit.CallbackContext) { + exceeded = true + } + + // test sleep is short enough + i := setupSecondaryLimitInjecter(t, every, sleep) + c, err := github_ratelimit.NewRateLimitWaiterClient(i, + github_ratelimit.WithLimitDetectedCallback(callback), + github_ratelimit.WithSingleSleepLimit(5*time.Second, onLimitExceeded)) + if err != nil { + t.Fatal(err) + } + + // initialize injecter timing + _, _ = c.Get("/") + + // prepare an override - force sleep time to be 0, + // so that it will not sleep at all regardless of the original config. + limit := github_ratelimit.WithSingleSleepLimit(0, onLimitExceeded) + ctx := github_ratelimit.WithOverrideConfig(context.Background(), limit) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, "/", nil) + if err != nil { + t.Fatal(err) + } + + // wait for next sleep to kick in, but issue the request with the override + waitForNextSleep(i) + + // attempt during rate limit + slept = false + exceeded = false + _, err = c.Do(req) + if err != nil { + t.Fatal(err) + } + // expect no sleep because the override is set to 0 + if slept || !exceeded { + t.Fatal(slept, exceeded) + } + + // prepare an override with a different nature (extra safety check) + exceeded = false + usedAltCallback := false + onSleepAlt := func(*github_ratelimit.CallbackContext) { + usedAltCallback = true + } + + limit = github_ratelimit.WithSingleSleepLimit(10*time.Second, onLimitExceeded) + sleepCB := github_ratelimit.WithLimitDetectedCallback(onSleepAlt) + ctx = github_ratelimit.WithOverrideConfig(context.Background(), limit, sleepCB) + req, err = http.NewRequestWithContext(ctx, http.MethodGet, "/", nil) + if err != nil { + t.Fatal(err) + } + + // attempt during rate limit + _, err = c.Do(req) + if err != nil { + t.Fatal(err) + } + if !usedAltCallback || exceeded { + t.Fatal(slept, exceeded) + } + +} From 6f205a0c49d270acddd502381d74e4da304e372d Mon Sep 17 00:00:00 2001 From: Gal Ofri Date: Sat, 13 Jan 2024 19:53:10 +0200 Subject: [PATCH 7/7] readme --- README.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/README.md b/README.md index 37987b5..b28ce47 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,11 @@ The RoundTripper accepts a set of optional callbacks: note: to detect secondary rate limits and take a custom action without sleeping, use SingleSleepLimit=0 and OnSingleLimitExceeded(). +## Per-Request Options +Use WithOverrideConfig() to override the configuration for a specific request using a context. +Per-request overrides may be useful for special-cases of user requests, +as well as fine-grained policy control (e.g., for a sophisticated pagination mechanism). + ## License This package is distributed under the MIT license found in the LICENSE file. Contribution and feedback is welcome.