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>
74 lines
1.4 KiB
Go
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
|
|
}
|