diff --git a/server/internal/pkg/ratelimit/ratelimit.go b/server/internal/pkg/ratelimit/ratelimit.go index 2e4435e..12bd93f 100644 --- a/server/internal/pkg/ratelimit/ratelimit.go +++ b/server/internal/pkg/ratelimit/ratelimit.go @@ -53,9 +53,44 @@ func (l *Limiter) AllowRequestDaily(keyID uint64, limit int) bool { if limit > 0 && c.n > int64(limit) { return false } + // 代理热路径顺带做回收:停用密钥的旧日期条目和过期的用户窗口不再堆积 + if len(l.reqDaily) > dailySweepThreshold { + l.sweepLocked(d) + } return true } +const dailySweepThreshold = 4096 + +// sweepLocked 回收不再活跃的条目(调用方持锁): +// - 日计数器:日期不是今天的条目下次使用本来就会清零,直接删除 +// - 用户速率窗口:窗口内无命中的条目删除 +func (l *Limiter) sweepLocked(today string) { + now := time.Now() + for k, c := range l.reqDaily { + if c.date != today { + delete(l.reqDaily, k) + } + } + for k, c := range l.tokDaily { + if c.date != today { + delete(l.tokDaily, k) + } + } + for k, w := range l.userHits { + cutoff := now.Add(-w.window) + alive := 0 + for _, t := range w.times { + if t.After(cutoff) { + alive++ + } + } + if alive == 0 { + delete(l.userHits, k) + } + } +} + // AddTokens 累计密钥今日 token 用量(请求结束后记账)。 func (l *Limiter) AddTokens(keyID uint64, n int64) { if n <= 0 { diff --git a/server/internal/pkg/ratelimit/ratelimit_test.go b/server/internal/pkg/ratelimit/ratelimit_test.go index 19dd39e..b6d3096 100644 --- a/server/internal/pkg/ratelimit/ratelimit_test.go +++ b/server/internal/pkg/ratelimit/ratelimit_test.go @@ -1,6 +1,9 @@ package ratelimit -import "testing" +import ( + "testing" + "time" +) func TestAllowRequestDaily(t *testing.T) { l := New() @@ -50,3 +53,40 @@ func TestAllowUserRate(t *testing.T) { t.Fatal("limit 0 should always allow") } } + +// 超过阈值时清扫:旧日期的日计数条目和窗口内无命中的用户条目被回收。 +func TestSweepRemovesStaleEntries(t *testing.T) { + l := New() + l.AllowRequestDaily(1, 0) + l.AddTokens(1, 100) + l.AllowUserRate(2, 5) + // 人为把日计数条目改成"昨天",模拟跨天后的陈旧条目 + l.mu.Lock() + for _, c := range l.reqDaily { + c.date = "2000-01-01" + } + for _, c := range l.tokDaily { + c.date = "2000-01-01" + } + // 用户速率条目:命中时间回溯到窗口外,模拟久未活动的用户 + if w, ok := l.userHits[2]; ok { + w.times = []time.Time{time.Now().Add(-time.Hour)} + } + l.mu.Unlock() + + for i := 0; i < dailySweepThreshold+1; i++ { + l.AllowRequestDaily(uint64(100+i), 0) + } + + l.mu.Lock() + _, reqStale := l.reqDaily[1] + _, tokStale := l.tokDaily[1] + _, userStale := l.userHits[2] + l.mu.Unlock() + if reqStale || tokStale { + t.Error("陈旧的日计数条目应被清扫") + } + if userStale { + t.Error("无命中的用户速率条目应被清扫") + } +}