feat: WS 每 IP 并发连接上限(默认10,(租户,IP) 键,finally 释放)

Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
浅念
2026-10-12 20:54:28 +08:00
parent fd0f49870c
commit 54a65ed7da
5 changed files with 173 additions and 8 deletions
+24 -7
View File
@@ -5,6 +5,8 @@ from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.engine.models import ConversationSource, ConversationStatus
from app.engine.routes.runtime import get_runtime
from app.engine.services.conn_limiter import DEFAULT_WS_CONN_LIMIT
from app.engine.services.conn_limiter import limiter as conn_limiter
from app.engine.services.conversation import ConversationManager
logger = logging.getLogger(__name__)
@@ -13,21 +15,32 @@ router = APIRouter()
@router.websocket("/ws/chat")
async def ws_chat(websocket: WebSocket):
await websocket.accept()
# Task 8: 每 (租户, IP) 并发连接上限——必须在 accept 前拒绝(close 代
# accept,客户端握手即断)。IP 提取与下方 client_info 同一来源(XFF 优先,
# scope client 兜底),握手期 headers/scope 已就绪、不依赖首条 init 消息;
# 客户端自报的 client_info.ip 仍按原逻辑只进会话记录,不参与限流。
runtime = get_runtime(websocket)
client_ip = websocket.client.host if websocket.client else ""
forwarded = websocket.headers.get("x-forwarded-for", "").split(",")[0].strip()
if forwarded:
client_ip = forwarded
limit = runtime.config.data["engine"].get(
"ws_conn_limit", DEFAULT_WS_CONN_LIMIT)
if not conn_limiter.try_acquire(runtime.config.slug, client_ip, limit=limit):
await websocket.close(code=1013, reason="too many connections")
return
session_id = None
try:
await websocket.accept()
# Task 5b: 租户运行时经 state/默认租户兜底解析(过渡期无中间件);
# manager 由装配线从 runtime 造(db 会话 + runtime 引用)
runtime = get_runtime(websocket)
async with runtime.session_factory() as db:
manager = ConversationManager(db, runtime=runtime)
# Collect client info
client_ip = websocket.client.host if websocket.client else ""
forwarded = websocket.headers.get("x-forwarded-for", "").split(",")[0].strip()
if forwarded:
client_ip = forwarded
# Collect client info(client_ip/forwarded 已在握手期提取,复用)
user_agent = websocket.headers.get("user-agent", "")
client_info = {"ip": client_ip, "user_agent": user_agent}
@@ -107,3 +120,7 @@ async def ws_chat(websocket: WebSocket):
logger.info("Client disconnected: %s", session_id)
except Exception:
logger.exception("WebSocket error for session %s", session_id)
finally:
# Task 8: 正常关闭与异常断开都归还名额——acquire 成功后所有路径
# 必经此处,不泄漏(accept 失败/循环中抛错同样覆盖)
conn_limiter.release(runtime.config.slug, client_ip)
+57
View File
@@ -0,0 +1,57 @@
"""WS 每 IP 并发连接上限(spec §6.1 P2)——进程内 (租户, IP) 计数闸。
动机: 平台化后单 IP 可对所有租户域名各开一圈 WS 耗尽单进程资源;
REST 消息接口已有的 rate_limit.py 管不到长连接握手。key 是
(tenant_slug, ip) 二元组——各租户名额独立,不共享预算。
约束与取舍(与 rate_limit.py 同一族):
- 单进程内存态, ``--workers 1`` 约束下准确; 多 worker 各自一份计数,
实际上限 = limit * workers。
- 纯同步 dict 操作、方法内无 await——单事件循环下 check-then-set
原子,无需锁;不得在多线程里使用。
- 计数归零即删 key, 长期运行不随历史 (租户, IP) 对积累内存;
release 未持有的 key 是安全 no-op(不出负数)。
- limit 可在 try_acquire 时按租户配置覆盖: 模块单例服务全部租户,而
ws_conn_limit 是 engine 组的租户级配置——按次传入让各租户差异实时
生效(含 registry_reload 后), 又不必维护"每租户一个 limiter 实例"
的第二层簿记(实例缓存还有配置改了不生效的滞留问题)。
"""
from typing import Optional
DEFAULT_WS_CONN_LIMIT = 10
class ConnectionLimiter:
"""每 (tenant, ip) 并发计数器: 满则 try_acquire 拒绝, release 归零清键。"""
def __init__(self, limit: int = DEFAULT_WS_CONN_LIMIT) -> None:
self._limit = limit
self._counts: dict[tuple[str, str], int] = {}
def try_acquire(self, tenant: str, ip: str, *,
limit: Optional[int] = None) -> bool:
"""未满则记账并返回 True; 已满返回 False(不记账)。"""
cap = self._limit if limit is None else limit
key = (tenant, ip)
n = self._counts.get(key, 0)
if n >= cap:
return False
self._counts[key] = n + 1
return True
def release(self, tenant: str, ip: str) -> None:
"""归还一个名额; 计数归零删 key; 未持有的 key 是安全 no-op。"""
key = (tenant, ip)
n = self._counts.get(key, 0)
if n <= 1:
self._counts.pop(key, None)
else:
self._counts[key] = n - 1
# app 级单例: chat.py 的 WS 路由共用这一份进程内计数(--workers 1 约束)
limiter = ConnectionLimiter()
def _reset_for_test() -> None:
limiter._counts.clear()