39 lines
1.6 KiB
Python
39 lines
1.6 KiB
Python
import logging
|
|
from typing import Dict, Set
|
|
from fastapi import WebSocket, WebSocketDisconnect
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ConnectionManager:
|
|
"""WebSocket 实时连接与订阅频道管理器"""
|
|
|
|
def __init__(self):
|
|
# 活跃连接池: market_address -> Set[WebSocket]
|
|
self.active_connections: Dict[str, Set[WebSocket]] = {}
|
|
|
|
async def connect(self, websocket: WebSocket, channel: str):
|
|
await websocket.accept()
|
|
if channel not in self.active_connections:
|
|
self.active_connections[channel] = set()
|
|
self.active_connections[channel].add(websocket)
|
|
logger.info(f"Client connected to channel {channel}. Total: {len(self.active_connections[channel])}")
|
|
|
|
def disconnect(self, websocket: WebSocket, channel: str):
|
|
if channel in self.active_connections:
|
|
self.active_connections[channel].discard(websocket)
|
|
if not self.active_connections[channel]:
|
|
del self.active_connections[channel]
|
|
logger.info(f"Client disconnected from channel {channel}")
|
|
|
|
async def broadcast_to_channel(self, channel: str, message: dict):
|
|
"""向特定频道所有订阅者推送 JSON 消息"""
|
|
if channel in self.active_connections:
|
|
for connection in list(self.active_connections[channel]):
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception as e:
|
|
logger.error(f"Error broadcasting message: {e}")
|
|
self.disconnect(connection, channel)
|
|
|
|
ws_manager = ConnectionManager()
|