From 06b9fd12fd8f1731d6bb3a824e427f78516eba25 Mon Sep 17 00:00:00 2001
From: openqj <openqj@gmail.com>
Date: Tue, 27 Oct 2026 00:15:16 +0800
Subject: [PATCH] HA stage 2: DB-shared rate limiting + second server unit +
 monitor both-servers check (#52)

Co-authored-by: host <host@test.local>
---
 cloud/server.py                     | 87 +++++++++++++++++++++++------
 cloud/tests/test_rate_limit_db.py   | 74 ++++++++++++++++++++++++
 deploy/codex-cloud-server-2.service | 20 +++++++
 deploy/monitor.sh                   |  8 ++-
 4 files changed, 170 insertions(+), 19 deletions(-)
 create mode 100644 cloud/tests/test_rate_limit_db.py
 create mode 100644 deploy/codex-cloud-server-2.service

diff --git a/cloud/server.py b/cloud/server.py
index fab43d3..c8a90d0 100644
--- a/cloud/server.py
+++ b/cloud/server.py
@@ -127,8 +127,8 @@ TERMINAL_STATUSES = frozenset(
 DEFAULT_SWEEP_THRESHOLD_SECONDS = 300
 # 限流默认值：每 IP / 每 API key 每分钟 120 次；环境变量
 # CLOUD_RATE_LIMIT_PER_MINUTE / CLOUD_RATE_LIMIT_WINDOW_SECONDS 可调。
-# 限流是内存滑动窗口，不跨进程共享（多 worker 部署时每进程独立计数，
-# 这是诚实边界，不是全局限流）。
+# 限流默认内存滑动窗口；Postgres 下改走 DB 共享计数（多 server
+# 共用，见 DbRateLimiter），并发下可能多放行少量请求（非原子计数）。
 DEFAULT_RATE_LIMIT_PER_MINUTE = 120
 DEFAULT_RATE_LIMIT_WINDOW_SECONDS = 60
 TASK_STATUS_KEYS = (
@@ -201,6 +201,69 @@ class RateLimiter:
             return True
 
 
+class DbRateLimiter:
+    """DB 共享滑动窗口限流（HA 阶段 2，多 server 共用计数）。
+
+    与 RateLimiter 同接口（allow(key)，per_minute/window_seconds
+    属性可改，测试注入小阈值验证 429 的路子不变）：每次 allow 写
+    一行 hit、清过期行、数窗口内行数。Postgres 下多 server 共享
+    同一张表，计数跨进程；SQLite 下与内存版等价（单测用）。
+    并发下计数与写入非原子，可能多放行少量请求——限流是 abuse
+    缓解，不是精确配额。表与索引懒建（workers 表同款模式），
+    不进 POSTGRES_SCHEMA。列名用 hit_key（key 在部分方言里是
+    保留字，避开它）。
+    """
+
+    def __init__(self, db_path, per_minute=DEFAULT_RATE_LIMIT_PER_MINUTE,
+                 window_seconds=DEFAULT_RATE_LIMIT_WINDOW_SECONDS,
+                 time_fn=None):
+        self.db_path = str(db_path)
+        self.per_minute = per_minute
+        self.window_seconds = window_seconds
+        self._time_fn = time_fn or time.time
+
+    def allow(self, key):
+        now = self._time_fn()
+        window = self.window_seconds
+        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,))
+            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()
+            conn.commit()
+            return row[0] <= limit
+        finally:
+            conn.close()
+
+
+def _make_rate_limiter(db_path):
+    """建限流器：Postgres 下用 DB 共享计数（多 server 共用），
+    否则用内存版（单机与单测）。后端只看环境变量，与其它分发一致。"""
+    per_minute = _rate_limit_from_env(
+        "CLOUD_RATE_LIMIT_PER_MINUTE", DEFAULT_RATE_LIMIT_PER_MINUTE)
+    window_seconds = _rate_limit_from_env(
+        "CLOUD_RATE_LIMIT_WINDOW_SECONDS",
+        DEFAULT_RATE_LIMIT_WINDOW_SECONDS)
+    if _db_backend_name() == "postgres":
+        return DbRateLimiter(db_path, per_minute=per_minute,
+                             window_seconds=window_seconds)
+    return RateLimiter(per_minute=per_minute,
+                       window_seconds=window_seconds)
+
+
 def _db_backend_name():
     """当前数据库后端名：环境变量，未设则 sqlite；绝不返回连接串。"""
     return (os.environ.get("CLOUD_DB_BACKEND") or "sqlite").strip().lower()
@@ -705,13 +768,7 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler):
         # 无 key 时按客户端 IP 计。
         limiter = getattr(self.server, "rate_limiter", None)
         if limiter is None:
-            limiter = RateLimiter(
-                per_minute=_rate_limit_from_env(
-                    "CLOUD_RATE_LIMIT_PER_MINUTE",
-                    DEFAULT_RATE_LIMIT_PER_MINUTE),
-                window_seconds=_rate_limit_from_env(
-                    "CLOUD_RATE_LIMIT_WINDOW_SECONDS",
-                    DEFAULT_RATE_LIMIT_WINDOW_SECONDS))
+            limiter = _make_rate_limiter(self.server.db_path)
             self.server.rate_limiter = limiter
         # 测试注入：server.rate_limit_per_minute / rate_limit_window_seconds
         # 改动后即时生效，不用重启 server。
