snowy-yang
a7725f8812
在引入 uv 依赖管理、ruff 静态检查与格式化、pyright 类型检查之前， 先把当前可正常运行的状态固定下来，作为后续大规模自动改动的回滚点。 此时状态：97 项测试全绿，HTTP + WebSocket 多牧场服务器可用。
251 lines
8.9 KiB
Python
251 lines
8.9 KiB
Python
"""实时通道：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 itertools
|
||
import threading
|
||
import time
|
||
from collections import defaultdict
|
||
from typing import Any, Dict, List, Optional, Set
|
||
|
||
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__ = ("cid", "farm_id", "loop", "queue", "name", "joined_at", "alive",
|
||
"_sent", "_dropped")
|
||
|
||
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: Optional[Client] = 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: Optional[Client] = 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
|
||
try:
|
||
self.sessions.autosave_all()
|
||
except Exception:
|
||
pass
|
||
|
||
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: Optional[Dict[str, Any]] = 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
|