From 156a45dabf83208db7d74159aaaf909b4c1d10b1 Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Thu, 25 Jun 2026 20:31:59 +0200 Subject: [PATCH 1/2] fast take --- ratelimit.go | 182 +++++++++++++++++++++---------------- ratelimit_coverage_test.go | 91 +++++++++++++++++++ ratelimit_perf_test.go | 96 +++++++++++++++++++ 3 files changed, 293 insertions(+), 76 deletions(-) create mode 100644 ratelimit_coverage_test.go create mode 100644 ratelimit_perf_test.go diff --git a/ratelimit.go b/ratelimit.go index 813f464..38f4c27 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -3,72 +3,111 @@ package ratelimit import ( "context" "math" + "sync" "sync/atomic" "time" "golang.org/x/time/rate" ) -// equals to -1 -var minusOne = ^uint32(0) - // Limiter allows a burst of request during the defined duration +// +// The default (None) strategy is a fixed-window limiter: up to maxCount tokens +// are available per interval and the budget is refilled in a burst when the +// window rolls over. It is served entirely from a small mutex-guarded counter, +// so Take is a couple of field operations on the hot path. The previous design +// delivered one token per call over an unbuffered channel fed by a dedicated +// background goroutine, which charged every caller a goroutine handoff and +// scheduler wakeup and spawned a goroutine per limiter. That overhead dominates +// once tools push tens of thousands of Take calls per second across many +// workers, so it has been removed. type Limiter struct { strategy Strategy + + // maxCount is the number of tokens granted per interval. It is atomic so + // GetLimit/SetLimit stay lock-free. maxCount atomic.Uint32 interval time.Duration - count atomic.Uint32 - ticker *time.Ticker - tokens chan struct{} - ctx context.Context - // internal - cancelFunc context.CancelFunc + + // mu guards the fixed-window state (count, next, interval) for the None + // strategy. The critical section is only a few field operations. + mu sync.Mutex + count uint32 // tokens remaining in the current window + next time.Time // end of the current window (when count refills) + + ctx context.Context + // done is closed by Stop to release any waiter blocked in Take. + done chan struct{} + stopOnce sync.Once // wraps uber's leaky bucket limiter sizing it to the desired tokens per duration leakyBucketLimiter *rate.Limiter } -func (limiter *Limiter) run(ctx context.Context) { - defer close(limiter.tokens) +// Take one token from the bucket +func (limiter *Limiter) Take() { + if limiter.strategy == LeakyBucket { + _ = limiter.leakyBucketLimiter.Wait(context.TODO()) + return + } + for { - if limiter.count.Load() == 0 { - <-limiter.ticker.C - limiter.count.Store(limiter.maxCount.Load()) + limiter.mu.Lock() + now := time.Now() + switch { + case limiter.next.IsZero(): + // first take in this window; the initial budget is already set + limiter.next = now.Add(limiter.interval) + case !now.Before(limiter.next): + // one or more windows elapsed, refill to the cap. Advance from the + // previous boundary to avoid drift, snapping forward if the limiter + // sat idle for longer than a full window. + limiter.count = limiter.maxCount.Load() + limiter.next = limiter.next.Add(limiter.interval) + if now.After(limiter.next) { + limiter.next = now.Add(limiter.interval) + } + } + if limiter.count > 0 { + limiter.count-- + limiter.mu.Unlock() + return + } + wait := time.Until(limiter.next) + limiter.mu.Unlock() + + if wait <= 0 { + continue + } + timer := time.NewTimer(wait) + var ctxDone <-chan struct{} + if limiter.ctx != nil { + ctxDone = limiter.ctx.Done() } select { - case <-ctx.Done(): - // Internal Context - limiter.ticker.Stop() + case <-timer.C: + case <-ctxDone: + timer.Stop() return - case <-limiter.ctx.Done(): - limiter.ticker.Stop() + case <-limiter.done: + timer.Stop() return - case limiter.tokens <- struct{}{}: - limiter.count.Add(minusOne) - case <-limiter.ticker.C: - limiter.count.Store(limiter.maxCount.Load()) } } } -// Take one token from the bucket -func (limiter *Limiter) Take() { - switch limiter.strategy { - case LeakyBucket: - _ = limiter.leakyBucketLimiter.Wait(context.TODO()) - default: - <-limiter.tokens - } -} - // CanTake checks if the rate limiter has any token func (limiter *Limiter) CanTake() bool { - switch limiter.strategy { - case LeakyBucket: + if limiter.strategy == LeakyBucket { return limiter.leakyBucketLimiter.Tokens() > 0 - default: - return limiter.count.Load() > 0 } + limiter.mu.Lock() + defer limiter.mu.Unlock() + // a rolled-over window will refill on the next take + if !limiter.next.IsZero() && !time.Now().Before(limiter.next) { + return true + } + return limiter.count > 0 } // GetLimit returns current rate limit per given duration @@ -76,76 +115,67 @@ func (limiter *Limiter) GetLimit() uint { return uint(limiter.maxCount.Load()) } -// GetLimit returns current rate limit per given duration +// SetLimit sets the current rate limit per given duration func (limiter *Limiter) SetLimit(max uint) { limiter.maxCount.Store(uint32(max)) - switch limiter.strategy { - case LeakyBucket: + if limiter.strategy == LeakyBucket { limiter.leakyBucketLimiter.SetBurst(int(max)) - default: } } -// GetLimit returns current rate limit per given duration +// SetDuration sets the current rate limit duration func (limiter *Limiter) SetDuration(d time.Duration) { - limiter.interval = d - switch limiter.strategy { - case LeakyBucket: + if limiter.strategy == LeakyBucket { + limiter.interval = d limiter.leakyBucketLimiter.SetLimit(rate.Every(d)) - default: - limiter.ticker.Reset(d) + return } + limiter.mu.Lock() + limiter.interval = d + // recompute the window boundary against the new duration on the next take + limiter.next = time.Time{} + limiter.mu.Unlock() } -// Stop the rate limiter canceling the internal context +// Stop the rate limiter releasing any waiter blocked in Take func (limiter *Limiter) Stop() { - switch limiter.strategy { - case LeakyBucket: // NOP - default: - if limiter.cancelFunc != nil { - limiter.cancelFunc() - } + if limiter.strategy == LeakyBucket { + return } + limiter.stopOnce.Do(func() { + if limiter.done != nil { + close(limiter.done) + } + }) } // New creates a new limiter instance with the tokens amount and the interval func New(ctx context.Context, max uint, duration time.Duration) *Limiter { - internalctx, cancel := context.WithCancel(context.TODO()) - - maxCount := &atomic.Uint32{} - maxCount.Store(uint32(max)) limiter := &Limiter{ - ticker: time.NewTicker(duration), - tokens: make(chan struct{}), - ctx: ctx, - cancelFunc: cancel, - strategy: None, - interval: duration, + strategy: None, + interval: duration, + ctx: ctx, + done: make(chan struct{}), } limiter.maxCount.Store(uint32(max)) - limiter.count.Store(uint32(max)) - go limiter.run(internalctx) - + limiter.count = uint32(max) return limiter } // NewUnlimited create a bucket with approximated unlimited tokens func NewUnlimited(ctx context.Context) *Limiter { - internalctx, cancel := context.WithCancel(context.TODO()) limiter := &Limiter{ - ticker: time.NewTicker(time.Millisecond), - tokens: make(chan struct{}), - ctx: ctx, - cancelFunc: cancel, + strategy: None, + interval: time.Millisecond, + ctx: ctx, + done: make(chan struct{}), } limiter.maxCount.Store(math.MaxUint32) - limiter.count.Store(math.MaxUint32) - go limiter.run(internalctx) - + limiter.count = math.MaxUint32 return limiter } -// NewUnlimited create a bucket with approximated unlimited tokens +// NewLeakyBucket create a bucket with a smooth (leaky bucket) token rate func NewLeakyBucket(ctx context.Context, max uint, duration time.Duration) *Limiter { limiter := &Limiter{ strategy: LeakyBucket, diff --git a/ratelimit_coverage_test.go b/ratelimit_coverage_test.go new file mode 100644 index 0000000..5b3e2a9 --- /dev/null +++ b/ratelimit_coverage_test.go @@ -0,0 +1,91 @@ +package ratelimit + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// drainAvailable consumes every currently available token without blocking, +// returning how many were taken (bounded to guard against a runaway loop). +func drainAvailable(l *Limiter) int { + n := 0 + for n <= 1000 && l.CanTake() { + l.Take() + n++ + } + return n +} + +func TestGetSetLimit(t *testing.T) { + l := New(context.Background(), 5, time.Hour) + require.Equal(t, uint(5), l.GetLimit()) + l.SetLimit(8) + require.Equal(t, uint(8), l.GetLimit()) +} + +// TestRefillBoundedToMax is the guard for the bounded-burst property: after a +// window rolls over the budget refills to exactly max, never accumulating. +func TestRefillBoundedToMax(t *testing.T) { + const max = 3 + l := New(context.Background(), max, 200*time.Millisecond) + + require.Equal(t, max, drainAvailable(l), "first window must grant exactly max") + require.False(t, l.CanTake(), "budget must be exhausted within the window") + + // let the window roll over while idle + time.Sleep(260 * time.Millisecond) + + require.True(t, l.CanTake(), "a rolled-over window must refill") + require.Equal(t, max, drainAvailable(l), "refill must cap at max, not accumulate idle windows") +} + +// TestSetLimitAppliesNextWindow ensures a new limit takes effect once the +// current window refills. +func TestSetLimitAppliesNextWindow(t *testing.T) { + l := New(context.Background(), 2, 200*time.Millisecond) + require.Equal(t, 2, drainAvailable(l)) + + l.SetLimit(5) + time.Sleep(260 * time.Millisecond) + + require.Equal(t, 5, drainAvailable(l), "new limit must apply after the window refills") +} + +// TestSetDuration ensures shortening the duration shortens the wait for the next +// token instead of leaving the caller blocked on the old interval. +func TestSetDuration(t *testing.T) { + l := New(context.Background(), 1, time.Hour) + l.Take() // drain; the next Take would otherwise block ~1h + + l.SetDuration(150 * time.Millisecond) + + start := time.Now() + l.Take() + elapsed := time.Since(start) + + require.Less(t, elapsed, time.Second, "SetDuration must shorten the wait") + require.GreaterOrEqual(t, elapsed, 100*time.Millisecond, "should still wait about the new duration") +} + +func TestCanTakeLeakyBucket(t *testing.T) { + // CanTake on the LeakyBucket strategy reflects the underlying token pool. + // With a full burst it must report available; the exhausted case is not + // asserted because rate.Limiter.Tokens() returns a float that creeps back + // above zero immediately after a take. + l := NewLeakyBucket(context.Background(), 5, time.Hour) + require.True(t, l.CanTake()) +} + +// TestLeakyBucketSetters exercises the LeakyBucket branches of the setters and +// Stop (a no-op for that strategy). +func TestLeakyBucketSetters(t *testing.T) { + l := NewLeakyBucket(context.Background(), 2, time.Second) + require.Equal(t, uint(2), l.GetLimit()) + l.SetLimit(5) + require.Equal(t, uint(5), l.GetLimit()) + l.SetDuration(2 * time.Second) + l.Stop() // must be a safe no-op for LeakyBucket +} diff --git a/ratelimit_perf_test.go b/ratelimit_perf_test.go new file mode 100644 index 0000000..20af042 --- /dev/null +++ b/ratelimit_perf_test.go @@ -0,0 +1,96 @@ +package ratelimit + +import ( + "context" + "runtime" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// TestNoBackgroundGoroutine guards the core of this change: the default limiter +// no longer spawns a background goroutine per instance to serve tokens over a +// channel. Creating many limiters must not grow the goroutine count. +func TestNoBackgroundGoroutine(t *testing.T) { + runtime.GC() + before := runtime.NumGoroutine() + + const n = 200 + limiters := make([]*Limiter, 0, n) + for i := 0; i < n; i++ { + l := New(context.Background(), 100, time.Second) + l.Take() + limiters = append(limiters, l) + } + // give any (unexpected) goroutines a chance to show up + time.Sleep(50 * time.Millisecond) + + after := runtime.NumGoroutine() + runtime.KeepAlive(limiters) + require.Less(t, after-before, 20, "creating limiters must not spawn a goroutine each") +} + +// TestStopUnblocksTake ensures Stop releases a caller blocked waiting for the +// next window, matching the previous channel-close behavior. +func TestStopUnblocksTake(t *testing.T) { + l := New(context.Background(), 1, time.Hour) + l.Take() // drain the only token so the next Take blocks + + done := make(chan struct{}) + go func() { + l.Take() + close(done) + }() + + time.Sleep(50 * time.Millisecond) + l.Stop() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Stop did not unblock a pending Take") + } +} + +// TestContextCancelUnblocksTake ensures cancelling the context handed to New +// releases a blocked Take. +func TestContextCancelUnblocksTake(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + l := New(ctx, 1, time.Hour) + l.Take() + + done := make(chan struct{}) + go func() { + l.Take() + close(done) + }() + + time.Sleep(50 * time.Millisecond) + cancel() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("context cancel did not unblock a pending Take") + } +} + +func BenchmarkTakeUnlimited(b *testing.B) { + l := NewUnlimited(context.Background()) + b.ResetTimer() + for i := 0; i < b.N; i++ { + l.Take() + } +} + +func BenchmarkTakeParallel(b *testing.B) { + // high rate so the window never blocks within the benchmark + l := New(context.Background(), 1<<30, time.Second) + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + l.Take() + } + }) +} From 3d8f01efa72a007e9208810175144b44136e18f6 Mon Sep 17 00:00:00 2001 From: Mzack9999 Date: Wed, 23 Sep 2026 21:58:14 +0200 Subject: [PATCH 2/2] fix benchmark --- unlimited_idle_linux_test.go | 6 ------ 1 file changed, 6 deletions(-) diff --git a/unlimited_idle_linux_test.go b/unlimited_idle_linux_test.go index 4ff57ef..eccba8f 100644 --- a/unlimited_idle_linux_test.go +++ b/unlimited_idle_linux_test.go @@ -22,12 +22,6 @@ func BenchmarkUnlimitedIdle(b *testing.B) { for _, limiter := range limiters { limiter.Stop() } - for _, limiter := range limiters { - if limiter.tokens != nil { - for range limiter.tokens { - } - } - } }() workers := runtime.NumGoroutine() - before cpuTime := func() time.Duration {