NeoBot/core/ws_pool.py

"""
WebSocket 连接池模块

该模块实现了 WebSocket 连接池功能，用于管理多个 WebSocket 连接，
提高并发处理能力和连接复用效率。
"""
import asyncio
import websockets
from websockets.legacy.client import WebSocketClientProtocol
from typing import Optional, Dict, Any, cast, Union, AsyncGenerator
import uuid
from loguru import logger
import contextlib

from .config_loader import global_config
from .utils.exceptions import WebSocketError, WebSocketConnectionError


class WSConnection:
    """
    WebSocket 连接包装类
    
    封装单个 WebSocket 连接的状态和操作
    """
    def __init__(self, conn: WebSocketClientProtocol, conn_id: str):
        self.conn = conn
        self.conn_id = conn_id
        self.last_used = asyncio.get_event_loop().time()
        self.is_active = True
        self._pending_requests: Dict[str, asyncio.Future] = {}
    
    async def send(self, data: Union[Dict[Any, Any], bytes]):
        """
        发送数据到 WebSocket 连接
        """
        if not self.is_active:
            raise WebSocketError(f"连接 {self.conn_id} 已关闭")
        
        try:
            await self.conn.send(data)
            self.last_used = asyncio.get_event_loop().time()
        except Exception as e:
            self.is_active = False
            raise WebSocketError(f"发送数据失败: {e}")
    
    async def recv(self):
        """
        从 WebSocket 连接接收数据
        """
        if not self.is_active:
            raise WebSocketError(f"连接 {self.conn_id} 已关闭")
        
        try:
            data = await self.conn.recv()
            self.last_used = asyncio.get_event_loop().time()
            return data
        except Exception as e:
            self.is_active = False
            raise WebSocketError(f"接收数据失败: {e}")
    
    async def ping(self, timeout: int = 5) -> bool:
        """
        对 WebSocket 连接执行 ping-pong 健康检查
        """
        if not self.is_active:
            return False
        try:
            # 使用 wait_for 包装 ping
            pong_waiter = await self.conn.ping()
            await asyncio.wait_for(pong_waiter, timeout=timeout)
            return True
        except (asyncio.TimeoutError, websockets.exceptions.ConnectionClosed, Exception):
            self.is_active = False
            return False

    async def close(self):
        """
        关闭 WebSocket 连接
        """
        if self.is_active:
            self.is_active = False
            try:
                await self.conn.close()
            except Exception:
                pass


