Sakurasan
9d3648ffe7
reqDaily/tokDaily/userHits 此前只增不减：停用密钥与久未活动用户的 条目长期堆积。超过阈值时在代理热路径清扫——旧日期的日计数条目 （下次使用本就清零）与窗口内无命中的用户条目直接删除。
93 lines
2.2 KiB
Go
93 lines
2.2 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestAllowRequestDaily(t *testing.T) {
|
|
l := New()
|
|
if !l.AllowRequestDaily(1, 2) {
|
|
t.Fatal("first request should be allowed")
|
|
}
|
|
if !l.AllowRequestDaily(1, 2) {
|
|
t.Fatal("second request should be allowed")
|
|
}
|
|
if l.AllowRequestDaily(1, 2) {
|
|
t.Fatal("third request should be blocked")
|
|
}
|
|
// 另一个 key 不受影响
|
|
if !l.AllowRequestDaily(2, 2) {
|
|
t.Fatal("other key should be allowed")
|
|
}
|
|
}
|
|
|
|
func TestTokensDaily(t *testing.T) {
|
|
l := New()
|
|
l.AddTokens(1, 100)
|
|
l.AddTokens(1, 50)
|
|
if got := l.TokensUsed(1); got != 150 {
|
|
t.Fatalf("tokens = %d, want 150", got)
|
|
}
|
|
if got := l.TokensUsed(2); got != 0 {
|
|
t.Fatalf("other key tokens = %d, want 0", got)
|
|
}
|
|
// 配额检查语义：已达 150 再设配额 150 应拒绝
|
|
l.AddTokens(3, 150)
|
|
if l.TokensUsed(3) >= 150 {
|
|
// 注意：>= 表示已达上限，Auth 层据此拒绝
|
|
}
|
|
}
|
|
|
|
func TestAllowUserRate(t *testing.T) {
|
|
l := New()
|
|
// limit=1：立即第二次应被拒绝
|
|
if !l.AllowUserRate(9, 1) {
|
|
t.Fatal("first should be allowed")
|
|
}
|
|
if l.AllowUserRate(9, 1) {
|
|
t.Fatal("second in same second should be blocked")
|
|
}
|
|
// limit<=0 不限制
|
|
if !l.AllowUserRate(9, 0) {
|
|
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("无命中的用户速率条目应被清扫")
|
|
}
|
|
}
|