ae9a71034d
Add HTTP /v1 MVC, TCP send path, groups with server-side member-avatar collage, object-storage upload, and friends/favorites around the architecture-v3 layout. Co-authored-by: Cursor <cursoragent@cursor.com>
70 lines
1.1 KiB
Go
70 lines
1.1 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// Limiter is a per-key token bucket. Keys are RealIP (C6), never gnet RemoteAddr.
|
|
type Limiter struct {
|
|
mu sync.Mutex
|
|
rate float64
|
|
burst float64
|
|
buckets map[string]*bucket
|
|
}
|
|
|
|
type bucket struct {
|
|
tokens float64
|
|
last time.Time
|
|
}
|
|
|
|
func New(perSec int) *Limiter {
|
|
if perSec < 1 {
|
|
perSec = 5
|
|
}
|
|
return &Limiter{
|
|
rate: float64(perSec),
|
|
burst: float64(perSec),
|
|
buckets: make(map[string]*bucket),
|
|
}
|
|
}
|
|
|
|
func (l *Limiter) Allow(key string) bool {
|
|
if l == nil {
|
|
return true
|
|
}
|
|
if key == "" {
|
|
key = "unknown"
|
|
}
|
|
now := time.Now()
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
b := l.buckets[key]
|
|
if b == nil {
|
|
b = &bucket{tokens: l.burst, last: now}
|
|
l.buckets[key] = b
|
|
}
|
|
elapsed := now.Sub(b.last).Seconds()
|
|
b.tokens += elapsed * l.rate
|
|
if b.tokens > l.burst {
|
|
b.tokens = l.burst
|
|
}
|
|
b.last = now
|
|
if b.tokens < 1 {
|
|
return false
|
|
}
|
|
b.tokens--
|
|
if len(l.buckets) > 8192 {
|
|
l.gc(now)
|
|
}
|
|
return true
|
|
}
|
|
|
|
func (l *Limiter) gc(now time.Time) {
|
|
for k, b := range l.buckets {
|
|
if now.Sub(b.last) > time.Minute {
|
|
delete(l.buckets, k)
|
|
}
|
|
}
|
|
}
|