fin_data_unify_plat/backend/app/websocket/connection_manager.py

# backend/app/websocket/connection_manager.py
"""
WebSocket 连接管理器
支持 1000+ 并发连接，心跳机制，订阅管理
"""

import asyncio
import json
import time
from typing import Dict, Set, Optional, List
from datetime import datetime
from fastapi import WebSocket, WebSocketDisconnect
from collections import defaultdict
import uuid


class ConnectionManager:
    """
    WebSocket 连接管理器
    
    功能:
    - 连接管理（存储、断开清理）
    - 认证验证
    - 心跳机制（30秒间隔）
    - 订阅管理（订阅/取消订阅）
    - 消息推送（广播、定向推送）
    
    性能优化:
    - 异步 IO
    - 连接池管理
    - 消息序列化优化
    """
    
    def __init__(self):
        # user_id -> Set[WebSocket] (支持多连接)
        self.active_connections: Dict[int, Set[WebSocket]] = defaultdict(set)
        
        # WebSocket -> user_id (反向映射)
        self.connection_users: Dict[WebSocket, int] = {}
        
        # user_id -> Set[symbols] (订阅管理)
        self.subscriptions: Dict[int, Set[str]] = defaultdict(set)
        
        # symbol -> Set[user_id] (反向映射，用于广播)
        self.symbol_subscribers: Dict[str, Set[int]] = defaultdict(set)
        
        # WebSocket -> connection_id (连接标识)
        self.connection_ids: Dict[WebSocket, str] = {}
        
        # 心跳时间记录 (WebSocket -> last_heartbeat)
        self.heartbeat_times: Dict[WebSocket, float] = {}
        
        # 连接统计
        self.total_connections = 0
        self.total_messages_sent = 0
        
        # 心跳超时时间（秒）
        self.heartbeat_timeout = 90
        
        # 锁（用于并发安全）
        self._lock = asyncio.Lock()
    
    async def connect(self, websocket: WebSocket, user_id: int, client_ip: str = None):
        """
        建立连接
        
        Args:
            websocket: WebSocket 连接对象
            user_id: 用户 ID
            client_ip: 客户端 IP
        
        Returns:
            str: connection_id
        """
        async with self._lock:
            # 接受连接
            await websocket.accept()
            
            # 生成连接 ID
            connection_id = str(uuid.uuid4())
            
            # 存储连接信息
            self.active_connections[user_id].add(websocket)
            self.connection_users[websocket] = user_id
            self.connection_ids[websocket] = connection_id
            self.heartbeat_times[websocket] = time.time()
            
            # 更新统计
            self.total_connections += 1
            
            # 发送连接成功消息
            await self.send_to_connection(websocket, {
                "type": "system",
                "event": "connected",
                "connection_id": connection_id,
                "time": datetime.now().isoformat(),
                "message": "WebSocket 连接成功"
            })
            
            return connection_id
    
    async def disconnect(self, websocket: WebSocket):
        """
        断开连接
        
        Args:
            websocket: WebSocket 连接对象
        """
        async with self._lock:
            user_id = self.connection_users.get(websocket)
            if user_id is None:
                return
            
            # 清理订阅
            subscribed_symbols = self.subscriptions.get(user_id, set())
            for symbol in subscribed_symbols:
                self.symbol_subscribers[symbol].discard(user_id)
            
            # 清理连接
            self.active_connections[user_id].discard(websocket)
            if not self.active_connections[user_id]:
                del self.active_connections[user_id]
                if user_id in self.subscriptions:
                    del self.subscriptions[user_id]
            
            # 清理反向映射
            del self.connection_users[websocket]
            del self.connection_ids[websocket]
            del self.heartbeat_times[websocket]
            
            # 更新统计
            self.total_connections -= 1
    
    async def subscribe(self, websocket: WebSocket, symbols: List[str]):
        """
        订阅品种
        
        Args:
            websocket: WebSocket 连接对象
            symbols: 品种代码列表
        """
        async with self._lock:
            user_id = self.connection_users.get(websocket)
            if user_id is None:
                return
            
            # 添加订阅
            for symbol in symbols:
                self.subscriptions[user_id].add(symbol)
                self.symbol_subscribers[symbol].add(user_id)
            
            # 发送订阅确认
            await self.send_to_connection(websocket, {
                "type": "system",
                "event": "subscribed",
                "symbols": symbols,
                "time": datetime.now().isoformat(),
                "message": f"已订阅 {len(symbols)} 个品种"
            })
    
    async def unsubscribe(self, websocket: WebSocket, symbols: List[str]):
        """
        取消订阅
        
        Args:
            websocket: WebSocket 连接对象
            symbols: 品种代码列表
        """
        async with self._lock:
            user_id = self.connection_users.get(websocket)
            if user_id is None:
                return
            
            # 取消订阅
            for symbol in symbols:
                self.subscriptions[user_id].discard(symbol)
                self.symbol_subscribers[symbol].discard(user_id)
            
            # 发送取消确认
            await self.send_to_connection(websocket, {
                "type": "system",
                "event": "unsubscribed",
                "symbols": symbols,
                "time": datetime.now().isoformat(),
                "message": f"已取消订阅 {len(symbols)} 个品种"
            })
    
    async def send_to_connection(self, websocket: WebSocket, message: dict):
        """
        向单个连接发送消息
        
        Args:
            websocket: WebSocket 连接对象
            message: 消息内容
        """
        try:
            await websocket.send_json(message)
            self.total_messages_sent += 1
        except Exception as e:
            # 连接已断开，清理
            await self.disconnect(websocket)
    
    async def send_to_user(self, user_id: int, message: dict):
        """
        向用户的所有连接发送消息
        
        Args:
            user_id: 用户 ID
            message: 消息内容
        """
        connections = self.active_connections.get(user_id, set())
        for websocket in list(connections):  # 使用 list 防止迭代时修改
            await self.send_to_connection(websocket, message)
    
    async def broadcast_to_symbol(self, symbol: str, message: dict):
        """
        向订阅该品种的所有用户广播
        
        Args:
            symbol: 品种代码
            message: 消息内容
        """
        subscribers = self.symbol_subscribers.get(symbol, set())
        for user_id in list(subscribers):
            await self.send_to_user(user_id, message)
    
    async def broadcast(self, message: dict):
        """
        向所有连接广播
        
        Args:
            message: 消息内容
        """
        for user_id, connections in self.active_connections.items():
            for websocket in list(connections):
                await self.send_to_connection(websocket, message)
    
    async def handle_heartbeat(self, websocket: WebSocket):
        """
        处理心跳
        
        Args:
            websocket: WebSocket 连接对象
        """
        self.heartbeat_times[websocket] = time.time()
        await self.send_to_connection(websocket, {
            "type": "system",
            "event": "heartbeat",
            "time": datetime.now().isoformat()
        })
    
    async def check_heartbeat_timeout(self):
        """
        检查心跳超时
        
        超过 heartbeat_timeout 秒无心跳的连接将被断开
        """
        current_time = time.time()
        timeout_connections = []
        
        async with self._lock:
            for websocket, last_time in self.heartbeat_times.items():
                if current_time - last_time > self.heartbeat_timeout:
                    timeout_connections.append(websocket)
        
        # 断开超时连接
        for websocket in timeout_connections:
            try:
                await websocket.close(code=4003, reason="心跳超时")
            except:
                pass
            await self.disconnect(websocket)
    
    def get_connection_count(self) -> int:
        """获取当前连接数"""
        return self.total_connections
    
    def get_user_count(self) -> int:
        """获取当前用户数"""
        return len(self.active_connections)
    
    def get_subscription_count(self) -> int:
        """获取订阅总数"""
        return sum(len(symbols) for symbols in self.symbol_subscribers.values())
    
    def get_statistics(self) -> dict:
        """获取统计信息"""
        return {
            "total_connections": self.total_connections,
            "active_users": self.get_user_count(),
            "total_subscriptions": self.get_subscription_count(),
            "total_messages_sent": self.total_messages_sent,
            "symbol_subscribers": {
                symbol: len(users) 
                for symbol, users in self.symbol_subscribers.items()
            }
        }
    
    def get_user_subscriptions(self, user_id: int) -> List[str]:
        """获取用户订阅列表"""
        return list(self.subscriptions.get(user_id, set()))


