From 101e2f313b4a9ebd3d1ab7e1b7b238caacdb1b97 Mon Sep 17 00:00:00 2001 From: Nakul Bharti Date: Thu, 10 Sep 2026 13:15:31 +0530 Subject: [PATCH 1/4] fix: apply configured leaky bucket throughput --- ratelimit.go | 15 ++++++++++++--- ratelimit_test.go | 29 +++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 3 deletions(-) diff --git a/ratelimit.go b/ratelimit.go index 7c09cff..c7283f7 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -92,7 +92,7 @@ func (limiter *Limiter) Take() { switch limiter.strategy { case LeakyBucket: - _ = limiter.leakyBucketLimiter.Wait(context.TODO()) + _ = limiter.leakyBucketLimiter.Wait(limiter.ctx) default: <-limiter.tokens } @@ -142,6 +142,7 @@ func (limiter *Limiter) SetLimit(max uint) { switch limiter.strategy { case LeakyBucket: + limiter.leakyBucketLimiter.SetLimit(leakyBucketRate(max, limiter.interval)) limiter.leakyBucketLimiter.SetBurst(int(max)) default: } @@ -169,7 +170,7 @@ func (limiter *Limiter) SetDuration(d time.Duration) { limiter.interval = d switch limiter.strategy { case LeakyBucket: - limiter.leakyBucketLimiter.SetLimit(rate.Every(d)) + limiter.leakyBucketLimiter.SetLimit(leakyBucketRate(limiter.GetLimit(), d)) default: limiter.ticker.Reset(d) } @@ -241,7 +242,8 @@ func NewUnlimited(ctx context.Context) *Limiter { func NewLeakyBucket(ctx context.Context, max uint, duration time.Duration) *Limiter { limiter := &Limiter{ strategy: LeakyBucket, - leakyBucketLimiter: rate.NewLimiter(rate.Every(duration), int(max)), + leakyBucketLimiter: rate.NewLimiter(leakyBucketRate(max, duration), int(max)), + ctx: ctx, } limiter.maxCount.Store(uint32(max)) @@ -249,3 +251,10 @@ func NewLeakyBucket(ctx context.Context, max uint, duration time.Duration) *Limi return limiter } + +func leakyBucketRate(max uint, duration time.Duration) rate.Limit { + if duration <= 0 { + return rate.Inf + } + return rate.Limit(float64(max) / duration.Seconds()) +} diff --git a/ratelimit_test.go b/ratelimit_test.go index f8fbc7e..8aa91db 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -108,4 +108,33 @@ func TestRateLimit(t *testing.T) { expected := 3 * time.Second require.True(t, took >= expected) }) + + t.Run("LeakyBucket applies max per duration", func(t *testing.T) { + limiter := NewLeakyBucket(context.Background(), 50, time.Second) + require.InDelta(t, 50, float64(limiter.leakyBucketLimiter.Limit()), 0.001) + + limiter.SetLimit(25) + require.InDelta(t, 25, float64(limiter.leakyBucketLimiter.Limit()), 0.001) + limiter.SetDuration(500 * time.Millisecond) + require.InDelta(t, 50, float64(limiter.leakyBucketLimiter.Limit()), 0.001) + }) + + t.Run("LeakyBucket observes cancellation", func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + limiter := NewLeakyBucket(ctx, 1, time.Hour) + limiter.Take() + + done := make(chan struct{}) + go func() { + limiter.Take() + close(done) + }() + cancel() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("Take remained blocked after limiter context cancellation") + } + }) } From 9a6bcf83b7fc2be9c1310dd7f80379d292a39e87 Mon Sep 17 00:00:00 2001 From: Nakul Bharti Date: Thu, 10 Sep 2026 13:42:35 +0530 Subject: [PATCH 2/4] test: wait for limiter accounting --- ratelimit_test.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ratelimit_test.go b/ratelimit_test.go index 8aa91db..82c6816 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -93,7 +93,10 @@ func TestRateLimit(t *testing.T) { limiter.Take() limiter.Take() limiter.Take() - require.False(t, limiter.CanTake()) + // The token producer decrements its atomic count immediately after the + // unbuffered handoff. Under the race detector the receiver can resume in + // that tiny window, so wait for the producer-side accounting to settle. + require.Eventually(t, func() bool { return !limiter.CanTake() }, time.Second, time.Millisecond) }) t.Run("LeakyBucket", func(t *testing.T) { From 7535caba304c26e4bf5b8137eb1ea1acca5bbe34 Mon Sep 17 00:00:00 2001 From: Nakul Bharti Date: Sat, 12 Sep 2026 22:04:55 +0530 Subject: [PATCH 3/4] test: tolerate scheduler timing on Windows --- auto_ratelimit_test.go | 5 ++++- ratelimit_test.go | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/auto_ratelimit_test.go b/auto_ratelimit_test.go index 98527bd..fe58292 100644 --- a/auto_ratelimit_test.go +++ b/auto_ratelimit_test.go @@ -187,7 +187,10 @@ func TestAutoLimiterCreateOrDefault(t *testing.T) { require.True(t, recreated.CanTake()) recreated.Take() } - require.False(t, recreated.CanTake()) + // Take returns at the unbuffered handoff; the producer decrements its + // accounting immediately afterwards. Wait for that scheduler-sized window + // instead of making the test depend on goroutine ordering. + require.Eventually(t, func() bool { return !recreated.CanTake() }, time.Second, time.Millisecond) } func TestAutoLimiterAddAndTake(t *testing.T) { diff --git a/ratelimit_test.go b/ratelimit_test.go index 82c6816..3064c4a 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -27,7 +27,10 @@ func TestRateLimit(t *testing.T) { // take another one above max limiter.Take() took = time.Since(start).Nanoseconds() - require.GreaterOrEqual(t, took, expected.Nanoseconds()) + // Runtime timers may wake fractionally before their nominal deadline on + // some platforms. A small tolerance still proves that the full refill + // window was enforced without making CI depend on timer granularity. + require.GreaterOrEqual(t, took, (expected - 10*time.Millisecond).Nanoseconds()) }) t.Run("Unlimited Rate Limit", func(t *testing.T) { From aeeca5645294bcae4abe2e676aac65b6b844085a Mon Sep 17 00:00:00 2001 From: Nakul Bharti Date: Sat, 12 Sep 2026 22:13:06 +0530 Subject: [PATCH 4/4] fix: stop leaky bucket waiters --- ratelimit.go | 9 +++++++-- ratelimit_test.go | 18 ++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/ratelimit.go b/ratelimit.go index c7283f7..cd83599 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -191,7 +191,10 @@ func (limiter *Limiter) Stop() { } switch limiter.strategy { - case LeakyBucket: // NOP + case LeakyBucket: + if limiter.cancelFunc != nil { + limiter.cancelFunc() + } default: if limiter.cancelFunc != nil { limiter.cancelFunc() @@ -240,10 +243,12 @@ func NewUnlimited(ctx context.Context) *Limiter { // NewLeakyBucket creates a limiter that uses golang.org/x/time/rate. func NewLeakyBucket(ctx context.Context, max uint, duration time.Duration) *Limiter { + internalctx, cancel := context.WithCancel(ctx) limiter := &Limiter{ strategy: LeakyBucket, leakyBucketLimiter: rate.NewLimiter(leakyBucketRate(max, duration), int(max)), - ctx: ctx, + ctx: internalctx, + cancelFunc: cancel, } limiter.maxCount.Store(uint32(max)) diff --git a/ratelimit_test.go b/ratelimit_test.go index 3064c4a..c82bcef 100644 --- a/ratelimit_test.go +++ b/ratelimit_test.go @@ -143,4 +143,22 @@ func TestRateLimit(t *testing.T) { t.Fatal("Take remained blocked after limiter context cancellation") } }) + + t.Run("LeakyBucket stop unblocks waiters", func(t *testing.T) { + limiter := NewLeakyBucket(context.Background(), 1, time.Hour) + limiter.Take() + + done := make(chan struct{}) + go func() { + limiter.Take() + close(done) + }() + limiter.Stop() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("Take remained blocked after limiter stop") + } + }) }