zTools2/app/services/whiteboard_hub.py

"""白板 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 as exc:  # pragma: no cover
            logger.debug("关闭 websocket 时出错 board=%s client=%s: %s",
                         conn.board_id, conn.client_id, exc)
        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 as exc:  # pragma: no cover
                logger.debug("关闭 websocket 时出错 board=%s client=%s: %s",
                             conn.board_id, conn.client_id, exc)
        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