SimpleFarm/farm/server/hub.py

"""实时通道：WebSocket 连接池 + 状态推送线程。

为什么选 WebSocket 而不是 SSE
-----------------------------
这个游戏是**双向**的：服务器要往下推状态（作物成熟、别人来串门），
客户端也要往上推动作（播种、喂食）。SSE 只能单向下行，上行还得再开一个
HTTP 通道，等于把一条逻辑连接拆成两套鉴权、两套重连、两套错误处理。
WebSocket 一次握手就能双向跑，还能带上 ``action`` 回执，所以这里选它。

线程模型（重点）
----------------
uvicorn 跑在 asyncio 事件循环里，而游戏引擎是同步的、跑在工作线程里。
两者之间不能直接调用对方的 API，所以：

* 每个连接持有一个 :class:`asyncio.Queue` 作为**出口缓冲区**；
* 任意线程想发消息，就 ``loop.call_soon_threadsafe()`` 把消息塞进队列；
* 每条连接有自己的写协程，只负责「从队列取出 → ws.send_json」。

这样引擎线程永远不会阻塞在网络上，事件循环也永远不会被磁盘/锁卡住。
"""

from __future__ import annotations

import asyncio
import contextlib
import itertools
import threading
import time
from collections import defaultdict
from typing import Any

STATE_INTERVAL = 0.25  # 状态推送的最小间隔（秒），天然把高频操作合并成一次
QUEUE_SIZE = 32  # 单连接出口缓冲区；满了就丢最旧的（状态类消息可丢）

_seq_counter = itertools.count(1)


def next_seq() -> int:
    """全局单调递增的消息序号。

    客户端靠它判断"这条状态比我看过的新"，从而丢弃乱序/重复的帧。
    放在 hub 里是因为推送线程和 HTTP 层都要用它。
    """
    return next(_seq_counter)


class Client:
    """一条 WebSocket 连接。"""

    _ids = itertools.count(1)

    __slots__ = (
        "_dropped",
        "_sent",
        "alive",
        "cid",
        "farm_id",
        "joined_at",
        "loop",
        "name",
        "queue",
    )

    def __init__(self, farm_id: str, loop: asyncio.AbstractEventLoop, name: str = "") -> None:
        self.cid = f"c{next(Client._ids)}"
        self.farm_id = farm_id
        self.loop = loop
        self.queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_SIZE)
        self.name = name
        self.joined_at = time.time()
        self.alive = True
        self._sent = 0
        self._dropped = 0

    # ------------------------------------------------------------ 线程安全投递
    def push(self, message: dict[str, Any]) -> None:
        """从任意线程把消息排队到该连接的写协程。"""
        if not self.alive:
            return
        try:
            self.loop.call_soon_threadsafe(self._enqueue, message)
        except RuntimeError:
            # 事件循环已关闭（客户端刚好断开）
            self.alive = False

    def _enqueue(self, message: dict[str, Any]) -> None:
        if self.queue.full():
            # 丢最旧的一条：实时状态宁可少一帧，也不能让连接被拖死
            try:
                self.queue.get_nowait()
                self._dropped += 1
            except asyncio.QueueEmpty:
                pass
        self.queue.put_nowait(message)

    async def next_message(self) -> dict[str, Any]:
        message = await self.queue.get()
        self._sent += 1
        return message

    def stats(self) -> dict[str, Any]:
        return {
            "client_id": self.cid,
            "farm": self.farm_id,
            "name": self.name,
            "sent": self._sent,
            "dropped": self._dropped,
            "online_seconds": round(time.time() - self.joined_at, 1),
        }


