42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
"""WebSocket 管理 - 实时推送"""
|
|
|
|
import logging
|
|
import json
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 活跃连接
|
|
_connections: list[WebSocket] = []
|
|
|
|
|
|
async def ws_endpoint(websocket: WebSocket):
|
|
"""WebSocket 连接处理"""
|
|
await websocket.accept()
|
|
_connections.append(websocket)
|
|
logger.info(f"WebSocket 连接建立，当前连接数: {len(_connections)}")
|
|
try:
|
|
while True:
|
|
# 保持连接，接收客户端心跳
|
|
data = await websocket.receive_text()
|
|
if data == "ping":
|
|
await websocket.send_text("pong")
|
|
except WebSocketDisconnect:
|
|
_connections.remove(websocket)
|
|
logger.info(f"WebSocket 断开，当前连接数: {len(_connections)}")
|
|
|
|
|
|
async def broadcast_update(data: dict):
|
|
"""向所有连接推送更新"""
|
|
if not _connections:
|
|
return
|
|
message = json.dumps(data, ensure_ascii=False, default=str)
|
|
dead = []
|
|
for ws in _connections:
|
|
try:
|
|
await ws.send_text(message)
|
|
except Exception:
|
|
dead.append(ws)
|
|
for ws in dead:
|
|
_connections.remove(ws)
|