diff --git a/cloud/server.py b/cloud/server.py index de042b5..7f00f4e 100644 --- a/cloud/server.py +++ b/cloud/server.py @@ -265,13 +265,21 @@ class DbRateLimiter: """DB 共享滑动窗口限流（HA 阶段 2，多 server 共用计数）。 与 RateLimiter 同接口（allow(key)，per_minute/window_seconds - 属性可改，测试注入小阈值验证 429 的路子不变）：每次 allow 写 - 一行 hit、清过期行、数窗口内行数。Postgres 下多 server 共享 - 同一张表，计数跨进程；SQLite 下与内存版等价（单测用）。 + 属性可改，测试注入小阈值验证 429 的路子不变）。Postgres 下多 + server 共享同一张表，计数跨进程；SQLite 下与内存版等价（单测用）。 并发下计数与写入非原子，可能多放行少量请求——限流是 abuse 缓解，不是精确配额。表与索引懒建（workers 表同款模式）， 不进 POSTGRES_SCHEMA。列名用 hit_key（key 在部分方言里是 保留字，避开它）。 + + 结构要点（2026-10-27 死锁修复）：allow() 只做 INSERT + 窗口内 + COUNT，不再 DELETE——旧实现每请求在同一事务里 DELETE 旧行 + + INSERT + COUNT，控制台并发加载 + 双 server 下两个事务在 + rate_limit_hits 表上死锁，未捕获的异常直接断连、nginx 报 + 502。过期行改由后台 sweep_rate_limit_hits() 定期清（见下）， + 请求路径上只剩 INSERT + COUNT，不存在 DELETE↔DELETE/INSERT + 的死锁形状；COUNT 本身按窗口过滤，过期行残留不影响计数 + 正确性，只影响表大小。 """ def __init__(self, db_path, per_minute=DEFAULT_RATE_LIMIT_PER_MINUTE, @@ -281,6 +289,25 @@ class DbRateLimiter: self.per_minute = per_minute self.window_seconds = window_seconds self._time_fn = time_fn or time.time + self._ensure_lock = threading.Lock() + self._ensured = False + + def _ensure_table(self, conn): + """建表与索引：每实例只做一次（首个 allow 时），不是每请求 + 都跑 DDL——DDL 也会拿锁，没必要让每个请求都经过它。""" + if self._ensured: + return + with self._ensure_lock: + if self._ensured: + return + conn.execute( + "CREATE TABLE IF NOT EXISTS rate_limit_hits(" + "hit_key TEXT, ts REAL)") + conn.execute( + "CREATE INDEX IF NOT EXISTS rate_limit_hits_key_ts" + " ON rate_limit_hits(hit_key, ts)") + conn.commit() + self._ensured = True def allow(self, key): now = self._time_fn() @@ -288,21 +315,14 @@ class DbRateLimiter: limit = self.per_minute conn = _db_connect(self.db_path) try: - conn.execute( - "CREATE TABLE IF NOT EXISTS rate_limit_hits(" - "hit_key TEXT, ts REAL)") - conn.execute( - "CREATE INDEX IF NOT EXISTS rate_limit_hits_key_ts" - " ON rate_limit_hits(hit_key, ts)") - conn.execute( - "DELETE FROM rate_limit_hits WHERE ts <= ?", - (now - window,)) + self._ensure_table(conn) conn.execute( "INSERT INTO rate_limit_hits(hit_key, ts) VALUES(?, ?)", (key, now)) row = conn.execute( - "SELECT COUNT(*) FROM rate_limit_hits WHERE hit_key = ?", - (key,)).fetchone() + "SELECT COUNT(*) FROM rate_limit_hits" + " WHERE hit_key = ? AND ts > ?", + (key, now - window)).fetchone() conn.commit() return row[0] <= limit finally: @@ -324,6 +344,70 @@ def _make_rate_limiter(db_path): window_seconds=window_seconds) +RATE_LIMIT_SWEEP_INTERVAL_SECONDS = 60 +# Postgres advisory lock key：两个 server 各起一个 sweeper 时，只 +# 有拿到锁的那个真正 DELETE，另一个直接跳过——清理者之间也不 +# 会互相死锁。本项目其它地方不用 advisory lock，不会撞 key。 +RATE_LIMIT_SWEEP_ADVISORY_KEY = 20261027 + + +def sweep_rate_limit_hits(db_path, window_seconds, now=None): + """删过期 hit 行：只由后台定期调用，不在请求路径上。 + + 返回 True 表示本次真正执行了删除（表为空时也算），False + 表示 Postgres 下没拿到 advisory 锁、另一台 server 正在清， + 本次跳过。清理是纯卫生（防表无限增长）：allow() 的 COUNT + 本身按窗口过滤，过期行残留不影响计数，所以清理失败或 + 跳过都不影响限流正确性，只影响表大小。 + """ + now = time.time() if now is None else now + conn = _db_connect(db_path) + try: + conn.execute( + "CREATE TABLE IF NOT EXISTS rate_limit_hits(" + "hit_key TEXT, ts REAL)") + conn.execute( + "CREATE INDEX IF NOT EXISTS rate_limit_hits_key_ts" + " ON rate_limit_hits(hit_key, ts)") + if _db_backend_name() == "postgres": + row = conn.execute( + "SELECT pg_try_advisory_xact_lock(?)", + (RATE_LIMIT_SWEEP_ADVISORY_KEY,)).fetchone() + if not row or not row[0]: + return False + conn.execute( + "DELETE FROM rate_limit_hits WHERE ts <= ?", + (now - window_seconds,)) + conn.commit() + return True + finally: + conn.close() + + +def start_rate_limit_sweeper(limiter, db_path): + """给 server 主进程起后台清理线程：每 60 秒清一次过期行。 + + 只对 DbRateLimiter 起（内存版在 allow() 里自行修剪，不需要）。 + 线程 daemon、异常自愈（记 warning 后继续下一轮），绝不因 + 清理失败拖垮 server。调用方只在 main() 里调一次，测试不调。 + """ + if not isinstance(limiter, DbRateLimiter): + return None + + def _loop(): + while True: + time.sleep(RATE_LIMIT_SWEEP_INTERVAL_SECONDS) + try: + sweep_rate_limit_hits(db_path, limiter.window_seconds) + except Exception as exc: # noqa: BLE001 -- 清理是卫生操作 + _server_log.warning("rate limit sweep failed: %s", exc) + + thread = threading.Thread(target=_loop, daemon=True, + name="rate-limit-sweeper") + thread.start() + return thread + + def _db_backend_name(): """当前数据库后端名：环境变量，未设则 sqlite；绝不返回连接串。""" return (os.environ.get("CLOUD_DB_BACKEND") or "sqlite").strip().lower() @@ -4274,6 +4358,7 @@ def main(argv=None): args = parser.parse_args(argv) server = create_server(args.db, args.host, args.port, backend=args.db_backend) + start_rate_limit_sweeper(server.rate_limiter, args.db) print( "cloud-control-plane listening on %s:%d db=%s" % (args.host, server.server_address[1], args.db), diff --git a/cloud/tests/test_rate_limit_db.py b/cloud/tests/test_rate_limit_db.py index 80aee23..f26d6e3 100644 --- a/cloud/tests/test_rate_limit_db.py +++ b/cloud/tests/test_rate_limit_db.py @@ -8,16 +8,19 @@ DbRateLimiter 与内存版同接口：窗口内达到阈值即 False、窗口 """ import os +import sqlite3 import sys import tempfile import unittest from pathlib import Path try: - from cloud.server import DbRateLimiter, RateLimiter + from cloud.server import (DbRateLimiter, RateLimiter, + sweep_rate_limit_hits) except ImportError: sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - from cloud.server import DbRateLimiter, RateLimiter + from cloud.server import (DbRateLimiter, RateLimiter, + sweep_rate_limit_hits) class DbRateLimiterTestCase(unittest.TestCase): @@ -72,3 +75,60 @@ class DbRateLimiterTestCase(unittest.TestCase): self.assertTrue(lim.allow("k")) self.assertTrue(lim.allow("k")) self.assertFalse(lim.allow("k")) + + def _rows(self): + conn = sqlite3.connect(self.db_path) + try: + return sorted(conn.execute( + "SELECT hit_key, ts FROM rate_limit_hits").fetchall()) + finally: + conn.close() + + def _seed(self, rows): + conn = sqlite3.connect(self.db_path) + try: + conn.executemany( + "INSERT INTO rate_limit_hits(hit_key, ts) VALUES(?, ?)", + rows) + conn.commit() + finally: + conn.close() + + def test_allow_does_not_delete_old_rows(self): + # 结构修复的核心断言：allow() 只 INSERT + 窗口内 COUNT， + # 不再 DELETE——过期行在 allow() 之后必须还在表里，只等 + # 后台 sweep 来清。旧实现里这一步 DELETE 会把它们清掉。 + lim = self._limiter(per_minute=10, window=60) + self.assertTrue(lim.allow("k")) + self.assertTrue(lim.allow("k")) + self._seed([("old", 500.0), ("k", 500.0)]) + self.now[0] = 2000.0 + self.assertTrue(lim.allow("k")) + rows = self._rows() + self.assertIn(("old", 500.0), rows) + self.assertIn(("k", 500.0), rows) + self.assertIn(("k", 2000.0), rows) + + def test_windowed_count_ignores_expired_rows(self): + # COUNT 按窗口过滤：即使过期行还残留在表里，窗口滑过去 + # 后计数也必须恢复，不依赖 sweep 先跑过。 + lim = self._limiter(per_minute=3, window=60) + for _ in range(3): + self.assertTrue(lim.allow("k1")) + self.assertFalse(lim.allow("k1")) + self.now[0] += 61 + self.assertTrue(lim.allow("k1")) + # 旧行仍在表里（allow 没删），但计数已不含它们。 + self.assertEqual(len(self._rows()), 5) + + def test_sweep_deletes_only_expired_rows(self): + lim = self._limiter(per_minute=10, window=60) + self.assertTrue(lim.allow("seed")) + self._seed([("old", 100.0), ("new", 1990.0)]) + self.assertTrue(sweep_rate_limit_hits(self.db_path, 60, now=2000.0)) + self.assertEqual(self._rows(), [("new", 1990.0)]) + + def test_sweep_on_fresh_db_is_noop_true(self): + # 表还不存在时 sweep 不应炸：建空表、返回 True。 + self.assertTrue(sweep_rate_limit_hits(self.db_path, 60, now=2000.0)) + self.assertEqual(self._rows(), [])