claude_dev
918bbed0fc
- 修复 deps.py 中 get_vn_service 的引用错误 (vn_service.vn_service -> vn_service) - 移除登录页面上的明文密码提示 - 改进前端错误处理，避免数据加载失败导致登录显示错误
394 lines
11 KiB
Python
394 lines
11 KiB
Python
"""
|
|
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"]
|