SimpleFarm/farm/server/hub.py

"""实时通道：SSE 订阅池 + 状态推送线程。

为什么是 SSE
------------
这个游戏的下行（作物成熟、别人来串门）需要服务器**主动**推给客户端，
而上行（播种、喂食）本来就是普通的 HTTP 写请求。SSE（Server-Sent Events）
正好是"一条 HTTP 长连接单向推流"，上行复用 REST，两边各用最自然的协议：

* 纯 HTTP 语义：`curl`、浏览器 `EventSource`、任何 HTTP 客户端都能读；
* 不用额外协议栈：客户端只需 urllib 流式读，连第三方库都不用装；
* 天然支持重连（浏览器 EventSource 自动重连，本项目的客户端自己重连）。

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

* 每条订阅持有一个 :class:`asyncio.Queue` 作为**出口缓冲区**；
* 任意线程想发消息，就 ``loop.call_soon_threadsafe()`` 把消息塞进队列；
* 每条订阅有一个生成器协程，只负责「从队列取出 → yield 成 SSE 帧」。

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

from __future__ import annotations

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

STATE_INTERVAL = 0.25  # 状态推送的最小间隔（秒），天然把高频操作合并成一次
QUEUE_SIZE = 32  # 单订阅出口缓冲区；满了就丢最旧的（状态类消息可丢）
HEARTBEAT_SECONDS = 15.0  # 无消息时的 SSE 心跳间隔，防中间代理掐断空闲连接
IDLE_TICK = 1.0  # 生成器空转的唤醒间隔：用于发现"服务器要停机了"

# 停机时推给每个订阅者的收尾消息：客户端据此知道是服务器主动断开的
CLOSE_MESSAGE: dict[str, Any] = {"type": "bye", "message": "服务器正在停机"}

_seq_counter = itertools.count(1)


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

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


def sse_frame(message: dict[str, Any]) -> str:
    """把一条消息编码成 SSE 帧。

    SSE 的格式很简单：``data: <一行文本>`` 加一个空行。
    JSON 本身不含裸换行（``json.dumps`` 会把换行转义），所以一帧一行就够。
    """
    return f"data: {json.dumps(message, ensure_ascii=False)}\n\n"


def sse_heartbeat() -> str:
    """SSE 注释行心跳：客户端会忽略，但能让代理/防火墙看到连接还活着。"""
    return ": keep-alive\n\n"


class Client:
    """一条 SSE 订阅（一个浏览器标签页 / 一个 TUI / 一个脚本）。"""

    _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:
    """按牧场分组的 SSE 订阅登记处。"""

    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 close_all(self) -> int:
        """让所有 SSE 流主动收尾，返回受影响的订阅数。

        停机时必须做这件事：SSE 是**永不自己结束**的长连接，只要还有一条挂着，
        uvicorn 的优雅停机就要一直等它，`stop` 会卡到超时。发一条 ``bye``
        让每个流的生成器自己 break，连接随之正常关闭。
        """
        with self._lock:
            targets = list(self._all)
        return self._push_many(targets, dict(CLOSE_MESSAGE), None)

    # ------------------------------------------------------------ 广播
    def _push_many(
        self, targets: list[Client], message: dict[str, Any], exclude: Client | None
    ) -> int:
        """给一批订阅者推送同一条消息，返回送出的条数。

        ``Client.push`` 只是把消息排进队列、不改动它，所以这里可以放心让所有
        订阅者共用同一个 dict（不必每人拷一份）。
        """
        sent = 0
        for client in targets:
            if client is exclude:
                continue
            client.push(message)
            sent += 1
        return sent

    def broadcast(
        self, farm_id: str, message: dict[str, Any], exclude: Client | None = None
    ) -> int:
        """推给某个牧场里的所有订阅者。"""
        return self._push_many(self.clients(farm_id), message, exclude)

    def broadcast_all(self, message: dict[str, Any], exclude: Client | None = None) -> int:
        """推给所有牧场的所有订阅者。"""
        with self._lock:
            targets = list(self._all)
        return self._push_many(targets, message, exclude)


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.errors = 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:
                # 推送线程绝不能因为单个牧场出错而退出；
                # 但也不能完全无声——计数进 self.errors，/api/health 可观测
                self.errors += 1

            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