class WSConnectionPool:
    """
    WebSocket 连接池
    
    管理多个 WebSocket 连接，提供连接的获取、释放和回收功能
    """
    def __init__(self, pool_size: int = 3, max_idle_time: int = 300):
        """
        初始化连接池
        
        :param pool_size: 连接池大小
        :param max_idle_time: 连接最大空闲时间（秒）
        """
        self.pool_size = pool_size
        self.max_idle_time = max_idle_time
        self.pool: asyncio.Queue[WSConnection] = asyncio.Queue(maxsize=pool_size)
        self._closed = False
        self._cleanup_task: Optional[asyncio.Task] = None
        self._current_size = 0  # 当前管理的连接数（包括池中和借出的）
        self._lock = asyncio.Lock() # 用于保护 _current_size 的修改
        
        # 从全局配置读取参数
        self.url = global_config.napcat_ws.uri
        self.token = global_config.napcat_ws.token
        self.reconnect_interval = global_config.napcat_ws.reconnect_interval
        
        logger.info(f"WebSocket 连接池初始化完成，大小: {pool_size}")
    
    async def initialize(self):
        """
        初始化连接池，创建初始连接
        """
        if self._closed:
            raise WebSocketError("连接池已关闭")
        
        # 启动连接清理任务
        self._cleanup_task = asyncio.create_task(self._cleanup_idle_connections())
        
        # 预热连接池
        for _ in range(self.pool_size):
            try:
                conn = await self._create_connection()
                await self.pool.put(conn)
                async with self._lock:
                    self._current_size += 1
                logger.info(f"WebSocket 连接 {conn.conn_id} 已创建并加入连接池")
            except Exception as e:
                logger.error(f"创建初始连接失败: {e}")
                # 初始连接失败不抛出异常，允许后续动态创建
    
    async def _create_connection(self) -> WSConnection:
        """
        创建新的 WebSocket 连接
        """
        headers = {"Authorization": f"Bearer {self.token}"} if self.token else {}
        
        try:
            conn_id = str(uuid.uuid4())
            websocket_raw = await websockets.connect(
                self.url, additional_headers=headers
            )
            websocket = cast(WebSocketClientProtocol, websocket_raw)
            
            conn = WSConnection(websocket, conn_id)
            logger.info(f"WebSocket 连接 {conn_id} 已建立")
            return conn
        except Exception as e:
            raise WebSocketConnectionError(f"创建 WebSocket 连接失败: {e}")
    
    @contextlib.asynccontextmanager
    async def connection(self) -> AsyncGenerator[WSConnection, None]:
        """
        获取连接的上下文管理器
        """
        conn = await self.get_connection()
        try:
            yield conn
        finally:
            await self.release_connection(conn)

    async def get_connection(self) -> WSConnection:
        """
        从连接池获取一个健康的连接，包含健康检查。
        """
        if self._closed:
            raise WebSocketError("连接池已关闭")
        
        start_time = asyncio.get_event_loop().time()
        timeout = 10  # 获取连接的总超时时间

        while True:
            if asyncio.get_event_loop().time() - start_time > timeout:
                raise WebSocketError("获取连接超时")

            try:
                # 1. 尝试从池中获取
                conn = self.pool.get_nowait()
                
                # 健康检查
                if await conn.ping():
                    logger.debug(f"连接 {conn.conn_id} 健康检查通过")
                    return conn
                else:
                    logger.warning(f"连接 {conn.conn_id} 健康检查失败，丢弃")
                    await conn.close()
                    async with self._lock:
                        self._current_size -= 1
                    # 继续循环，尝试获取下一个或创建新的
                    continue

            except asyncio.QueueEmpty:
                # 池为空，检查是否可以创建新连接
                async with self._lock:
                    if self._current_size < self.pool_size:
                        # 有配额，创建新连接
                        self._current_size += 1 # 先占位
                        create_new = True
                    else:
                        create_new = False
                
                if create_new:
                    try:
                        conn = await self._create_connection()
                        return conn
                    except Exception as e:
                        async with self._lock:
                            self._current_size -= 1 # 回滚占位
                        logger.error(f"创建新连接失败: {e}")
                        await asyncio.sleep(1) # 避免快速失败循环
                        continue
                else:
                    # 没有配额，等待池中有可用连接
                    try:
                        conn = await asyncio.wait_for(self.pool.get(), timeout=1.0)
                        # 获取到了，进行健康检查（在下一次循环中处理，或者这里直接处理）
                        # 为了代码复用，我们把 conn 放回去（或者直接用），这里直接用
                        if await conn.ping():
                            return conn
                        else:
                            await conn.close()
                            async with self._lock:
                                self._current_size -= 1
                            continue
                    except asyncio.TimeoutError:
                        continue
    
    async def release_connection(self, conn: WSConnection):
        """
        释放连接回连接池
        """
        if self._closed:
            await conn.close()
            return
        
        if not conn.is_active:
            logger.warning(f"连接 {conn.conn_id} 已失效，不返回连接池")
            await conn.close()
            async with self._lock:
                self._current_size -= 1
            return
        
        try:
            # 尝试放回池中
            self.pool.put_nowait(conn)
            logger.debug(f"连接 {conn.conn_id} 已返回连接池")
        except asyncio.QueueFull:
            # 理论上不应该发生，除非 _current_size 逻辑有误
            logger.warning(f"连接池已满，关闭多余连接 {conn.conn_id}")
            await conn.close()
            async with self._lock:
                self._current_size -= 1
        except Exception as e:
            logger.error(f"释放连接失败: {e}")
            await conn.close()
            async with self._lock:
                self._current_size -= 1
    
    async def _cleanup_idle_connections(self):
        """
        清理空闲连接任务
        """
        while not self._closed:
            await asyncio.sleep(60)  # 每分钟检查一次
            
            try:
                # 我们不替换队列，而是取出检查再放回
                # 这样比较安全，但可能会暂时清空池子
                # 更好的做法是只检查队头的连接
                
                # 获取当前队列大小
                qsize = self.pool.qsize()
                for _ in range(qsize):
                    try:
                        conn = self.pool.get_nowait()
                    except asyncio.QueueEmpty:
                        break
                    
                    current_time = asyncio.get_event_loop().time()
                    if current_time - conn.last_used > self.max_idle_time:
                        logger.info(f"清理空闲连接 {conn.conn_id}")
                        await conn.close()
                        async with self._lock:
                            self._current_size -= 1
                    else:
                        # 还没过期，放回去
                        try:
                            self.pool.put_nowait(conn)
                        except asyncio.QueueFull:
                            # 竞争条件下可能满了
                            await conn.close()
                            async with self._lock:
                                self._current_size -= 1
            except Exception as e:
                logger.error(f"清理空闲连接失败: {e}")
    
    async def close(self):
        """
        关闭连接池
        """
        if self._closed:
            return
        
        self._closed = True
        
        # 停止清理任务
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
        
        # 关闭所有连接
        while not self.pool.empty():
            try:
                conn = self.pool.get_nowait()
                await conn.close()
            except asyncio.QueueEmpty:
                break
        
        logger.info("WebSocket 连接池已关闭")