Files
openteam/server/internal/pkg/ratelimit/ratelimit.go
T
SakurasanandClaude Sonnet 5 d0e31b198f feat(server): API relay gateway backend M0-M4
Gin + GORM + pure-Go SQLite. Users/auth (JWT), API key management with
quotas, proxy gateway with weighted channel failover and health checks,
usage/billing ledger, cross-protocol conversion (Anthropic Messages /
OpenAI Chat Completions / OpenAI Responses), and channel/model admin API.
Channels declare native API formats and auto-convert the rest.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-15 21:05:02 +08:00

74 lines
1.4 KiB
Go

package ratelimit
import (
"sync"
"time"
)
// Limiter is a token-bucket rate limiter keyed by string.
type Limiter struct {
mu sync.Mutex
rate float64 // tokens per second
burst float64
tokens map[string]*bucket
}
type bucket struct {
tokens float64
lastFill time.Time
}
// New creates a limiter refilling `rate` tokens/sec with `burst` capacity.
func New(rate float64, burst int) *Limiter {
return &Limiter{
rate: rate,
burst: float64(burst),
tokens: map[string]*bucket{},
}
}
// Allow checks whether `key` may take one token now.
func (l *Limiter) Allow(key string) bool {
return l.Take(key, 1)
}
// Take checks whether `key` may take n tokens now.
func (l *Limiter) Take(key string, n float64) bool {
l.mu.Lock()
defer l.mu.Unlock()
now := time.Now()
b, ok := l.tokens[key]
if !ok {
b = &bucket{tokens: l.burst, lastFill: now}
l.tokens[key] = b
}
// Refill based on elapsed time.
elapsed := now.Sub(b.lastFill).Seconds()
b.tokens = minF(l.burst, b.tokens+elapsed*l.rate)
b.lastFill = now
if b.tokens >= n {
b.tokens -= n
return true
}
return false
}
// Sweep removes idle buckets to bound memory. Call periodically.
func (l *Limiter) Sweep(olderThan time.Duration) {
l.mu.Lock()
defer l.mu.Unlock()
cutoff := time.Now().Add(-olderThan)
for k, b := range l.tokens {
if b.lastFill.Before(cutoff) {
delete(l.tokens, k)
}
}
}
func minF(a, b float64) float64 {
if a < b {
return a
}
return b
}