NotAShelf
f46697bd21
The TokenBucket ratelimiter accumulated time drift over multiple refills because I'm an idiot. We were using 'now' as base for lastFill calc. but this could case rate limiting to become inaccurate over time. Now we advance lastFill by *exact* periods from previous value. Signed-off-by: NotAShelf <raf@notashelf.dev> Change-Id: Ia3990b441ab6072f51dfdfa4a2511b5f6a6a6964
116 lines
2.4 KiB
Go
116 lines
2.4 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestTokenBucket_Allow(t *testing.T) {
|
|
// Create bucket with 10 tokens, refill 10 per second
|
|
tb := NewTokenBucket(10, 10, time.Second)
|
|
|
|
// Should allow first 10 requests
|
|
for i := 0; i < 10; i++ {
|
|
if !tb.Allow() {
|
|
t.Errorf("request %d should be allowed", i)
|
|
}
|
|
}
|
|
|
|
// 11th request should be denied (no tokens left)
|
|
if tb.Allow() {
|
|
t.Error("request 11 should be denied (bucket empty)")
|
|
}
|
|
}
|
|
|
|
func TestTokenBucket_Refill(t *testing.T) {
|
|
// Create bucket with 5 tokens, refill 5 per 100ms
|
|
tb := NewTokenBucket(5, 5, 100*time.Millisecond)
|
|
|
|
// Consume all tokens
|
|
for range 5 {
|
|
tb.Allow()
|
|
}
|
|
|
|
// Should be denied
|
|
if tb.Allow() {
|
|
t.Error("should be denied before refill")
|
|
}
|
|
|
|
// Wait for refill
|
|
time.Sleep(150 * time.Millisecond)
|
|
|
|
// Should be allowed again
|
|
if !tb.Allow() {
|
|
t.Error("should be allowed after refill")
|
|
}
|
|
}
|
|
|
|
func TestTokenBucket_NoDrift(t *testing.T) {
|
|
// Test that lastFill advances by exact periods with no drift
|
|
tb := NewTokenBucket(50, 10, 100*time.Millisecond)
|
|
|
|
// Consume most tokens to avoid hitting capacity
|
|
for range 45 {
|
|
tb.Allow()
|
|
}
|
|
|
|
// Wait for refill and establish baseline
|
|
time.Sleep(120 * time.Millisecond)
|
|
tb.Allow() // Triggers refill
|
|
|
|
// Record baseline
|
|
tb.mu.Lock()
|
|
baseline := tb.lastFill
|
|
tb.mu.Unlock()
|
|
|
|
// Wait for exactly 10 refill periods (1 second)
|
|
time.Sleep(time.Second + 10*time.Millisecond)
|
|
|
|
// Trigger refill calculation
|
|
tb.Allow()
|
|
|
|
// Verify lastFill advanced by exactly 10 periods
|
|
tb.mu.Lock()
|
|
expectedFill := baseline.Add(10 * 100 * time.Millisecond)
|
|
actualFill := tb.lastFill
|
|
tb.mu.Unlock()
|
|
|
|
// Time should match exactly (within 1ms for timing jitter)
|
|
drift := actualFill.Sub(expectedFill)
|
|
if drift < 0 {
|
|
drift = -drift
|
|
}
|
|
if drift > time.Millisecond {
|
|
t.Errorf(
|
|
"time drift detected: expected %v, got %v (drift: %v)",
|
|
expectedFill,
|
|
actualFill,
|
|
drift,
|
|
)
|
|
}
|
|
}
|
|
|
|
func TestTokenBucket_MultipleRefills(t *testing.T) {
|
|
// Create bucket with 10 tokens, refill 5 per 50ms
|
|
tb := NewTokenBucket(10, 5, 50*time.Millisecond)
|
|
|
|
// Consume all tokens
|
|
for range 10 {
|
|
tb.Allow()
|
|
}
|
|
|
|
// Wait for 2 refill periods (should add 10 tokens)
|
|
time.Sleep(120 * time.Millisecond)
|
|
|
|
// Should be able to consume 10 tokens (capped at capacity)
|
|
allowed := 0
|
|
for range 15 {
|
|
if tb.Allow() {
|
|
allowed++
|
|
}
|
|
}
|
|
|
|
if allowed != 10 {
|
|
t.Errorf("expected 10 tokens after refill, got %d", allowed)
|
|
}
|
|
}
|