perf(jupyter): 5s (workspace_id, user_id) validation cache

Jupyter 一次会话会拉几十次 auth_request(HTML shell / static /
WebSocket / api/contents / autosave / kernels),每次都跑 JWT
verify + WorkspaceMembers JOIN + Scripts.is_locked 查 + runtime
RPC,重复开销大。

新增 module-level (workspace_id, user_id) -> payload 缓存:
* 只缓存 membership 校验通过 + 拿到 runtime 信息的成功结果
  (x-upstream-addr、x-jupyter-internal-token)
* JWT 验签、lock check 仍每请求执行(前者是信任边界,后者
  per-URI 状态易变)
* TTL 5s,time.monotonic(),threading.Lock 保护
* 失败结果(lock 403 / runtime 500)不写缓存

折衷:被踢出 workspace 后最坏 5s 仍返 200;Runtime 单实例下
无需 Redis。新增 8 个 case 覆盖 hit/miss/TTL/lock-every-request/
jwt-every-request/failure-does-not-populate。
This commit is contained in:
tao.chen
2026-08-21 13:02:49 +08:00
parent 69a9a48a0b
commit 2b1f6b5303
2 changed files with 298 additions and 3 deletions
+62 -3
View File
@@ -8,6 +8,8 @@
"""
import re
import threading
import time
from common.auth.jwt import JwtError, verify_jwt_token
from common.auth.membership import MembershipError, load_active_membership
@@ -20,6 +22,26 @@ from sqlalchemy.ext.asyncio import AsyncSession
from backend.dependencies import database_session
from backend.runtime_client import RuntimeClientError
# ---------------------------------------------------------------------------
# (workspace_id, user_id) -> (expires_at_monotonic, payload) 的 5 秒验证结果缓存。
#
# payload 至少包含 x-upstream-addr 与 x-jupyter-internal-token,只缓存
# "membership 校验通过 + 拿到 runtime 信息"后的成功结果;403/500 不写缓存。
#
# 设计说明:
# * TTL 只有 5 秒,且 Runtime 容器单实例(CLAUDE.md "Service rules"
# "Runtime must stay single-replica while file leases and Jupyter tickets
# use the simplified implementation"),module-level 内存缓存是安全的,
# 无需 Redis 之类的外部存储。
# * JWT 验签与 lock check 不进缓存:前者是每请求必须的信任边界;后者是
# per-URI 且 5s 内可能解锁/加锁,跨用户/跨 notebook 不应共享缓存。
# * 折衷:用户被踢出 workspace / membership 撤销后,最坏 5 秒内本接口仍会
# 对已缓存的 (workspace_id, user_id) 返回 200,这是可接受的折衷。
_JUPYTER_AUTH_CACHE: dict[tuple[str, str], tuple[float, dict[str, str]]] = {}
_JUPYTER_AUTH_CACHE_LOCK = threading.Lock()
_JUPYTER_AUTH_CACHE_TTL_SECONDS = 5.0
router = APIRouter(tags=["jupyter"])
security = HTTPBearer(auto_error=False)
@@ -89,6 +111,24 @@ async def load_active_membership_or_403(
) from exc
def _jupyter_auth_cache_get(workspace_id: str, user_id: str) -> dict[str, str] | None:
with _JUPYTER_AUTH_CACHE_LOCK:
entry = _JUPYTER_AUTH_CACHE.get((workspace_id, user_id))
if entry is None:
return None
expires_at, payload = entry
if time.monotonic() >= expires_at:
_JUPYTER_AUTH_CACHE.pop((workspace_id, user_id), None)
return None
return payload
def _jupyter_auth_cache_put(workspace_id: str, user_id: str, payload: dict[str, str]) -> None:
expires_at = time.monotonic() + _JUPYTER_AUTH_CACHE_TTL_SECONDS
with _JUPYTER_AUTH_CACHE_LOCK:
_JUPYTER_AUTH_CACHE[(workspace_id, user_id)] = (expires_at, payload)
# 供 Nginx auth_request 调用:验证访问 Jupyter 的身份、成员关系和文件锁,
# 再返回应转发到的 Jupyter 地址及内部令牌。
@router.get("/api/v1/auth/jupyter")
@@ -132,8 +172,14 @@ async def verify_jupyter_access(
detail="Invalid Authentication Token",
)
await load_active_membership_or_403(session, user_id, workspace_id)
# JWT 验签之后、昂贵的 membership/runtime 查找之前先查缓存。命中时跳过
# membership 与 runtime,但仍要跑下面的 lock check(per-URI,缓存不含它)。
cached = _jupyter_auth_cache_get(workspace_id, user_id)
if cached is None:
await load_active_membership_or_403(session, user_id, workspace_id)
# lock check 永远执行、不进缓存:同一 (workspace_id, user_id) 的不同 URI
# 状态不同,且 5s 内可能解锁/加锁。
notebook_path = extract_notebook_path(original_uri, workspace_id)
if notebook_path and await check_notebook_is_locked(
session,
@@ -146,6 +192,11 @@ async def verify_jupyter_access(
detail=f"Notebook '{notebook_path}' is currently locked",
)
if cached is not None:
response.headers["x-upstream-addr"] = cached["x_upstream_addr"]
response.headers["x-jupyter-internal-token"] = cached["x_jupyter_internal_token"]
return {"status": "ok"}
runtime_client = request.app.state.runtime_client
ws_info = await runtime_client.get_workspace(workspace_id)
if not ws_info or ws_info.get("status") != "running":
@@ -166,6 +217,14 @@ async def verify_jupyter_access(
detail="Jupyter instance returned no port",
)
response.headers["x-upstream-addr"] = f"{jupyter_base_url}:{target_port}"
response.headers["x-jupyter-internal-token"] = jupyter_token or ""
headers_payload = {
"x_upstream_addr": f"{jupyter_base_url}:{target_port}",
"x_jupyter_internal_token": jupyter_token or "",
}
# 只缓存成功结果;lock check 失败(403)或 runtime 启动失败(500)在上方
# 已提前返回,不会走到这里污染缓存。
_jupyter_auth_cache_put(workspace_id, user_id, headers_payload)
response.headers["x-upstream-addr"] = headers_payload["x_upstream_addr"]
response.headers["x-jupyter-internal-token"] = headers_payload["x_jupyter_internal_token"]
return {"status": "ok"}