fix: 限流清理移出请求路径，根治 rate_limit_hits 死锁
allow() 原先每请求在同一事务里 DELETE 旧行 + INSERT + COUNT， 控制台并发加载 + 双 server 下两个事务在 rate_limit_hits 表上 死锁，未捕获异常直接断连、nginx 报 502（2026-10-27 实测， 环境与链路两区同时 502 均源于此，与各自端点逻辑无关）。 结构修复： - allow() 只做 INSERT + 窗口内 COUNT（COUNT 按 ts > now-window 过滤，过期行残留不影响计数正确性）；建表/索引每实例只做一次。 - 过期行改由后台线程 sweep_rate_limit_hits() 每 60 秒清一次； Postgres 下拿事务级 advisory lock，另一台 server 正在清时 直接跳过，清理者之间也不死锁；清理失败只记 warning。 - main() 里对 DbRateLimiter 起 sweeper 线程（内存版不需要）。 单测：allow 不再删旧行、窗口计数忽略过期行、sweep 只删过期 行、空库 sweep 不炸。
This commit is contained in:
host committed
1 parent
5ba4d35839
commit
8bda29eac3
2 files changed
+161
-16
No files matched your search
+99
-14
@@ -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),
|
||||
|
||||
@@ -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(), [])
|
||||
Reference in new issue
Block a user