"""白板 WebSocket 连接管理器(实时同步 + 心跳 + 失活清理)。 设计要点: - 进程内单例 ``WhiteboardHub``,维护 ``{board_id: set[Connection]}``。 - 每个 Connection 封装 websocket / board_id / client_id / last_heartbeat。 - 心跳:客户端每 ``heartbeat_interval_seconds``(默认 3s)发一次 ping,服务端回 pong 并 刷新 last_heartbeat。reaper 每秒扫描,超过 ``interval * threshold``(默认 15s)未心跳 的连接判为失活,关闭并从 hub 移除。 - 内存安全:disconnect 幂等;空 set 从 dict 删除;broadcast 对单连接异常立即 disconnect; close_board 关闭并清理整个 board 的连接集合。 - 并发:用一个 asyncio.Lock 保护 ``_boards`` 的结构变更(add/remove board key), 集合内的连接增删用 set 原子操作(Python 单线程事件循环下安全)。 注意:hub 是进程内存,多 worker 下不互通。生产部署需单 worker 或后续接 Redis pub/sub。 """ from __future__ import annotations import asyncio import logging import time from dataclasses import dataclass, field from fastapi import WebSocket from ..config import get_settings logger = logging.getLogger("zikai.whiteboard") @dataclass(eq=False) class Connection: """一个白板在线连接。eq=False 使其按对象 identity 哈希/比较,可放入 set。""" websocket: WebSocket board_id: str client_id: str last_heartbeat: float = field(default_factory=time.monotonic) async def send_json(self, msg: dict) -> bool: """发送一条消息;失败返回 False(调用方据此 disconnect)。""" try: await self.websocket.send_json(msg) return True except Exception as exc: # WebSocketDisconnect / 已关闭 / 编码失败 logger.debug("发送失败 board=%s client=%s: %s", self.board_id, self.client_id, exc) return False def touch(self) -> None: self.last_heartbeat = time.monotonic() class WhiteboardHub: """白板连接管理器单例。""" def __init__(self) -> None: cfg = get_settings().whiteboard self.heartbeat_interval = cfg.heartbeat_interval_seconds self.heartbeat_miss_threshold = cfg.heartbeat_miss_threshold self.timeout_seconds = self.heartbeat_interval * self.heartbeat_miss_threshold # {board_id: set[Connection]} self._boards: dict[str, set[Connection]] = {} self._lock = asyncio.Lock() # ---------------- 连接生命周期 ---------------- async def register(self, conn: Connection) -> None: """把已 accept 的连接加入 board 集合(WebSocket accept 由 controller 负责)。""" async with self._lock: conns = self._boards.setdefault(conn.board_id, set()) conns.add(conn) logger.info("连接接入 board=%s client=%s(当前 %d 人)", conn.board_id, conn.client_id, self.connection_count(conn.board_id)) async def disconnect(self, conn: Connection) -> None: """幂等移除连接;空 set 从 dict 删除以防内存泄漏。""" async with self._lock: conns = self._boards.get(conn.board_id) if conns is None: return conns.discard(conn) if not conns: self._boards.pop(conn.board_id, None) # 尽力关闭 websocket(可能已关闭) try: await conn.websocket.close() except Exception: # pragma: no cover pass logger.info("连接移除 board=%s client=%s(剩余 %d 人)", conn.board_id, conn.client_id, self.connection_count(conn.board_id)) def connection_count(self, board_id: str) -> int: """调试/监控用:某 board 当前在线人数。""" return len(self._boards.get(board_id, ())) # ---------------- 广播 ---------------- async def broadcast(self, board_id: str, msg: dict, exclude: Connection | None = None) -> None: """把 msg 发给 board 内所有在线连接(可排除发送者)。单连接失败不影响其他。""" async with self._lock: conns = list(self._boards.get(board_id, ())) dead: list[Connection] = [] for conn in conns: if exclude is not None and conn is exclude: continue ok = await conn.send_json(msg) if not ok: dead.append(conn) # 发送失败的连接统一清理 for conn in dead: await self.disconnect(conn) # ---------------- 心跳 reaper ---------------- async def reap_loop(self, stop: asyncio.Event) -> None: """后台循环:扫描失活连接。每秒一次,粒度细于 timeout。""" logger.info("白板心跳 reaper 已启动(间隔 1s,超时 %ds)", self.timeout_seconds) while not stop.is_set(): try: await self._reap_once() except Exception as exc: # pragma: no cover logger.warning("reaper 循环异常:%s", exc) try: await asyncio.wait_for(stop.wait(), timeout=1.0) except asyncio.TimeoutError: pass async def _reap_once(self) -> None: now = time.monotonic() async with self._lock: # 快照待检查连接,避免持锁时 await stale: list[Connection] = [] for board_id, conns in self._boards.items(): for conn in conns: if now - conn.last_heartbeat > self.timeout_seconds: stale.append(conn) for conn in stale: logger.warning("心跳失活,移除 board=%s client=%s(静默 %ds)", conn.board_id, conn.client_id, int(now - conn.last_heartbeat)) await self.disconnect(conn) async def close_board(self, board_id: str) -> None: """关闭并清理某 board 的所有连接(删除白板时调用)。""" async with self._lock: conns = self._boards.pop(board_id, None) if not conns: return await asyncio.gather( *(c.send_json({"type": "error", "msg": "白板已被删除"}) for c in conns), return_exceptions=True, ) for conn in conns: try: await conn.websocket.close() except Exception: # pragma: no cover pass logger.info("关闭白板 board=%s,踢出 %d 个连接", board_id, len(conns)) # 进程内单例(由 main.py lifespan / controller 共享) _hub: WhiteboardHub | None = None def get_hub() -> WhiteboardHub: global _hub if _hub is None: _hub = WhiteboardHub() return _hub def reset_hub() -> None: """测试用:重置单例。""" global _hub _hub = None