Merge pull request #94 from openqj/feat-regperf

perf: 注册开通复用进程内管理员会话

cloud/sub2api.py

@@ -16,6 +16,7 @@
 import json
 import os
 import secrets
+import threading
 import time
 import urllib.error
 import urllib.request
@@ -112,8 +113,63 @@ def get_user_gateway_key(conn, user_id):
         return None
 
 
+# 进程内共享管理员会话：(origin, email) -> (token, refresh, expires_at)。
+# 只复用 token 元组，不共用实例（实例不跨线程）；注册与首任务开通
+# 每次新建实例都要登录一次（实测约 111ms，占开通链路约 1/3），
+# 共享后同一进程内第二次起免登录。401/刷新失败会回落到重登录，
+# 密码轮换后下一次登录即用新密码刷新缓存，不会卡死。
+_SHARED_ADMIN = {}
+_SHARED_ADMIN_LOCK = threading.RLock()
+
+
+def _store_shared_session(admin):
+    key = getattr(admin, "_shared_key", None)
+    if key is None or not admin._token:
+        return
+    with _SHARED_ADMIN_LOCK:
+        _SHARED_ADMIN[key] = (admin._token, admin._refresh,
+                              admin._expires_at)
+
+
+def _evict_shared_session(admin):
+    key = getattr(admin, "_shared_key", None)
+    if key is None:
+        return
+    with _SHARED_ADMIN_LOCK:
+        _SHARED_ADMIN.pop(key, None)
+
+
+def shared_admin_session(request=None):
+    """带进程内 token 复用的管理员会话（生产开通路径默认用它）。
+
+    缓存命中且未过期时直接播种 token，免一次管理员登录；
+    未命中/已过期时在锁内登录一次再返回：突发冷缓存下
+    N 个并发开通也只登录一次（其余线程等锁后直接播种），
+    不会各登各的把网关刷出 N 次 bcrypt。调用方显式传
+    admin 的单测不受影响（不走这条工厂）。
+    """
+    admin = Sub2apiAdmin.from_env(request=request)
+    key = (admin.origin, admin.email)
+    admin._shared_key = key
+    with _SHARED_ADMIN_LOCK:
+        cached = _SHARED_ADMIN.get(key)
+        if cached and cached[2] > time.time():
+            admin._token, admin._refresh, admin._expires_at = cached
+            return admin
+        # 缓存未命中：在锁内登录（_login 成功后经
+        # _store_shared_session 写回同一缓存，RLock 可重入）。
+        admin._login()
+        return admin
+
+
 class Sub2apiAdmin:
-    """管理员会话：登录、建/删用户、查组。token 自动刷新一次。"""
+    """管理员会话：登录、建/删用户、查组。token 自动刷新一次。
+
+    每次实例只在进程内复用 token 元组（见 shared_admin_session）：
+    实例本身不跨线程共用，只复用 (token, refresh, expires_at)，
+    401 时走已有刷新->重登录链路自愈，密码轮换后下一次登录
+    即用新密码并刷新缓存，不会卡在旧 token 上。
+    """
 
     def __init__(self, origin, email, password, request=None):
         if not origin:
@@ -127,6 +183,7 @@ class Sub2apiAdmin:
         self._token = None
         self._refresh = None
         self._expires_at = 0
+        self._shared_key = None
 
     @classmethod
     def from_env(cls, request=None):
@@ -141,6 +198,7 @@ class Sub2apiAdmin:
             {"email": self.email, "password": self.password})
         payload = (data or {}).get("data") or {}
         if st != 200 or not payload.get("access_token"):
+            _evict_shared_session(self)
             raise Sub2apiError(
                 "sub2api 管理员登录失败: %s"
                 % ((data or {}).get("message") or st), status=st)
@@ -151,6 +209,7 @@ class Sub2apiAdmin:
         except (TypeError, ValueError):
             ttl = 3600
         self._expires_at = time.time() + max(60, ttl - 60)
+        _store_shared_session(self)
 
     def _refresh_token(self):
         if not self._refresh:
@@ -169,6 +228,7 @@ class Sub2apiAdmin:
             except (TypeError, ValueError):
                 ttl = 3600
             self._expires_at = time.time() + max(60, ttl - 60)
+            _store_shared_session(self)
             return
         self._login()
 
@@ -310,7 +370,9 @@ def ensure_user_gateway_key(conn, user_id, admin=None, request=None):
     入库失败时尽力清理刚建的 key 与账号（避免孤儿 key），再抛原错；
     授余额失败同样走这条清理（余额为 0 的 key 在网关侧直接 403，
     留着没有意义）。同一用户并发开通时建号撞 email 已存在，
