diff --git a/ratelimit.go b/ratelimit.go index cd83599..e4212c1 100644 --- a/ratelimit.go +++ b/ratelimit.go @@ -10,20 +10,36 @@ import ( "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 { unlimited *unlimitedLimiter strategy Strategy - maxCount atomic.Uint32 - interval time.Duration - count atomic.Uint32 - ticker *time.Ticker - tokens chan struct{} - ctx context.Context - // internal + + // maxCount is the number of tokens granted per interval. It is atomic so + // GetLimit/SetLimit stay lock-free. + maxCount atomic.Uint32 + interval time.Duration + + // 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 cancelFunc context.CancelFunc // wraps uber's leaky bucket limiter sizing it to the desired tokens per duration @@ -56,30 +72,6 @@ func (limiter *Limiter) initUnlimitedBucket() *Limiter { return u.finite } -func (limiter *Limiter) run(ctx context.Context) { - defer close(limiter.tokens) - for { - if limiter.count.Load() == 0 { - <-limiter.ticker.C - limiter.count.Store(limiter.maxCount.Load()) - } - - select { - case <-ctx.Done(): - // Internal Context - limiter.ticker.Stop() - return - case <-limiter.ctx.Done(): - limiter.ticker.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() { if limiter.unlimited != nil { @@ -90,11 +82,53 @@ func (limiter *Limiter) Take() { return } - switch limiter.strategy { - case LeakyBucket: + if limiter.strategy == LeakyBucket { _ = limiter.leakyBucketLimiter.Wait(limiter.ctx) - default: - <-limiter.tokens + return + } + + for { + 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 <-timer.C: + case <-ctxDone: + timer.Stop() + return + case <-limiter.done: + timer.Stop() + return + } } } @@ -108,12 +142,16 @@ func (limiter *Limiter) CanTake() bool { return true } - 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 @@ -139,12 +177,9 @@ func (limiter *Limiter) SetLimit(max uint) { } limiter.maxCount.Store(uint32(max)) - - switch limiter.strategy { - case LeakyBucket: + if limiter.strategy == LeakyBucket { limiter.leakyBucketLimiter.SetLimit(leakyBucketRate(max, limiter.interval)) limiter.leakyBucketLimiter.SetBurst(int(max)) - default: } } @@ -167,16 +202,18 @@ func (limiter *Limiter) SetDuration(d time.Duration) { return } - limiter.interval = d - switch limiter.strategy { - case LeakyBucket: + if limiter.strategy == LeakyBucket { + limiter.interval = d limiter.leakyBucketLimiter.SetLimit(leakyBucketRate(limiter.GetLimit(), d)) - default: - limiter.ticker.Reset(d) + return } + limiter.mu.Lock() + limiter.interval = d + limiter.next = time.Now().Add(d) + 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() { if limiter.unlimited != nil { limiter.unlimited.mu.Lock() @@ -190,38 +227,31 @@ func (limiter *Limiter) Stop() { return } - switch limiter.strategy { - case LeakyBucket: - if limiter.cancelFunc != nil { - limiter.cancelFunc() - } - default: + if limiter.strategy == LeakyBucket { if limiter.cancelFunc != nil { limiter.cancelFunc() } + 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) + limiter.next = time.Now().Add(duration) return limiter } 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() + } + }) +} diff --git a/unlimited_benchmark_test.go b/unlimited_benchmark_test.go index 43dff95..e6a04a3 100644 --- a/unlimited_benchmark_test.go +++ b/unlimited_benchmark_test.go @@ -31,11 +31,6 @@ func BenchmarkUnlimitedLifecycle(b *testing.B) { for b.Loop() { limiter := NewUnlimited(context.Background()) limiter.Stop() - // Wait for the old implementation to exit so workers do not accumulate. - if limiter.tokens != nil { - for range limiter.tokens { - } - } } } 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 { diff --git a/unlimited_test.go b/unlimited_test.go index 5cd2a41..3f0993b 100644 --- a/unlimited_test.go +++ b/unlimited_test.go @@ -16,9 +16,8 @@ func TestUnlimitedNoResources(t *testing.T) { synctest.Test(t, func(t *testing.T) { limiter := NewUnlimited(context.Background()) defer limiter.Stop() - require.Nil(t, limiter.ticker, "unlimited mode must not schedule refills") - require.Nil(t, limiter.tokens, "unlimited mode must not exchange tokens") - require.Nil(t, limiter.cancelFunc, "unlimited mode must not start a worker") + require.NotNil(t, limiter.unlimited, "unlimited mode must use a lazy bucket") + require.Nil(t, limiter.unlimited.finite, "unlimited mode must not create a finite bucket yet") require.Equal(t, uint(math.MaxUint32), limiter.GetLimit()) for i := 0; i < 1000; i++ { limiter.Take() @@ -147,9 +146,8 @@ func TestUnlimitedWrappers(t *testing.T) { automatic, err := auto.get("unlimited") require.NoError(t, err) for _, limiter := range []*Limiter{direct, automatic} { - require.Nil(t, limiter.ticker) - require.Nil(t, limiter.tokens) - require.Nil(t, limiter.cancelFunc) + require.NotNil(t, limiter.unlimited) + require.Nil(t, limiter.unlimited.finite) require.Equal(t, uint(math.MaxUint32), limiter.GetLimit()) } auto.Stop()