# 全局连接管理器实例
connection_manager = ConnectionManager()


# ============== WebSocket 路由处理 ==============

async def websocket_handler(websocket: WebSocket, user_id: int):
    """
    WebSocket 消息处理
    
    Args:
        websocket: WebSocket 连接对象
        user_id: 用户 ID
    """
    # 建立连接
    connection_id = await connection_manager.connect(websocket, user_id)
    
    try:
        while True:
            # 接收消息
            data = await websocket.receive_json()
            
            # 处理消息
            action = data.get("action")
            
            if action == "subscribe":
                # 订阅品种
                symbols = data.get("symbols", [])
                if symbols:
                    await connection_manager.subscribe(websocket, symbols)
            
            elif action == "unsubscribe":
                # 取消订阅
                symbols = data.get("symbols", [])
                if symbols:
                    await connection_manager.unsubscribe(websocket, symbols)
            
            elif action == "heartbeat":
                # 心跳
                await connection_manager.handle_heartbeat(websocket)
            
            elif action == "query":
                # 查询订阅
                subscriptions = connection_manager.get_user_subscriptions(user_id)
                await connection_manager.send_to_connection(websocket, {
                    "type": "system",
                    "event": "query_result",
                    "subscriptions": subscriptions,
                    "time": datetime.now().isoformat()
                })
            
            else:
                # 未知操作
                await connection_manager.send_to_connection(websocket, {
                    "type": "system",
                    "event": "error",
                    "message": f"未知操作: {action}",
                    "time": datetime.now().isoformat()
                })
    
    except WebSocketDisconnect:
        # 客户端断开
        await connection_manager.disconnect(websocket)
    
    except Exception as e:
        # 其他异常
        await connection_manager.disconnect(websocket)


# ============== 心跳检查任务 ==============

async def heartbeat_checker():
    """
    心跳检查后台任务
    
    每 30 秒检查一次心跳超时
    """
    while True:
        await asyncio.sleep(30)
        await connection_manager.check_heartbeat_timeout()