69 lines
1.4 KiB
Go
69 lines
1.4 KiB
Go
package main
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// RateLimiter implements a per-user token bucket rate limiter.
|
|
type RateLimiter struct {
|
|
mu sync.Mutex
|
|
users map[string]*bucket
|
|
rate int // requests per minute
|
|
}
|
|
|
|
type bucket struct {
|
|
tokens float64
|
|
lastRefill time.Time
|
|
}
|
|
|
|
// NewRateLimiter creates a rate limiter with the given rate (requests/minute).
|
|
func NewRateLimiter(ratePerMinute int) *RateLimiter {
|
|
if ratePerMinute < 0 {
|
|
ratePerMinute = 0
|
|
}
|
|
return &RateLimiter{
|
|
users: make(map[string]*bucket),
|
|
rate: ratePerMinute,
|
|
}
|
|
}
|
|
|
|
// Allow checks if a request from the given user ID should be allowed.
|
|
func (rl *RateLimiter) Allow(userID string) bool {
|
|
// 0 = unlimited.
|
|
if rl.rate == 0 {
|
|
return true
|
|
}
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
|
|
b, exists := rl.users[userID]
|
|
now := time.Now()
|
|
|
|
if !exists {
|
|
b = &bucket{tokens: 1, lastRefill: now}
|
|
rl.users[userID] = b
|
|
}
|
|
|
|
elapsed := now.Sub(b.lastRefill).Seconds()
|
|
refill := elapsed * float64(rl.rate) / 60.0
|
|
b.tokens += refill
|
|
if b.tokens > 1 {
|
|
b.tokens = 1
|
|
}
|
|
b.lastRefill = now
|
|
|
|
if b.tokens >= 1 {
|
|
b.tokens--
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// Reset clears all rate limiter state.
|
|
func (rl *RateLimiter) Reset() {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
rl.users = make(map[string]*bucket)
|
|
}
|