73 lines
1.6 KiB
Go
73 lines
1.6 KiB
Go
// FILE: logwisp/src/internal/limit/token_bucket.go
|
|
package limit
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// TokenBucket implements a thread-safe token bucket rate limiter.
|
|
type TokenBucket struct {
|
|
capacity float64
|
|
tokens float64
|
|
refillRate float64 // tokens per second
|
|
lastRefill time.Time
|
|
mu sync.Mutex
|
|
}
|
|
|
|
// NewTokenBucket creates a new token bucket with a given capacity and refill rate.
|
|
func NewTokenBucket(capacity float64, refillRate float64) *TokenBucket {
|
|
return &TokenBucket{
|
|
capacity: capacity,
|
|
tokens: capacity, // Start full
|
|
refillRate: refillRate,
|
|
lastRefill: time.Now(),
|
|
}
|
|
}
|
|
|
|
// Allow attempts to consume one token, returning true if successful.
|
|
func (tb *TokenBucket) Allow() bool {
|
|
return tb.AllowN(1)
|
|
}
|
|
|
|
// AllowN attempts to consume n tokens, returning true if successful.
|
|
func (tb *TokenBucket) AllowN(n float64) bool {
|
|
tb.mu.Lock()
|
|
defer tb.mu.Unlock()
|
|
|
|
tb.refill()
|
|
|
|
if tb.tokens >= n {
|
|
tb.tokens -= n
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// Tokens returns the current number of available tokens in the bucket.
|
|
func (tb *TokenBucket) Tokens() float64 {
|
|
tb.mu.Lock()
|
|
defer tb.mu.Unlock()
|
|
|
|
tb.refill()
|
|
return tb.tokens
|
|
}
|
|
|
|
// refill adds new tokens to the bucket based on the elapsed time.
|
|
func (tb *TokenBucket) refill() {
|
|
now := time.Now()
|
|
elapsed := now.Sub(tb.lastRefill).Seconds()
|
|
|
|
// Handle time sync issues causing negative elapsed time
|
|
if elapsed < 0 {
|
|
// Clock went backwards, reset to current time but don't add tokens
|
|
tb.lastRefill = now
|
|
elapsed = 0
|
|
}
|
|
|
|
tb.tokens += elapsed * tb.refillRate
|
|
if tb.tokens > tb.capacity {
|
|
tb.tokens = tb.capacity
|
|
}
|
|
tb.lastRefill = now
|
|
} |