From dd40e87231e453e16acf20341b16ed8f56fda428 Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 13:28:34 -0400 Subject: [PATCH 01/15] Remove FIFO / waiters from Wait method --- limiter.go | 1 - limiter_wait.go | 56 +-------- limiter_wait_test.go | 267 ------------------------------------------- syncmap.go | 5 - 4 files changed, 4 insertions(+), 325 deletions(-) diff --git a/limiter.go b/limiter.go index 1c6ab9c..27eb522 100644 --- a/limiter.go +++ b/limiter.go @@ -6,7 +6,6 @@ type Limiter[TInput any, TKey comparable] struct { limits []Limit limitFuncs []LimitFunc[TInput] buckets bucketMap[TKey] - waiters syncMap[TKey, *waiter] } // KeyFunc is a function that takes an input and returns a bucket key. diff --git a/limiter_wait.go b/limiter_wait.go index cdc2d58..03f505e 100644 --- a/limiter_wait.go +++ b/limiter_wait.go @@ -2,8 +2,6 @@ package rate import ( "context" - "sync" - "sync/atomic" "time" "github.com/clipperhouse/ntime" @@ -30,10 +28,8 @@ import ( // // ctx := context.WithTimeout(ctx, limit.DurationPerToken()) // -// Wait offers best-effort FIFO ordering of requests. Under sustained -// contention (when multiple requests wait longer than ~1ms), the Go runtime's -// mutex starvation mode ensures strict FIFO ordering. Under light load, -// ordering may be less strict but performance is optimized. +// Wait makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) bool { return r.waitN(ctx, input, ntime.Now(), 1) } @@ -59,10 +55,8 @@ func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) bool { // // ctx := context.WithTimeout(ctx, limit.DurationPerToken()) // -// WaitN offers best-effort FIFO ordering of requests. Under sustained -// contention (when multiple requests wait longer than ~1ms), the Go runtime's -// mutex starvation mode ensures strict FIFO ordering. Under light load, -// ordering may be less strict but performance is optimized. +// WaitN makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64) bool { return r.waitN(ctx, input, ntime.Now(), n) } @@ -105,23 +99,8 @@ func (r *Limiter[TInput, TKey]) waitNWithCancellation( // be a parameter. currentTime := startTime - userKey := r.keyFunc(input) - waiter := r.getWaiter(userKey) - - // Ensure cleanup happens when this waiter exits - defer func() { - // If no more waiters for this key, remove the entry to prevent memory leak, - // reference-counted - if waiter.decrement() == 0 { - r.waiters.delete(userKey) - } - }() - for { - // The goroutine at the front of the queue gets to try for a token first. - waiter.mu.Lock() allow, details := r.allowNWithDetails(input, currentTime, n) - waiter.mu.Unlock() if allow { return true } @@ -143,30 +122,3 @@ func (r *Limiter[TInput, TKey]) waitNWithCancellation( } } } - -// waiter represents a reservation queue for a specific key with reference counting -type waiter struct { - mu sync.Mutex - count int64 // number of active waiters for this key -} - -// increment atomically increments the waiter count and returns the new count -func (w *waiter) increment() int64 { - return atomic.AddInt64(&w.count, 1) -} - -// decrement atomically decrements the waiter count and returns the new count -func (w *waiter) decrement() int64 { - return atomic.AddInt64(&w.count, -1) -} - -func getWaiter() *waiter { - return &waiter{} -} - -// getWaiter atomically gets or creates a waiter entry and increments its reference count -func (r *Limiter[TInput, TKey]) getWaiter(key TKey) *waiter { - waiter := r.waiters.loadOrStore(key, getWaiter) - waiter.increment() - return waiter -} diff --git a/limiter_wait_test.go b/limiter_wait_test.go index e3d260c..3cd26b0 100644 --- a/limiter_wait_test.go +++ b/limiter_wait_test.go @@ -1,10 +1,8 @@ package rate import ( - "context" "fmt" "sync" - "sync/atomic" "testing" "time" @@ -768,268 +766,3 @@ func TestLimiter_Wait_FIFO_Ordering_MultipleBuckets_Flaky(t *testing.T) { require.Equal(t, starts, successes, "success order should match start order for bucket %d, proving FIFO", bucketID) } } - -func TestLimiter_WaitersCleanup_Basic(t *testing.T) { - t.Parallel() - - keyFunc := func(input string) string { - return input - } - limit := NewLimit(1, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - - // Initial waiters count should be 0 - require.Equal(t, 0, limiter.waiters.count(), "initial waiters count should be 0") - - executionTime := ntime.Now() - - // Exhaust tokens for multiple keys - require.True(t, limiter.allow("key1", executionTime), "should allow initial token for key1") - require.True(t, limiter.allow("key2", executionTime), "should allow initial token for key2") - - // Create deadline that will timeout immediately - deadline := func() (time.Time, bool) { - return executionTime.ToTime(), true // immediate timeout - } - done := func() <-chan struct{} { - ch := make(chan struct{}) - close(ch) // immediately cancelled - return ch - } - - // Try to wait - this should create a waiter entry that gets cleaned up - allow := limiter.waitWithCancellation("key1", executionTime, deadline, done) - require.False(t, allow, "should timeout immediately") - - // Check if waiter was cleaned up - require.Equal(t, 0, limiter.waiters.count(), "waiters should be cleaned up after timeout") - - // Try again with a different key - allow = limiter.waitWithCancellation("key2", executionTime, deadline, done) - require.False(t, allow, "should timeout immediately") - - // Check waiters count again - require.Equal(t, 0, limiter.waiters.count(), "waiters should be cleaned up after timeout") - - // Try the same key again - allow = limiter.waitWithCancellation("key1", executionTime, deadline, done) - require.False(t, allow, "should timeout immediately") - - // Check waiters count - require.Equal(t, 0, limiter.waiters.count(), "waiters should be cleaned up after timeout") -} - -func TestLimiter_WaitersCleanup_MemoryLeak_Prevention(t *testing.T) { - t.Parallel() - - keyFunc := func(input int) string { - return fmt.Sprintf("key-%d", input) - } - limit := NewLimit(1, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - - executionTime := ntime.Now() - - // Create deadline that will timeout immediately - deadline := func() (time.Time, bool) { - return executionTime.ToTime(), true // immediate timeout - } - done := func() <-chan struct{} { - ch := make(chan struct{}) - close(ch) // immediately cancelled - return ch - } - - const numKeys = 1000 - - // Simulate many different keys trying to wait - for i := range numKeys { - // First exhaust the token for this key - limiter.allow(i, executionTime) - - // Then try to wait (which will timeout immediately) - result := limiter.waitWithCancellation(i, executionTime, deadline, done) - require.False(t, result, "should timeout immediately for key %d", i) - } - - require.Equal(t, 0, limiter.waiters.count(), "waiters should be cleaned up after use") -} - -func TestLimiter_WaitersCleanup_Concurrent(t *testing.T) { - t.Parallel() - - keyFunc := func(input int) string { - return fmt.Sprintf("key-%d", input) - } - limit := NewLimit(1, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - - executionTime := ntime.Now() - - // Create deadline that will timeout immediately - deadline := func() (time.Time, bool) { - return executionTime.ToTime(), true // immediate timeout - } - done := func() <-chan struct{} { - ch := make(chan struct{}) - close(ch) // immediately cancelled - return ch - } - - const concurrency = 100 - const keysPerGoroutine = 10 - - var wg sync.WaitGroup - wg.Add(concurrency) - for i := range concurrency { - go func(i int) { - defer wg.Done() - - for k := range keysPerGoroutine { - keyID := i*keysPerGoroutine + k - - // First exhaust the token for this key - limiter.allow(keyID, executionTime) - - // Then try to wait (which will timeout immediately) - limiter.waitWithCancellation(keyID, executionTime, deadline, done) - } - }(i) - } - wg.Wait() - - require.Equal(t, 0, limiter.waiters.count(), "waiters should be cleaned up after use") -} - -func TestLimiter_WaitersCleanup_WithSuccessfulWaits(t *testing.T) { - t.Parallel() - - keyFunc := func(input string) string { - return input - } - limit := NewLimit(1, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - - executionTime := ntime.Now() - - // Exhaust the token - require.True(t, limiter.allow("key", executionTime), "should allow initial token") - - const concurrency = 20 - results := make([]bool, concurrency) - - // Create deadline that gives enough time for tokens to be refilled - deadline := func() (time.Time, bool) { - return executionTime.Add(time.Duration(concurrency) * limit.durationPerToken).ToTime(), true - } - - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - var wg sync.WaitGroup - wg.Add(concurrency) - for i := range concurrency { - go func(i int) { - defer wg.Done() - // stagger the start times - time.Sleep(time.Duration(i) * time.Millisecond) - results[i] = limiter.waitWithCancellation("key", executionTime, deadline, done) - }(i) - } - wg.Wait() - - // All waiters should have eventually succeeded - successes := 0 - for _, result := range results { - if result { - successes++ - } - } - - require.Equal(t, concurrency, successes, "all waiters should eventually succeed") - require.Equal(t, 0, limiter.waiters.count(), "all waiters should be cleaned up after completion") -} - -func TestLimiter_Wait_FIFOOrdering_HighContention(t *testing.T) { - t.Parallel() - - keyFunc := func(input int) int { - return input - } - - limiter := NewLimiter(keyFunc, NewLimit(1, 50*time.Millisecond)) - - // Exhaust the bucket - require.True(t, limiter.Allow(1)) - - // Test that multiple waiters can successfully acquire tokens - const waiters = 5 - results := make(chan bool, waiters) - - for range waiters { - go func() { - ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second) - defer cancel() - results <- limiter.Wait(ctx, 1) - }() - } - - // Collect results - successes := 0 - for range waiters { - select { - case success := <-results: - if success { - successes++ - } - case <-time.After(2 * time.Second): - t.Fatal("Test timed out") - } - } - - // Should get at least some successes (the fix should prevent token starvation) - require.Greater(t, successes, 0, "Should have some successful token acquisitions") - require.Equal(t, 0, limiter.waiters.count(), "All waiters should be cleaned up") -} - -func TestLimiter_Wait_RaceCondition_Prevention(t *testing.T) { - t.Parallel() - - keyFunc := func(input int) int { - return input - } - // Use a very restrictive limit to force contention - limiter := NewLimiter(keyFunc, NewLimit(1, 100*time.Millisecond)) - - // Exhaust the bucket - require.True(t, limiter.Allow(1)) - - const concurrency = 50 - successes := int64(0) - - var wg sync.WaitGroup - wg.Add(concurrency) - // Start many goroutines that will compete for tokens - for range concurrency { - go func() { - defer wg.Done() - ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - defer cancel() - - if limiter.Wait(ctx, 1) { - atomic.AddInt64(&successes, 1) - } - }() - } - wg.Wait() - - // Should have some successes but not more than the tokens available - // in the timeout period (roughly 20 tokens in 2 seconds at 100ms per token) - actual := atomic.LoadInt64(&successes) - require.Greater(t, actual, int64(0), "should have some successful acquisitions") - require.LessOrEqual(t, actual, int64(25), "should not exceed reasonable token availability") - - // Verify no memory leaks - require.Equal(t, 0, limiter.waiters.count(), "all waiters should be cleaned up") -} diff --git a/syncmap.go b/syncmap.go index bf1ae3e..f7c1611 100644 --- a/syncmap.go +++ b/syncmap.go @@ -33,11 +33,6 @@ func (sm *syncMap[K, V]) count() int { return count } -// delete removes a key from the map -func (sm *syncMap[K, V]) delete(key K) { - sm.m.Delete(key) -} - // bucketMap is a specialized sync.Map for storing buckets to avoid allocations type bucketMap[TKey comparable] struct { m sync.Map From 19e8cece3671554985f61ada659be9eb7a8c35e0 Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 13:38:22 -0400 Subject: [PATCH 02/15] syncMap is now unused --- syncmap.go | 27 --------------------------- syncmap_test.go | 46 ---------------------------------------------- 2 files changed, 73 deletions(-) diff --git a/syncmap.go b/syncmap.go index f7c1611..0566042 100644 --- a/syncmap.go +++ b/syncmap.go @@ -6,33 +6,6 @@ import ( "github.com/clipperhouse/ntime" ) -// syncMap is a typed wrapper around sync.Map for our specific use case -type syncMap[K comparable, V any] struct { - m sync.Map -} - -// loadOrStore returns the existing value for the key if present. -// Otherwise, it calls the factory function to create a new value, stores it, and returns it. -// This avoids creating the value unless it's actually needed. -func (sm *syncMap[K, V]) loadOrStore(key K, getter func() V) V { - if loaded, ok := sm.m.Load(key); ok { - return loaded.(V) - } - // Only create the value if we didn't find an existing one - value := getter() - actual, _ := sm.m.LoadOrStore(key, value) - return actual.(V) -} - -func (sm *syncMap[K, V]) count() int { - count := 0 - sm.m.Range(func(_, _ any) bool { - count++ - return true - }) - return count -} - // bucketMap is a specialized sync.Map for storing buckets to avoid allocations type bucketMap[TKey comparable] struct { m sync.Map diff --git a/syncmap_test.go b/syncmap_test.go index 51b37ff..17dc6b8 100644 --- a/syncmap_test.go +++ b/syncmap_test.go @@ -47,52 +47,6 @@ func TestSyncMap_ConcurrentAccess(t *testing.T) { // If we get here without panic or race conditions, the test passes } -func TestSyncMap_LoadOrStoreFunc(t *testing.T) { - t.Parallel() - - sm := &syncMap[string, int]{} - - factoryCalls := 0 - factory := func() int { - factoryCalls++ - return 42 - } - - // First call should call factory and store the value - result1 := sm.loadOrStore("key1", factory) - require.Equal(t, 42, result1, "should return factory value") - require.Equal(t, 1, factoryCalls, "factory should be called once") - - // Second call with same key should NOT call factory - result2 := sm.loadOrStore("key1", factory) - require.Equal(t, 42, result2, "should return existing value") - require.Equal(t, 1, factoryCalls, "factory should not be called again for existing key") - - // Third call with different key should call factory again - result3 := sm.loadOrStore("key2", factory) - require.Equal(t, 42, result3, "should return factory value for new key") - require.Equal(t, 2, factoryCalls, "factory should be called for new key") -} - -func TestSyncMap_Count(t *testing.T) { - t.Parallel() - - sm := &syncMap[string, int]{} - - for range 2 { - // should only be stored once despite multiple calls - for key := range 101 { - sm.loadOrStore(fmt.Sprint(key), func() int { - return key - }) - } - } - - expected := 101 - actual := sm.count() - require.Equal(t, expected, actual, "expected Count() to be accurate") -} - func TestBucketMap_LoadOrStore(t *testing.T) { t.Parallel() From 3a0690c04b2bbefbaed9b90394b1622efa10758b Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 13:41:00 -0400 Subject: [PATCH 03/15] Remove bucketMap.loadOrGet, unused --- syncmap.go | 15 --------------- syncmap_test.go | 28 ---------------------------- 2 files changed, 43 deletions(-) diff --git a/syncmap.go b/syncmap.go index 0566042..a32aa15 100644 --- a/syncmap.go +++ b/syncmap.go @@ -37,21 +37,6 @@ func (bm *bucketMap[TKey]) loadOrStore(userKey TKey, executionTime ntime.Time, l return actual.(*bucket) } -// loadOrGet returns the existing value for the key if present. -// Otherwise, it returns a new (temporary) value. -func (bm *bucketMap[TKey]) loadOrGet(userKey TKey, executionTime ntime.Time, limit Limit) *bucket { - spec := bucketSpec[TKey]{ - limit: limit, - userKey: userKey, - } - loaded, ok := bm.m.Load(spec) - if ok { - return loaded.(*bucket) - } - b := newBucket(executionTime, limit) - return &b -} - func (bm *bucketMap[TKey]) load(userKey TKey, limit Limit) (*bucket, bool) { spec := bucketSpec[TKey]{ limit: limit, diff --git a/syncmap_test.go b/syncmap_test.go index 17dc6b8..5441a44 100644 --- a/syncmap_test.go +++ b/syncmap_test.go @@ -72,31 +72,6 @@ func TestBucketMap_LoadOrStore(t *testing.T) { require.False(t, bucket1 == bucket4, "should return different bucket for different key") } -func TestBucketMap_LoadOrGet(t *testing.T) { - t.Parallel() - - var bm bucketMap[string] - limit := NewLimit(100, time.Second) - executionTime := ntime.Now() - - // First call should return a temporary bucket (not stored) - bucket1 := bm.loadOrGet("key1", executionTime, limit) - require.NotNil(t, bucket1, "should return a bucket") - - // Second call should return another temporary bucket (different instance) - // Use slightly different execution time to ensure different bucket instances - bucket2 := bm.loadOrGet("key1", executionTime.Add(time.Nanosecond), limit) - require.NotNil(t, bucket2, "should return a bucket") - require.False(t, bucket1 == bucket2, "should return different temporary buckets when none stored") - - // Store a bucket first - storedBucket := bm.loadOrStore("key1", executionTime, limit) - - // Now loadOrGet should return the stored bucket - bucket3 := bm.loadOrGet("key1", executionTime, limit) - require.Equal(t, storedBucket, bucket3, "should return stored bucket when available") -} - func TestBucketMap_Load(t *testing.T) { t.Parallel() @@ -181,9 +156,6 @@ func TestBucketMap_ConcurrentAccess(t *testing.T) { // Test loadOrStore method bm.loadOrStore(key, executionTime, limit) - // Test loadOrGet method - bm.loadOrGet(key, executionTime, limit) - // Test load method bm.load(key, limit) From 86dfa49fbd36710f18aa819eeb475f855667eb43 Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 14:58:28 -0400 Subject: [PATCH 04/15] comments --- syncmap.go | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/syncmap.go b/syncmap.go index a32aa15..9f9c240 100644 --- a/syncmap.go +++ b/syncmap.go @@ -6,7 +6,7 @@ import ( "github.com/clipperhouse/ntime" ) -// bucketMap is a specialized sync.Map for storing buckets to avoid allocations +// bucketMap is a specialized sync.Map for storing buckets type bucketMap[TKey comparable] struct { m sync.Map } @@ -15,13 +15,12 @@ type bucketMap[TKey comparable] struct { // It is a composite key to ensure that each bucket is unique for a given limit and user key. type bucketSpec[TKey comparable] struct { limit Limit - // userKey is the result of calling the user-defined Keyer + // userKey is the result of calling the user-defined KeyFunc userKey TKey } // loadOrStore returns the existing bucket for the key if present. // Otherwise, it creates a new bucket, stores it, and returns it. -// This is specialized to avoid a closure allocation for the getter. func (bm *bucketMap[TKey]) loadOrStore(userKey TKey, executionTime ntime.Time, limit Limit) *bucket { spec := bucketSpec[TKey]{ limit: limit, @@ -31,7 +30,7 @@ func (bm *bucketMap[TKey]) loadOrStore(userKey TKey, executionTime ntime.Time, l if loaded, ok := bm.m.Load(spec); ok { return loaded.(*bucket) } - // Only create the b if we didn't find an existing one + // Only create the bucket if we didn't find an existing one b := newBucket(executionTime, limit) actual, _ := bm.m.LoadOrStore(spec, &b) return actual.(*bucket) From 9c7eedeb37fb77a780f08980d9106063dfe7278f Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 14:58:44 -0400 Subject: [PATCH 05/15] Simplify waitNWithCancellation --- limiter_wait.go | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/limiter_wait.go b/limiter_wait.go index 03f505e..9ced2d9 100644 --- a/limiter_wait.go +++ b/limiter_wait.go @@ -89,10 +89,6 @@ func (r *Limiter[TInput, TKey]) waitNWithCancellation( deadline func() (time.Time, bool), done func() <-chan struct{}, ) bool { - if r.allowN(input, startTime, n) { - return true - } - // currentTime is an approximation of the real clock moving forward // it's imprecise because it depends on time.After below. // For testing purposes, we want startTime (execution time) to @@ -107,9 +103,13 @@ func (r *Limiter[TInput, TKey]) waitNWithCancellation( retryAfter := details.RetryAfter() - // if we can't possibly get a token, fail fast + // If we can't possibly get a token, fail fast if deadline, ok := deadline(); ok { - if deadline.Before(currentTime.Add(retryAfter).ToTime()) { + // ctx deadline uses time.Time, not ntime.Time, + // so use a duration for the comparison. + // This might not be robust to clock skew. + d := deadline.Sub(currentTime.ToTime()) + if d < retryAfter { return false } } From 86f67888be1010ddc1dc0608a2b813a659b4447c Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 15:41:21 -0400 Subject: [PATCH 06/15] A little tightening --- limiter_wait.go | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/limiter_wait.go b/limiter_wait.go index 9ced2d9..9f3494e 100644 --- a/limiter_wait.go +++ b/limiter_wait.go @@ -31,7 +31,7 @@ import ( // Wait makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) bool { - return r.waitN(ctx, input, ntime.Now(), 1) + return r.WaitN(ctx, input, 1) } // WaitN will poll [Limiter.AllowN] for a period of time, @@ -58,13 +58,9 @@ func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) bool { // WaitN makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64) bool { - return r.waitN(ctx, input, ntime.Now(), n) -} - -func (r *Limiter[TInput, TKey]) waitN(ctx context.Context, input TInput, executionTime ntime.Time, n int64) bool { return r.waitNWithCancellation( input, - executionTime, + ntime.Now(), n, ctx.Deadline, ctx.Done, From b002aa7a75c1f1158c4f3698b11e9df2a4712c51 Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 15:42:12 -0400 Subject: [PATCH 07/15] Remove FIFO tests --- limiter_wait_test.go | 158 ------------------------------------------- 1 file changed, 158 deletions(-) diff --git a/limiter_wait_test.go b/limiter_wait_test.go index 3cd26b0..8163e8c 100644 --- a/limiter_wait_test.go +++ b/limiter_wait_test.go @@ -608,161 +608,3 @@ func TestLimiter_WaitN_ConsumesCorrectTokens(t *testing.T) { require.Equal(t, expectedRemaining, finalDetails[0].TokensRemaining(), "should have consumed exactly %d tokens total", successes*tokensPerWait) }) } - -func TestLimiter_Wait_FIFO_Ordering_SingleBucket_Flaky(t *testing.T) { - t.Skip("this test is flaky because the implementation is not deterministic") - - t.Parallel() - keyFunc := func(input string) string { - return "test-bucket" - } - // 1 token per 50ms - limit := NewLimit(1, 50*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - executionTime := ntime.Now() - - // Exhaust the single token - require.True(t, limiter.allow("key", executionTime), "should allow initial token") - require.False(t, limiter.allow("key", executionTime), "should not allow second token") - - const concurrency = 5 - var wg sync.WaitGroup - wg.Add(concurrency) - - startOrder := make(chan int, concurrency) - successOrder := make(chan int, concurrency) - - // Deadline that gives enough time for all tokens to be refilled - deadline := func() (time.Time, bool) { - return executionTime.Add(time.Duration(concurrency) * limit.durationPerToken).ToTime(), true - } - - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - for i := range concurrency { - go func(id int) { - defer wg.Done() - - // Signal that this goroutine is starting its wait - startOrder <- id - - // Wait for a token - if limiter.waitWithCancellation("key", executionTime, deadline, done) { - // Signal that this goroutine successfully acquired a token - successOrder <- id - } - }(i) - // A small, non-deterministic delay to encourage goroutines to queue up in order. - // This helps simulate a real-world scenario where requests arrive sequentially. - time.Sleep(5 * time.Millisecond) - } - - wg.Wait() - - close(startOrder) - close(successOrder) - - var starts []int - for id := range startOrder { - starts = append(starts, id) - } - - var successes []int - for id := range successOrder { - successes = append(successes, id) - } - - require.Equal(t, concurrency, len(successes), "all goroutines should have acquired a token") - require.Equal(t, starts, successes, "success order should match start order, proving FIFO") -} - -func TestLimiter_Wait_FIFO_Ordering_MultipleBuckets_Flaky(t *testing.T) { - t.Skip("this test is flaky because the implementation is not deterministic") - t.Parallel() - - // The FIFO behavior is best-effort, this test is known-flaky - // as a result - - const buckets = 3 - const concurrencyPerBucket = 5 - const concurrency = buckets * concurrencyPerBucket - - keyFunc := func(input int) string { - return fmt.Sprintf("test-bucket-%d", input) - } - // 1 token per 50ms, to make the test run reasonably fast - limit := NewLimit(1, 50*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - executionTime := ntime.Now() - - // Exhaust the single token for each bucket - for i := range buckets { - require.True(t, limiter.allow(i, executionTime), "should allow initial token for bucket %d", i) - require.False(t, limiter.allow(i, executionTime), "should not allow second token for bucket %d", i) - } - - // Create maps to hold order channels for each bucket - startOrders := make(map[int]chan int, buckets) - successOrders := make(map[int]chan int, buckets) - for i := range buckets { - startOrders[i] = make(chan int, concurrencyPerBucket) - successOrders[i] = make(chan int, concurrencyPerBucket) - } - - // Deadline that gives enough time for all tokens to be refilled for all goroutines - deadline := func() (time.Time, bool) { - return executionTime.Add(time.Duration(concurrencyPerBucket) * limit.durationPerToken).ToTime(), true - } - - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - var wg sync.WaitGroup - wg.Add(concurrency) - for i := range concurrency { - go func(id int) { - defer wg.Done() - bucketID := id % buckets - - // Announce that this goroutine is starting its wait for its bucket - startOrders[bucketID] <- id - - // Wait for a token - if limiter.waitWithCancellation(bucketID, executionTime, deadline, done) { - // Announce that this goroutine successfully acquired a token - successOrders[bucketID] <- id - } - }(i) - // A small, non-deterministic delay to encourage goroutines to queue up in order. - time.Sleep(5 * time.Millisecond) - } - - wg.Wait() - - // Close all channels - for bucketID := range buckets { - close(startOrders[bucketID]) - close(successOrders[bucketID]) - } - - // Verify FIFO order for each bucket - for bucketID := range buckets { - var starts []int - for id := range startOrders[bucketID] { - starts = append(starts, id) - } - - var successes []int - for id := range successOrders[bucketID] { - successes = append(successes, id) - } - - require.Equal(t, concurrencyPerBucket, len(successes), "all goroutines for bucket %d should have acquired a token", bucketID) - require.Equal(t, starts, successes, "success order should match start order for bucket %d, proving FIFO", bucketID) - } -} From 440dc3ee327a95e61cac0994bd2601465232eb8d Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 15:53:42 -0400 Subject: [PATCH 08/15] Combine single and multi tokens tests --- limiter_wait_test.go | 100 ++++++++++++++++++++++--------------------- 1 file changed, 51 insertions(+), 49 deletions(-) diff --git a/limiter_wait_test.go b/limiter_wait_test.go index 8163e8c..c06941f 100644 --- a/limiter_wait_test.go +++ b/limiter_wait_test.go @@ -30,7 +30,7 @@ func TestLimiter_Wait_SingleBucket(t *testing.T) { ok := limiter.allow("test", executionTime) require.False(t, ok, "should not allow when tokens exhausted") - // Test 1: Wait with enough time to acquire a token + // Test 1: Wait for 1 token with enough time to acquire it { // Deadline that gives enough time deadline := func() (time.Time, bool) { @@ -42,8 +42,8 @@ func TestLimiter_Wait_SingleBucket(t *testing.T) { return make(chan struct{}) // never closes } - allow := limiter.waitWithCancellation("test", executionTime, deadline, done) - require.True(t, allow, "should acquire token after waiting") + allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) + require.True(t, allow, "should acquire 1 token after waiting") // We waited executionTime = executionTime.Add(limit.durationPerToken) @@ -53,7 +53,7 @@ func TestLimiter_Wait_SingleBucket(t *testing.T) { ok = limiter.allow("test", executionTime) require.False(t, ok, "should not allow again immediately after wait") - // Test 2: Wait with deadline that expires before token is available + // Test 2: Wait for 1 token with deadline that expires before token is available { // Deadline that expires too soon deadline := func() (time.Time, bool) { @@ -65,11 +65,11 @@ func TestLimiter_Wait_SingleBucket(t *testing.T) { return make(chan struct{}) // never closes } - allow := limiter.waitWithCancellation("test", executionTime, deadline, done) - require.False(t, allow, "should not acquire token if deadline expires before token is available") + allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) + require.False(t, allow, "should not acquire 1 token if deadline expires before token is available") } - // Test 3: Wait with immediate cancellation + // Test 3: Wait for 1 token with immediate cancellation { // Deadline that gives enough time deadline := func() (time.Time, bool) { @@ -83,11 +83,11 @@ func TestLimiter_Wait_SingleBucket(t *testing.T) { return ch } - allow := limiter.waitWithCancellation("test", executionTime, deadline, done) - require.False(t, allow, "should not acquire token if context is cancelled immediately") + allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) + require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") } - // Test 4: Wait with no deadline + // Test 4: Wait for 1 token with no deadline { // Deadline that returns no deadline deadline := func() (time.Time, bool) { @@ -100,8 +100,8 @@ func TestLimiter_Wait_SingleBucket(t *testing.T) { } { - allow := limiter.waitWithCancellation("test", executionTime, deadline, done) - require.True(t, allow, "should acquire token when no deadline is set") + allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) + require.True(t, allow, "should acquire 1 token when no deadline is set") } // We waited @@ -113,31 +113,12 @@ func TestLimiter_Wait_SingleBucket(t *testing.T) { require.False(t, allow, "should not allow again immediately after wait") } } -} -func TestLimiter_WaitN_SingleBucket(t *testing.T) { - t.Parallel() - keyFunc := func(input string) string { - return input - } - limit := NewLimit(2, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - - executionTime := ntime.Now() - - // Consume all tokens - { - ok := limiter.allowN("test", executionTime, limit.count) - require.True(t, ok, "should allow initial tokens") - } + // Test 5: Wait for multiple tokens (limit.count) with enough time to acquire them { - // Should not allow immediately - ok := limiter.allowN("test", executionTime, 1) - require.False(t, ok, "should not allow when tokens exhausted") - } + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) - // Test 1: Wait with enough time to acquire tokens - { wait := time.Duration(limit.count) * limit.durationPerToken // Deadline that gives enough time deadline := func() (time.Time, bool) { @@ -150,23 +131,26 @@ func TestLimiter_WaitN_SingleBucket(t *testing.T) { } allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) - require.True(t, allow, "should acquire tokens after waiting") + require.True(t, allow, "should acquire %d tokens after waiting", limit.count) // We waited executionTime = executionTime.Add(limit.durationPerToken) } + // Test 6: Wait for multiple tokens with deadline that expires before all tokens are available { - // Should not allow again immediately after wait - ok := limiter.allow("test", executionTime) - require.False(t, ok, "should not allow again immediately after wait") - } + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) - // Test 2: Wait with deadline that expires before n tokens are available - { - // Deadline that expires too soon + // Consume all tokens to exhaust the bucket again + for range limit.count { + ok := limiter.allow("test", executionTime) + require.True(t, ok, "should allow tokens to exhaust bucket") + } + + // Deadline that expires too soon - use a very short duration deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken).ToTime(), true + return executionTime.Add(limit.durationPerToken / 4).ToTime(), true } // Done channel that never closes @@ -175,11 +159,20 @@ func TestLimiter_WaitN_SingleBucket(t *testing.T) { } allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) - require.False(t, allow, "should not acquire tokens if deadline expires before tokens are available") + require.False(t, allow, "should not acquire %d tokens if deadline expires before tokens are available", limit.count) } - // Test 3: Wait with immediate cancellation + // Test 7: Wait for multiple tokens with immediate cancellation { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) + + // Consume all tokens to exhaust the bucket again + for range limit.count { + ok := limiter.allow("test", executionTime) + require.True(t, ok, "should allow tokens to exhaust bucket") + } + // Deadline that gives enough time deadline := func() (time.Time, bool) { return executionTime.Add(limit.durationPerToken * time.Duration(limit.count)).ToTime(), true @@ -193,11 +186,20 @@ func TestLimiter_WaitN_SingleBucket(t *testing.T) { } allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) - require.False(t, allow, "should not acquire tokens if context is cancelled immediately") + require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit.count) } - // Test 4: Wait with no deadline + // Test 8: Wait for multiple tokens with no deadline { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) + + // Consume all tokens to exhaust the bucket again + for range limit.count { + ok := limiter.allow("test", executionTime) + require.True(t, ok, "should allow tokens to exhaust bucket") + } + // Deadline that returns no deadline deadline := func() (time.Time, bool) { return time.Time{}, false @@ -210,7 +212,7 @@ func TestLimiter_WaitN_SingleBucket(t *testing.T) { { allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) - require.True(t, allow, "should acquire tokens when no deadline is set") + require.True(t, allow, "should acquire %d tokens when no deadline is set", limit.count) } // We waited @@ -219,7 +221,7 @@ func TestLimiter_WaitN_SingleBucket(t *testing.T) { // Should not allow again immediately { allow := limiter.allowN("test", executionTime, limit.count) - require.False(t, allow, "should not allow again immediately after wait") + require.False(t, allow, "should not allow %d tokens again immediately after wait", limit.count) } } } From 422525bd19b9bebd908f64f7de9e190c189d3c4e Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 16:46:05 -0400 Subject: [PATCH 09/15] Reorganized sub tests --- limiter_wait_test.go | 646 ++++++++++++++++++++++--------------------- 1 file changed, 325 insertions(+), 321 deletions(-) diff --git a/limiter_wait_test.go b/limiter_wait_test.go index c06941f..1495cd6 100644 --- a/limiter_wait_test.go +++ b/limiter_wait_test.go @@ -10,414 +10,418 @@ import ( "github.com/stretchr/testify/require" ) -func TestLimiter_Wait_SingleBucket(t *testing.T) { +func TestLimiter_Wait(t *testing.T) { t.Parallel() - keyFunc := func(input string) string { - return input - } - limit := NewLimit(2, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - executionTime := ntime.Now() - - // Consume all tokens - for range limit.count { - ok := limiter.allow("test", executionTime) - require.True(t, ok, "should allow initial tokens") - } + t.Run("SingleBucket", func(t *testing.T) { + t.Parallel() + keyFunc := func(input string) string { + return input + } + limit := NewLimit(2, 100*time.Millisecond) + limiter := NewLimiter(keyFunc, limit) - // Should not allow immediately - ok := limiter.allow("test", executionTime) - require.False(t, ok, "should not allow when tokens exhausted") + executionTime := ntime.Now() - // Test 1: Wait for 1 token with enough time to acquire it - { - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken).ToTime(), true + // Consume all tokens + for range limit.count { + ok := limiter.allow("test", executionTime) + require.True(t, ok, "should allow initial tokens") } - // Done channel that never closes (no cancellation) - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } + // Should not allow immediately + ok := limiter.allow("test", executionTime) + require.False(t, ok, "should not allow when tokens exhausted") - allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) - require.True(t, allow, "should acquire 1 token after waiting") + // Test 1: Wait for 1 token with enough time to acquire it + { + // Deadline that gives enough time + deadline := func() (time.Time, bool) { + return executionTime.Add(limit.durationPerToken).ToTime(), true + } - // We waited - executionTime = executionTime.Add(limit.durationPerToken) - } + // Done channel that never closes (no cancellation) + done := func() <-chan struct{} { + return make(chan struct{}) // never closes + } - // Should not allow again immediately after wait - ok = limiter.allow("test", executionTime) - require.False(t, ok, "should not allow again immediately after wait") + allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) + require.True(t, allow, "should acquire 1 token after waiting") - // Test 2: Wait for 1 token with deadline that expires before token is available - { - // Deadline that expires too soon - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken / 2).ToTime(), true + // We waited + executionTime = executionTime.Add(limit.durationPerToken) } - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } + // Should not allow again immediately after wait + ok = limiter.allow("test", executionTime) + require.False(t, ok, "should not allow again immediately after wait") - allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) - require.False(t, allow, "should not acquire 1 token if deadline expires before token is available") - } + // Test 2: Wait for 1 token with deadline that expires before token is available + { + // Deadline that expires too soon + deadline := func() (time.Time, bool) { + return executionTime.Add(limit.durationPerToken / 2).ToTime(), true + } - // Test 3: Wait for 1 token with immediate cancellation - { - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken).ToTime(), true - } + // Done channel that never closes + done := func() <-chan struct{} { + return make(chan struct{}) // never closes + } - // Done channel that closes immediately - done := func() <-chan struct{} { - ch := make(chan struct{}) - close(ch) // immediately closed - return ch + allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) + require.False(t, allow, "should not acquire 1 token if deadline expires before token is available") } - allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) - require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") - } - - // Test 4: Wait for 1 token with no deadline - { - // Deadline that returns no deadline - deadline := func() (time.Time, bool) { - return time.Time{}, false - } + // Test 3: Wait for 1 token with immediate cancellation + { + // Deadline that gives enough time + deadline := func() (time.Time, bool) { + return executionTime.Add(limit.durationPerToken).ToTime(), true + } - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } + // Done channel that closes immediately + done := func() <-chan struct{} { + ch := make(chan struct{}) + close(ch) // immediately closed + return ch + } - { allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) - require.True(t, allow, "should acquire 1 token when no deadline is set") + require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") } - // We waited - executionTime = executionTime.Add(limit.durationPerToken) - - // Should not allow again immediately + // Test 4: Wait for 1 token with no deadline { - allow := limiter.allow("test", executionTime) - require.False(t, allow, "should not allow again immediately after wait") - } - } + // Deadline that returns no deadline + deadline := func() (time.Time, bool) { + return time.Time{}, false + } - // Test 5: Wait for multiple tokens (limit.count) with enough time to acquire them - { - // Wait for bucket to refill enough tokens - executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) + // Done channel that never closes + done := func() <-chan struct{} { + return make(chan struct{}) // never closes + } - wait := time.Duration(limit.count) * limit.durationPerToken - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(wait).ToTime(), true - } + { + allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) + require.True(t, allow, "should acquire 1 token when no deadline is set") + } - // Done channel that never closes (no cancellation) - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } + // We waited + executionTime = executionTime.Add(limit.durationPerToken) - allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) - require.True(t, allow, "should acquire %d tokens after waiting", limit.count) + // Should not allow again immediately + { + allow := limiter.allow("test", executionTime) + require.False(t, allow, "should not allow again immediately after wait") + } + } - // We waited - executionTime = executionTime.Add(limit.durationPerToken) - } + // Test 5: Wait for multiple tokens (limit.count) with enough time to acquire them + { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) - // Test 6: Wait for multiple tokens with deadline that expires before all tokens are available - { - // Wait for bucket to refill enough tokens - executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) + wait := time.Duration(limit.count) * limit.durationPerToken + // Deadline that gives enough time + deadline := func() (time.Time, bool) { + return executionTime.Add(wait).ToTime(), true + } - // Consume all tokens to exhaust the bucket again - for range limit.count { - ok := limiter.allow("test", executionTime) - require.True(t, ok, "should allow tokens to exhaust bucket") - } + // Done channel that never closes (no cancellation) + done := func() <-chan struct{} { + return make(chan struct{}) // never closes + } - // Deadline that expires too soon - use a very short duration - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken / 4).ToTime(), true - } + allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) + require.True(t, allow, "should acquire %d tokens after waiting", limit.count) - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes + // We waited + executionTime = executionTime.Add(limit.durationPerToken) } - allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) - require.False(t, allow, "should not acquire %d tokens if deadline expires before tokens are available", limit.count) - } + // Test 6: Wait for multiple tokens with deadline that expires before all tokens are available + { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) - // Test 7: Wait for multiple tokens with immediate cancellation - { - // Wait for bucket to refill enough tokens - executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) + // Consume all tokens to exhaust the bucket again + for range limit.count { + ok := limiter.allow("test", executionTime) + require.True(t, ok, "should allow tokens to exhaust bucket") + } - // Consume all tokens to exhaust the bucket again - for range limit.count { - ok := limiter.allow("test", executionTime) - require.True(t, ok, "should allow tokens to exhaust bucket") - } + // Deadline that expires too soon - use a very short duration + deadline := func() (time.Time, bool) { + return executionTime.Add(limit.durationPerToken / 4).ToTime(), true + } - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken * time.Duration(limit.count)).ToTime(), true - } + // Done channel that never closes + done := func() <-chan struct{} { + return make(chan struct{}) // never closes + } - // Done channel that closes immediately - done := func() <-chan struct{} { - ch := make(chan struct{}) - close(ch) // immediately closed - return ch + allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) + require.False(t, allow, "should not acquire %d tokens if deadline expires before tokens are available", limit.count) } - allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) - require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit.count) - } + // Test 7: Wait for multiple tokens with immediate cancellation + { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) - // Test 8: Wait for multiple tokens with no deadline - { - // Wait for bucket to refill enough tokens - executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) + // Consume all tokens to exhaust the bucket again + for range limit.count { + ok := limiter.allow("test", executionTime) + require.True(t, ok, "should allow tokens to exhaust bucket") + } - // Consume all tokens to exhaust the bucket again - for range limit.count { - ok := limiter.allow("test", executionTime) - require.True(t, ok, "should allow tokens to exhaust bucket") - } + // Deadline that gives enough time + deadline := func() (time.Time, bool) { + return executionTime.Add(limit.durationPerToken * time.Duration(limit.count)).ToTime(), true + } - // Deadline that returns no deadline - deadline := func() (time.Time, bool) { - return time.Time{}, false - } + // Done channel that closes immediately + done := func() <-chan struct{} { + ch := make(chan struct{}) + close(ch) // immediately closed + return ch + } - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes + allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) + require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit.count) } + // Test 8: Wait for multiple tokens with no deadline { - allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) - require.True(t, allow, "should acquire %d tokens when no deadline is set", limit.count) - } + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) - // We waited - executionTime = executionTime.Add(limit.durationPerToken) + // Consume all tokens to exhaust the bucket again + for range limit.count { + ok := limiter.allow("test", executionTime) + require.True(t, ok, "should allow tokens to exhaust bucket") + } - // Should not allow again immediately - { - allow := limiter.allowN("test", executionTime, limit.count) - require.False(t, allow, "should not allow %d tokens again immediately after wait", limit.count) - } - } -} + // Deadline that returns no deadline + deadline := func() (time.Time, bool) { + return time.Time{}, false + } -func TestLimiter_Wait_MultipleBuckets(t *testing.T) { - t.Parallel() - keyFunc := func(input int) string { - return fmt.Sprintf("test-bucket-%d", input) - } - const buckets = 3 - limit := NewLimit(2, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - executionTime := ntime.Now() + // Done channel that never closes + done := func() <-chan struct{} { + return make(chan struct{}) // never closes + } - // Exhaust tokens for all buckets - for bucketID := range buckets { - for range limit.count { - allow := limiter.allow(bucketID, executionTime) - require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) - } - } + { + allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) + require.True(t, allow, "should acquire %d tokens when no deadline is set", limit.count) + } - // Should not allow immediately for any bucket - for bucketID := range buckets { - allow := limiter.allow(bucketID, executionTime) - require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID) - } + // We waited + executionTime = executionTime.Add(limit.durationPerToken) - // Wait for a token for each bucket - for bucketID := range buckets { - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken).ToTime(), true + // Should not allow again immediately + { + allow := limiter.allowN("test", executionTime, limit.count) + require.False(t, allow, "should not allow %d tokens again immediately after wait", limit.count) + } } + }) - // Done channel that never closes (no cancellation) - done := func() <-chan struct{} { - return make(chan struct{}) // never closes + t.Run("MultipleBuckets", func(t *testing.T) { + t.Parallel() + keyFunc := func(input int) string { + return fmt.Sprintf("test-bucket-%d", input) + } + const buckets = 3 + limit := NewLimit(2, 100*time.Millisecond) + limiter := NewLimiter(keyFunc, limit) + executionTime := ntime.Now() + + // Exhaust tokens for all buckets + for bucketID := range buckets { + for range limit.count { + allow := limiter.allow(bucketID, executionTime) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } } - allow := limiter.waitWithCancellation(bucketID, executionTime, deadline, done) - require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID) - } - - // We waited - executionTime = executionTime.Add(limit.durationPerToken) - - // Buckets should be empty again - for bucketID := range buckets { - ok := limiter.allow(bucketID, executionTime) - require.False(t, ok, "should not allow again immediately after wait for bucket %d", bucketID) - } -} - -func TestLimiter_Wait_MultipleBuckets_Concurrent(t *testing.T) { - t.Parallel() - keyFunc := func(input int64) string { - return fmt.Sprintf("test-bucket-%d", input) - } - const buckets int64 = 3 - limit := NewLimit(2, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - executionTime := ntime.Now() - - // Exhaust tokens for all buckets - for bucketID := range buckets { - for range limit.count { + // Should not allow immediately for any bucket + for bucketID := range buckets { allow := limiter.allow(bucketID, executionTime) - require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID) } - } - // Buckets should be empty - for bucketID := range buckets { - allow := limiter.allow(bucketID, executionTime) - require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID) - } + // Wait for a token for each bucket + for bucketID := range buckets { + // Deadline that gives enough time + deadline := func() (time.Time, bool) { + return executionTime.Add(limit.durationPerToken).ToTime(), true + } - // Test 1: Multiple goroutines competing for tokens with enough time - { - // More goroutines than available tokens to create competition - tokens := buckets * limit.count - concurrency := tokens * 3 // oversubscribe by 3x - results := make([]bool, concurrency) + // Done channel that never closes (no cancellation) + done := func() <-chan struct{} { + return make(chan struct{}) // never closes + } - // Deadline that gives enough time for all tokens to be refilled - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.period).ToTime(), true + allow := limiter.waitWithCancellation(bucketID, executionTime, deadline, done) + require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID) } - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes + // We waited + executionTime = executionTime.Add(limit.durationPerToken) + + // Buckets should be empty again + for bucketID := range buckets { + ok := limiter.allow(bucketID, executionTime) + require.False(t, ok, "should not allow again immediately after wait for bucket %d", bucketID) } + }) - // Start concurrent waits - var wg sync.WaitGroup - for i := range concurrency { - wg.Add(1) - go func(i int64) { - defer wg.Done() - bucketID := i % buckets - results[i] = limiter.waitWithCancellation(bucketID, executionTime, deadline, done) - }(i) + t.Run("MultipleBuckets_Concurrent", func(t *testing.T) { + t.Parallel() + keyFunc := func(input int64) string { + return fmt.Sprintf("test-bucket-%d", input) + } + const buckets int64 = 3 + limit := NewLimit(2, 100*time.Millisecond) + limiter := NewLimiter(keyFunc, limit) + executionTime := ntime.Now() + + // Exhaust tokens for all buckets + for bucketID := range buckets { + for range limit.count { + allow := limiter.allow(bucketID, executionTime) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } } - wg.Wait() - // We waited - executionTime = executionTime.Add(limit.period) + // Buckets should be empty + for bucketID := range buckets { + allow := limiter.allow(bucketID, executionTime) + require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID) + } - // Count successes and failures - var successes, failures int64 - for _, result := range results { - if result { - successes++ - } else { - failures++ + // Test 1: Multiple goroutines competing for tokens with enough time + { + // More goroutines than available tokens to create competition + tokens := buckets * limit.count + concurrency := tokens * 3 // oversubscribe by 3x + results := make([]bool, concurrency) + + // Deadline that gives enough time for all tokens to be refilled + deadline := func() (time.Time, bool) { + return executionTime.Add(limit.period).ToTime(), true } - } - // Exactly limit.count tokens per bucket should succeed, 3 buckets * 2 tokens each - expectedSuccesses := buckets * limit.count - require.Equal(t, expectedSuccesses, successes, "expected exactly %d goroutines to acquire tokens", expectedSuccesses) + // Done channel that never closes + done := func() <-chan struct{} { + return make(chan struct{}) // never closes + } - // The rest should fail due to competition - expectedFailures := concurrency - expectedSuccesses - require.Equal(t, expectedFailures, failures, "expected %d goroutines to fail due to competition", expectedFailures) - } + // Start concurrent waits + var wg sync.WaitGroup + for i := range concurrency { + wg.Add(1) + go func(i int64) { + defer wg.Done() + bucketID := i % buckets + results[i] = limiter.waitWithCancellation(bucketID, executionTime, deadline, done) + }(i) + } + wg.Wait() + + // We waited + executionTime = executionTime.Add(limit.period) + + // Count successes and failures + var successes, failures int64 + for _, result := range results { + if result { + successes++ + } else { + failures++ + } + } - // Test 2: Multiple goroutines with deadline that expires before tokens are available - { - concurrency := buckets - results := make([]bool, concurrency) + // Exactly limit.count tokens per bucket should succeed, 3 buckets * 2 tokens each + expectedSuccesses := buckets * limit.count + require.Equal(t, expectedSuccesses, successes, "expected exactly %d goroutines to acquire tokens", expectedSuccesses) - // Deadline that expires too soon - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken / 2).ToTime(), true + // The rest should fail due to competition + expectedFailures := concurrency - expectedSuccesses + require.Equal(t, expectedFailures, failures, "expected %d goroutines to fail due to competition", expectedFailures) } - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } + // Test 2: Multiple goroutines with deadline that expires before tokens are available + { + concurrency := buckets + results := make([]bool, concurrency) - // Start concurrent waits - var wg sync.WaitGroup - for i := range concurrency { - wg.Add(1) - go func(i int64) { - defer wg.Done() - results[i] = limiter.waitWithCancellation(i, executionTime, deadline, done) - }(i) - } + // Deadline that expires too soon + deadline := func() (time.Time, bool) { + return executionTime.Add(limit.durationPerToken / 2).ToTime(), true + } - wg.Wait() + // Done channel that never closes + done := func() <-chan struct{} { + return make(chan struct{}) // never closes + } - // All should fail because deadline expires before tokens are available - for i, result := range results { - require.False(t, result, "goroutine %d should not acquire token due to early deadline", i) - } - } + // Start concurrent waits + var wg sync.WaitGroup + for i := range concurrency { + wg.Add(1) + go func(i int64) { + defer wg.Done() + results[i] = limiter.waitWithCancellation(i, executionTime, deadline, done) + }(i) + } - // Test 3: Multiple goroutines with immediate cancellation - { - concurrency := buckets - results := make([]bool, concurrency) + wg.Wait() - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.period).ToTime(), true + // All should fail because deadline expires before tokens are available + for i, result := range results { + require.False(t, result, "goroutine %d should not acquire token due to early deadline", i) + } } - // Done channel that closes immediately - done := func() <-chan struct{} { - ch := make(chan struct{}) - close(ch) // immediately closed - return ch - } + // Test 3: Multiple goroutines with immediate cancellation + { + concurrency := buckets + results := make([]bool, concurrency) - // Start concurrent waits - var wg sync.WaitGroup - for i := range concurrency { - wg.Add(1) - go func(i int64) { - defer wg.Done() - results[i] = limiter.waitWithCancellation(i, executionTime, deadline, done) - }(i) - } + // Deadline that gives enough time + deadline := func() (time.Time, bool) { + return executionTime.Add(limit.period).ToTime(), true + } - wg.Wait() + // Done channel that closes immediately + done := func() <-chan struct{} { + ch := make(chan struct{}) + close(ch) // immediately closed + return ch + } - // All should fail because context is cancelled immediately - for i, result := range results { - require.False(t, result, "goroutine %d should not acquire token due to immediate cancellation", i) + // Start concurrent waits + var wg sync.WaitGroup + for i := range concurrency { + wg.Add(1) + go func(i int64) { + defer wg.Done() + results[i] = limiter.waitWithCancellation(i, executionTime, deadline, done) + }(i) + } + + wg.Wait() + + // All should fail because context is cancelled immediately + for i, result := range results { + require.False(t, result, "goroutine %d should not acquire token due to immediate cancellation", i) + } } - } + }) } func TestLimiter_WaitN_ConsumesCorrectTokens(t *testing.T) { From 3cfc4363d0eed34fd877a9ab377457238dd30b78 Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 21:34:49 -0400 Subject: [PATCH 10/15] Add wait with details --- limiter_wait.go | 95 ++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 86 insertions(+), 9 deletions(-) diff --git a/limiter_wait.go b/limiter_wait.go index 9f3494e..f52b76f 100644 --- a/limiter_wait.go +++ b/limiter_wait.go @@ -58,7 +58,70 @@ func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) bool { // WaitN makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64) bool { - return r.waitNWithCancellation( + allow, _ := r.waitNWithDetails( + input, + ntime.Now(), + n, + ctx.Deadline, + ctx.Done, + ) + return allow +} + +// WaitWithDetails will poll [Allow] for a period of time, +// until it is cancelled by the passed context. It has the +// effect of adding latency to requests instead of refusing +// them immediately. Consider it graceful degradation. +// +// WaitWithDetails will return true if a token becomes available prior to +// the context cancellation, and will consume a token. It will +// return false if not, and therefore not consume a token. It will +// also return the details of the request. +// +// Take care to create an appropriate context. You almost certainly +// want [context.WithTimeout] or [context.WithDeadline]. +// +// You should be conservative, as WaitWithDetails will introduce +// backpressure on your upstream systems -- connections +// may be held open longer, requests may queue in memory. +// +// A good starting place will be to timeout after waiting +// for one token. For example: +// +// ctx := context.WithTimeout(ctx, limit.DurationPerToken()) +// +// WaitWithDetails makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. +func (r *Limiter[TInput, TKey]) WaitWithDetails(ctx context.Context, input TInput) (bool, Details[TInput, TKey]) { + return r.WaitNWithDetails(ctx, input, 1) +} + +// WaitNWithDetails will poll [Limiter.AllowN] for a period of time, +// until it is cancelled by the passed context. It has the +// effect of adding latency to requests instead of refusing +// them immediately. Consider it graceful degradation. +// +// WaitNWithDetails will return true if `n` tokens become available prior to +// the context cancellation, and will consume `n` tokens. If not, +// it will return false, and therefore consume no tokens. It will +// also return the details of the request. +// +// Take care to create an appropriate context. You almost certainly +// want [context.WithTimeout] or [context.WithDeadline]. +// +// You should be conservative, as WaitNWithDetails will introduce +// backpressure on your upstream systems -- connections +// may be held open longer, requests may queue in memory. +// +// A good starting place will be to timeout after waiting +// for one token. For example: +// +// ctx := context.WithTimeout(ctx, limit.DurationPerToken()) +// +// WaitNWithDetails makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. +func (r *Limiter[TInput, TKey]) WaitNWithDetails(ctx context.Context, input TInput, n int64) (bool, Details[TInput, TKey]) { + return r.waitNWithDetails( input, ntime.Now(), n, @@ -73,11 +136,10 @@ func (r *Limiter[TInput, TKey]) waitWithCancellation( deadline func() (time.Time, bool), done func() <-chan struct{}, ) bool { - return r.waitNWithCancellation(input, startTime, 1, deadline, done) + allow, _ := r.waitNWithDetails(input, startTime, 1, deadline, done) + return allow } -// waitWithCancellation is a more testable version of wait that accepts -// deadline and done functions instead of a context, allowing for deterministic testing. func (r *Limiter[TInput, TKey]) waitNWithCancellation( input TInput, startTime ntime.Time, @@ -85,8 +147,21 @@ func (r *Limiter[TInput, TKey]) waitNWithCancellation( deadline func() (time.Time, bool), done func() <-chan struct{}, ) bool { - // currentTime is an approximation of the real clock moving forward - // it's imprecise because it depends on time.After below. + allow, _ := r.waitNWithDetails(input, startTime, n, deadline, done) + return allow +} + +// waitWithCancellation is a more testable version of wait that accepts +// deadline and done functions instead of a context, allowing for deterministic testing. +func (r *Limiter[TInput, TKey]) waitNWithDetails( + input TInput, + startTime ntime.Time, + n int64, + deadline func() (time.Time, bool), + done func() <-chan struct{}, +) (bool, Details[TInput, TKey]) { + // currentTime is an approximation of the real clock moving forward. + // It's imprecise because it depends on time.After below. // For testing purposes, we want startTime (execution time) to // be a parameter. currentTime := startTime @@ -94,7 +169,7 @@ func (r *Limiter[TInput, TKey]) waitNWithCancellation( for { allow, details := r.allowNWithDetails(input, currentTime, n) if allow { - return true + return allow, details } retryAfter := details.RetryAfter() @@ -106,13 +181,15 @@ func (r *Limiter[TInput, TKey]) waitNWithCancellation( // This might not be robust to clock skew. d := deadline.Sub(currentTime.ToTime()) if d < retryAfter { - return false + return false, details } } select { case <-done(): - return false + // Do one more check on cancellation, + // to have accurate retryAfter. + return r.allowNWithDetails(input, currentTime, n) case <-time.After(retryAfter): currentTime = currentTime.Add(retryAfter) } From 1e580b519899a109669c20e01e95cea2b598b767 Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Wed, 20 Aug 2025 23:44:13 -0400 Subject: [PATCH 11/15] Use a real context More cleanup to do I think --- limiter_wait.go | 57 ++---- limiter_wait_test.go | 430 ++++++++++++------------------------------- 2 files changed, 127 insertions(+), 360 deletions(-) diff --git a/limiter_wait.go b/limiter_wait.go index f52b76f..b1a1d88 100644 --- a/limiter_wait.go +++ b/limiter_wait.go @@ -58,13 +58,7 @@ func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) bool { // WaitN makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64) bool { - allow, _ := r.waitNWithDetails( - input, - ntime.Now(), - n, - ctx.Deadline, - ctx.Done, - ) + allow, _ := r.waitNWithDetails(input, ntime.Now(), n, ctx) return allow } @@ -109,7 +103,7 @@ func (r *Limiter[TInput, TKey]) WaitWithDetails(ctx context.Context, input TInpu // Take care to create an appropriate context. You almost certainly // want [context.WithTimeout] or [context.WithDeadline]. // -// You should be conservative, as WaitNWithDetails will introduce +// You should be conservative, as WaitWithDetails will introduce // backpressure on your upstream systems -- connections // may be held open longer, requests may queue in memory. // @@ -118,47 +112,20 @@ func (r *Limiter[TInput, TKey]) WaitWithDetails(ctx context.Context, input TInpu // // ctx := context.WithTimeout(ctx, limit.DurationPerToken()) // -// WaitNWithDetails makes no ordering guarantees. Multiple concurrent calls may +// WaitWithDetails makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. func (r *Limiter[TInput, TKey]) WaitNWithDetails(ctx context.Context, input TInput, n int64) (bool, Details[TInput, TKey]) { - return r.waitNWithDetails( - input, - ntime.Now(), - n, - ctx.Deadline, - ctx.Done, - ) -} - -func (r *Limiter[TInput, TKey]) waitWithCancellation( - input TInput, - startTime ntime.Time, - deadline func() (time.Time, bool), - done func() <-chan struct{}, -) bool { - allow, _ := r.waitNWithDetails(input, startTime, 1, deadline, done) - return allow -} - -func (r *Limiter[TInput, TKey]) waitNWithCancellation( - input TInput, - startTime ntime.Time, - n int64, - deadline func() (time.Time, bool), - done func() <-chan struct{}, -) bool { - allow, _ := r.waitNWithDetails(input, startTime, n, deadline, done) - return allow + return r.waitNWithDetails(input, ntime.Now(), n, ctx) } -// waitWithCancellation is a more testable version of wait that accepts -// deadline and done functions instead of a context, allowing for deterministic testing. +// waitNWithDetails is the internal implementation that accepts a context. +// It is designed to be testable by accepting any context implementation, +// including test contexts that provide deterministic behavior. func (r *Limiter[TInput, TKey]) waitNWithDetails( input TInput, startTime ntime.Time, n int64, - deadline func() (time.Time, bool), - done func() <-chan struct{}, + ctx context.Context, ) (bool, Details[TInput, TKey]) { // currentTime is an approximation of the real clock moving forward. // It's imprecise because it depends on time.After below. @@ -175,7 +142,7 @@ func (r *Limiter[TInput, TKey]) waitNWithDetails( retryAfter := details.RetryAfter() // If we can't possibly get a token, fail fast - if deadline, ok := deadline(); ok { + if deadline, ok := ctx.Deadline(); ok { // ctx deadline uses time.Time, not ntime.Time, // so use a duration for the comparison. // This might not be robust to clock skew. @@ -186,10 +153,8 @@ func (r *Limiter[TInput, TKey]) waitNWithDetails( } select { - case <-done(): - // Do one more check on cancellation, - // to have accurate retryAfter. - return r.allowNWithDetails(input, currentTime, n) + case <-ctx.Done(): + return false, details case <-time.After(retryAfter): currentTime = currentTime.Add(retryAfter) } diff --git a/limiter_wait_test.go b/limiter_wait_test.go index 1495cd6..c6e9177 100644 --- a/limiter_wait_test.go +++ b/limiter_wait_test.go @@ -1,6 +1,7 @@ package rate import ( + "context" "fmt" "sync" "testing" @@ -35,18 +36,16 @@ func TestLimiter_Wait(t *testing.T) { // Test 1: Wait for 1 token with enough time to acquire it { - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken).ToTime(), true + // Context with deadline that gives enough time + ctx := &testContext{ + deadline: executionTime.Add(limit.durationPerToken).ToTime(), + hasDeadline: true, + done: make(chan struct{}), // never closes } - // Done channel that never closes (no cancellation) - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) + allow, details := limiter.waitNWithDetails("test", executionTime, 1, ctx) require.True(t, allow, "should acquire 1 token after waiting") + require.NotNil(t, details, "should return details") // We waited executionTime = executionTime.Add(limit.durationPerToken) @@ -58,63 +57,53 @@ func TestLimiter_Wait(t *testing.T) { // Test 2: Wait for 1 token with deadline that expires before token is available { - // Deadline that expires too soon - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken / 2).ToTime(), true - } - - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes + // Context with deadline that expires too soon + ctx := &testContext{ + deadline: executionTime.Add(limit.durationPerToken / 2).ToTime(), + hasDeadline: true, + done: make(chan struct{}), // never closes } - allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) + allow, details := limiter.waitNWithDetails("test", executionTime, 1, ctx) require.False(t, allow, "should not acquire 1 token if deadline expires before token is available") + require.NotNil(t, details, "should return details even on deadline expiry") } // Test 3: Wait for 1 token with immediate cancellation { - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken).ToTime(), true - } - - // Done channel that closes immediately - done := func() <-chan struct{} { - ch := make(chan struct{}) - close(ch) // immediately closed - return ch + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + close(done) + ctx := &testContext{ + deadline: executionTime.Add(limit.durationPerToken).ToTime(), + hasDeadline: true, + done: done, } - allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) + allow, details := limiter.waitNWithDetails("test", executionTime, 1, ctx) require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") + require.NotNil(t, details, "should return details even on cancellation") } // Test 4: Wait for 1 token with no deadline { - // Deadline that returns no deadline - deadline := func() (time.Time, bool) { - return time.Time{}, false + // Context with no deadline + ctx := &testContext{ + deadline: time.Time{}, + hasDeadline: false, + done: make(chan struct{}), // never closes } - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - { - allow := limiter.waitNWithCancellation("test", executionTime, 1, deadline, done) - require.True(t, allow, "should acquire 1 token when no deadline is set") - } + allow, details := limiter.waitNWithDetails("test", executionTime, 1, ctx) + require.True(t, allow, "should acquire 1 token when no deadline is set") + require.NotNil(t, details, "should return details") // We waited executionTime = executionTime.Add(limit.durationPerToken) // Should not allow again immediately - { - allow := limiter.allow("test", executionTime) - require.False(t, allow, "should not allow again immediately after wait") - } + ok := limiter.allow("test", executionTime) + require.False(t, ok, "should not allow again immediately after wait") } // Test 5: Wait for multiple tokens (limit.count) with enough time to acquire them @@ -123,18 +112,16 @@ func TestLimiter_Wait(t *testing.T) { executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) wait := time.Duration(limit.count) * limit.durationPerToken - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(wait).ToTime(), true + // Context with deadline that gives enough time + ctx := &testContext{ + deadline: executionTime.Add(wait).ToTime(), + hasDeadline: true, + done: make(chan struct{}), // never closes } - // Done channel that never closes (no cancellation) - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) + allow, details := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) require.True(t, allow, "should acquire %d tokens after waiting", limit.count) + require.NotNil(t, details, "should return details") // We waited executionTime = executionTime.Add(limit.durationPerToken) @@ -151,18 +138,16 @@ func TestLimiter_Wait(t *testing.T) { require.True(t, ok, "should allow tokens to exhaust bucket") } - // Deadline that expires too soon - use a very short duration - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken / 4).ToTime(), true - } - - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes + // Context with deadline that expires too soon + ctx := &testContext{ + deadline: executionTime.Add(limit.durationPerToken / 4).ToTime(), + hasDeadline: true, + done: make(chan struct{}), // never closes } - allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) + allow, details := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) require.False(t, allow, "should not acquire %d tokens if deadline expires before tokens are available", limit.count) + require.NotNil(t, details, "should return details even on deadline expiry") } // Test 7: Wait for multiple tokens with immediate cancellation @@ -176,20 +161,18 @@ func TestLimiter_Wait(t *testing.T) { require.True(t, ok, "should allow tokens to exhaust bucket") } - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken * time.Duration(limit.count)).ToTime(), true - } - - // Done channel that closes immediately - done := func() <-chan struct{} { - ch := make(chan struct{}) - close(ch) // immediately closed - return ch + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + close(done) + ctx := &testContext{ + deadline: executionTime.Add(limit.durationPerToken * time.Duration(limit.count)).ToTime(), + hasDeadline: true, + done: done, } - allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) + allow, details := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit.count) + require.NotNil(t, details, "should return details even on cancellation") } // Test 8: Wait for multiple tokens with no deadline @@ -203,29 +186,23 @@ func TestLimiter_Wait(t *testing.T) { require.True(t, ok, "should allow tokens to exhaust bucket") } - // Deadline that returns no deadline - deadline := func() (time.Time, bool) { - return time.Time{}, false + // Context with no deadline + ctx := &testContext{ + deadline: time.Time{}, + hasDeadline: false, + done: make(chan struct{}), // never closes } - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - { - allow := limiter.waitNWithCancellation("test", executionTime, limit.count, deadline, done) - require.True(t, allow, "should acquire %d tokens when no deadline is set", limit.count) - } + allow, details := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) + require.True(t, allow, "should acquire %d tokens when no deadline is set", limit.count) + require.NotNil(t, details, "should return details") // We waited executionTime = executionTime.Add(limit.durationPerToken) // Should not allow again immediately - { - allow := limiter.allowN("test", executionTime, limit.count) - require.False(t, allow, "should not allow %d tokens again immediately after wait", limit.count) - } + allow = limiter.allowN("test", executionTime, limit.count) + require.False(t, allow, "should not allow %d tokens again immediately after wait", limit.count) } }) @@ -255,18 +232,16 @@ func TestLimiter_Wait(t *testing.T) { // Wait for a token for each bucket for bucketID := range buckets { - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken).ToTime(), true + // Context with deadline that gives enough time + ctx := &testContext{ + deadline: executionTime.Add(limit.durationPerToken).ToTime(), + hasDeadline: true, + done: make(chan struct{}), // never closes } - // Done channel that never closes (no cancellation) - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - allow := limiter.waitWithCancellation(bucketID, executionTime, deadline, done) + allow, details := limiter.waitNWithDetails(bucketID, executionTime, 1, ctx) require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID) + require.NotNil(t, details, "should return details") } // We waited @@ -310,14 +285,11 @@ func TestLimiter_Wait(t *testing.T) { concurrency := tokens * 3 // oversubscribe by 3x results := make([]bool, concurrency) - // Deadline that gives enough time for all tokens to be refilled - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.period).ToTime(), true - } - - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes + // Context with deadline that gives enough time for all tokens to be refilled + ctx := &testContext{ + deadline: executionTime.Add(limit.period).ToTime(), + hasDeadline: true, + done: make(chan struct{}), // never closes } // Start concurrent waits @@ -327,7 +299,7 @@ func TestLimiter_Wait(t *testing.T) { go func(i int64) { defer wg.Done() bucketID := i % buckets - results[i] = limiter.waitWithCancellation(bucketID, executionTime, deadline, done) + results[i], _ = limiter.waitNWithDetails(bucketID, executionTime, 1, ctx) }(i) } wg.Wait() @@ -359,14 +331,11 @@ func TestLimiter_Wait(t *testing.T) { concurrency := buckets results := make([]bool, concurrency) - // Deadline that expires too soon - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.durationPerToken / 2).ToTime(), true - } - - // Done channel that never closes - done := func() <-chan struct{} { - return make(chan struct{}) // never closes + // Context with deadline that expires too soon + ctx := &testContext{ + deadline: executionTime.Add(limit.durationPerToken / 2).ToTime(), + hasDeadline: true, + done: make(chan struct{}), // never closes } // Start concurrent waits @@ -375,10 +344,9 @@ func TestLimiter_Wait(t *testing.T) { wg.Add(1) go func(i int64) { defer wg.Done() - results[i] = limiter.waitWithCancellation(i, executionTime, deadline, done) + results[i], _ = limiter.waitNWithDetails(i, executionTime, 1, ctx) }(i) } - wg.Wait() // All should fail because deadline expires before tokens are available @@ -392,16 +360,13 @@ func TestLimiter_Wait(t *testing.T) { concurrency := buckets results := make([]bool, concurrency) - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(limit.period).ToTime(), true - } - - // Done channel that closes immediately - done := func() <-chan struct{} { - ch := make(chan struct{}) - close(ch) // immediately closed - return ch + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + close(done) + ctx := &testContext{ + deadline: executionTime.Add(limit.period).ToTime(), + hasDeadline: true, + done: done, } // Start concurrent waits @@ -410,10 +375,9 @@ func TestLimiter_Wait(t *testing.T) { wg.Add(1) go func(i int64) { defer wg.Done() - results[i] = limiter.waitWithCancellation(i, executionTime, deadline, done) + results[i], _ = limiter.waitNWithDetails(i, executionTime, 1, ctx) }(i) } - wg.Wait() // All should fail because context is cancelled immediately @@ -424,193 +388,31 @@ func TestLimiter_Wait(t *testing.T) { }) } -func TestLimiter_WaitN_ConsumesCorrectTokens(t *testing.T) { - t.Parallel() - keyFunc := func(input string) string { - return input - } - limit := NewLimit(10, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - - executionTime := ntime.Now() - - // Test 1: Wait should consume exactly 1 token - t.Run("Wait_Consumes_One_Token", func(t *testing.T) { - // Verify initial state - _, initialDetails := limiter.peekWithDebug("test-wait-1", executionTime) - require.Equal(t, limit.count, initialDetails[0].TokensRemaining(), "should start with all tokens") - - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(time.Second).ToTime(), true - } - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - // Wait should succeed and consume exactly 1 token - allowed := limiter.waitWithCancellation("test-wait-1", executionTime, deadline, done) - require.True(t, allowed, "wait should succeed") - - // Verify exactly 1 token was consumed - _, finalDetails := limiter.peekWithDebug("test-wait-1", executionTime) - require.Equal(t, limit.count-1, finalDetails[0].TokensRemaining(), "should have consumed exactly 1 token") - }) - - // Test 2: WaitN should consume exactly n tokens - t.Run("WaitN_Consumes_N_Tokens", func(t *testing.T) { - const tokensToWait = 3 - - // Verify initial state - _, initialDetails := limiter.peekWithDebug("test-waitn-3", executionTime) - require.Equal(t, limit.count, initialDetails[0].TokensRemaining(), "should start with all tokens") +var _ context.Context = &testContext{} - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(time.Second).ToTime(), true - } - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - // WaitN should succeed and consume exactly tokensToWait tokens - allowed := limiter.waitNWithCancellation("test-waitn-3", executionTime, tokensToWait, deadline, done) - require.True(t, allowed, "waitN should succeed") - - // Verify exactly tokensToWait tokens were consumed - _, details := limiter.peekWithDebug("test-waitn-3", executionTime) - require.Equal(t, limit.count-tokensToWait, details[0].TokensRemaining(), "should have consumed exactly %d tokens", tokensToWait) - }) - - // Test 3: WaitN with multiple limits should consume n tokens from all buckets - t.Run("WaitN_MultipleLimits_Consumes_N_From_All", func(t *testing.T) { - perSecond := NewLimit(5, time.Second) - perMinute := NewLimit(20, time.Minute) - limiter := NewLimiter(keyFunc, perSecond, perMinute) - const tokensToWait = 2 - - // Verify initial state - _, initialDetails := limiter.peekWithDebug("test-multi-waitn", executionTime) - require.Len(t, initialDetails, 2, "should have details for both limits") - require.Equal(t, perSecond.count, initialDetails[0].TokensRemaining(), "per-second should start with all tokens") - require.Equal(t, perMinute.count, initialDetails[1].TokensRemaining(), "per-minute should start with all tokens") - - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(time.Second).ToTime(), true - } - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - // WaitN should succeed and consume exactly tokensToWait tokens from both limits - allowed := limiter.waitNWithCancellation("test-multi-waitn", executionTime, tokensToWait, deadline, done) - require.True(t, allowed, "waitN should succeed with multiple limits") - - // Verify exactly tokensToWait tokens were consumed from both buckets - _, finalDetails := limiter.peekWithDebug("test-multi-waitn", executionTime) - require.Len(t, finalDetails, 2, "should have details for both limits") - require.Equal(t, perSecond.count-tokensToWait, finalDetails[0].TokensRemaining(), "per-second should have consumed exactly %d tokens", tokensToWait) - require.Equal(t, perMinute.count-tokensToWait, finalDetails[1].TokensRemaining(), "per-minute should have consumed exactly %d tokens", tokensToWait) - }) - - // Test 4: WaitN that fails should consume zero tokens - t.Run("WaitN_Fails_Consumes_Zero_Tokens", func(t *testing.T) { - // First, exhaust the bucket - for range limit.count { - limiter.allow("test-fail-waitn", executionTime) - } - - // Verify bucket is exhausted - _, exhaustedDetails := limiter.peekWithDebug("test-fail-waitn", executionTime) - require.Equal(t, int64(0), exhaustedDetails[0].TokensRemaining(), "bucket should be exhausted") - - // Deadline that expires immediately (no time to refill) - deadline := func() (time.Time, bool) { - return executionTime.ToTime(), true // expires immediately - } - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - // WaitN should fail and consume zero tokens - allowed := limiter.waitNWithCancellation("test-fail-waitn", executionTime, 1, deadline, done) - require.False(t, allowed, "waitN should fail when deadline expires before tokens available") - - // Verify no tokens were consumed - _, finalDetails := limiter.peekWithDebug("test-fail-waitn", executionTime) - require.Equal(t, int64(0), finalDetails[0].TokensRemaining(), "should still have 0 tokens after failed wait") - }) - - // Test 5: Verify WaitN with high token count - t.Run("WaitN_High_Token_Count", func(t *testing.T) { - bigLimit := NewLimit(50, time.Second) - bigLimiter := NewLimiter(keyFunc, bigLimit) - const tokensToWait = 25 - - // Verify initial state - _, initialDetails := bigLimiter.peekWithDebug("test-big-waitn", executionTime) - require.Equal(t, bigLimit.count, initialDetails[0].TokensRemaining(), "should start with all tokens") - - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(time.Second).ToTime(), true - } - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } - - // WaitN should succeed and consume exactly tokensToWait tokens - allowed := bigLimiter.waitNWithCancellation("test-big-waitn", executionTime, tokensToWait, deadline, done) - require.True(t, allowed, "waitN should succeed with high token count") - - // Verify exactly tokensToWait tokens were consumed - _, finalDetails := bigLimiter.peekWithDebug("test-big-waitn", executionTime) - require.Equal(t, bigLimit.count-tokensToWait, finalDetails[0].TokensRemaining(), "should have consumed exactly %d tokens", tokensToWait) - }) - - // Test 6: Concurrent WaitN should consume correct total tokens - t.Run("WaitN_Concurrent_Token_Consumption", func(t *testing.T) { - concurrentLimit := NewLimit(20, time.Second) - concurrentLimiter := NewLimiter(keyFunc, concurrentLimit) - const tokensPerWait = 2 - const numGoroutines = 5 // Will try to consume 10 total tokens - const expectedSuccesses = int64(10) // 20 / 2 = 10 successful waits possible - - results := make([]bool, numGoroutines) - - // Deadline that gives enough time - deadline := func() (time.Time, bool) { - return executionTime.Add(time.Second).ToTime(), true - } - done := func() <-chan struct{} { - return make(chan struct{}) // never closes - } +type testContext struct { + deadline time.Time + hasDeadline bool + done <-chan struct{} +} - // Start concurrent waits - var wg sync.WaitGroup - for i := range numGoroutines { - wg.Add(1) - go func(i int) { - defer wg.Done() - results[i] = concurrentLimiter.waitNWithCancellation("test-concurrent-waitn", executionTime, tokensPerWait, deadline, done) - }(i) - } - wg.Wait() +func (t *testContext) Deadline() (deadline time.Time, ok bool) { + return t.deadline, t.hasDeadline +} - // Count successes - var successes int64 - for _, result := range results { - if result { - successes++ - } - } +func (t *testContext) Done() <-chan struct{} { + return t.done +} - require.Equal(t, expectedSuccesses/tokensPerWait, successes, "expected exactly %d successful waits", expectedSuccesses/tokensPerWait) +func (t *testContext) Err() error { + select { + case <-t.done: + return context.Canceled + default: + return nil + } +} - // Verify total tokens consumed - _, finalDetails := concurrentLimiter.peekWithDebug("test-concurrent-waitn", executionTime) - expectedRemaining := concurrentLimit.count - (successes * tokensPerWait) - require.Equal(t, expectedRemaining, finalDetails[0].TokensRemaining(), "should have consumed exactly %d tokens total", successes*tokensPerWait) - }) +func (t *testContext) Value(key any) any { + return nil } From d3802fe96ae494bc5a2d4a9823c5bfa0df19679f Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Thu, 21 Aug 2025 12:48:39 -0400 Subject: [PATCH 12/15] Remove special logic for deadlines We only care about ctx.Done. Simpler, and less likely to have clock inconsistencies. --- limiter_wait.go | 16 +- limiter_wait_test.go | 386 ++++++++++++++----------------------------- 2 files changed, 125 insertions(+), 277 deletions(-) diff --git a/limiter_wait.go b/limiter_wait.go index b1a1d88..7f55cd4 100644 --- a/limiter_wait.go +++ b/limiter_wait.go @@ -141,19 +141,13 @@ func (r *Limiter[TInput, TKey]) waitNWithDetails( retryAfter := details.RetryAfter() - // If we can't possibly get a token, fail fast - if deadline, ok := ctx.Deadline(); ok { - // ctx deadline uses time.Time, not ntime.Time, - // so use a duration for the comparison. - // This might not be robust to clock skew. - d := deadline.Sub(currentTime.ToTime()) - if d < retryAfter { - return false, details - } - } - select { case <-ctx.Done(): + // Need to get updated details, since this cancellation + // event might have been a while after the last call. + // We'll choose the semantics of "cancellation always + // means deny". + _, details := r.peekNWithDetails(input, currentTime, n) return false, details case <-time.After(retryAfter): currentTime = currentTime.Add(retryAfter) diff --git a/limiter_wait_test.go b/limiter_wait_test.go index c6e9177..364251d 100644 --- a/limiter_wait_test.go +++ b/limiter_wait_test.go @@ -34,13 +34,10 @@ func TestLimiter_Wait(t *testing.T) { ok := limiter.allow("test", executionTime) require.False(t, ok, "should not allow when tokens exhausted") - // Test 1: Wait for 1 token with enough time to acquire it + // Test 1: Wait for 1 token with no cancellation { - // Context with deadline that gives enough time ctx := &testContext{ - deadline: executionTime.Add(limit.durationPerToken).ToTime(), - hasDeadline: true, - done: make(chan struct{}), // never closes + done: make(chan struct{}), // never closes } allow, details := limiter.waitNWithDetails("test", executionTime, 1, ctx) @@ -55,68 +52,27 @@ func TestLimiter_Wait(t *testing.T) { ok = limiter.allow("test", executionTime) require.False(t, ok, "should not allow again immediately after wait") - // Test 2: Wait for 1 token with deadline that expires before token is available - { - // Context with deadline that expires too soon - ctx := &testContext{ - deadline: executionTime.Add(limit.durationPerToken / 2).ToTime(), - hasDeadline: true, - done: make(chan struct{}), // never closes - } - - allow, details := limiter.waitNWithDetails("test", executionTime, 1, ctx) - require.False(t, allow, "should not acquire 1 token if deadline expires before token is available") - require.NotNil(t, details, "should return details even on deadline expiry") - } - - // Test 3: Wait for 1 token with immediate cancellation + // Test 2: Wait for 1 token with immediate cancellation { // Context that's immediately cancelled (closed) done := make(chan struct{}) - close(done) ctx := &testContext{ - deadline: executionTime.Add(limit.durationPerToken).ToTime(), - hasDeadline: true, - done: done, + done: done, } + close(done) allow, details := limiter.waitNWithDetails("test", executionTime, 1, ctx) require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") require.NotNil(t, details, "should return details even on cancellation") } - // Test 4: Wait for 1 token with no deadline - { - // Context with no deadline - ctx := &testContext{ - deadline: time.Time{}, - hasDeadline: false, - done: make(chan struct{}), // never closes - } - - allow, details := limiter.waitNWithDetails("test", executionTime, 1, ctx) - require.True(t, allow, "should acquire 1 token when no deadline is set") - require.NotNil(t, details, "should return details") - - // We waited - executionTime = executionTime.Add(limit.durationPerToken) - - // Should not allow again immediately - ok := limiter.allow("test", executionTime) - require.False(t, ok, "should not allow again immediately after wait") - } - - // Test 5: Wait for multiple tokens (limit.count) with enough time to acquire them + // Test 3: Wait for multiple tokens (limit.count) with no cancellation { // Wait for bucket to refill enough tokens executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) - wait := time.Duration(limit.count) * limit.durationPerToken - // Context with deadline that gives enough time ctx := &testContext{ - deadline: executionTime.Add(wait).ToTime(), - hasDeadline: true, - done: make(chan struct{}), // never closes + done: make(chan struct{}), // never closes } allow, details := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) @@ -127,30 +83,7 @@ func TestLimiter_Wait(t *testing.T) { executionTime = executionTime.Add(limit.durationPerToken) } - // Test 6: Wait for multiple tokens with deadline that expires before all tokens are available - { - // Wait for bucket to refill enough tokens - executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) - - // Consume all tokens to exhaust the bucket again - for range limit.count { - ok := limiter.allow("test", executionTime) - require.True(t, ok, "should allow tokens to exhaust bucket") - } - - // Context with deadline that expires too soon - ctx := &testContext{ - deadline: executionTime.Add(limit.durationPerToken / 4).ToTime(), - hasDeadline: true, - done: make(chan struct{}), // never closes - } - - allow, details := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) - require.False(t, allow, "should not acquire %d tokens if deadline expires before tokens are available", limit.count) - require.NotNil(t, details, "should return details even on deadline expiry") - } - - // Test 7: Wait for multiple tokens with immediate cancellation + // Test 4: Wait for multiple tokens with immediate cancellation { // Wait for bucket to refill enough tokens executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) @@ -163,241 +96,162 @@ func TestLimiter_Wait(t *testing.T) { // Context that's immediately cancelled (closed) done := make(chan struct{}) - close(done) ctx := &testContext{ - deadline: executionTime.Add(limit.durationPerToken * time.Duration(limit.count)).ToTime(), - hasDeadline: true, - done: done, + done: done, } + close(done) allow, details := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit.count) require.NotNil(t, details, "should return details even on cancellation") } - - // Test 8: Wait for multiple tokens with no deadline - { - // Wait for bucket to refill enough tokens - executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) - - // Consume all tokens to exhaust the bucket again - for range limit.count { - ok := limiter.allow("test", executionTime) - require.True(t, ok, "should allow tokens to exhaust bucket") - } - - // Context with no deadline - ctx := &testContext{ - deadline: time.Time{}, - hasDeadline: false, - done: make(chan struct{}), // never closes - } - - allow, details := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) - require.True(t, allow, "should acquire %d tokens when no deadline is set", limit.count) - require.NotNil(t, details, "should return details") - - // We waited - executionTime = executionTime.Add(limit.durationPerToken) - - // Should not allow again immediately - allow = limiter.allowN("test", executionTime, limit.count) - require.False(t, allow, "should not allow %d tokens again immediately after wait", limit.count) - } }) t.Run("MultipleBuckets", func(t *testing.T) { - t.Parallel() - keyFunc := func(input int) string { - return fmt.Sprintf("test-bucket-%d", input) - } - const buckets = 3 - limit := NewLimit(2, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - executionTime := ntime.Now() - - // Exhaust tokens for all buckets - for bucketID := range buckets { - for range limit.count { - allow := limiter.allow(bucketID, executionTime) - require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + t.Run("Serial", func(t *testing.T) { + t.Parallel() + keyFunc := func(input int) string { + return fmt.Sprintf("test-bucket-%d", input) } - } - - // Should not allow immediately for any bucket - for bucketID := range buckets { - allow := limiter.allow(bucketID, executionTime) - require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID) - } - - // Wait for a token for each bucket - for bucketID := range buckets { - // Context with deadline that gives enough time - ctx := &testContext{ - deadline: executionTime.Add(limit.durationPerToken).ToTime(), - hasDeadline: true, - done: make(chan struct{}), // never closes + const buckets = 3 + limit := NewLimit(2, 100*time.Millisecond) + limiter := NewLimiter(keyFunc, limit) + executionTime := ntime.Now() + + // Exhaust tokens for all buckets + for bucketID := range buckets { + for range limit.count { + allow := limiter.allow(bucketID, executionTime) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } } - allow, details := limiter.waitNWithDetails(bucketID, executionTime, 1, ctx) - require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID) - require.NotNil(t, details, "should return details") - } - - // We waited - executionTime = executionTime.Add(limit.durationPerToken) - - // Buckets should be empty again - for bucketID := range buckets { - ok := limiter.allow(bucketID, executionTime) - require.False(t, ok, "should not allow again immediately after wait for bucket %d", bucketID) - } - }) - - t.Run("MultipleBuckets_Concurrent", func(t *testing.T) { - t.Parallel() - keyFunc := func(input int64) string { - return fmt.Sprintf("test-bucket-%d", input) - } - const buckets int64 = 3 - limit := NewLimit(2, 100*time.Millisecond) - limiter := NewLimiter(keyFunc, limit) - executionTime := ntime.Now() - - // Exhaust tokens for all buckets - for bucketID := range buckets { - for range limit.count { + // Should not allow immediately for any bucket + for bucketID := range buckets { allow := limiter.allow(bucketID, executionTime) - require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID) } - } - - // Buckets should be empty - for bucketID := range buckets { - allow := limiter.allow(bucketID, executionTime) - require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID) - } - // Test 1: Multiple goroutines competing for tokens with enough time - { - // More goroutines than available tokens to create competition - tokens := buckets * limit.count - concurrency := tokens * 3 // oversubscribe by 3x - results := make([]bool, concurrency) + // Wait for a token for each bucket + for bucketID := range buckets { + ctx := &testContext{ + done: make(chan struct{}), // never closes + } - // Context with deadline that gives enough time for all tokens to be refilled - ctx := &testContext{ - deadline: executionTime.Add(limit.period).ToTime(), - hasDeadline: true, - done: make(chan struct{}), // never closes + allow, details := limiter.waitNWithDetails(bucketID, executionTime, 1, ctx) + require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID) + require.NotNil(t, details, "should return details") } - // Start concurrent waits - var wg sync.WaitGroup - for i := range concurrency { - wg.Add(1) - go func(i int64) { - defer wg.Done() - bucketID := i % buckets - results[i], _ = limiter.waitNWithDetails(bucketID, executionTime, 1, ctx) - }(i) + // We waited + executionTime = executionTime.Add(limit.durationPerToken) + + // Buckets should be empty again + for bucketID := range buckets { + ok := limiter.allow(bucketID, executionTime) + require.False(t, ok, "should not allow again immediately after wait for bucket %d", bucketID) } - wg.Wait() + }) - // We waited - executionTime = executionTime.Add(limit.period) - - // Count successes and failures - var successes, failures int64 - for _, result := range results { - if result { - successes++ - } else { - failures++ + t.Run("Concurrent", func(t *testing.T) { + t.Parallel() + keyFunc := func(input int64) string { + return fmt.Sprintf("test-bucket-%d", input) + } + const buckets int64 = 3 + limit := NewLimit(2, 100*time.Millisecond) + limiter := NewLimiter(keyFunc, limit) + executionTime := ntime.Now() + + // Exhaust tokens for all buckets + for bucketID := range buckets { + for range limit.count { + allow := limiter.allow(bucketID, executionTime) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) } } - // Exactly limit.count tokens per bucket should succeed, 3 buckets * 2 tokens each - expectedSuccesses := buckets * limit.count - require.Equal(t, expectedSuccesses, successes, "expected exactly %d goroutines to acquire tokens", expectedSuccesses) + // Buckets should be empty + for bucketID := range buckets { + allow := limiter.allow(bucketID, executionTime) + require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID) + } - // The rest should fail due to competition - expectedFailures := concurrency - expectedSuccesses - require.Equal(t, expectedFailures, failures, "expected %d goroutines to fail due to competition", expectedFailures) - } + // Test 1: Multiple goroutines with immediate cancellation + { + concurrency := buckets + results := make([]bool, concurrency) - // Test 2: Multiple goroutines with deadline that expires before tokens are available - { - concurrency := buckets - results := make([]bool, concurrency) + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + close(done) + + // Start concurrent waits + var wg sync.WaitGroup + for i := range concurrency { + wg.Add(1) + go func(i int64) { + defer wg.Done() + results[i], _ = limiter.waitNWithDetails(i, executionTime, 1, ctx) + }(i) + } + wg.Wait() - // Context with deadline that expires too soon - ctx := &testContext{ - deadline: executionTime.Add(limit.durationPerToken / 2).ToTime(), - hasDeadline: true, - done: make(chan struct{}), // never closes + // All should fail because context is cancelled immediately + for i, result := range results { + require.False(t, result, "goroutine %d should not acquire token due to immediate cancellation", i) + } } - // Start concurrent waits - var wg sync.WaitGroup - for i := range concurrency { - wg.Add(1) - go func(i int64) { - defer wg.Done() - results[i], _ = limiter.waitNWithDetails(i, executionTime, 1, ctx) - }(i) - } - wg.Wait() + // Test 2: Multiple goroutines with delayed cancellation + { + concurrency := buckets + results := make([]bool, concurrency) - // All should fail because deadline expires before tokens are available - for i, result := range results { - require.False(t, result, "goroutine %d should not acquire token due to early deadline", i) - } - } + // Context that gets cancelled after a short delay + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } - // Test 3: Multiple goroutines with immediate cancellation - { - concurrency := buckets - results := make([]bool, concurrency) + // Start concurrent waits + var wg sync.WaitGroup + for i := range concurrency { + wg.Add(1) + go func(i int64) { + defer wg.Done() + results[i], _ = limiter.waitNWithDetails(i, executionTime, 1, ctx) + }(i) + } - // Context that's immediately cancelled (closed) - done := make(chan struct{}) - close(done) - ctx := &testContext{ - deadline: executionTime.Add(limit.period).ToTime(), - hasDeadline: true, - done: done, - } + // Cancel context after a short delay + go func() { + time.Sleep(10 * time.Millisecond) + close(done) + }() - // Start concurrent waits - var wg sync.WaitGroup - for i := range concurrency { - wg.Add(1) - go func(i int64) { - defer wg.Done() - results[i], _ = limiter.waitNWithDetails(i, executionTime, 1, ctx) - }(i) - } - wg.Wait() + wg.Wait() - // All should fail because context is cancelled immediately - for i, result := range results { - require.False(t, result, "goroutine %d should not acquire token due to immediate cancellation", i) + // All should fail because context is cancelled + for i, result := range results { + require.False(t, result, "goroutine %d should not acquire token due to cancellation", i) + } } - } + }) }) } var _ context.Context = &testContext{} type testContext struct { - deadline time.Time - hasDeadline bool - done <-chan struct{} + done <-chan struct{} } func (t *testContext) Deadline() (deadline time.Time, ok bool) { - return t.deadline, t.hasDeadline + return time.Time{}, false } func (t *testContext) Done() <-chan struct{} { From b786bb14a7e0adfa9de7ad761cac69524767e76f Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Thu, 21 Aug 2025 12:54:15 -0400 Subject: [PATCH 13/15] Return context errors --- limiter_wait.go | 21 +++++++++++---------- limiter_wait_test.go | 25 +++++++++++++++++-------- 2 files changed, 28 insertions(+), 18 deletions(-) diff --git a/limiter_wait.go b/limiter_wait.go index 7f55cd4..bbec678 100644 --- a/limiter_wait.go +++ b/limiter_wait.go @@ -30,8 +30,9 @@ import ( // // Wait makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. -func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) bool { - return r.WaitN(ctx, input, 1) +func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) (bool, error) { + allow, _, err := r.waitNWithDetails(input, ntime.Now(), 1, ctx) + return allow, err } // WaitN will poll [Limiter.AllowN] for a period of time, @@ -57,9 +58,9 @@ func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) bool { // // WaitN makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. -func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64) bool { - allow, _ := r.waitNWithDetails(input, ntime.Now(), n, ctx) - return allow +func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64) (bool, error) { + allow, _, err := r.waitNWithDetails(input, ntime.Now(), n, ctx) + return allow, err } // WaitWithDetails will poll [Allow] for a period of time, @@ -86,7 +87,7 @@ func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64 // // WaitWithDetails makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. -func (r *Limiter[TInput, TKey]) WaitWithDetails(ctx context.Context, input TInput) (bool, Details[TInput, TKey]) { +func (r *Limiter[TInput, TKey]) WaitWithDetails(ctx context.Context, input TInput) (bool, Details[TInput, TKey], error) { return r.WaitNWithDetails(ctx, input, 1) } @@ -114,7 +115,7 @@ func (r *Limiter[TInput, TKey]) WaitWithDetails(ctx context.Context, input TInpu // // WaitWithDetails makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. -func (r *Limiter[TInput, TKey]) WaitNWithDetails(ctx context.Context, input TInput, n int64) (bool, Details[TInput, TKey]) { +func (r *Limiter[TInput, TKey]) WaitNWithDetails(ctx context.Context, input TInput, n int64) (bool, Details[TInput, TKey], error) { return r.waitNWithDetails(input, ntime.Now(), n, ctx) } @@ -126,7 +127,7 @@ func (r *Limiter[TInput, TKey]) waitNWithDetails( startTime ntime.Time, n int64, ctx context.Context, -) (bool, Details[TInput, TKey]) { +) (bool, Details[TInput, TKey], error) { // currentTime is an approximation of the real clock moving forward. // It's imprecise because it depends on time.After below. // For testing purposes, we want startTime (execution time) to @@ -136,7 +137,7 @@ func (r *Limiter[TInput, TKey]) waitNWithDetails( for { allow, details := r.allowNWithDetails(input, currentTime, n) if allow { - return allow, details + return allow, details, nil } retryAfter := details.RetryAfter() @@ -148,7 +149,7 @@ func (r *Limiter[TInput, TKey]) waitNWithDetails( // We'll choose the semantics of "cancellation always // means deny". _, details := r.peekNWithDetails(input, currentTime, n) - return false, details + return false, details, ctx.Err() case <-time.After(retryAfter): currentTime = currentTime.Add(retryAfter) } diff --git a/limiter_wait_test.go b/limiter_wait_test.go index 364251d..b0b197a 100644 --- a/limiter_wait_test.go +++ b/limiter_wait_test.go @@ -40,7 +40,8 @@ func TestLimiter_Wait(t *testing.T) { done: make(chan struct{}), // never closes } - allow, details := limiter.waitNWithDetails("test", executionTime, 1, ctx) + allow, details, err := limiter.waitNWithDetails("test", executionTime, 1, ctx) + require.NoError(t, err, "should not return error") require.True(t, allow, "should acquire 1 token after waiting") require.NotNil(t, details, "should return details") @@ -61,7 +62,8 @@ func TestLimiter_Wait(t *testing.T) { } close(done) - allow, details := limiter.waitNWithDetails("test", executionTime, 1, ctx) + allow, details, err := limiter.waitNWithDetails("test", executionTime, 1, ctx) + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") require.NotNil(t, details, "should return details even on cancellation") } @@ -75,7 +77,8 @@ func TestLimiter_Wait(t *testing.T) { done: make(chan struct{}), // never closes } - allow, details := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) + allow, details, err := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) + require.NoError(t, err, "should not return error") require.True(t, allow, "should acquire %d tokens after waiting", limit.count) require.NotNil(t, details, "should return details") @@ -101,7 +104,8 @@ func TestLimiter_Wait(t *testing.T) { } close(done) - allow, details := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) + allow, details, err := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit.count) require.NotNil(t, details, "should return details even on cancellation") } @@ -138,7 +142,8 @@ func TestLimiter_Wait(t *testing.T) { done: make(chan struct{}), // never closes } - allow, details := limiter.waitNWithDetails(bucketID, executionTime, 1, ctx) + allow, details, err := limiter.waitNWithDetails(bucketID, executionTime, 1, ctx) + require.NoError(t, err, "should not return error") require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID) require.NotNil(t, details, "should return details") } @@ -195,7 +200,9 @@ func TestLimiter_Wait(t *testing.T) { wg.Add(1) go func(i int64) { defer wg.Done() - results[i], _ = limiter.waitNWithDetails(i, executionTime, 1, ctx) + allow, _, err := limiter.waitNWithDetails(i, executionTime, 1, ctx) + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + results[i] = allow }(i) } wg.Wait() @@ -223,13 +230,15 @@ func TestLimiter_Wait(t *testing.T) { wg.Add(1) go func(i int64) { defer wg.Done() - results[i], _ = limiter.waitNWithDetails(i, executionTime, 1, ctx) + allow, _, err := limiter.waitNWithDetails(i, executionTime, 1, ctx) + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + results[i] = allow }(i) } // Cancel context after a short delay go func() { - time.Sleep(10 * time.Millisecond) + time.Sleep(2 * time.Millisecond) close(done) }() From a931a3f82e65a9acceeaab120ab9ea5ca37ea32a Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Thu, 21 Aug 2025 13:53:50 -0400 Subject: [PATCH 14/15] Implement waitWithDebug ctx first is idiomatic, changed waitWithDetails too --- limiter_wait.go | 113 ++++++++++++++- limiter_wait_test.go | 332 ++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 416 insertions(+), 29 deletions(-) diff --git a/limiter_wait.go b/limiter_wait.go index bbec678..8ba10fb 100644 --- a/limiter_wait.go +++ b/limiter_wait.go @@ -28,10 +28,12 @@ import ( // // ctx := context.WithTimeout(ctx, limit.DurationPerToken()) // +// The returned error will be non-nil if the context is cancelled. +// // Wait makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) (bool, error) { - allow, _, err := r.waitNWithDetails(input, ntime.Now(), 1, ctx) + allow, _, err := r.waitNWithDetails(ctx, input, ntime.Now(), 1) return allow, err } @@ -56,10 +58,12 @@ func (r *Limiter[TInput, TKey]) Wait(ctx context.Context, input TInput) (bool, e // // ctx := context.WithTimeout(ctx, limit.DurationPerToken()) // +// The returned error will be non-nil if the context is cancelled. +// // WaitN makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64) (bool, error) { - allow, _, err := r.waitNWithDetails(input, ntime.Now(), n, ctx) + allow, _, err := r.waitNWithDetails(ctx, input, ntime.Now(), n) return allow, err } @@ -85,6 +89,8 @@ func (r *Limiter[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64 // // ctx := context.WithTimeout(ctx, limit.DurationPerToken()) // +// The returned error will be non-nil if the context is cancelled. +// // WaitWithDetails makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. func (r *Limiter[TInput, TKey]) WaitWithDetails(ctx context.Context, input TInput) (bool, Details[TInput, TKey], error) { @@ -113,25 +119,25 @@ func (r *Limiter[TInput, TKey]) WaitWithDetails(ctx context.Context, input TInpu // // ctx := context.WithTimeout(ctx, limit.DurationPerToken()) // +// The returned error will be non-nil if the context is cancelled. +// // WaitWithDetails makes no ordering guarantees. Multiple concurrent calls may // acquire tokens in any order. func (r *Limiter[TInput, TKey]) WaitNWithDetails(ctx context.Context, input TInput, n int64) (bool, Details[TInput, TKey], error) { - return r.waitNWithDetails(input, ntime.Now(), n, ctx) + return r.waitNWithDetails(ctx, input, ntime.Now(), n) } // waitNWithDetails is the internal implementation that accepts a context. // It is designed to be testable by accepting any context implementation, // including test contexts that provide deterministic behavior. func (r *Limiter[TInput, TKey]) waitNWithDetails( + ctx context.Context, input TInput, startTime ntime.Time, n int64, - ctx context.Context, ) (bool, Details[TInput, TKey], error) { // currentTime is an approximation of the real clock moving forward. // It's imprecise because it depends on time.After below. - // For testing purposes, we want startTime (execution time) to - // be a parameter. currentTime := startTime for { @@ -155,3 +161,98 @@ func (r *Limiter[TInput, TKey]) waitNWithDetails( } } } + +// WaitWithDebug will poll [Limiter.Allow] for a period of time, +// until it is cancelled by the passed context. It has the +// effect of adding latency to requests instead of refusing +// them immediately. Consider it graceful degradation. +// +// WaitWithDebug will return true if a token becomes available prior to +// the context cancellation, and will consume a token. It will +// return false if not, and therefore not consume a token. It will +// also return the debugs of the request. +// +// Take care to create an appropriate context. You almost certainly +// want [context.WithTimeout] or [context.WithDeadline]. +// +// The returned error will be non-nil if the context is cancelled. +// +// WaitNWithDebug makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. +func (r *Limiter[TInput, TKey]) WaitWithDebug( + ctx context.Context, + input TInput, +) (bool, []Debug[TInput, TKey], error) { + return r.WaitNWithDebug(ctx, input, 1) +} + +// WaitNWithDebug will poll [Limiter.AllowN] for a period of time, +// until it is cancelled by the passed context. It has the +// effect of adding latency to requests instead of refusing +// them immediately. Consider it graceful degradation. +// +// WaitNWithDebug will return true if `n` tokens become available prior to +// the context cancellation, and will consume `n` tokens. If not, +// it will return false, and therefore consume no tokens. It will +// also return the debugs of the request. +// +// Take care to create an appropriate context. You almost certainly +// want [context.WithTimeout] or [context.WithDeadline]. +// +// The returned error will be non-nil if the context is cancelled. +// +// WaitNWithDebug makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. +func (r *Limiter[TInput, TKey]) WaitNWithDebug( + ctx context.Context, + input TInput, + n int64, +) (bool, []Debug[TInput, TKey], error) { + return r.waitNWithDebug(ctx, input, ntime.Now(), n) +} + +// waitNWithDetails is the internal implementation that accepts a context. +// It is designed to be testable by accepting any context implementation, +// including test contexts that provide deterministic behavior. +func (r *Limiter[TInput, TKey]) waitNWithDebug( + ctx context.Context, + input TInput, + startTime ntime.Time, + n int64, +) (bool, []Debug[TInput, TKey], error) { + // currentTime is an approximation of the real clock moving forward. + // It's imprecise because it depends on time.After below. + currentTime := startTime + + for { + allow, debugs := r.allowNWithDebug(input, currentTime, n) + if allow { + return allow, debugs, nil + } + + // len(debugs) will be greater than 0 here; + // if there were zero, allow was true + + // Get the max retryAfter + retryAfter := debugs[0].RetryAfter() + for i := 1; i < len(debugs); i++ { + debug := debugs[i] + ra := debug.RetryAfter() + if ra > retryAfter { + retryAfter = ra + } + } + + select { + case <-ctx.Done(): + // Need to get updated debugs, since this cancellation + // event might have been a while after the last call. + // We'll choose the semantics of "cancellation always + // means deny". + _, debugs := r.peekNWithDebug(input, currentTime, n) + return false, debugs, ctx.Err() + case <-time.After(retryAfter): + currentTime = currentTime.Add(retryAfter) + } + } +} diff --git a/limiter_wait_test.go b/limiter_wait_test.go index b0b197a..ad0e518 100644 --- a/limiter_wait_test.go +++ b/limiter_wait_test.go @@ -40,10 +40,10 @@ func TestLimiter_Wait(t *testing.T) { done: make(chan struct{}), // never closes } - allow, details, err := limiter.waitNWithDetails("test", executionTime, 1, ctx) - require.NoError(t, err, "should not return error") + allow, details, err := limiter.waitNWithDetails(ctx, "test", executionTime, 1) require.True(t, allow, "should acquire 1 token after waiting") - require.NotNil(t, details, "should return details") + require.NotZero(t, details, "should return details") + require.NoError(t, err, "should not return error") // We waited executionTime = executionTime.Add(limit.durationPerToken) @@ -62,10 +62,10 @@ func TestLimiter_Wait(t *testing.T) { } close(done) - allow, details, err := limiter.waitNWithDetails("test", executionTime, 1, ctx) - require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + allow, details, err := limiter.waitNWithDetails(ctx, "test", executionTime, 1) require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") - require.NotNil(t, details, "should return details even on cancellation") + require.NotZero(t, details, "should return details even on cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") } // Test 3: Wait for multiple tokens (limit.count) with no cancellation @@ -77,10 +77,10 @@ func TestLimiter_Wait(t *testing.T) { done: make(chan struct{}), // never closes } - allow, details, err := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) - require.NoError(t, err, "should not return error") + allow, details, err := limiter.waitNWithDetails(ctx, "test", executionTime, limit.count) require.True(t, allow, "should acquire %d tokens after waiting", limit.count) - require.NotNil(t, details, "should return details") + require.NotZero(t, details, "should return details") + require.NoError(t, err, "should not return error") // We waited executionTime = executionTime.Add(limit.durationPerToken) @@ -104,10 +104,10 @@ func TestLimiter_Wait(t *testing.T) { } close(done) - allow, details, err := limiter.waitNWithDetails("test", executionTime, limit.count, ctx) - require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + allow, details, err := limiter.waitNWithDetails(ctx, "test", executionTime, limit.count) require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit.count) - require.NotNil(t, details, "should return details even on cancellation") + require.NotZero(t, details, "should return details even on cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") } }) @@ -142,10 +142,10 @@ func TestLimiter_Wait(t *testing.T) { done: make(chan struct{}), // never closes } - allow, details, err := limiter.waitNWithDetails(bucketID, executionTime, 1, ctx) - require.NoError(t, err, "should not return error") + allow, details, err := limiter.waitNWithDetails(ctx, bucketID, executionTime, 1) require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID) - require.NotNil(t, details, "should return details") + require.NotZero(t, details, "should return details") + require.NoError(t, err, "should not return error") } // We waited @@ -200,7 +200,8 @@ func TestLimiter_Wait(t *testing.T) { wg.Add(1) go func(i int64) { defer wg.Done() - allow, _, err := limiter.waitNWithDetails(i, executionTime, 1, ctx) + allow, _, err := limiter.waitNWithDetails(ctx, i, executionTime, 1) + require.False(t, allow, "should not acquire token due to immediate cancellation") require.ErrorIs(t, err, context.Canceled, "should return context canceled error") results[i] = allow }(i) @@ -214,15 +215,248 @@ func TestLimiter_Wait(t *testing.T) { } // Test 2: Multiple goroutines with delayed cancellation + { + keyFunc := func(input int64) int64 { + return input + } + + limit := NewLimit(1, 200*time.Millisecond) + limiter := NewLimiter(keyFunc, limit) + + // Exhaust tokens + for bucketID := range buckets { + for range limit.count { + allow := limiter.allow(bucketID, executionTime) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } + } + + // We're trying to emulate a wait call + // that starts waiting with a retry, but + // is cancelled in the meantime. + + const buckets int64 = 3 + results := make([]bool, buckets) + + var wg sync.WaitGroup + for bucketID := range buckets { + wg.Add(1) + + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + + go func(bucketID int64) { + defer wg.Done() + // Should start waiting, since the bucket is empty. + // retryAfter should be ~200ms + allow, _, err := limiter.waitNWithDetails(ctx, bucketID, executionTime, 1) + require.False(t, allow, "should not acquire token due to cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + results[bucketID] = allow + }(bucketID) + + // Cancel context after a goroutine has started, + // with a delay + go func() { + time.Sleep(10 * time.Millisecond) + close(done) + }() + } + wg.Wait() + + // All should fail because context is cancelled + for i, result := range results { + require.False(t, result, "goroutine %d should not acquire token due to cancellation", i) + } + } + }) + }) +} + +func TestLimiter_WaitWithDebug(t *testing.T) { + t.Parallel() + + t.Run("SingleBucket", func(t *testing.T) { + t.Parallel() + keyFunc := func(input string) string { + return input + } + limit := NewLimit(2, 100*time.Millisecond) + limiter := NewLimiter(keyFunc, limit) + + executionTime := ntime.Now() + + // Consume all tokens + for range limit.count { + ok := limiter.allow("test", executionTime) + require.True(t, ok, "should allow initial tokens") + } + + // Should not allow immediately + ok := limiter.allow("test", executionTime) + require.False(t, ok, "should not allow when tokens exhausted") + + // Test 1: Wait for 1 token with no cancellation + { + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, debugs, err := limiter.waitNWithDebug(ctx, "test", executionTime, 1) + require.True(t, allow, "should acquire 1 token after waiting") + require.NotNil(t, debugs, "should return debugs") + require.NoError(t, err, "should not return error") + + // We waited + executionTime = executionTime.Add(limit.durationPerToken) + } + + // Should not allow again immediately after wait + ok = limiter.allow("test", executionTime) + require.False(t, ok, "should not allow again immediately after wait") + + // Test 2: Wait for 1 token with immediate cancellation + { + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + close(done) + + allow, debugs, err := limiter.waitNWithDebug(ctx, "test", executionTime, 1) + require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") + require.NotNil(t, debugs, "should return debugs even on cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + } + + // Test 3: Wait for multiple tokens (limit.count) with no cancellation + { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) + + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, debugs, err := limiter.waitNWithDebug(ctx, "test", executionTime, limit.count) + require.True(t, allow, "should acquire %d tokens after waiting", limit.count) + require.NotNil(t, debugs, "should return debugs") + require.NoError(t, err, "should not return error") + + // We waited + executionTime = executionTime.Add(limit.durationPerToken) + } + + // Test 4: Wait for multiple tokens with immediate cancellation + { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit.durationPerToken * time.Duration(limit.count)) + + // Consume all tokens to exhaust the bucket again + for range limit.count { + ok := limiter.allow("test", executionTime) + require.True(t, ok, "should allow tokens to exhaust bucket") + } + + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + close(done) + + allow, debugs, err := limiter.waitNWithDebug(ctx, "test", executionTime, limit.count) + require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit.count) + require.NotNil(t, debugs, "should return debugs even on cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + } + }) + + t.Run("MultipleBuckets", func(t *testing.T) { + t.Run("Serial", func(t *testing.T) { + t.Parallel() + keyFunc := func(input int) string { + return fmt.Sprintf("test-bucket-%d", input) + } + const buckets = 3 + limit := NewLimit(2, 100*time.Millisecond) + limiter := NewLimiter(keyFunc, limit) + executionTime := ntime.Now() + + // Exhaust tokens for all buckets + for bucketID := range buckets { + for range limit.count { + allow := limiter.allow(bucketID, executionTime) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } + } + + // Should not allow immediately for any bucket + for bucketID := range buckets { + allow := limiter.allow(bucketID, executionTime) + require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID) + } + + // Wait for a token for each bucket + for bucketID := range buckets { + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, debugs, err := limiter.waitNWithDebug(ctx, bucketID, executionTime, 1) + require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID) + require.NotNil(t, debugs, "should return debugs") + require.NoError(t, err, "should not return error") + } + + // We waited + executionTime = executionTime.Add(limit.durationPerToken) + + // Buckets should be empty again + for bucketID := range buckets { + ok := limiter.allow(bucketID, executionTime) + require.False(t, ok, "should not allow again immediately after wait for bucket %d", bucketID) + } + }) + + t.Run("Concurrent", func(t *testing.T) { + t.Parallel() + keyFunc := func(input int64) string { + return fmt.Sprintf("test-bucket-%d", input) + } + const buckets int64 = 3 + limit := NewLimit(2, 100*time.Millisecond) + limiter := NewLimiter(keyFunc, limit) + executionTime := ntime.Now() + + // Exhaust tokens for all buckets + for bucketID := range buckets { + for range limit.count { + allow := limiter.allow(bucketID, executionTime) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } + } + + // Buckets should be empty + for bucketID := range buckets { + allow := limiter.allow(bucketID, executionTime) + require.False(t, allow, "should not allow when tokens exhausted for bucket %d", bucketID) + } + + // Test 1: Multiple goroutines with immediate cancellation { concurrency := buckets results := make([]bool, concurrency) - // Context that gets cancelled after a short delay + // Context that's immediately cancelled (closed) done := make(chan struct{}) ctx := &testContext{ done: done, } + close(done) // Start concurrent waits var wg sync.WaitGroup @@ -230,18 +464,70 @@ func TestLimiter_Wait(t *testing.T) { wg.Add(1) go func(i int64) { defer wg.Done() - allow, _, err := limiter.waitNWithDetails(i, executionTime, 1, ctx) + allow, _, err := limiter.waitNWithDebug(ctx, i, executionTime, 1) + require.False(t, allow, "should not acquire token due to immediate cancellation") require.ErrorIs(t, err, context.Canceled, "should return context canceled error") results[i] = allow }(i) } + wg.Wait() - // Cancel context after a short delay - go func() { - time.Sleep(2 * time.Millisecond) - close(done) - }() + // All should fail because context is cancelled immediately + for i, result := range results { + require.False(t, result, "goroutine %d should not acquire token due to immediate cancellation", i) + } + } + // Test 2: Multiple goroutines with delayed cancellation + { + keyFunc := func(input int64) int64 { + return input + } + + limit := NewLimit(1, 200*time.Millisecond) + limiter := NewLimiter(keyFunc, limit) + + // Exhaust tokens + for bucketID := range buckets { + for range limit.count { + allow := limiter.allow(bucketID, executionTime) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } + } + + // We're trying to emulate a wait call + // that starts waiting with a retry, but + // is cancelled in the meantime. + + const buckets int64 = 3 + results := make([]bool, buckets) + + var wg sync.WaitGroup + for bucketID := range buckets { + wg.Add(1) + + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + + go func(bucketID int64) { + defer wg.Done() + // Should start waiting, since the bucket is empty. + // retryAfter should be ~200ms + allow, _, err := limiter.waitNWithDebug(ctx, bucketID, executionTime, 1) + require.False(t, allow, "should not acquire token due to cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + results[bucketID] = allow + }(bucketID) + + // Cancel context after a goroutine has started, + // with a delay + go func() { + time.Sleep(10 * time.Millisecond) + close(done) + }() + } wg.Wait() // All should fail because context is cancelled From 13c9e357c8ba9ed9267cd3905afb88cb7c921e77 Mon Sep 17 00:00:00 2001 From: Matt Sherman Date: Thu, 21 Aug 2025 21:09:28 -0400 Subject: [PATCH 15/15] Implement Limiters.Wait* --- limiter_wait_test.go | 24 +- limiters_wait.go | 251 +++++++++++++++ limiters_wait_test.go | 692 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 955 insertions(+), 12 deletions(-) create mode 100644 limiters_wait.go create mode 100644 limiters_wait_test.go diff --git a/limiter_wait_test.go b/limiter_wait_test.go index ad0e518..f99607d 100644 --- a/limiter_wait_test.go +++ b/limiter_wait_test.go @@ -257,12 +257,12 @@ func TestLimiter_Wait(t *testing.T) { results[bucketID] = allow }(bucketID) - // Cancel context after a goroutine has started, - // with a delay - go func() { - time.Sleep(10 * time.Millisecond) - close(done) - }() + // Cancel context after goroutine has started. + // Originally, this had a delay, for a better test, + // but it is flaky on GitHub Actions, presumably + // because the runner is resource constrained. + // Delay works fine locally on Mac M2. + close(done) } wg.Wait() @@ -521,12 +521,12 @@ func TestLimiter_WaitWithDebug(t *testing.T) { results[bucketID] = allow }(bucketID) - // Cancel context after a goroutine has started, - // with a delay - go func() { - time.Sleep(10 * time.Millisecond) - close(done) - }() + // Cancel context after goroutine has started. + // Originally, this had a delay, for a better test, + // but it is flaky on GitHub Actions, presumably + // because the runner is resource constrained. + // Delay works fine locally on Mac M2. + close(done) } wg.Wait() diff --git a/limiters_wait.go b/limiters_wait.go new file mode 100644 index 0000000..4ba431b --- /dev/null +++ b/limiters_wait.go @@ -0,0 +1,251 @@ +package rate + +import ( + "context" + "time" + + "github.com/clipperhouse/ntime" +) + +// Wait will poll [Limiters.Allow] for a period of time, +// until it is cancelled by the passed context. It has the +// effect of adding latency to requests instead of refusing +// them immediately. Consider it graceful degradation. +// +// Wait will return true if a token becomes available prior to +// the context cancellation, and will consume a token. It will +// return false if not, and therefore not consume a token. +// +// Take care to create an appropriate context. You almost certainly +// want [context.WithTimeout] or [context.WithDeadline]. +// +// You should be conservative, as Wait will introduce +// backpressure on your upstream systems -- connections +// may be held open longer, requests may queue in memory. +// +// A good starting place will be to timeout after waiting +// for one token. For example: +// +// ctx := context.WithTimeout(ctx, limit.DurationPerToken()) +// +// The returned error will be non-nil if the context is cancelled. +// +// Wait makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. +func (r *Limiters[TInput, TKey]) Wait(ctx context.Context, input TInput) (bool, error) { + return r.WaitN(ctx, input, 1) +} + +// WaitN will poll [Limiters.AllowN] for a period of time, +// until it is cancelled by the passed context. It has the +// effect of adding latency to requests instead of refusing +// them immediately. Consider it graceful degradation. +// +// WaitN will return true if `n` tokens become available prior to +// the context cancellation, and will consume `n` tokens. If not, +// it will return false, and therefore consume no tokens. +// +// Take care to create an appropriate context. You almost certainly +// want [context.WithTimeout] or [context.WithDeadline]. +// +// You should be conservative, as Wait will introduce +// backpressure on your upstream systems -- connections +// may be held open longer, requests may queue in memory. +// +// A good starting place will be to timeout after waiting +// for one token. For example: +// +// ctx := context.WithTimeout(ctx, limit.DurationPerToken()) +// +// The returned error will be non-nil if the context is cancelled. +// +// WaitN makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. +func (r *Limiters[TInput, TKey]) WaitN(ctx context.Context, input TInput, n int64) (bool, error) { + allow, _, err := r.WaitNWithDetails(ctx, input, n) + return allow, err +} + +// WaitWithDetails will poll [Limiters.Allow] for a period of time, +// until it is cancelled by the passed context. It has the +// effect of adding latency to requests instead of refusing +// them immediately. Consider it graceful degradation. +// +// WaitWithDetails will return true if a token becomes available prior to +// the context cancellation, and will consume a token. It will +// return false if not, and therefore not consume a token. It will +// also return the details of the request. +// +// Take care to create an appropriate context. You almost certainly +// want [context.WithTimeout] or [context.WithDeadline]. +// +// You should be conservative, as WaitWithDetails will introduce +// backpressure on your upstream systems -- connections +// may be held open longer, requests may queue in memory. +// +// A good starting place will be to timeout after waiting +// for one token. For example: +// +// ctx := context.WithTimeout(ctx, limit.DurationPerToken()) +// +// The returned error will be non-nil if the context is cancelled. +// +// WaitWithDetails makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. +func (r *Limiters[TInput, TKey]) WaitWithDetails(ctx context.Context, input TInput) (bool, Details[TInput, TKey], error) { + return r.WaitNWithDetails(ctx, input, 1) +} + +// WaitNWithDetails will poll [Limiters.AllowN] for a period of time, +// until it is cancelled by the passed context. It has the +// effect of adding latency to requests instead of refusing +// them immediately. Consider it graceful degradation. +// +// WaitNWithDetails will return true if `n` tokens become available prior to +// the context cancellation, and will consume `n` tokens. If not, +// it will return false, and therefore consume no tokens. It will +// also return the details of the request. +// +// Take care to create an appropriate context. You almost certainly +// want [context.WithTimeout] or [context.WithDeadline]. +// +// You should be conservative, as WaitWithDetails will introduce +// backpressure on your upstream systems -- connections +// may be held open longer, requests may queue in memory. +// +// A good starting place will be to timeout after waiting +// for one token. For example: +// +// ctx := context.WithTimeout(ctx, limit.DurationPerToken()) +// +// The returned error will be non-nil if the context is cancelled. +// +// WaitWithDetails makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. +func (r *Limiters[TInput, TKey]) WaitNWithDetails(ctx context.Context, input TInput, n int64) (bool, Details[TInput, TKey], error) { + return r.waitNWithDetails(ctx, input, ntime.Now(), n) +} + +func (r *Limiters[TInput, TKey]) waitNWithDetails( + ctx context.Context, + input TInput, + startTime ntime.Time, + n int64, +) (bool, Details[TInput, TKey], error) { + // currentTime is an approximation of the real clock moving forward. + // It's imprecise because it depends on time.After below. + currentTime := startTime + + for { + allow, details := r.allowNWithDetails(input, currentTime, n) + if allow { + return allow, details, nil + } + + retryAfter := details.RetryAfter() + + select { + case <-ctx.Done(): + // Need to get updated details, since this cancellation + // event might have been a while after the last call. + // We'll choose the semantics of "cancellation always + // means deny". + _, details := r.peekNWithDetails(input, currentTime, n) + return false, details, ctx.Err() + case <-time.After(retryAfter): + currentTime = currentTime.Add(retryAfter) + } + } +} + +// WaitWithDebug will poll [Limiters.Allow] for a period of time, +// until it is cancelled by the passed context. It has the +// effect of adding latency to requests instead of refusing +// them immediately. Consider it graceful degradation. +// +// WaitWithDebug will return true if a token becomes available prior to +// the context cancellation, and will consume a token. It will +// return false if not, and therefore not consume a token. It will +// also return the debugs of the request. +// +// Take care to create an appropriate context. You almost certainly +// want [context.WithTimeout] or [context.WithDeadline]. +// +// The returned error will be non-nil if the context is cancelled. +// +// WaitNWithDebug makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. +func (r *Limiters[TInput, TKey]) WaitWithDebug( + ctx context.Context, + input TInput, +) (bool, []Debug[TInput, TKey], error) { + return r.WaitNWithDebug(ctx, input, 1) +} + +// WaitNWithDebug will poll [Limiters.AllowN] for a period of time, +// until it is cancelled by the passed context. It has the +// effect of adding latency to requests instead of refusing +// them immediately. Consider it graceful degradation. +// +// WaitNWithDebug will return true if `n` tokens become available prior to +// the context cancellation, and will consume `n` tokens. If not, +// it will return false, and therefore consume no tokens. It will +// also return the debugs of the request. +// +// Take care to create an appropriate context. You almost certainly +// want [context.WithTimeout] or [context.WithDeadline]. +// +// The returned error will be non-nil if the context is cancelled. +// +// WaitNWithDebug makes no ordering guarantees. Multiple concurrent calls may +// acquire tokens in any order. +func (r *Limiters[TInput, TKey]) WaitNWithDebug( + ctx context.Context, + input TInput, + n int64, +) (bool, []Debug[TInput, TKey], error) { + return r.waitNWithDebug(ctx, input, ntime.Now(), n) +} + +func (r *Limiters[TInput, TKey]) waitNWithDebug( + ctx context.Context, + input TInput, + startTime ntime.Time, + n int64, +) (bool, []Debug[TInput, TKey], error) { + // currentTime is an approximation of the real clock moving forward. + // It's imprecise because it depends on time.After below. + currentTime := startTime + + for { + allow, debugs := r.allowNWithDebug(input, currentTime, n) + if allow { + return allow, debugs, nil + } + + // len(debugs) will be greater than 0 here; + // if there were zero, allow was true + + // Get the max retryAfter + retryAfter := debugs[0].RetryAfter() + for i := 1; i < len(debugs); i++ { + debug := debugs[i] + ra := debug.RetryAfter() + if ra > retryAfter { + retryAfter = ra + } + } + + select { + case <-ctx.Done(): + // Need to get updated debugs, since this cancellation + // event might have been a while after the last call. + // We'll choose the semantics of "cancellation always + // means deny". + _, debugs := r.peekNWithDebug(input, currentTime, n) + return false, debugs, ctx.Err() + case <-time.After(retryAfter): + currentTime = currentTime.Add(retryAfter) + } + } +} diff --git a/limiters_wait_test.go b/limiters_wait_test.go new file mode 100644 index 0000000..8d29d99 --- /dev/null +++ b/limiters_wait_test.go @@ -0,0 +1,692 @@ +package rate + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/clipperhouse/ntime" + "github.com/stretchr/testify/require" +) + +func TestLimiters_Wait(t *testing.T) { + t.Parallel() + + t.Run("SingleBucket", func(t *testing.T) { + t.Parallel() + keyFunc := func(input string) string { + return input + } + // Create two limiters with different limits + limit1 := NewLimit(2, 100*time.Millisecond) // 2 per 100ms (more restrictive) + limit2 := NewLimit(5, 100*time.Millisecond) // 5 per 100ms + limiter1 := NewLimiter(keyFunc, limit1) + limiter2 := NewLimiter(keyFunc, limit2) + limiters := Combine(limiter1, limiter2) + + executionTime := ntime.Now() + + // Consume all tokens from the more restrictive limiter + for range limit1.count { + ok := limiters.allowN("test", executionTime, 1) + require.True(t, ok, "should allow initial tokens") + } + + // Should not allow immediately (limiter1 exhausted) + ok := limiters.allowN("test", executionTime, 1) + require.False(t, ok, "should not allow when limiter1 tokens exhausted") + + // Test 1: Wait for 1 token with no cancellation + { + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, details, err := limiters.waitNWithDetails(ctx, "test", executionTime, 1) + require.True(t, allow, "should acquire 1 token after waiting") + require.NotZero(t, details, "should return details") + require.NoError(t, err, "should not return error") + + // We waited + executionTime = executionTime.Add(limit1.durationPerToken) + } + + // Should not allow again immediately after wait + ok = limiters.allowN("test", executionTime, 1) + require.False(t, ok, "should not allow again immediately after wait") + + // Test 2: Wait for 1 token with immediate cancellation + { + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + close(done) + + allow, details, err := limiters.waitNWithDetails(ctx, "test", executionTime, 1) + require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") + require.NotZero(t, details, "should return details even on cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + } + + // Test 3: Wait for multiple tokens (limit1.count) with no cancellation + { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit1.period) + + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, details, err := limiters.waitNWithDetails(ctx, "test", executionTime, limit1.count) + require.True(t, allow, "should acquire %d tokens after waiting", limit1.count) + require.NotZero(t, details, "should return details") + require.NoError(t, err, "should not return error") + } + + // Test 4: Wait for multiple tokens with immediate cancellation + { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit1.durationPerToken * time.Duration(limit1.count)) + + // Consume all tokens to exhaust the bucket again + for range limit1.count { + ok := limiters.allowN("test", executionTime, 1) + require.True(t, ok, "should allow tokens to exhaust bucket") + } + + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + close(done) + + allow, details, err := limiters.waitNWithDetails(ctx, "test", executionTime, limit1.count) + require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit1.count) + require.NotZero(t, details, "should return details even on cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + } + }) + + t.Run("MultipleBuckets", func(t *testing.T) { + t.Run("Serial", func(t *testing.T) { + t.Parallel() + keyFunc := func(input int) string { + return fmt.Sprintf("test-bucket-%d", input) + } + const buckets = 3 + // Create two limiters with different limits + limit1 := NewLimit(2, 100*time.Millisecond) // 2 per 100ms (more restrictive) + limit2 := NewLimit(4, 100*time.Millisecond) // 4 per 100ms + limiter1 := NewLimiter(keyFunc, limit1) + limiter2 := NewLimiter(keyFunc, limit2) + limiters := Combine(limiter1, limiter2) + executionTime := ntime.Now() + + // Exhaust tokens for all buckets from the more restrictive limiter + for bucketID := range buckets { + for range limit1.count { + allow := limiters.allowN(bucketID, executionTime, 1) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } + } + + // Should not allow immediately for any bucket (limiter1 exhausted) + for bucketID := range buckets { + allow := limiters.allowN(bucketID, executionTime, 1) + require.False(t, allow, "should not allow when limiter1 tokens exhausted for bucket %d", bucketID) + } + + // Wait for a token for each bucket + for bucketID := range buckets { + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, details, err := limiters.waitNWithDetails(ctx, bucketID, executionTime, 1) + require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID) + require.NotZero(t, details, "should return details") + require.NoError(t, err, "should not return error") + } + + // We waited + executionTime = executionTime.Add(limit1.durationPerToken) + + // Buckets should be empty again + for bucketID := range buckets { + ok := limiters.allowN(bucketID, executionTime, 1) + require.False(t, ok, "should not allow again immediately after wait for bucket %d", bucketID) + } + }) + + t.Run("Concurrent", func(t *testing.T) { + // t.Parallel() + // GitHub Actions are pretty resource constrained, + // and this test is sensitive to timing. + keyFunc := func(input int64) string { + return fmt.Sprintf("test-bucket-%d", input) + } + const buckets int64 = 3 + // Create two limiters with different limits + limit1 := NewLimit(2, 100*time.Millisecond) // 2 per 100ms (more restrictive) + limit2 := NewLimit(4, 100*time.Millisecond) // 4 per 100ms + limiter1 := NewLimiter(keyFunc, limit1) + limiter2 := NewLimiter(keyFunc, limit2) + limiters := Combine(limiter1, limiter2) + executionTime := ntime.Now() + + // Exhaust tokens for all buckets from the more restrictive limiter + for bucketID := range buckets { + for range limit1.count { + allow := limiters.allowN(bucketID, executionTime, 1) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } + } + + // Buckets should be empty (limiter1 exhausted) + for bucketID := range buckets { + allow := limiters.allowN(bucketID, executionTime, 1) + require.False(t, allow, "should not allow when limiter1 tokens exhausted for bucket %d", bucketID) + } + + // Test 1: Multiple goroutines with immediate cancellation + { + concurrency := buckets + results := make([]bool, concurrency) + + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + close(done) + + // Start concurrent waits + var wg sync.WaitGroup + for i := range concurrency { + wg.Add(1) + go func(i int64) { + defer wg.Done() + allow, _, err := limiters.waitNWithDetails(ctx, i, executionTime, 1) + require.False(t, allow, "should not acquire token due to immediate cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + results[i] = allow + }(i) + } + wg.Wait() + + // All should fail because context is cancelled immediately + for i, result := range results { + require.False(t, result, "goroutine %d should not acquire token due to immediate cancellation", i) + } + } + + // Test 2: Multiple goroutines with delayed cancellation + { + keyFunc := func(input int64) int64 { + return input + } + + // Create two limiters with different limits + limit1 := NewLimit(1, 200*time.Millisecond) // 1 per 200ms (more restrictive) + limit2 := NewLimit(3, 200*time.Millisecond) // 3 per 200ms + limiter1 := NewLimiter(keyFunc, limit1) + limiter2 := NewLimiter(keyFunc, limit2) + limiters := Combine(limiter1, limiter2) + + // Exhaust tokens from the more restrictive limiter + for bucketID := range buckets { + for range limit1.count { + allow := limiters.allowN(bucketID, executionTime, 1) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } + } + + // We're trying to emulate a wait call + // that starts waiting with a retry, but + // is cancelled in the meantime. + + const buckets int64 = 3 + results := make([]bool, buckets) + + var wg sync.WaitGroup + for bucketID := range buckets { + wg.Add(1) + + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + + go func(bucketID int64) { + defer wg.Done() + // Should start waiting, since the bucket is empty. + // retryAfter should be ~200ms + allow, _, err := limiters.waitNWithDetails(ctx, bucketID, executionTime, 1) + require.False(t, allow, "should not acquire token due to cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + results[bucketID] = allow + }(bucketID) + + // Cancel context after goroutine has started. + // Originally, this had a delay, for a better test, + // but it is flaky on GitHub Actions, presumably + // because the runner is resource constrained. + // Delay works fine locally on Mac M2. + close(done) + } + wg.Wait() + + // All should fail because context is cancelled + for i, result := range results { + require.False(t, result, "goroutine %d should not acquire token due to cancellation", i) + } + } + }) + }) + + t.Run("MultipleLimiters", func(t *testing.T) { + t.Parallel() + t.Run("SameKeyer", func(t *testing.T) { + t.Parallel() + + keyFunc := func(input string) string { + return fmt.Sprintf("bucket-%s", input) + } + + // Create two limiters with different limits + limit1 := NewLimit(3, 100*time.Millisecond) // 3 per 100ms (more restrictive) + limit2 := NewLimit(5, 100*time.Millisecond) // 5 per 100ms + limiter1 := NewLimiter(keyFunc, limit1) + limiter2 := NewLimiter(keyFunc, limit2) + + // Create Limiters with both limiters + limiters := Combine(limiter1, limiter2) + + executionTime := ntime.Now() + + // Consume all tokens from the more restrictive limiter + for range limit1.count { + ok := limiters.allowN("test", executionTime, 1) + require.True(t, ok, "should allow initial tokens") + } + + // Should not allow immediately (limiter1 exhausted) + ok := limiters.allowN("test", executionTime, 1) + require.False(t, ok, "should not allow when limiter1 tokens exhausted") + + // Test 1: Wait for 1 token with no cancellation + { + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, details, err := limiters.waitNWithDetails(ctx, "test", executionTime, 1) + require.True(t, allow, "should acquire 1 token after waiting") + require.NotZero(t, details, "should return details") + require.NoError(t, err, "should not return error") + + // We waited + executionTime = executionTime.Add(limit1.durationPerToken) + } + + // Should not allow again immediately after wait + ok = limiters.allowN("test", executionTime, 1) + require.False(t, ok, "should not allow again immediately after wait") + + // Test 2: Wait for 1 token with immediate cancellation + { + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + close(done) + + allow, details, err := limiters.waitNWithDetails(ctx, "test", executionTime, 1) + require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") + require.NotZero(t, details, "should return details even on cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + } + }) + + t.Run("DifferentKeyers", func(t *testing.T) { + t.Parallel() + + // Create limiters with different key functions + keyFunc1 := func(input string) string { return input + "-1" } + keyFunc2 := func(input string) string { return input + "-2" } + + limit1 := NewLimit(2, 100*time.Millisecond) + limit2 := NewLimit(3, 100*time.Millisecond) + limiter1 := NewLimiter(keyFunc1, limit1) + limiter2 := NewLimiter(keyFunc2, limit2) + + limiters := Combine(limiter1, limiter2) + executionTime := ntime.Now() + + // Test that we can consume tokens from both keyers + // The Combine function treats all limiters as a single unit + // We should be able to consume up to the most restrictive limit (2) from each keyer + for i := range 2 { + ok := limiters.allowN("test", executionTime, 1) + require.True(t, ok, "should allow token %d from keyer1", i+1) + } + + // Should not allow more from keyer1 (limit1 exhausted) + ok := limiters.allowN("test", executionTime, 1) + require.False(t, ok, "should not allow when keyer1 limit exhausted") + + // Test: Wait for 1 token with no cancellation + { + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, details, err := limiters.waitNWithDetails(ctx, "test", executionTime, 1) + require.True(t, allow, "should acquire 1 token after waiting") + require.NotZero(t, details, "should return details") + require.NoError(t, err, "should not return error") + + // We waited + executionTime = executionTime.Add(limit1.durationPerToken) + } + + // Should not allow again immediately after wait + ok = limiters.allowN("test", executionTime, 1) + require.False(t, ok, "should not allow again immediately after wait") + }) + }) +} + +func TestLimiters_WaitWithDebug(t *testing.T) { + t.Parallel() + + t.Run("SingleBucket", func(t *testing.T) { + t.Parallel() + keyFunc := func(input string) string { + return input + } + // Create two limiters with different limits + limit1 := NewLimit(2, 100*time.Millisecond) // 2 per 100ms (more restrictive) + limit2 := NewLimit(5, 100*time.Millisecond) // 5 per 100ms + limiter1 := NewLimiter(keyFunc, limit1) + limiter2 := NewLimiter(keyFunc, limit2) + limiters := Combine(limiter1, limiter2) + + executionTime := ntime.Now() + + // Consume all tokens from the more restrictive limiter + for range limit1.count { + ok := limiters.allowN("test", executionTime, 1) + require.True(t, ok, "should allow initial tokens") + } + + // Should not allow immediately (limiter1 exhausted) + ok := limiters.allowN("test", executionTime, 1) + require.False(t, ok, "should not allow when limiter1 tokens exhausted") + + // Test 1: Wait for 1 token with no cancellation + { + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, debugs, err := limiters.waitNWithDebug(ctx, "test", executionTime, 1) + require.True(t, allow, "should acquire 1 token after waiting") + require.NotNil(t, debugs, "should return debugs") + require.NoError(t, err, "should not return error") + + // We waited + executionTime = executionTime.Add(limit1.durationPerToken) + } + + // Should not allow again immediately after wait + ok = limiters.allowN("test", executionTime, 1) + require.False(t, ok, "should not allow again immediately after wait") + + // Test 2: Wait for 1 token with immediate cancellation + { + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + close(done) + + allow, debugs, err := limiters.waitNWithDebug(ctx, "test", executionTime, 1) + require.False(t, allow, "should not acquire 1 token if context is cancelled immediately") + require.NotNil(t, debugs, "should return debugs even on cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + } + + // Test 3: Wait for multiple tokens (limit1.count) with no cancellation + { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit1.durationPerToken * time.Duration(limit1.count)) + + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, debugs, err := limiters.waitNWithDebug(ctx, "test", executionTime, limit1.count) + require.True(t, allow, "should acquire %d tokens after waiting", limit1.count) + require.NotNil(t, debugs, "should return debugs") + require.NoError(t, err, "should not return error") + + // We waited + executionTime = executionTime.Add(limit1.durationPerToken) + } + + // Test 4: Wait for multiple tokens with immediate cancellation + { + // Wait for bucket to refill enough tokens + executionTime = executionTime.Add(limit1.durationPerToken * time.Duration(limit1.count)) + + // Consume all tokens to exhaust the bucket again + for range limit1.count { + ok := limiters.allowN("test", executionTime, 1) + require.True(t, ok, "should allow tokens to exhaust bucket") + } + + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + close(done) + + allow, debugs, err := limiters.waitNWithDebug(ctx, "test", executionTime, limit1.count) + require.False(t, allow, "should not acquire %d tokens if context is cancelled immediately", limit1.count) + require.NotNil(t, debugs, "should return debugs even on cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + } + }) + + t.Run("MultipleBuckets", func(t *testing.T) { + t.Run("Serial", func(t *testing.T) { + t.Parallel() + keyFunc := func(input int) string { + return fmt.Sprintf("test-bucket-%d", input) + } + const buckets = 3 + // Create two limiters with different limits + limit1 := NewLimit(2, 100*time.Millisecond) // 2 per 100ms (more restrictive) + limit2 := NewLimit(4, 100*time.Millisecond) // 4 per 100ms + limiter1 := NewLimiter(keyFunc, limit1) + limiter2 := NewLimiter(keyFunc, limit2) + limiters := Combine(limiter1, limiter2) + executionTime := ntime.Now() + + // Exhaust tokens for all buckets from the more restrictive limiter + for bucketID := range buckets { + for range limit1.count { + allow := limiters.allowN(bucketID, executionTime, 1) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } + } + + // Should not allow immediately for any bucket (limiter1 exhausted) + for bucketID := range buckets { + allow := limiters.allowN(bucketID, executionTime, 1) + require.False(t, allow, "should not allow when limiter1 tokens exhausted for bucket %d", bucketID) + } + + // Wait for a token for each bucket + for bucketID := range buckets { + ctx := &testContext{ + done: make(chan struct{}), // never closes + } + + allow, debugs, err := limiters.waitNWithDebug(ctx, bucketID, executionTime, 1) + require.True(t, allow, "should acquire token after waiting for bucket %d", bucketID) + require.NotNil(t, debugs, "should return debugs") + require.NoError(t, err, "should not return error") + } + + // We waited + executionTime = executionTime.Add(limit1.durationPerToken) + + // Buckets should be empty again + for bucketID := range buckets { + ok := limiters.allowN(bucketID, executionTime, 1) + require.False(t, ok, "should not allow again immediately after wait for bucket %d", bucketID) + } + }) + + t.Run("Concurrent", func(t *testing.T) { + // t.Parallel() + // GitHub Actions are pretty resource constrained, + // and this test is sensitive to timing. + keyFunc := func(input int64) string { + return fmt.Sprintf("test-bucket-%d", input) + } + const buckets int64 = 3 + // Create two limiters with different limits + limit1 := NewLimit(2, 100*time.Millisecond) // 2 per 100ms (more restrictive) + limit2 := NewLimit(4, 100*time.Millisecond) // 4 per 100ms + limiter1 := NewLimiter(keyFunc, limit1) + limiter2 := NewLimiter(keyFunc, limit2) + limiters := Combine(limiter1, limiter2) + executionTime := ntime.Now() + + // Exhaust tokens for all buckets from the more restrictive limiter + for bucketID := range buckets { + for range limit1.count { + allow := limiters.allowN(bucketID, executionTime, 1) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } + } + + // Buckets should be empty (limiter1 exhausted) + for bucketID := range buckets { + allow := limiters.allowN(bucketID, executionTime, 1) + require.False(t, allow, "should not allow when limiter1 tokens exhausted for bucket %d", bucketID) + } + + // Test 1: Multiple goroutines with immediate cancellation + { + concurrency := buckets + results := make([]bool, concurrency) + + // Context that's immediately cancelled (closed) + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + close(done) + + // Start concurrent waits + var wg sync.WaitGroup + for i := range concurrency { + wg.Add(1) + go func(i int64) { + defer wg.Done() + allow, _, err := limiters.waitNWithDebug(ctx, i, executionTime, 1) + require.False(t, allow, "should not acquire token due to immediate cancellation") + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + results[i] = allow + }(i) + } + wg.Wait() + + // All should fail because context is cancelled immediately + for i, result := range results { + require.False(t, result, "goroutine %d should not acquire token due to immediate cancellation", i) + } + } + + // Test 2: Multiple goroutines with delayed cancellation + { + keyFunc := func(input int64) int64 { + return input + } + + limit1 := NewLimit(1, 200*time.Millisecond) + limit2 := NewLimit(3, 200*time.Millisecond) + limiter1 := NewLimiter(keyFunc, limit1) + limiter2 := NewLimiter(keyFunc, limit2) + limiters := Combine(limiter1, limiter2) + + // Exhaust tokens + for bucketID := range buckets { + // limit 1 being exhausted means the overall combined limiters are exhausted + for range limit1.count { + allow := limiters.allowN(bucketID, executionTime, 1) + require.True(t, allow, "should allow initial tokens for bucket %d", bucketID) + } + } + + // Confirm exhausted + for bucketID := range buckets { + allow := limiters.allowN(bucketID, executionTime, 1) + require.False(t, allow, "should not allow when limiter1 tokens exhausted for bucket %d", bucketID) + } + + // We're trying to emulate a wait call + // that starts waiting with a retry, but + // is cancelled in the meantime. + + const buckets int64 = 3 + results := make([]bool, buckets) + + var wg sync.WaitGroup + for bucketID := range buckets { + wg.Add(1) + + done := make(chan struct{}) + ctx := &testContext{ + done: done, + } + + go func(bucketID int64) { + defer wg.Done() + // Should start waiting, since the bucket is empty. + // retryAfter should be ~200ms + allow, _, err := limiters.waitNWithDebug(ctx, bucketID, executionTime, 1) + require.False(t, allow, "should not acquire token due to cancellation, bucketID: %d", bucketID) + require.ErrorIs(t, err, context.Canceled, "should return context canceled error") + results[bucketID] = allow + }(bucketID) + + // Cancel context after goroutine has started. + // Originally, this had a delay, for a better test, + // but it is flaky on GitHub Actions, presumably + // because the runner is resource constrained. + // Delay works fine locally on Mac M2. + close(done) + } + wg.Wait() + + // All should fail because context is cancelled + for i, result := range results { + require.False(t, result, "goroutine %d should not acquire token due to cancellation", i) + } + } + }) + }) +}