Merge pull request #97 from openqj/feat-login-retry

登录限流下开通突发不再硬失败（节流 + 429 重试）

cloud/sub2api.py

@@ -13,8 +13,10 @@
   便于单测（见 cloud/tests/test_sub2api.py）。
 """
 
+import collections
 import json
 import os
+import random
 import secrets
 import threading
 import time
@@ -162,6 +164,67 @@ def shared_admin_session(request=None):
         return admin
 
 
+def _is_rate_limited_error(exc):
+    if getattr(exc, "status", None) == 429:
+        return True
+    return "too many requests" in str(exc).lower()
+
+
+# 登录重试：网关登录端点 20 次/分钟/IP（固定窗口），突发注册
+# 必然撞限；退避 2/4/8/16/32 秒（+0-2 秒抖动）可跨过 60 秒窗口。
+_LOGIN_RETRY_ATTEMPTS = 6
+_LOGIN_RETRY_BASE_SECONDS = 2.0
+
+
+def _call_with_login_retry(fn):
+    """登录类调用：先取进程内节流名额，再遇 429 退避重试。
+
+    节流名额保证单进程内登录不超过预算（网关 20 次/分钟/IP，
+    留 5 次余量给其它进程与人工登录）；重试兜底跨进程竞争。
+    其它错误与耗尽次数后原错抛出。
+    """
+    _acquire_login_slot()
+    delay = _LOGIN_RETRY_BASE_SECONDS
+    for attempt in range(_LOGIN_RETRY_ATTEMPTS):
+        try:
+            return fn()
+        except Sub2apiError as exc:
+            if not _is_rate_limited_error(exc) \
+                    or attempt == _LOGIN_RETRY_ATTEMPTS - 1:
+                raise
+            time.sleep(delay + random.uniform(0, 2))
+            delay *= 2
+
+
+# 进程内登录节流：60 秒窗口内最多 _LOGIN_BUDGET_PER_MINUTE 次。
+# 只协调本进程；多进程突发仍靠 429 重试兜底（硬保证需共享
+# Redis 节流，控制面目前没有 Redis 客户端，未做）。
+_LOGIN_BUDGET_PER_MINUTE = 15
+_LOGIN_WINDOW_SECONDS = 60.0
+_login_slots = collections.deque()
+_login_slots_lock = threading.Lock()
+
+
+def _acquire_login_slot():
+    """等一个登录名额：窗口满了就等最早一次滑出窗口再进。
+
+    等待时长由排队位置决定（突发 100 路约 7 分钟）；注册是
+    best-effort 语义，失败仍有 worker 首任务懒开通兜底，故不设
+    提前放弃——放弃会把有序排队打回同步重试波。
+    """
+    while True:
+        with _login_slots_lock:
+            now = time.time()
+            while _login_slots \
+                    and _login_slots[0] <= now - _LOGIN_WINDOW_SECONDS:
+                _login_slots.popleft()
+            if len(_login_slots) < _LOGIN_BUDGET_PER_MINUTE:
+                _login_slots.append(now)
+                return
+            wait = _login_slots[0] + _LOGIN_WINDOW_SECONDS - now
+        time.sleep(max(wait, 0))
+
+
 class Sub2apiAdmin:
     """管理员会话：登录、建/删用户、查组。token 自动刷新一次。
 
@@ -193,15 +256,22 @@ class Sub2apiAdmin:
                    request=request)
 
     def _login(self):
-        st, data = self.request(
-            "POST", self.origin + "/api/v1/auth/login",
-            {"email": self.email, "password": self.password})
-        payload = (data or {}).get("data") or {}
-        if st != 200 or not payload.get("access_token"):
+        def _once():
+            st, data = self.request(
+                "POST", self.origin + "/api/v1/auth/login",
+                {"email": self.email, "password": self.password})
+            payload = (data or {}).get("data") or {}
+            if st != 200 or not payload.get("access_token"):
+                raise Sub2apiError(
+                    "sub2api 管理员登录失败: %s"
+                    % ((data or {}).get("message") or st), status=st)
+            return payload
+
+        try:
+            payload = _call_with_login_retry(_once)
+        except Sub2apiError:
             _evict_shared_session(self)