-    则轮询复用胜出者入库的行。调用方负责 conn 的提交语义：本函数在入库后
+    则轮询复用胜出者入库的行。管理员工单默认走进程内共享会话
+    （同一进程内第二次起免管理员登录，见 shared_admin_session）。
+    调用方负责 conn 的提交语义：本函数在入库后
     commit（开通是独立事务，不跟调用方的业务事务绑在一起）。
     """
     ensure_gateway_keys_table(conn)
@@ -320,7 +382,7 @@ def ensure_user_gateway_key(conn, user_id, admin=None, request=None):
     origin = gateway_origin()
     if not origin:
         raise Sub2apiError("SUB2API_BASE_URL 未配置，无法开通网关账号")
-    admin = admin or Sub2apiAdmin.from_env(request=request)
+    admin = admin or shared_admin_session(request=request)
     req = request or _default_request
     email = "codex-%s@codex.local" % user_id
     username = "codex-" + str(user_id)[:12]

cloud/tests/test_sub2api.py

@@ -7,6 +7,7 @@
 import os
 import sqlite3
 import sys
+import threading
 import unittest
 from pathlib import Path
 from unittest import mock
@@ -31,6 +32,7 @@ class FakeApi:
 
     def __init__(self):
         self.calls = []
+        self._lock = threading.Lock()
         self.admin_token = "adm-tok"
         self.user_token = "usr-tok"
         self.groups = [{"id": 2, "name": "codex-system"}]
@@ -43,6 +45,10 @@ class FakeApi:
         self.create_user_hook = None
 
     def __call__(self, method, url, body=None, token=None):
+        with self._lock:
+            return self._handle(method, url, body, token)
+
+    def _handle(self, method, url, body=None, token=None):
         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":
@@ -216,6 +222,89 @@ class TestEnsureGatewayKey(unittest.TestCase):
             " WHERE user_id = ?", ("user-race2",)).fetchone()
         self.assertEqual(row[0], 0)
 
+    def _admin_env(self):
+        old = {k: os.environ.get(k) for k in
+               ("SUB2API_ADMIN_EMAIL", "SUB2API_ADMIN_PASSWORD")}
+        os.environ["SUB2API_ADMIN_EMAIL"] = "adm@x"
+        os.environ["SUB2API_ADMIN_PASSWORD"] = "pw"
+        sub2api._SHARED_ADMIN.clear()
+        return old
+
+    def _restore_admin_env(self, old):
+        for k, v in old.items():
+            if v is None:
+                os.environ.pop(k, None)
+            else:
+                os.environ[k] = v
+        sub2api._SHARED_ADMIN.clear()
+
+    def _admin_login_count(self):
+        return len([
+            c for c in self.api.calls
+            if c[0] == "POST" and c[1] == "/api/v1/auth/login"
+            and (c[2] or {}).get("email") == "adm@x"])
+
+    def test_shared_admin_session_reuses_login(self):
+        old = self._admin_env()
+        try:
+            conn1 = _conn()
+            key1 = ensure_user_gateway_key(conn1, "user-s1",
+                                           request=self.api)
+            self.assertTrue(key1.startswith("sk-probe-"))
+            conn2 = _conn()
+            key2 = ensure_user_gateway_key(conn2, "user-s2",
+                                           request=self.api)
+            self.assertTrue(key2.startswith("sk-probe-"))
+            # 同一进程内第二次开通不再管理员登录（只剩首次那一次）。
+            self.assertEqual(self._admin_login_count(), 1)
+        finally:
+            self._restore_admin_env(old)
+
+    def test_shared_admin_session_expired_refetches(self):
+        old = self._admin_env()
+        sub2api._SHARED_ADMIN[("http://x:28084", "adm@x")] = (
+            "old-tok", "old-ref", 0)
+        try:
+            conn = _conn()
+            key = ensure_user_gateway_key(conn, "user-s3",
+                                          request=self.api)
+            self.assertTrue(key.startswith("sk-probe-"))
+            self.assertEqual(self._admin_login_count(), 1)
+        finally:
+            self._restore_admin_env(old)
+
+    def test_shared_admin_session_concurrent_burst_logs_in_once(self):
+        old = self._admin_env()
+        barrier = threading.Barrier(10)
+        results = []
+        errors = []
+
+        def run(i):
+            try:
+                barrier.wait(timeout=10)
+                conn = _conn()
+                key = ensure_user_gateway_key(
+                    conn, "user-c%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(10)]
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join(timeout=30)
+            self.assertEqual(errors, [])
+            self.assertEqual(len(results), 10)
+            self.assertTrue(all(k.startswith("sk-probe-")
+                                for k in results))
+            # 突发冷缓存：10 路并发也只登录一次。
+            self.assertEqual(self._admin_login_count(), 1)
+        finally:
+            self._restore_admin_env(old)
+
 
 class TestProviderCommand(unittest.TestCase):
     def test_server_command_sub2api(self):