diff --git a/cloud/sub2api.py b/cloud/sub2api.py index 4583de4..3fb4343 100644 --- a/cloud/sub2api.py +++ b/cloud/sub2api.py @@ -119,7 +119,7 @@ def get_user_gateway_key(conn, user_id): # 共享后同一进程内第二次起免登录。401/刷新失败会回落到重登录, # 密码轮换后下一次登录即用新密码刷新缓存,不会卡死。 _SHARED_ADMIN = {} -_SHARED_ADMIN_LOCK = threading.Lock() +_SHARED_ADMIN_LOCK = threading.RLock() def _store_shared_session(admin): @@ -143,7 +143,9 @@ def shared_admin_session(request=None): """带进程内 token 复用的管理员会话(生产开通路径默认用它)。 缓存命中且未过期时直接播种 token,免一次管理员登录; - 未命中/已过期时照常登录,成功后写回缓存。调用方显式传 + 未命中/已过期时在锁内登录一次再返回:突发冷缓存下 + N 个并发开通也只登录一次(其余线程等锁后直接播种), + 不会各登各的把网关刷出 N 次 bcrypt。调用方显式传 admin 的单测不受影响(不走这条工厂)。 """ admin = Sub2apiAdmin.from_env(request=request) @@ -151,9 +153,13 @@ def shared_admin_session(request=None): 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 + 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: diff --git a/cloud/tests/test_sub2api.py b/cloud/tests/test_sub2api.py index e89c1d2..cb8a217 100644 --- a/cloud/tests/test_sub2api.py +++ b/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": @@ -267,6 +273,38 @@ class TestEnsureGatewayKey(unittest.TestCase): 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):