-            raise Sub2apiError(
-                "sub2api 管理员登录失败: %s"
-                % ((data or {}).get("message") or st), status=st)
+            raise
         self._token = payload["access_token"]
         self._refresh = payload.get("refresh_token")
         try:
@@ -298,16 +368,23 @@ class Sub2apiAdmin:
 
 
 def user_login(origin, email, password, request=None):
-    """以 sub2api 用户身份登录：返回 (access_token, refresh_token)。"""
+    """以 sub2api 用户身份登录：返回 (access_token, refresh_token)。
+
+    遇网关登录限流（429）时退避重试，见 _call_with_login_retry。
+    """
     req = request or _default_request
-    st, data = req("POST", origin + "/api/v1/auth/login",
-                   {"email": email, "password": password})
-    payload = (data or {}).get("data") or {}
-    if st != 200 or not payload.get("access_token"):
-        raise Sub2apiError(
-            "sub2api 用户登录失败: %s"
-            % ((data or {}).get("message") or st), status=st)
-    return payload["access_token"], payload.get("refresh_token")
+
+    def _once():
+        st, data = req("POST", origin + "/api/v1/auth/login",
+                       {"email": email, "password": password})
+        payload = (data or {}).get("data") or {}
+        if st != 200 or not payload.get("access_token"):
+            raise Sub2apiError(
+                "sub2api 用户登录失败: %s"
+                % ((data or {}).get("message") or st), status=st)
+        return payload["access_token"], payload.get("refresh_token")
+
+    return _call_with_login_retry(_once)
 
 
 def user_create_key(origin, user_token, name, group_id, request=None):

cloud/tests/test_sub2api.py

@@ -43,6 +43,8 @@ class FakeApi:
         self.fail_balance = False
         self.fail_create_user = False
         self.create_user_hook = None
+        self.login_429_times = 0
+        self.admin_login_429_times = 0
 
     def __call__(self, method, url, body=None, token=None):
         with self._lock:
@@ -52,10 +54,18 @@ class FakeApi:
         path = url.split("28084", 1)[-1] if "28084" in url else url
         self.calls.append((method, path, body, token))
         if path == "/api/v1/auth/login" and body.get("email") == "adm@x":
+            if self.admin_login_429_times > 0:
+                self.admin_login_429_times -= 1
+                return 429, {"message":
+                             "Too many requests, please try again later"}
             return 200, {"data": {"access_token": self.admin_token,
                                   "refresh_token": "adm-ref",
                                   "expires_in": 3600}}
         if path == "/api/v1/auth/login":
+            if self.login_429_times > 0:
+                self.login_429_times -= 1
+                return 429, {"message":
+                             "Too many requests, please try again later"}
             return 200, {"data": {"access_token": self.user_token,
                                   "refresh_token": "usr-ref",
                                   "expires_in": 3600}}
@@ -122,12 +132,15 @@ class TestEnsureGatewayKey(unittest.TestCase):
         self.api = FakeApi()
         self.admin = Sub2apiAdmin("http://x:28084", "adm@x", "pw",
                                   request=self.api)
+        # 节流名额是进程级共享状态：每个单测从空开始，避免跨测累积。
+        sub2api._login_slots.clear()
         self.env = {"SUB2API_BASE_URL": "http://x:28084",
                     "SUB2API_GROUP_NAME": "codex-system"}
         self.old = {k: os.environ.get(k) for k in self.env}
         os.environ.update(self.env)
 
     def tearDown(self):
+        sub2api._login_slots.clear()
         for k, v in self.old.items():
             if v is None:
                 os.environ.pop(k, None)
@@ -305,6 +318,104 @@ class TestEnsureGatewayKey(unittest.TestCase):
         finally:
             self._restore_admin_env(old)
 