class ConnectionHub:
    """按牧场分组的连接登记处。"""

    def __init__(self) -> None:
        self._lock = threading.RLock()
        self._by_farm: dict[str, set[Client]] = defaultdict(set)
        self._all: set[Client] = set()

    # ------------------------------------------------------------ 增删
    def register(self, client: Client) -> None:
        with self._lock:
            self._by_farm[client.farm_id].add(client)
            self._all.add(client)

    def unregister(self, client: Client) -> None:
        with self._lock:
            client.alive = False
            bucket = self._by_farm.get(client.farm_id)
            if bucket is not None:
                bucket.discard(client)
                if not bucket:
                    self._by_farm.pop(client.farm_id, None)
            self._all.discard(client)

    # ------------------------------------------------------------ 查询
    def clients(self, farm_id: str) -> list[Client]:
        with self._lock:
            return list(self._by_farm.get(farm_id, ()))

    def online(self, farm_id: str) -> int:
        with self._lock:
            return len(self._by_farm.get(farm_id, ()))

    def online_map(self) -> dict[str, int]:
        with self._lock:
            return {farm: len(clients) for farm, clients in self._by_farm.items() if clients}

    def total(self) -> int:
        with self._lock:
            return len(self._all)

    def snapshot(self) -> list[dict[str, Any]]:
        with self._lock:
            return [c.stats() for c in sorted(self._all, key=lambda c: c.cid)]

    # ------------------------------------------------------------ 广播
    def broadcast(
        self, farm_id: str, message: dict[str, Any], exclude: Client | None = None
    ) -> int:
        sent = 0
        for client in self.clients(farm_id):
            if client is exclude:
                continue
            client.push(message)
            sent += 1
        return sent

    def broadcast_all(self, message: dict[str, Any], exclude: Client | None = None) -> int:
        with self._lock:
            targets = list(self._all)
        sent = 0
        for client in targets:
            if client is exclude:
                continue
            client.push(message)
            sent += 1
        return sent


class Publisher(threading.Thread):
    """后台推送线程：把「状态变了」这件事节流后送给在线的客户端。

    引擎的 ``_notify()`` 只在 session 上打一个 dirty 标记，真正的推送由这里
    按 :data:`STATE_INTERVAL` 合并发送。好处是玩家连点十次浇水也只推一帧，
    而 AI 在一秒内调二十次接口也不会把网页刷爆。
    """

    def __init__(self, sessions, hub: ConnectionHub, autosave_interval: float = 60.0) -> None:
        super().__init__(name="farm-publisher", daemon=True)
        self.sessions = sessions
        self.hub = hub
        self.autosave_interval = autosave_interval
        self._stop = threading.Event()
        self.pushes = 0
        self._last_autosave = time.time()
        self._started_at = time.time()

    def stop(self) -> None:
        self._stop.set()

    def run(self) -> None:
        while not self._stop.wait(STATE_INTERVAL):
            now = time.time()
            try:
                for session in self.sessions.all():
                    if self.hub.online(session.farm_id) == 0:
                        # 没人在看就不做序列化，省 CPU
                        session.dirty = False
                        continue
                    if not session.dirty:
                        continue
                    session.dirty = False
                    self._push_state(session, now)
            except Exception:
                # 推送线程绝不能因为单个牧场出错而退出
                pass

            if now - self._last_autosave >= self.autosave_interval:
                self._last_autosave = now
                # 自动存档失败不能把推送线程带走
                with contextlib.suppress(Exception):
                    self.sessions.autosave_all()

    def _push_state(self, session, now: float) -> None:
        payload = session.state_message(seq=next_seq())
        # 在线人数放在信封上而不是 state 里：它是连接属性，不是游戏状态
        payload["online"] = self.hub.online(session.farm_id)
        day = session.engine.state.day
        if day != session.published_day:
            session.published_day = day
            self.hub.broadcast(
                session.farm_id,
                {
                    "type": "day",
                    "farm": session.farm_id,
                    "day": day,
                    "events": session.last_events,
                    "at": now,
                },
            )
        self.hub.broadcast(session.farm_id, payload)
        self.pushes += 1

    def uptime(self) -> float:
        return round(time.time() - self._started_at, 1)


def build_state_message(
    session, seq: int, note: str = "", extra: dict[str, Any] | None = None
) -> dict[str, Any]:
    """构造给客户端的状态消息（HTTP 与 WS 共用同一份结构）。"""
    message: dict[str, Any] = {
        "type": "state",
        "seq": seq,
        "farm": session.farm_id,
        "state": session.snapshot(),
        "log": session.engine.recent_log(limit=8),
        "at": time.time(),
    }
    if note:
        message["note"] = note
    if extra:
        message.update(extra)
    return message