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:
zzhyyyyy
2026-07-15 19:57:41 +08:00
parent be8154c7a5
commit 0ccbcff2f7
2 changed files with 80 additions and 15 deletions
+23 -15
View File
@@ -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: