Merge pull request #22 from gofri/gofri/per_req

feat: per-request overrides

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.
@@ -44,6 +44,13 @@ 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().
+
+## 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.

github_ratelimit/callback.go

@@ -1,15 +1,13 @@
 package github_ratelimit
 
 import (
-	"context"
 	"net/http"
 	"time"
 )
 
 // 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

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)
+}

github_ratelimit/github_ratelimit_test/ratelimit_test.go

@@ -403,19 +403,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 +425,6 @@ func TestCallbackContext(t *testing.T) {
 	}
 
 	r, err := github_ratelimit.NewRateLimitWaiter(i,
-		github_ratelimit.WithUserContext(userContext),
 		github_ratelimit.WithLimitDetectedCallback(callback),
 	)
 	if err != nil {
@@ -470,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)
+	}
+
+}

github_ratelimit/options.go

@@ -1,49 +1,32 @@
 package github_ratelimit
 
 import (
-	"context"
 	"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
-	}
-}
-
-// 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 {
-			continue
-		}
-		o(w)
+	return func(c *SecondaryRateLimitConfig) {
+		c.totalSleepLimit = &limit
+		c.onTotalLimitExceeded = callback
 	}
 }