Sakurasan
902ecaeacc
Reorganize project structure: - backend/cmd/openteam/ — entry point - backend/internal/ — core packages - backend/middleware/ — HTTP middleware - backend/router/ — route setup - backend/wire/ — dependency injection - backend/pkg/ — shared utilities - backend/go.mod, go.sum — Go module files Updated Makefile to work from backend/ directory. Removed old lowercase makefile.
123 lines
2.4 KiB
Go
123 lines
2.4 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// Limiter 内存限流器
|
|
type Limiter struct {
|
|
mu sync.Mutex
|
|
|
|
// 每用户每秒请求数
|
|
userRPS map[uint64]*tokenBucket
|
|
|
|
// 密钥每日请求计数
|
|
keyDailyReq map[uint64]*dailyCounter
|
|
|
|
// 密钥每日 token 计数
|
|
keyDailyTokens map[uint64]*dailyCounter
|
|
}
|
|
|
|
type tokenBucket struct {
|
|
tokens float64
|
|
maxTokens float64
|
|
refillRate float64
|
|
lastRefill time.Time
|
|
}
|
|
|
|
type dailyCounter struct {
|
|
date string
|
|
count int64
|
|
}
|
|
|
|
func New() *Limiter {
|
|
return &Limiter{
|
|
userRPS: make(map[uint64]*tokenBucket),
|
|
keyDailyReq: make(map[uint64]*dailyCounter),
|
|
keyDailyTokens: make(map[uint64]*dailyCounter),
|
|
}
|
|
}
|
|
|
|
// AllowRequest 检查用户级每秒请求限制
|
|
func (l *Limiter) AllowRequest(userID uint64, rps int) bool {
|
|
if rps <= 0 {
|
|
return true
|
|
}
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
bucket, ok := l.userRPS[userID]
|
|
if !ok {
|
|
bucket = &tokenBucket{
|
|
tokens: float64(rps),
|
|
maxTokens: float64(rps),
|
|
refillRate: float64(rps),
|
|
lastRefill: time.Now(),
|
|
}
|
|
l.userRPS[userID] = bucket
|
|
}
|
|
|
|
now := time.Now()
|
|
elapsed := now.Sub(bucket.lastRefill).Seconds()
|
|
bucket.tokens += elapsed * bucket.refillRate
|
|
if bucket.tokens > bucket.maxTokens {
|
|
bucket.tokens = bucket.maxTokens
|
|
}
|
|
bucket.lastRefill = now
|
|
|
|
if bucket.tokens < 1 {
|
|
return false
|
|
}
|
|
bucket.tokens--
|
|
return true
|
|
}
|
|
|
|
// AllowRequestDaily 检查密钥每日请求配额
|
|
func (l *Limiter) AllowRequestDaily(keyID uint64, quota int) bool {
|
|
if quota <= 0 {
|
|
return true
|
|
}
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
today := time.Now().UTC().Format("2006-01-02")
|
|
counter, ok := l.keyDailyReq[keyID]
|
|
if !ok || counter.date != today {
|
|
l.keyDailyReq[keyID] = &dailyCounter{date: today, count: 1}
|
|
return true
|
|
}
|
|
if counter.count >= int64(quota) {
|
|
return false
|
|
}
|
|
counter.count++
|
|
return true
|
|
}
|
|
|
|
// TokensUsed 返回密钥今日 token 用量
|
|
func (l *Limiter) TokensUsed(keyID uint64) int64 {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
today := time.Now().UTC().Format("2006-01-02")
|
|
counter, ok := l.keyDailyTokens[keyID]
|
|
if !ok || counter.date != today {
|
|
return 0
|
|
}
|
|
return counter.count
|
|
}
|
|
|
|
// AddTokens 累加密钥今日 token 用量
|
|
func (l *Limiter) AddTokens(keyID uint64, tokens int64) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
|
|
today := time.Now().UTC().Format("2006-01-02")
|
|
counter, ok := l.keyDailyTokens[keyID]
|
|
if !ok || counter.date != today {
|
|
l.keyDailyTokens[keyID] = &dailyCounter{date: today, count: tokens}
|
|
return
|
|
}
|
|
counter.count += tokens
|
|
}
|