fix(ratelimit): GC 按每个 key 自身窗口判过期,避免短窗口端点误删日闸/登录计数
- 根因:_buckets 全局共享、混着 60s(广告)/3600s(登录)/86400s(日闸)不同窗口的 key; 原 GC 用「当前调用方的 window_sec」判所有 key 过期,导致高频 60s 广告端点触发 GC 时, 把本该活 3600s/86400s 的登录/日闸计数一并删掉 → 规模上(超阈值才触发 GC)这些限流被 反复清零而失效(本系列新增的每日发码上限首当其冲)。 - 修:桶值改存 (start, count, window_sec),GC(_purge_expired)按每个 key 自身窗口判过期; 顺带把 _hit/_commit 里重复的 GC 抽成一处、阈值提为 _GC_THRESHOLD 常量(仿 sms.py,可测)。 - 测:新增 tests/test_ratelimit.py 覆盖「短窗口 GC 不误删仍在窗口内的长窗口 key」等 3 例。 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+23
-15
@@ -15,24 +15,35 @@ from fastapi import HTTPException, Request, status
|
||||
|
||||
from app.core.config import settings
|
||||
|
||||
# key -> (window_start_ts, count)
|
||||
_buckets: dict[str, tuple[float, int]] = {}
|
||||
# key -> (window_start_ts, count, window_sec)
|
||||
# 存每个 key 自己的 window_sec:_buckets 混着不同窗口(60s 广告 / 3600s 登录 / 86400s 日闸)的 key,
|
||||
# GC 必须按各 key 自己的窗口判过期(见 [_purge_expired]),否则短窗口调用触发的 GC 会误删长窗口 key。
|
||||
_buckets: dict[str, tuple[float, int, float]] = {}
|
||||
_lock = threading.Lock()
|
||||
_GC_THRESHOLD = 10000 # _buckets 超此阈值才顺手清过期 key(仿 sms.py;测试可 monkeypatch 调小强制每次扫)
|
||||
|
||||
|
||||
def _purge_expired(now: float) -> None:
|
||||
"""清过期 key(**仅在持有 _lock 时调用**)。按每个 key 自己存的 window_sec 判过期,而非调用方的窗口
|
||||
—— _buckets 是全局共享、混着 60s(广告)/3600s(登录)/86400s(日闸)不同窗口的 key;若用调用方窗口,
|
||||
高频的 60s 广告端点触发 GC 时会把本该活 3600s/86400s 的登录/日闸计数一并删掉,使其在规模上(超阈值才
|
||||
触发本清理)被反复清零而失效。仅在超阈值时扫,低频、开销可忽略。"""
|
||||
if len(_buckets) <= _GC_THRESHOLD:
|
||||
return
|
||||
for k in [k for k, (s, _, w) in _buckets.items() if now - s >= w]:
|
||||
_buckets.pop(k, None)
|
||||
|
||||
|
||||
def _hit(key: str, limit: int, window_sec: float) -> bool:
|
||||
"""记一次访问。返回 True=放行,False=超限。"""
|
||||
now = time.monotonic()
|
||||
with _lock:
|
||||
start, count = _buckets.get(key, (now, 0))
|
||||
start, count, _ = _buckets.get(key, (now, 0, window_sec))
|
||||
if now - start >= window_sec: # 窗口过期,重置
|
||||
start, count = now, 0
|
||||
count += 1
|
||||
_buckets[key] = (start, count)
|
||||
# 顺手清理过期 key,防内存无限涨(低频访问足够)
|
||||
if len(_buckets) > 10000:
|
||||
for k in [k for k, (s, _) in _buckets.items() if now - s >= window_sec]:
|
||||
_buckets.pop(k, None)
|
||||
_buckets[key] = (start, count, window_sec)
|
||||
_purge_expired(now) # 顺手清过期 key(按各自窗口),防内存无限涨
|
||||
return count <= limit
|
||||
|
||||
|
||||
@@ -111,7 +122,7 @@ def _peek(key: str, limit: int, window_sec: float) -> bool:
|
||||
与 [_commit] 配对实现「先判后记」——只在动作成功后才 _commit。"""
|
||||
now = time.monotonic()
|
||||
with _lock:
|
||||
start, count = _buckets.get(key, (now, 0))
|
||||
start, count, _ = _buckets.get(key, (now, 0, window_sec))
|
||||
if now - start >= window_sec: # 窗口已过期 → 视作已重置(count 归零)
|
||||
count = 0
|
||||
return count < limit
|
||||
@@ -121,14 +132,11 @@ def _commit(key: str, window_sec: float) -> None:
|
||||
"""记一次访问(+1)。窗口过期则以本次为起点重置。仅在动作成功后调用。"""
|
||||
now = time.monotonic()
|
||||
with _lock:
|
||||
start, count = _buckets.get(key, (now, 0))
|
||||
start, count, _ = _buckets.get(key, (now, 0, window_sec))
|
||||
if now - start >= window_sec: # 窗口过期,重置
|
||||
start, count = now, 0
|
||||
_buckets[key] = (start, count + 1)
|
||||
# 顺手清理过期 key,防内存无限涨(同 [_hit])
|
||||
if len(_buckets) > 10000:
|
||||
for k in [k for k, (s, _) in _buckets.items() if now - s >= window_sec]:
|
||||
_buckets.pop(k, None)
|
||||
_buckets[key] = (start, count + 1, window_sec)
|
||||
_purge_expired(now) # 顺手清过期 key(按各自窗口,同 [_hit])
|
||||
|
||||
|
||||
def check_rate_limits(request: Request, subject: str, rules: list[RateLimitRule]) -> None:
|
||||
|
||||
Reference in New Issue
Block a user