fast take

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
 }
 

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