@@ -2796,14 +2853,10 @@ def create_server(db_path, host="127.0.0.1", port=8765, backend=None):
     server = ThreadingHTTPServer((host, port), CloudControlPlaneHandler)
     server.db_path = db_path
     server.daemon_threads = True
-    # 限流器：内存滑动窗口，默认值与环境变量覆盖见 RateLimiter；
-    # 测试可直接改 server.rate_limit_per_minute / window_seconds。
-    server.rate_limiter = RateLimiter(
-        per_minute=_rate_limit_from_env(
-            "CLOUD_RATE_LIMIT_PER_MINUTE", DEFAULT_RATE_LIMIT_PER_MINUTE),
-        window_seconds=_rate_limit_from_env(
-            "CLOUD_RATE_LIMIT_WINDOW_SECONDS",
-            DEFAULT_RATE_LIMIT_WINDOW_SECONDS))
+    # 限流器：Postgres 下 DB 共享计数（多 server 共用），否则内存
+    # 滑动窗口；默认值与环境变量覆盖见 _make_rate_limiter；测试可
+    # 直接改 server.rate_limit_per_minute / window_seconds。
+    server.rate_limiter = _make_rate_limiter(db_path)
     server.rate_limit_per_minute = None
     server.rate_limit_window_seconds = None
     # OAuth 回调注入点：生产为 None（走真实换 token + 真实 userinfo），
diff --git a/cloud/tests/test_rate_limit_db.py b/cloud/tests/test_rate_limit_db.py
new file mode 100644
index 0000000..80aee23
--- /dev/null
+++ b/cloud/tests/test_rate_limit_db.py
@@ -0,0 +1,74 @@
+"""HA 阶段 2 单测：DB 共享限流。
+
+DbRateLimiter 与内存版同接口：窗口内达到阈值即 False、窗口
+滑过去后恢复、不同 key 互不干扰；关键一条是多实例共享计数
+（两个 server 进程各建一个实例，计数仍合流——多 server 限流
+的前提）。SQLite 跑同款 SQL，Postgres 方言差异由 db.py 适配层
+覆盖（本文件不连真实 Postgres）。
+"""
+
+import os
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+try:
+    from cloud.server import DbRateLimiter, RateLimiter
+except ImportError:
+    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+    from cloud.server import DbRateLimiter, RateLimiter
+
+
+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"))
diff --git a/deploy/codex-cloud-server-2.service b/deploy/codex-cloud-server-2.service
new file mode 100644
index 0000000..48f34db
--- /dev/null
+++ b/deploy/codex-cloud-server-2.service
@@ -0,0 +1,20 @@
+# 第二台控制面（HA 阶段 2）。
+#
+# 注意：unit 文件不在 deploy tar 里，改动后须单独 scp 到
+# /etc/systemd/system/ 再 daemon-reload + enable + start。
+# 与 codex-cloud-server.service 完全相同，只改端口（8766）：
+# 同代码、同 env、同 DB 连接串；nginx upstream 把两个端口都
+# 配上，任一挂了另一台继续接（webhook 去重在 DB 里，重复投递
+# 幂等）。WorkingDirectory 与 User 与一号机一致。
+[Unit]
+Description=Codex cloud control plane (second instance)
+After=network.target
+[Service]
+User=codexsvc
+WorkingDirectory=/opt/codex-app-server
+EnvironmentFile=/etc/codex-app-server/env
+ExecStart=/usr/bin/python3 -m cloud.server --host 127.0.0.1 --port 8766 --db /var/lib/codex-app-server/control.db
+Restart=on-failure
+RestartSec=10
+[Install]
+WantedBy=multi-user.target
diff --git a/deploy/monitor.sh b/deploy/monitor.sh
index d7d2de7..d015d28 100755
--- a/deploy/monitor.sh
+++ b/deploy/monitor.sh
@@ -6,7 +6,8 @@
 #   sh /opt/codex-app-server/deploy/monitor.sh --self-test
 #
 # 检查项（全部不需要密码、不读 env 文件）：
-#   server：systemctl is-active codex-cloud-server 为 active
+#   server：systemctl is-active codex-cloud-server 与
+#          codex-cloud-server-2 均为 active（HA 阶段 2 双机）
 #   worker：codex-autoscaler 或 codex-cloud-worker 任一为 active
 #          （autoscaler 接管后静态 worker 故意停用，只查静态会误报：
 #           2026-10-26 07:02:55 静态 worker 停用后每 5 分钟 worker fail，
@@ -41,7 +42,10 @@ DISK_PATH=${MONITOR_DISK_PATH:-/var/lib/codex-app-server}
 if [ ! -d "$DISK_PATH" ]; then DISK_PATH="/"; fi
 
 check_server() {
-  if systemctl is-active codex-cloud-server >/dev/null 2>&1; then
+  # HA 阶段 2：两台 server 都必须 active，任一挂了都算 fail
+  # （单台 active 时 API 还能用，但冗余已丢，必须可见）。
+  if systemctl is-active codex-cloud-server >/dev/null 2>&1 \
+    && systemctl is-active codex-cloud-server-2 >/dev/null 2>&1; then
     echo pass
   else
     echo fail