73 lines
1.2 KiB
Go
73 lines
1.2 KiB
Go
package api
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// Simple token-bucket rate limiter used by Server.cors middleware.
|
|
|
|
type tokenBucket struct {
|
|
tokens float64
|
|
lastFill time.Time
|
|
}
|
|
|
|
type rateLimiter struct {
|
|
rate float64 // tokens per second
|
|
burst float64
|
|
mu sync.Mutex
|
|
bk map[string]*tokenBucket
|
|
evictDur time.Duration
|
|
lastGC time.Time
|
|
}
|
|
|
|
func newRateLimiter(rate float64, burst int, evict time.Duration) *rateLimiter {
|
|
return &rateLimiter{
|
|
rate: rate,
|
|
burst: float64(burst),
|
|
bk: make(map[string]*tokenBucket),
|
|
evictDur: evict,
|
|
lastGC: time.Now(),
|
|
}
|
|
}
|
|
|
|
func (rl *rateLimiter) allow(key string) bool {
|
|
now := time.Now()
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
// GC old buckets occasionally
|
|
if now.Sub(rl.lastGC) > rl.evictDur {
|
|
for k, b := range rl.bk {
|
|
if now.Sub(b.lastFill) > rl.evictDur {
|
|
delete(rl.bk, k)
|
|
}
|
|
}
|
|
rl.lastGC = now
|
|
}
|
|
|
|
b, ok := rl.bk[key]
|
|
if !ok {
|
|
b = &tokenBucket{tokens: rl.burst, lastFill: now}
|
|
rl.bk[key] = b
|
|
}
|
|
|
|
// Refill
|
|
elapsed := now.Sub(b.lastFill).Seconds()
|
|
b.tokens = minf(rl.burst, b.tokens+elapsed*rl.rate)
|
|
b.lastFill = now
|
|
|
|
if b.tokens >= 1.0 {
|
|
b.tokens -= 1.0
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
func minf(a, b float64) float64 {
|
|
if a < b {
|
|
return a
|
|
}
|
|
return b
|
|
}
|