host
8bda29eac3
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 不炸。
135 lines
5.0 KiB
Python
135 lines
5.0 KiB
Python
"""HA 阶段 2 单测：DB 共享限流。
|
||
|
||
DbRateLimiter 与内存版同接口：窗口内达到阈值即 False、窗口
|
||
滑过去后恢复、不同 key 互不干扰；关键一条是多实例共享计数
|
||
（两个 server 进程各建一个实例，计数仍合流——多 server 限流
|
||
的前提）。SQLite 跑同款 SQL，Postgres 方言差异由 db.py 适配层
|
||
覆盖（本文件不连真实 Postgres）。
|
||
"""
|
||
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import unittest
|
||
from pathlib import Path
|
||
|
||
try:
|
||
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,
|
||
sweep_rate_limit_hits)
|
||
|
||
|
||
class DbRateLimiterTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
self.now = [1000.0]
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _limiter(self, per_minute=3, window=60):
|
||
return DbRateLimiter(self.db_path, per_minute=per_minute,
|
||
window_seconds=window,
|
||
time_fn=lambda: self.now[0])
|
||
|
||
def test_allow_up_to_limit_then_false(self):
|
||
lim = self._limiter()
|
||
self.assertTrue(lim.allow("k1"))
|
||
self.assertTrue(lim.allow("k1"))
|
||
self.assertTrue(lim.allow("k1"))
|
||
self.assertFalse(lim.allow("k1"))
|
||
|
||
def test_window_slides_and_recovers(self):
|
||
lim = self._limiter()
|
||
for _ in range(3):
|
||
self.assertTrue(lim.allow("k1"))
|
||
self.assertFalse(lim.allow("k1"))
|
||
self.now[0] += 61
|
||
self.assertTrue(lim.allow("k1"))
|
||
|
||
def test_keys_are_independent(self):
|
||
lim = self._limiter()
|
||
for _ in range(3):
|
||
self.assertTrue(lim.allow("a"))
|
||
self.assertFalse(lim.allow("a"))
|
||
self.assertTrue(lim.allow("b"))
|
||
|
||
def test_two_instances_share_counts(self):
|
||
# 多 server 的核心性质：两个实例（两个进程）计数合流。
|
||
a = self._limiter()
|
||
b = self._limiter()
|
||
self.assertTrue(a.allow("shared"))
|
||
self.assertTrue(a.allow("shared"))
|
||
self.assertTrue(b.allow("shared"))
|
||
self.assertFalse(b.allow("shared"))
|
||
self.assertFalse(a.allow("shared"))
|
||
|
||
def test_memory_limiter_unchanged(self):
|
||
lim = RateLimiter(per_minute=2, window_seconds=60,
|
||
time_fn=lambda: self.now[0])
|
||
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(), [])
|