fix: 限流日计数与用户速率窗口的条目回收

reqDaily/tokDaily/userHits 此前只增不减:停用密钥与久未活动用户的
条目长期堆积。超过阈值时在代理热路径清扫——旧日期的日计数条目
(下次使用本就清零)与窗口内无命中的用户条目直接删除。
This commit is contained in:
Sakurasan committed 2026-10-30 01:42:22 +08:00
1 parent bb20ebbed0
commit 9d3648ffe7
2 files changed
+76 -1

No files matched your search

@@ -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 {
@@ -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("无命中的用户速率条目应被清扫")
}
}