265 lines
8.0 KiB
Python
265 lines
8.0 KiB
Python
"""
|
|
Neural Sync - WebSocket Manager
|
|
|
|
实现实时事件广播和 WebSocket 连接管理。
|
|
"""
|
|
|
|
from typing import Dict, List, Set, Optional, Any
|
|
from datetime import datetime
|
|
import json
|
|
import logging
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
from enum import Enum
|
|
|
|
from ..utils.timezone import now_cn_aware, now_cn_naive
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ConnectionEventType(str, Enum):
|
|
"""连接事件类型"""
|
|
CONNECTED = "connected"
|
|
DISCONNECTED = "disconnected"
|
|
SUBSCRIBED = "subscribed"
|
|
UNSUBSCRIBED = "unsubscribed"
|
|
|
|
|
|
class WebSocketConnection:
|
|
"""WebSocket 连接包装器"""
|
|
|
|
def __init__(self, websocket: WebSocket, client_id: str):
|
|
self.websocket = websocket
|
|
self.client_id = client_id
|
|
self.subscriptions: Set[str] = set()
|
|
self.connected_at = now_cn_naive()
|
|
self.last_ping: Optional[datetime] = None
|
|
|
|
async def send(self, data: Dict[str, Any]) -> bool:
|
|
"""发送数据到客户端"""
|
|
try:
|
|
await self.websocket.send_json(data)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error sending to client {self.client_id}: {e}")
|
|
return False
|
|
|
|
async def ping(self) -> bool:
|
|
"""发送 ping"""
|
|
try:
|
|
await self.websocket.send_json({"type": "ping", "timestamp": now_cn_aware().timestamp()})
|
|
self.last_ping = now_cn_naive()
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error pinging client {self.client_id}: {e}")
|
|
return False
|
|
|
|
def subscribe(self, event_type: str):
|
|
"""订阅事件类型"""
|
|
self.subscriptions.add(event_type)
|
|
|
|
def unsubscribe(self, event_type: str):
|
|
"""取消订阅事件类型"""
|
|
self.subscriptions.discard(event_type)
|
|
|
|
def is_subscribed(self, event_type: str) -> bool:
|
|
"""检查是否订阅了事件类型"""
|
|
return event_type in self.subscriptions or "*" in self.subscriptions
|
|
|
|
|
|
class WebSocketManager:
|
|
"""
|
|
WebSocket 连接管理器
|
|
|
|
特性:
|
|
- 连接管理
|
|
- 事件订阅过滤
|
|
- 心跳检测
|
|
- 广播和单播
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._connections: Dict[str, WebSocketConnection] = {}
|
|
self._client_counter = 0
|
|
self._heartbeat_interval = 30 # seconds
|
|
self._heartbeat_task: Optional[Any] = None
|
|
|
|
@property
|
|
def connections(self) -> Dict[str, WebSocketConnection]:
|
|
"""获取所有连接"""
|
|
return self._connections.copy()
|
|
|
|
@property
|
|
def connection_count(self) -> int:
|
|
"""获取连接数"""
|
|
return len(self._connections)
|
|
|
|
async def connect(self, websocket: WebSocket) -> str:
|
|
"""
|
|
接受新的 WebSocket 连接
|
|
|
|
Args:
|
|
websocket: WebSocket 连接
|
|
|
|
Returns:
|
|
客户端 ID
|
|
"""
|
|
await websocket.accept()
|
|
|
|
self._client_counter += 1
|
|
client_id = f"client_{self._client_counter}"
|
|
|
|
connection = WebSocketConnection(websocket, client_id)
|
|
self._connections[client_id] = connection
|
|
|
|
# 发送连接确认
|
|
await connection.send({
|
|
"type": ConnectionEventType.CONNECTED,
|
|
"client_id": client_id,
|
|
"timestamp": now_cn_aware().timestamp(),
|
|
})
|
|
|
|
logger.info(f"WebSocket client connected: {client_id}")
|
|
return client_id
|
|
|
|
async def disconnect(self, client_id: str):
|
|
"""断开连接"""
|
|
if client_id in self._connections:
|
|
del self._connections[client_id]
|
|
logger.info(f"WebSocket client disconnected: {client_id}")
|
|
|
|
async def subscribe(self, client_id: str, event_types: List[str]):
|
|
"""订阅事件类型"""
|
|
if client_id in self._connections:
|
|
connection = self._connections[client_id]
|
|
for event_type in event_types:
|
|
connection.subscribe(event_type)
|
|
|
|
await connection.send({
|
|
"type": ConnectionEventType.SUBSCRIBED,
|
|
"event_types": list(connection.subscriptions),
|
|
"timestamp": now_cn_aware().timestamp(),
|
|
})
|
|
|
|
async def unsubscribe(self, client_id: str, event_types: List[str]):
|
|
"""取消订阅事件类型"""
|
|
if client_id in self._connections:
|
|
connection = self._connections[client_id]
|
|
for event_type in event_types:
|
|
connection.unsubscribe(event_type)
|
|
|
|
await connection.send({
|
|
"type": ConnectionEventType.UNSUBSCRIBED,
|
|
"event_types": list(connection.subscriptions),
|
|
"timestamp": now_cn_aware().timestamp(),
|
|
})
|
|
|
|
async def broadcast(self, event_type: str, data: Dict[str, Any]) -> int:
|
|
"""
|
|
广播事件到所有订阅的客户端
|
|
|
|
Args:
|
|
event_type: 事件类型
|
|
data: 事件数据
|
|
|
|
Returns:
|
|
发送成功的客户端数量
|
|
"""
|
|
message = {
|
|
"type": "event",
|
|
"event_type": event_type,
|
|
"data": data,
|
|
"timestamp": now_cn_aware().timestamp(),
|
|
}
|
|
|
|
success_count = 0
|
|
failed_clients = []
|
|
|
|
for client_id, connection in self._connections.items():
|
|
if connection.is_subscribed(event_type):
|
|
if await connection.send(message):
|
|
success_count += 1
|
|
else:
|
|
failed_clients.append(client_id)
|
|
|
|
# 清理失败的连接
|
|
for client_id in failed_clients:
|
|
await self.disconnect(client_id)
|
|
|
|
if success_count > 0:
|
|
logger.debug(f"Broadcast {event_type} to {success_count} clients")
|
|
|
|
return success_count
|
|
|
|
async def send_to_client(self, client_id: str, data: Dict[str, Any]) -> bool:
|
|
"""发送数据到指定客户端"""
|
|
if client_id in self._connections:
|
|
return await self._connections[client_id].send(data)
|
|
return False
|
|
|
|
async def start_heartbeat(self):
|
|
"""启动心跳检测"""
|
|
import asyncio
|
|
|
|
async def heartbeat_loop():
|
|
while True:
|
|
await asyncio.sleep(self._heartbeat_interval)
|
|
await self._send_heartbeat()
|
|
|
|
self._heartbeat_task = asyncio.create_task(heartbeat_loop())
|
|
logger.info("WebSocket heartbeat started")
|
|
|
|
async def stop_heartbeat(self):
|
|
"""停止心跳检测"""
|
|
if self._heartbeat_task:
|
|
self._heartbeat_task.cancel()
|
|
self._heartbeat_task = None
|
|
logger.info("WebSocket heartbeat stopped")
|
|
|
|
async def _send_heartbeat(self):
|
|
"""发送心跳到所有连接"""
|
|
failed_clients = []
|
|
|
|
for client_id, connection in self._connections.items():
|
|
if not await connection.ping():
|
|
failed_clients.append(client_id)
|
|
|
|
# 清理失败的连接
|
|
for client_id in failed_clients:
|
|
await self.disconnect(client_id)
|
|
|
|
async def close_all(self):
|
|
"""关闭所有连接"""
|
|
for client_id in list(self._connections.keys()):
|
|
await self.disconnect(client_id)
|
|
|
|
await self.stop_heartbeat()
|
|
|
|
|
|
# 全局 WebSocket 管理器实例
|
|
_global_manager: Optional[WebSocketManager] = None
|
|
|
|
|
|
def get_websocket_manager() -> WebSocketManager:
|
|
"""获取全局 WebSocket 管理器实例"""
|
|
global _global_manager
|
|
if _global_manager is None:
|
|
_global_manager = WebSocketManager()
|
|
return _global_manager
|
|
|
|
|
|
# 事件总线桥接 - 自动将事件广播到 WebSocket 客户端
|
|
async def setup_event_bridge(event_bus):
|
|
"""设置事件总线与 WebSocket 的桥接"""
|
|
from .events import EventType
|
|
|
|
async def broadcast_event(event):
|
|
"""广播事件到 WebSocket 客户端"""
|
|
manager = get_websocket_manager()
|
|
await manager.broadcast(
|
|
event_type=event.type.value,
|
|
data=event.data,
|
|
)
|
|
|
|
# 订阅所有事件
|
|
event_bus.subscribe_wildcard(broadcast_event)
|
|
logger.info("Event bus to WebSocket bridge established")
|