fix: 限流日计数与用户速率窗口的条目回收
reqDaily/tokDaily/userHits 此前只增不减:停用密钥与久未活动用户的 条目长期堆积。超过阈值时在代理热路径清扫——旧日期的日计数条目 (下次使用本就清零)与窗口内无命中的用户条目直接删除。
This commit is contained in:
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("无命中的用户速率条目应被清扫")
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user