mirror of
https://github.com/fastapiadmin/FastapiAdmin.git
synced 2026-09-23 13:13:09 +00:00
zhangtao
3f472d1e89
1. 迁移存储模块功能到工作流模块，合并冗余代码 2. 调整环境配置加载路径与初始化脚本目录 3. 更新文档与前端组件代码适配重构 4. 新增工作流相关CRUD、模型与迁移文件 5. 移除过期模块与冗余代码，优化日志配置
191 lines
7.9 KiB
Python
191 lines
7.9 KiB
Python
"""WebSocket 连接管理（进程内连接表 + Redis pub/sub 跨进程广播）。
|
||
|
||
单 worker 部署时所有连接都在同一进程；多 worker 部署时，推送消息必须先经 Redis
|
||
分发给各进程的订阅者，否则收件人连在其它 worker 上时消息会静默丢失。
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
import uuid
|
||
from typing import Any
|
||
|
||
from fastapi import WebSocket
|
||
from redis.asyncio import Redis
|
||
|
||
from app.core.logger import logger
|
||
|
||
# 本进程实例令牌：消息经 Redis 广播后会回到发送方进程，用它识别并跳过，避免重复投递
|
||
_PROCESS_TOKEN = uuid.uuid4().hex
|
||
|
||
# 全部 manager 注册表：应用关闭时统一停止订阅协程
|
||
_managers: list["WSConnectionManager"] = []
|
||
|
||
|
||
def stop_ws_relays() -> None:
|
||
"""应用关闭时停止所有跨进程订阅协程（幂等）。"""
|
||
for manager in _managers:
|
||
manager.stop_listener()
|
||
|
||
|
||
class WSConnectionManager:
|
||
"""维护 user_id -> 连接集合，支持同一用户多标签页；发送失败的连接立即剔除。
|
||
|
||
推送策略 = 本进程直发 + Redis publish；订阅协程负责把其它 worker 广播来的消息
|
||
投递给本进程命中目标（all / user_id 列表）的连接。Redis 不可用时自动退化为
|
||
仅本进程直发（等价于旧的单机行为）。
|
||
|
||
is_online / online_count 反映的是本进程连接——在线人数跨 worker 的全局统计
|
||
不在本类职责内，由各业务方按需自行汇总。
|
||
"""
|
||
|
||
def __init__(self, channel: str) -> None:
|
||
"""
|
||
参数:
|
||
- channel (str): 通道名，同时用于日志区分与 Redis 频道名（chat / transfer）。
|
||
"""
|
||
self._channel = channel
|
||
self._connections: dict[int, set[WebSocket]] = {}
|
||
self._redis: Redis | None = None
|
||
self._listener: asyncio.Task | None = None
|
||
_managers.append(self)
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# 连接生命周期
|
||
# ------------------------------------------------------------------ #
|
||
async def connect(self, user_id: int, ws: WebSocket, subprotocol: str | None = None) -> None:
|
||
"""接受握手并登记连接。
|
||
|
||
客户端通过 Sec-WebSocket-Protocol 携带令牌时，必须回显其提供的子协议，否则浏览器会判定握手失败。
|
||
首个连接到达时懒启动跨进程订阅（幂等），此后本进程才能收到其它 worker 的推送。
|
||
"""
|
||
await ws.accept(subprotocol=subprotocol)
|
||
self._connections.setdefault(user_id, set()).add(ws)
|
||
redis = getattr(ws.app.state, "redis", None)
|
||
if redis is not None:
|
||
self._start_listener(redis)
|
||
|
||
def disconnect(self, user_id: int, ws: WebSocket) -> None:
|
||
"""注销连接（幂等）"""
|
||
conns = self._connections.get(user_id)
|
||
if conns is None:
|
||
return
|
||
conns.discard(ws)
|
||
if not conns:
|
||
self._connections.pop(user_id, None)
|
||
|
||
def is_online(self, user_id: int) -> bool:
|
||
"""用户是否在线（本进程视角）"""
|
||
return bool(self._connections.get(user_id))
|
||
|
||
def online_count(self) -> int:
|
||
"""在线用户数（本进程视角）"""
|
||
return len(self._connections)
|
||
|
||
def all_connections(self) -> list[WebSocket]:
|
||
"""全部连接快照"""
|
||
return [ws for conns in self._connections.values() for ws in conns]
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# 推送
|
||
# ------------------------------------------------------------------ #
|
||
async def send_to_user(self, user_id: int | None, data: dict[str, Any]) -> None:
|
||
"""向指定用户的所有连接推送；单个连接异常不影响其余连接。"""
|
||
if user_id is None:
|
||
return
|
||
await self._relay({"users": [user_id]}, data)
|
||
|
||
async def send_to_users(self, user_ids: list[int], data: dict[str, Any]) -> None:
|
||
"""向多个用户推送"""
|
||
ids = sorted(set(user_ids))
|
||
if not ids:
|
||
return
|
||
await self._relay({"users": ids}, data)
|
||
|
||
async def broadcast(self, data: dict[str, Any]) -> None:
|
||
"""向全部连接广播"""
|
||
await self._relay({"all": True}, data)
|
||
|
||
@property
|
||
def _channel_name(self) -> str:
|
||
return f"fastapiadmin:ws:{self._channel}"
|
||
|
||
async def _relay(self, target: dict[str, Any], data: dict[str, Any]) -> None:
|
||
"""推送一条消息：先投本进程命中连接，再发布到 Redis 供其它 worker 投递。
|
||
|
||
回环到本进程的那份由订阅协程依据 _PROCESS_TOKEN 跳过，因此不会重复。
|
||
"""
|
||
await self._dispatch_local(target, data)
|
||
if self._redis is None:
|
||
return
|
||
try:
|
||
message = json.dumps({"sender": _PROCESS_TOKEN, "target": target, "data": data}, ensure_ascii=False)
|
||
await self._redis.publish(self._channel_name, message)
|
||
except Exception as e:
|
||
logger.warning("{} 通道跨进程广播失败（本进程已尽力投递）: {}", self._channel, e)
|
||
|
||
async def _dispatch_local(self, target: dict[str, Any], data: Any) -> None:
|
||
"""把消息投给本进程命中 target（{"all": true} 或 {"users": [...]}）的连接。"""
|
||
if target.get("all"):
|
||
pairs = [(None, ws) for ws in self.all_connections()]
|
||
else:
|
||
pairs = [
|
||
(uid, ws)
|
||
for uid in target.get("users", [])
|
||
for ws in list(self._connections.get(uid, ()))
|
||
]
|
||
for user_id, ws in pairs:
|
||
await self._send(user_id, ws, data)
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# 跨进程订阅
|
||
# ------------------------------------------------------------------ #
|
||
def _start_listener(self, redis: Redis) -> None:
|
||
"""启动跨进程订阅协程（幂等；异常退出后由下一次 connect 重启）。"""
|
||
if self._redis is None:
|
||
self._redis = redis
|
||
if self._listener is not None and not self._listener.done():
|
||
return
|
||
self._listener = asyncio.create_task(self._listen(redis), name=f"ws-relay-{self._channel}")
|
||
|
||
def stop_listener(self) -> None:
|
||
"""停止本 manager 的跨进程订阅协程。"""
|
||
if self._listener is not None and not self._listener.done():
|
||
self._listener.cancel()
|
||
self._listener = None
|
||
|
||
async def _listen(self, redis: Redis) -> None:
|
||
"""订阅 Redis 频道：把其它 worker 广播的消息投递给本进程命中目标的连接。"""
|
||
pubsub = redis.pubsub()
|
||
channel = self._channel_name
|
||
try:
|
||
await pubsub.subscribe(channel)
|
||
logger.info("{} 通道跨进程监听已启动: {}", self._channel, channel)
|
||
async for raw in pubsub.listen():
|
||
if raw.get("type") != "message":
|
||
continue
|
||
try:
|
||
payload = json.loads(raw["data"])
|
||
except (TypeError, ValueError):
|
||
logger.warning("{} 频道收到无法解析的消息，已跳过", channel)
|
||
continue
|
||
if payload.get("sender") == _PROCESS_TOKEN:
|
||
continue
|
||
await self._dispatch_local(payload.get("target") or {}, payload.get("data"))
|
||
except asyncio.CancelledError:
|
||
raise
|
||
except Exception as e:
|
||
logger.error("{} 通道跨进程监听中断: {}", channel, e)
|
||
finally:
|
||
try:
|
||
await pubsub.unsubscribe(channel)
|
||
finally:
|
||
await pubsub.close()
|
||
|
||
async def _send(self, user_id: int | None, ws: WebSocket, data: dict[str, Any]) -> None:
|
||
try:
|
||
await ws.send_json(data)
|
||
except Exception as e:
|
||
logger.warning("{} 通道推送失败，已剔除连接: user={}, err={}", self._channel, user_id, e)
|
||
if user_id is not None:
|
||
self.disconnect(user_id, ws)
|