30a263ed50
后端修复： - 白板删除踢人失效：delete_whiteboard 改 async def，删除后直接 await hub.close_board()。原实现用 asyncio.get_running_loop() 在同步 REST handler （threadpool）里调用必抛 RuntimeError 被 except 吞掉，close_board 从不执行。 同时移除 service 的 hub 依赖（close_board 改由 controller 调用，service 只管 DB）。 - delete_file 去重复查询：原先 get_out_with_disk_path + get_by_id 查两次， 合并为一次；磁盘 unlink 失败加 logger.warning（原静默吞掉致磁盘泄漏无记录）。 - get_hub 单例加 threading.Lock 双重检查（防 REST threadpool 与 WS 事件循环 并发首访各建一个 hub）。 - file_controller 公开 /api/files list 加 Query(ge=1, le=10000) 约束（原无上限可 DoS）。 前端修复： - applyRemoteUpdate 有未发送编辑时重发：合并远端更新后若本地有 pending 编辑 （editor.value !== lastSentText）重新 scheduleSend，避免被 lastSentText 短路丢弃。 - init 不覆盖未发送编辑：断线重连后若本地有未发送内容，作为新版本发上去而非被 init 覆盖。 - applyRemoteUpdate 仅在编辑器已有焦点时恢复焦点，避免抢按钮焦点。 - api() 401 时 location.reload() 触发浏览器 Basic Auth 弹窗（原只 toast 卡死）。 README： - 精简重写，补全 Ubuntu 从 0 安装、Apache 反代（含 WS）、配置项表格、防火墙说明。
189 lines
7.3 KiB
Python
189 lines
7.3 KiB
Python
"""白板 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 threading
|
||
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
|
||
self.max_connections_per_board = cfg.max_connections_per_board
|
||
# {board_id: set[Connection]}
|
||
self._boards: dict[str, set[Connection]] = {}
|
||
self._lock = asyncio.Lock()
|
||
|
||
# ---------------- 连接生命周期 ----------------
|
||
|
||
async def register(self, conn: Connection) -> bool:
|
||
"""把已 accept 的连接加入 board 集合。
|
||
|
||
返回 False 表示该 board 连接数已达上限（调用方应关闭连接）。
|
||
"""
|
||
async with self._lock:
|
||
conns = self._boards.setdefault(conn.board_id, set())
|
||
if len(conns) >= self.max_connections_per_board:
|
||
return False
|
||
conns.add(conn)
|
||
logger.info("连接接入 board=%s client=%s（当前 %d 人）",
|
||
conn.board_id, conn.client_id, self.connection_count(conn.board_id))
|
||
return True
|
||
|
||
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
|
||
_hub_lock = threading.Lock()
|
||
|
||
|
||
def get_hub() -> WhiteboardHub:
|
||
"""获取/创建进程内单例 hub。线程安全（lifespan 预热后通常不再进锁）。"""
|
||
global _hub
|
||
if _hub is None:
|
||
with _hub_lock:
|
||
if _hub is None:
|
||
_hub = WhiteboardHub()
|
||
return _hub
|
||
|
||
|
||
def reset_hub() -> None:
|
||
"""测试用：重置单例。"""
|
||
global _hub
|
||
with _hub_lock:
|
||
_hub = None
|