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>
2026-07-15 19:57:41 +08:00
parent be8154c7a5
commit 0ccbcff2f7
2 changed files with 80 additions and 15 deletions
@@ -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:
@@ -0,0 +1,57 @@
+"""ratelimit 内存桶过期清理(GC)测试。
+
+回归重点:_buckets 是**全局共享**、混着不同窗口(60s 广告 / 3600s 登录 / 86400s 日闸)的 key。
+GC 必须按【每个 key 自己存的 window_sec】判过期,而不是当前调用方的窗口 —— 否则高频的 60s 端点
+触发 GC 时会把本该存活更久的 3600s/86400s 计数(如短信日闸)一并删掉,使其被反复清零、限流失效。
+用 monkeypatch 把 _GC_THRESHOLD 调 0 强制每次都扫,免造上万条(仿 test_auth 里对 sms._GC_THRESHOLD 的做法)。
+"""
+from __future__ import annotations
+
+from app.core import ratelimit
+
+
+def test_purge_expired_respects_each_key_own_window(monkeypatch) -> None:
+    """短窗口(60s)触发的 GC 只删真正过期的 key,不得删掉仍在自身窗口内的长窗口 key。"""
+    monkeypatch.setattr(ratelimit, "_GC_THRESHOLD", 0)  # 强制每次都扫
+    ratelimit._buckets.clear()
+
+    now = 1_000_000.0
+    # 日闸:100s 前开窗、window=86400 → 远未过期,必须保留
+    ratelimit._buckets["sms-send-device-daily:D:IP"] = (now - 100, 7, 86400.0)
+    # 登录:1800s、window=3600 → 未过期,保留
+    ratelimit._buckets["sms-login-device:D:IP"] = (now - 1800, 2, 3600.0)
+    # 广告:120s、window=60 → 已过期,应删
+    ratelimit._buckets["ad-watch-report:IP2"] = (now - 120, 3, 60.0)
+
+    ratelimit._purge_expired(now)
+
+    assert "sms-send-device-daily:D:IP" in ratelimit._buckets
+    assert "sms-login-device:D:IP" in ratelimit._buckets
+    assert "ad-watch-report:IP2" not in ratelimit._buckets
+
+
+def test_purge_expired_keeps_long_window_key_older_than_short_window(monkeypatch) -> None:
+    """反证旧 bug:日闸 key 已老于 3600s,旧代码在 60s/3600s 端点触发 GC 时会误删它;
+    现在按自身 86400s 窗口判 → 未过期 → 必须保留。"""
+    monkeypatch.setattr(ratelimit, "_GC_THRESHOLD", 0)
+    ratelimit._buckets.clear()
+
+    now = 2_000_000.0
+    # 3700s 前开窗(> 1 小时),但 window=86400 → 未过期
+    ratelimit._buckets["sms-send-device-daily:D:IP"] = (now - 3700, 20, 86400.0)
+
+    ratelimit._purge_expired(now)
+
+    assert "sms-send-device-daily:D:IP" in ratelimit._buckets
+
+
+def test_purge_expired_noop_below_threshold(monkeypatch) -> None:
+    """未超阈值时不扫(即便有过期 key 也不动),避免每次请求都 O(n) 扫全表。"""
+    monkeypatch.setattr(ratelimit, "_GC_THRESHOLD", 10)
+    ratelimit._buckets.clear()
+
+    now = 3_000_000.0
+    ratelimit._buckets["stale:IP"] = (now - 999, 1, 60.0)  # 早过期,但没超阈值
+    ratelimit._purge_expired(now)
+
+    assert "stale:IP" in ratelimit._buckets  # 桶数没超阈值 → 不清理