fin_data_unify_plat/backend/app/services/realtime_service.py

"""
实时行情服务
"""
import json
import logging
from datetime import datetime
from typing import Dict, Set, Optional
from collections import defaultdict

import redis.asyncio as redis

from app.config import settings

logger = logging.getLogger(__name__)

# WebSocket 连接限流配置
MAX_CONNECTIONS_PER_USER = 5  # 每个用户最大连接数
MAX_CONNECTIONS_PER_SYMBOL = 100  # 每个品种最大连接数
MAX_TOTAL_CONNECTIONS = 100  # 总连接数限制


class RealtimeService:
    """实时行情服务"""
    
    # WebSocket 连接管理
    _active_connections: Dict[str, Set] = {}  # symbol -> set of websockets
    _user_connections: Dict[int, Set] = defaultdict(set)  # user_id -> set of websockets
    _websocket_user_map: Dict[int, int] = {}  # websocket id -> user_id
    _anonymous_connections: Set = set()  # 匿名连接集合
    
    def __init__(self):
        self.redis: Optional[redis.Redis] = None
    
    async def connect_redis(self):
        """连接 Redis"""
        self.redis = redis.from_url(settings.REDIS_URL, decode_responses=True)
        logger.info("Redis connected for realtime service")
    
    async def disconnect_redis(self):
        """断开 Redis 连接"""
        if self.redis:
            await self.redis.close()
    
    def register_connection(self, websocket, user_id: Optional[int] = None):
        """注册 WebSocket 连接"""
        if user_id is not None:
            self._websocket_user_map[id(websocket)] = user_id
            self._user_connections[user_id].add(websocket)
        else:
            self._anonymous_connections.add(websocket)
        logger.info(f"WebSocket registered (user={user_id})")
    
    def unregister_connection(self, websocket, user_id: int):
        """注销用户 WebSocket 连接"""
        ws_id = id(websocket)
        if ws_id in self._websocket_user_map:
            del self._websocket_user_map[ws_id]
        self._user_connections[user_id].discard(websocket)
        # 从所有品种订阅中移除
        for symbol in list(self._active_connections.keys()):
            self._active_connections[symbol].discard(websocket)
        logger.info(f"WebSocket unregistered (user={user_id})")
    
    def unregister_anonymous_connection(self, websocket):
        """注销匿名 WebSocket 连接"""
        self._anonymous_connections.discard(websocket)
        # 从所有品种订阅中移除
        for symbol in list(self._active_connections.keys()):
            self._active_connections[symbol].discard(websocket)
        logger.info("Anonymous WebSocket unregistered")
    
    def get_total_connections(self) -> int:
        """获取总连接数"""
        total_user_connections = sum(len(conns) for conns in self._user_connections.values())
        return total_user_connections + len(self._anonymous_connections)
    
    def get_user_connections(self, user_id: int) -> Set:
        """获取用户的连接集合"""
        return self._user_connections.get(user_id, set())
    
    async def subscribe_symbol(self, symbol: str, websocket, user_id: Optional[int] = None) -> bool:
        """
        订阅品种行情
        
        Args:
            symbol: 品种代码
            websocket: WebSocket 连接对象
            user_id: 用户 ID（用于限流）
            
        Returns:
            bool: 订阅是否成功
        """
        # 检查品种连接数限制
        if symbol in self._active_connections:
            if len(self._active_connections[symbol]) >= MAX_CONNECTIONS_PER_SYMBOL:
                logger.warning(f"Symbol {symbol} reached max connections limit ({MAX_CONNECTIONS_PER_SYMBOL})")
                return False
        
        # 检查用户连接数限制
        if user_id is not None:
            if len(self._user_connections[user_id]) >= MAX_CONNECTIONS_PER_USER:
                logger.warning(f"User {user_id} reached max connections limit ({MAX_CONNECTIONS_PER_USER})")
                return False
        
        if symbol not in self._active_connections:
            self._active_connections[symbol] = set()
        
        self._active_connections[symbol].add(websocket)
        
        logger.info(f"Client subscribed to {symbol}, total: {len(self._active_connections[symbol])}")
        return True
    
    async def unsubscribe_symbol(self, symbol: str, websocket, user_id: Optional[int] = None):
        """取消订阅品种行情"""
        if symbol in self._active_connections:
            self._active_connections[symbol].discard(websocket)
            if not self._active_connections[symbol]:
                del self._active_connections[symbol]
            
            logger.info(f"Client unsubscribed from {symbol}")
    
    async def broadcast_quote(self, symbol: str, quote: dict):
        """广播行情数据给所有订阅者"""
        if symbol in self._active_connections:
            message = json.dumps({
                "type": "quote",
                "symbol": symbol,
                "data": quote,
                "timestamp": datetime.utcnow().isoformat()
            })
            
            disconnected = set()
            for websocket in self._active_connections[symbol]:
                try:
                    await websocket.send_text(message)
                except Exception as e:
                    logger.error(f"Failed to send to websocket: {e}")
                    disconnected.add(websocket)
            
            # 清理断开的连接
            for ws in disconnected:
                self._active_connections[symbol].discard(ws)
    
    async def get_latest_quote(self, symbol: str) -> Optional[dict]:
        """从 Redis 获取最新行情"""
        if not self.redis:
            return None
        
        try:
            data = await self.redis.get(f"quote:{symbol}")
            if data:
                return json.loads(data)
        except Exception as e:
            logger.error(f"Failed to get quote from Redis: {e}")
        
        return None
    
    async def update_quote(self, symbol: str, quote: dict):
        """更新行情数据到 Redis"""
        if not self.redis:
            return
        
        try:
            quote["timestamp"] = datetime.utcnow().isoformat()
            await self.redis.set(
                f"quote:{symbol}",
                json.dumps(quote),
                ex=300  # 5 分钟过期
            )
            
            # 发布到 Redis Pub/Sub
            await self.redis.publish(
                f"quotes:{symbol}",
                json.dumps(quote)
            )
            
            # 广播给 WebSocket 客户端
            await self.broadcast_quote(symbol, quote)
            
        except Exception as e:
            logger.error(f"Failed to update quote: {e}")
    
    def get_active_subscriptions(self) -> Dict[str, int]:
        """获取活跃订阅统计"""
        return {
            symbol: len(connections)
            for symbol, connections in self._active_connections.items()
        }


# 全局实时行情服务实例
realtime_service = RealtimeService()