110 lines
2.8 KiB
Go
110 lines
2.8 KiB
Go
// Package ratelimit implements a per-key token-bucket rate limiter — no
|
|
// third-party rate-limiting library. A token bucket (rather than a hard
|
|
// fixed-window reset) is used deliberately: it smooths out bursts at
|
|
// window boundaries that a naive "reset every 60s" counter would allow
|
|
// (e.g. 20 requests at 0:59 plus another 20 at 1:01 both passing a
|
|
// "20/min" limit reset at the minute boundary). Safe for concurrent use.
|
|
package ratelimit
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type bucket struct {
|
|
tokens float64
|
|
lastRefill time.Time
|
|
}
|
|
|
|
// Limiter enforces "at most ratePerMinute events per key, per minute" with
|
|
// burst tolerance up to ratePerMinute tokens banked at once (i.e. a key
|
|
// that's been idle can burst up to the full per-minute allowance instantly,
|
|
// then is throttled to the steady-state rate — standard token-bucket
|
|
// behavior, not a stricter "evenly spaced" enforcement).
|
|
type Limiter struct {
|
|
ratePerMinute float64
|
|
mu sync.Mutex
|
|
buckets map[string]*bucket
|
|
|
|
stopCleanup chan struct{}
|
|
}
|
|
|
|
// New creates a limiter allowing ratePerMinute events per key. Pass 0 to
|
|
// disable limiting entirely (Allow always returns true) — this is how a
|
|
// zero/unset config value opts a listener out of rate limiting rather than
|
|
// silently blocking everything.
|
|
func New(ratePerMinute int) *Limiter {
|
|
l := &Limiter{
|
|
ratePerMinute: float64(ratePerMinute),
|
|
buckets: make(map[string]*bucket),
|
|
stopCleanup: make(chan struct{}),
|
|
}
|
|
if ratePerMinute > 0 {
|
|
go l.cleanupLoop()
|
|
}
|
|
return l
|
|
}
|
|
|
|
// Allow reports whether an event for key is permitted right now, consuming
|
|
// one token if so. Safe to call from many goroutines concurrently.
|
|
func (l *Limiter) Allow(key string) bool {
|
|
if l.ratePerMinute <= 0 {
|
|
return true
|
|
}
|
|
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
now := time.Now()
|
|
b, ok := l.buckets[key]
|
|
if !ok {
|
|
b = &bucket{tokens: l.ratePerMinute - 1, lastRefill: now}
|
|
l.buckets[key] = b
|
|
return true
|
|
}
|
|
|
|
elapsed := now.Sub(b.lastRefill).Seconds()
|
|
refill := elapsed * (l.ratePerMinute / 60.0)
|
|
b.tokens += refill
|
|
if b.tokens > l.ratePerMinute {
|
|
b.tokens = l.ratePerMinute
|
|
}
|
|
b.lastRefill = now
|
|
|
|
if b.tokens < 1 {
|
|
return false
|
|
}
|
|
b.tokens--
|
|
return true
|
|
}
|
|
|
|
// cleanupLoop periodically evicts buckets idle long enough to have fully
|
|
// refilled, so a limiter tracking many distinct one-off IPs doesn't grow
|
|
// unboundedly over a long-running server's lifetime.
|
|
func (l *Limiter) cleanupLoop() {
|
|
ticker := time.NewTicker(10 * time.Minute)
|
|
defer ticker.Stop()
|
|
for {
|
|
select {
|
|
case <-l.stopCleanup:
|
|
return
|
|
case <-ticker.C:
|
|
l.mu.Lock()
|
|
now := time.Now()
|
|
for key, b := range l.buckets {
|
|
if now.Sub(b.lastRefill) > 30*time.Minute {
|
|
delete(l.buckets, key)
|
|
}
|
|
}
|
|
l.mu.Unlock()
|
|
}
|
|
}
|
|
}
|
|
|
|
// Stop releases the background cleanup goroutine.
|
|
func (l *Limiter) Stop() {
|
|
if l.ratePerMinute > 0 {
|
|
close(l.stopCleanup)
|
|
}
|
|
}
|