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
54 lines
1.1 KiB
Go
54 lines
1.1 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// Implements a simple token bucket rate limiter
|
|
type TokenBucket struct {
|
|
mu sync.Mutex
|
|
tokens int
|
|
capacity int
|
|
refill int
|
|
interval time.Duration
|
|
lastFill time.Time
|
|
}
|
|
|
|
// Creates a rate limiter with specified capacity and refill rate capacity
|
|
func NewTokenBucket(capacity, refillPerInterval int, interval time.Duration) *TokenBucket {
|
|
return &TokenBucket{
|
|
tokens: capacity,
|
|
capacity: capacity,
|
|
refill: refillPerInterval,
|
|
interval: interval,
|
|
lastFill: time.Now(),
|
|
}
|
|
}
|
|
|
|
// Allow checks if a request should be allowed
|
|
func (tb *TokenBucket) Allow() bool {
|
|
tb.mu.Lock()
|
|
defer tb.mu.Unlock()
|
|
|
|
// Refill tokens based on elapsed time
|
|
now := time.Now()
|
|
elapsed := now.Sub(tb.lastFill)
|
|
if elapsed >= tb.interval {
|
|
periods := int(elapsed / tb.interval)
|
|
tb.tokens += periods * tb.refill
|
|
if tb.tokens > tb.capacity {
|
|
tb.tokens = tb.capacity
|
|
}
|
|
// Advance lastFill by exact periods to prevent drift
|
|
tb.lastFill = tb.lastFill.Add(time.Duration(periods) * tb.interval)
|
|
}
|
|
|
|
// Check if we have tokens available
|
|
if tb.tokens > 0 {
|
|
tb.tokens--
|
|
return true
|
|
}
|
|
|
|
return false
|
|
}
|