+    def test_shared_admin_session_concurrent_100(self):
+        # 本测只验并发正确性（各 key 唯一、管理员登录 1 次），不验节流：
+        # 把预算临时抬高，免得 101 次登录被 15 次/分钟的节流拖到 7 分钟。
+        orig_budget = sub2api._LOGIN_BUDGET_PER_MINUTE
+        sub2api._LOGIN_BUDGET_PER_MINUTE = 1000
+        old = self._admin_env()
+        barrier = threading.Barrier(100)
+        results = []
+        errors = []
+
+        def run(i):
+            try:
+                barrier.wait(timeout=30)
+                conn = _conn()
+                key = ensure_user_gateway_key(
+                    conn, "user-h%d" % i, request=self.api)
+                results.append(key)
+            except Exception as exc:  # noqa: BLE001 - 汇总后断言
+                errors.append(exc)
+
+        try:
+            threads = [threading.Thread(target=run, args=(i,))
+                       for i in range(100)]
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join(timeout=120)
+            self.assertEqual(errors, [])
+            self.assertEqual(len(results), 100)
+            self.assertTrue(all(k.startswith("sk-probe-")
+                                for k in results))
+            self.assertEqual(len(set(results)), 100)
+            # 100 路突发冷缓存也只登录一次。
+            self.assertEqual(self._admin_login_count(), 1)
+        finally:
+            sub2api._LOGIN_BUDGET_PER_MINUTE = orig_budget
+            self._restore_admin_env(old)
+
+    def test_user_login_retries_on_429_then_succeeds(self):
+        self.api.login_429_times = 2
+        with mock.patch("time.sleep", return_value=None):
+            token, refresh = sub2api.user_login(
+                "http://x:28084", "u@x", "pw", request=self.api)
+        self.assertEqual(token, "usr-tok")
+        self.assertEqual(refresh, "usr-ref")
+        logins = [c for c in self.api.calls
+                  if c[1] == "/api/v1/auth/login"
+                  and (c[2] or {}).get("email") != "adm@x"]
+        self.assertEqual(len(logins), 3)
+
+    def test_user_login_429_exhausted_raises(self):
+        self.api.login_429_times = 100
+        with mock.patch("time.sleep", return_value=None):
+            with self.assertRaises(sub2api.Sub2apiError):
+                sub2api.user_login("http://x:28084", "u@x", "pw",
+                                   request=self.api)
+        logins = [c for c in self.api.calls
+                  if c[1] == "/api/v1/auth/login"
+                  and (c[2] or {}).get("email") != "adm@x"]
+        self.assertEqual(len(logins), 6)
+
+    def test_admin_login_retries_on_429(self):
+        self.api.admin_login_429_times = 2
+        admin = Sub2apiAdmin("http://x:28084", "adm@x", "pw",
+                             request=self.api)
+        with mock.patch("time.sleep", return_value=None):
+            uid = admin.create_user("u@x", "pw", "u1")
+        self.assertEqual(uid, 11)
+        self.assertEqual(self._admin_login_count(), 3)
+
+    def test_login_slot_throttle_waits_for_window(self):
+        orig_budget = sub2api._LOGIN_BUDGET_PER_MINUTE
+        orig_window = sub2api._LOGIN_WINDOW_SECONDS
+        sub2api._LOGIN_BUDGET_PER_MINUTE = 2
+        sub2api._LOGIN_WINDOW_SECONDS = 60.0
+        sub2api._login_slots.clear()
+        now = [1000.0]
+        sleeps = []
+
+        def fake_time():
+            return now[0]
+
+        def fake_sleep(seconds):
+            sleeps.append(seconds)
+            now[0] += seconds
+
+        try:
+            with mock.patch("time.time", side_effect=fake_time), \
+                    mock.patch("time.sleep", side_effect=fake_sleep):
+                sub2api._acquire_login_slot()
+                sub2api._acquire_login_slot()
+                sub2api._acquire_login_slot()
+            self.assertEqual(sleeps, [60.0])
+        finally:
+            sub2api._LOGIN_BUDGET_PER_MINUTE = orig_budget
+            sub2api._LOGIN_WINDOW_SECONDS = orig_window
+            sub2api._login_slots.clear()
+
 
 class TestProviderCommand(unittest.TestCase):
     def test_server_command_sub2api(self):