""" WebSocket 连接管理器 管理实时数据推送的 WebSocket 连接 """ from typing import Dict, Set, List, Any from fastapi import WebSocket from datetime import datetime import json import logging logger = logging.getLogger(__name__) class ConnectionManager: """ WebSocket 连接管理器 管理 WebSocket 连接，支持消息广播和定向推送 """ def __init__(self): # active_connections: 连接ID -> WebSocket self.active_connections: Dict[str, WebSocket] = {} # user_connections: 用户名 -> 连接ID集合 self.user_connections: Dict[str, Set[str]] = {} # subscriptions: 连接ID -> 订阅类型集合 self.subscriptions: Dict[str, Set[str]] = {} # symbol_subscriptions: 品种代码 -> 连接ID集合 self.symbol_subscriptions: Dict[str, Set[str]] = {} self._connection_counter = 0 async def connect(self, websocket: WebSocket, user: str = "anonymous") -> str: """ 接受新的 WebSocket 连接 - **websocket**: WebSocket 实例 - **user**: 用户名（可选） """ await websocket.accept() # 生成连接ID connection_id = f"conn_{self._connection_counter}" self._connection_counter += 1 # 保存连接 self.active_connections[connection_id] = websocket # 绑定用户 if user not in self.user_connections: self.user_connections[user] = set() self.user_connections[user].add(connection_id) # 初始化订阅 self.subscriptions[connection_id] = set() logger.info(f"WebSocket connected: {connection_id} (user: {user})") return connection_id async def disconnect(self, connection_id: str) -> None: """ 断开 WebSocket 连接 - **connection_id**: 连接ID """ if connection_id not in self.active_connections: return # 获取连接 websocket = self.active_connections[connection_id] # 从用户绑定中移除 for user, conn_set in self.user_connections.items(): if connection_id in conn_set: conn_set.remove(connection_id) if not conn_set: del self.user_connections[user] break # 从品种订阅中移除 for symbol, conn_set in self.symbol_subscriptions.items(): if connection_id in conn_set: conn_set.remove(connection_id) if not conn_set: del self.symbol_subscriptions[symbol] # 移除连接和订阅 del self.active_connections[connection_id] del self.subscriptions[connection_id] logger.info(f"WebSocket disconnected: {connection_id}") async def send_personal_message(self, message: dict, connection_id: str) -> bool: """ 向特定连接发送消息 - **message**: 消息内容 - **connection_id**: 连接ID """ if connection_id not in self.active_connections: logger.warning(f"Connection not found: {connection_id}") return False try: websocket = self.active_connections[connection_id] await websocket.send_json(message) return True except Exception as e: logger.error(f"Error sending message to {connection_id}: {e}") await self.disconnect(connection_id) return False async def broadcast(self, message: dict) -> int: """ 向所有连接广播消息 - **message**: 消息内容 """ if not self.active_connections: return 0 # 复制连接列表，避免异步修改 disconnected = [] count = 0 for connection_id, websocket in list(self.active_connections.items()): try: await websocket.send_json(message) count += 1 except Exception as e: logger.error(f"Error broadcasting to {connection_id}: {e}") disconnected.append(connection_id) # 清理断开的连接 for conn_id in disconnected: await self.disconnect(conn_id) return count async def broadcast_to_user(self, message: dict, user: str) -> int: """ 向特定用户的所有连接广播消息 - **message**: 消息内容 - **user**: 用户名 """ if user not in self.user_connections: return 0 count = 0 for connection_id in list(self.user_connections[user]): if await self.send_personal_message(message, connection_id): count += 1 return count async def broadcast_to_symbol(self, message: dict, symbol: str) -> int: """ 向订阅了特定品种的连接广播消息 - **message**: 消息内容 - **symbol**: 品种代码 """ if symbol not in self.symbol_subscriptions: return 0 count = 0 for connection_id in list(self.symbol_subscriptions[symbol]): if await self.send_personal_message(message, connection_id): count += 1 return count def subscribe(self, connection_id: str, subscription_type: str) -> bool: """ 订阅消息类型 - **connection_id**: 连接ID - **subscription_type**: 订阅类型 (tick, order, position, trade) """ if connection_id not in self.subscriptions: return False self.subscriptions[connection_id].add(subscription_type) logger.info(f"Connection {connection_id} subscribed to {subscription_type}") return True def unsubscribe(self, connection_id: str, subscription_type: str) -> bool: """ 取消订阅 - **connection_id**: 连接ID - **subscription_type**: 订阅类型 """ if connection_id not in self.subscriptions: return False self.subscriptions[connection_id].discard(subscription_type) logger.info(f"Connection {connection_id} unsubscribed from {subscription_type}") return True def subscribe_symbol(self, connection_id: str, symbol: str) -> bool: """ 订阅品种数据 - **connection_id**: 连接ID - **symbol**: 品种代码 """ if connection_id not in self.active_connections: return False if symbol not in self.symbol_subscriptions: self.symbol_subscriptions[symbol] = set() self.symbol_subscriptions[symbol].add(connection_id) logger.info(f"Connection {connection_id} subscribed to symbol {symbol}") return True def unsubscribe_symbol(self, connection_id: str, symbol: str) -> bool: """ 取消品种订阅 - **connection_id**: 连接ID - **symbol**: 品种代码 """ if symbol not in self.symbol_subscriptions: return False self.symbol_subscriptions[symbol].discard(connection_id) if not self.symbol_subscriptions[symbol]: del self.symbol_subscriptions[symbol] logger.info(f"Connection {connection_id} unsubscribed from symbol {symbol}") return True def get_connection_count(self) -> int: """获取当前连接数""" return len(self.active_connections) def get_user_connection_count(self, user: str) -> int: """获取用户的连接数""" return len(self.user_connections.get(user, set())) async def broadcast_tick(self, tick_data: dict) -> int: """ 广播行情数据 - **tick_data**: 行情数据字典 """ message = { "type": "tick", "data": tick_data } # 如果有品种订阅，向订阅该品种的连接推送 symbol = tick_data.get("vt_symbol") if symbol and symbol in self.symbol_subscriptions: return await self.broadcast_to_symbol(message, symbol) # 否则向所有订阅 tick 的连接推送 count = 0 for conn_id, subscriptions in self.subscriptions.items(): if "tick" in subscriptions: if await self.send_personal_message(message, conn_id): count += 1 return count async def broadcast_order(self, order_data: dict) -> int: """ 广播订单数据 - **order_data**: 订单数据字典 """ message = { "type": "order", "data": order_data } count = 0 for conn_id, subscriptions in self.subscriptions.items(): if "order" in subscriptions: if await self.send_personal_message(message, conn_id): count += 1 return count async def broadcast_trade(self, trade_data: dict) -> int: """ 广播成交数据 - **trade_data**: 成交数据字典 """ message = { "type": "trade", "data": trade_data } count = 0 for conn_id, subscriptions in self.subscriptions.items(): if "trade" in subscriptions: if await self.send_personal_message(message, conn_id): count += 1 return count async def broadcast_position(self, position_data: dict) -> int: """ 广播持仓数据 - **position_data**: 持仓数据字典 """ message = { "type": "position", "data": position_data } count = 0 for conn_id, subscriptions in self.subscriptions.items(): if "position" in subscriptions: if await self.send_personal_message(message, conn_id): count += 1 return count async def broadcast_account(self, account_data: dict) -> int: """ 广播账户数据 - **account_data**: 账户数据字典 """ message = { "type": "account", "data": account_data } count = 0 for conn_id, subscriptions in self.subscriptions.items(): if "account" in subscriptions: if await self.send_personal_message(message, conn_id): count += 1 return count async def broadcast_log(self, log_data: dict) -> int: """ 广播日志数据 - **log_data**: 日志数据字典 """ message = { "type": "log", "data": log_data } count = 0 for conn_id, subscriptions in self.subscriptions.items(): if "log" in subscriptions: if await self.send_personal_message(message, conn_id): count += 1 return count async def broadcast_contract(self, contract_data: dict) -> int: """ 广播合约数据 - **contract_data**: 合约数据字典 """ message = { "type": "contract", "data": contract_data } count = 0 for conn_id, subscriptions in self.subscriptions.items(): if "contract" in subscriptions: if await self.send_personal_message(message, conn_id): count += 1 return count # 全局连接管理器实例 manager = ConnectionManager() __all__ = ["ConnectionManager", "manager"]