Tolan/backend/connection_pool.py

"""
豆包API连接池模块
"""
import asyncio
import time
from typing import Optional
from asyncio import Queue
from loguru import logger

from doubao_api.realtime_dialog_client import RealtimeDialogClient
from config import settings


class DoubaoConnectionPool:
    """豆包API连接池"""
    
    def __init__(
        self, 
        pool_size: int = None,
        max_idle_time: int = None
    ):
        self.pool_size = pool_size or settings.CONNECTION_POOL_SIZE
        self.max_idle_time = max_idle_time or settings.CONNECTION_MAX_IDLE_TIME
        
        # 空闲连接队列
        self.idle_connections: Queue = Queue(maxsize=self.pool_size)
        # 活跃连接集合
        self.active_connections: set = set()
        # 连接健康状态
        self.connection_health: dict = {}
        # 初始化标志
        self.is_initialized = False
    
    async def initialize(self):
        """初始化连接池 - 预创建所有连接"""
        if self.is_initialized:
            logger.warning("连接池已经初始化")
            return
        
        logger.info(f"开始初始化连接池，创建 {self.pool_size} 个连接...")
        
        tasks = []
        for i in range(self.pool_size):
            task = asyncio.create_task(self._create_connection(i))
            tasks.append(task)
        
        connections = await asyncio.gather(*tasks, return_exceptions=True)
        
        success_count = 0
        for i, conn in enumerate(connections):
            if isinstance(conn, Exception):
                logger.error(f"连接 #{i} 创建失败: {conn}")
            elif conn:
                await self.idle_connections.put(conn)
                self.connection_health[id(conn)] = {
                    'created_at': time.time(),
                    'last_used': time.time(),
                    'request_count': 0
                }
                success_count += 1
        
        logger.info(f"连接池初始化完成，成功创建 {success_count}/{self.pool_size} 个连接")
        self.is_initialized = True
        
        # 启动健康检查任务
        asyncio.create_task(self._health_check_loop())
    
    async def _create_connection(self, index: int):
        """创建单个豆包API连接"""
        try:
            logger.info(f"创建豆包API连接 #{index}...")
            
            # 检查API配置
            if not settings.DOUBAO_API_KEY or not settings.DOUBAO_APP_ID:
                logger.warning(f"连接 #{index}: 豆包API密钥未配置，创建模拟连接")
                return await self._create_mock_connection(index)
            
            # 创建真实的豆包API客户端
            client = RealtimeDialogClient(
                api_key=settings.DOUBAO_API_KEY,
                app_id=settings.DOUBAO_APP_ID,
                base_url=settings.DOUBAO_API_URL or "wss://openspeech.bytedance.com/api/v3/realtime/dialogue",
                output_audio_format="ogg_opus",  # 返回OGG/Opus格式
                recv_timeout=settings.DOUBAO_RECV_TIMEOUT,
                connect_timeout=settings.DOUBAO_CONNECT_TIMEOUT
            )
            
            # 连接并初始化
            await client.connect()
            
            logger.success(f"连接 #{index} 创建成功 (真实API)")
            return client
            
        except Exception as e:
            logger.error(f"连接 #{index} 创建失败: {e}")
            logger.warning(f"连接 #{index}: 回退到模拟连接")
            return await self._create_mock_connection(index)
    
    async def _create_mock_connection(self, index: int):
        """创建模拟连接（用于测试）"""
        logger.info(f"创建模拟连接 #{index}...")
        
        class MockConnection:
            def __init__(self, conn_id):
                self.conn_id = conn_id
                self.connected = True
                self.session_id = f"mock_session_{conn_id}"
            
            def is_connected(self):
                return self.connected
            
            async def send_audio(self, audio_data):
                """发送音频数据"""
                logger.debug(f"[模拟] 发送音频: {len(audio_data)} bytes")
                await asyncio.sleep(0.01)
            
            async def receive_audio(self):
                """接收音频数据"""
                await asyncio.sleep(0.1)
                # 返回模拟的OGG音频数据
                mock_audio = b"OggS" + b"\x00" * 100  # 模拟OGG头
                logger.debug(f"[模拟] 收到音频: {len(mock_audio)} bytes")
                return mock_audio
            
            async def close(self):
                self.connected = False
                logger.debug(f"[模拟] 连接已关闭")
        
        return MockConnection(index)
    
    async def acquire(self, timeout: int = 5) -> Optional:
        """获取一个可用连接"""
        try:
            # 从空闲队列获取连接
            connection = await asyncio.wait_for(
                self.idle_connections.get(), 
                timeout=timeout
            )
            
            # 移动到活跃集合
            self.active_connections.add(connection)
            
            # 更新健康状态
            conn_id = id(connection)
            if conn_id in self.connection_health:
                self.connection_health[conn_id]['last_used'] = time.time()
                self.connection_health[conn_id]['request_count'] += 1
            
            logger.debug(f"获取连接成功，当前活跃连接: {len(self.active_connections)}")
            return connection
            
        except asyncio.TimeoutError:
            logger.warning("连接池已耗尽，等待超时")
            return None
        except Exception as e:
            logger.error(f"获取连接失败: {e}")
            return None
    
    async def release(self, connection):
        """归还连接到池中"""
        if connection in self.active_connections:
            self.active_connections.remove(connection)
            
            # 检查连接是否健康
            if await self._is_connection_healthy(connection):
                await self.idle_connections.put(connection)
                logger.debug(f"连接已归还，当前空闲连接: {self.idle_connections.qsize()}")
            else:
                # 连接不健康，先关闭旧连接
                reason = "致命错误" if hasattr(connection, 'has_fatal_error') and connection.has_fatal_error else "健康检查失败"
                logger.warning(f"连接不健康 ({reason})，销毁并重新创建...")
                try:
                    await connection.close()
                except:
                    pass
                
                # 重新创建连接
                new_conn = await self._create_connection(-1)
                if new_conn:
                    await self.idle_connections.put(new_conn)
                    self.connection_health[id(new_conn)] = {
                        'created_at': time.time(),
                        'last_used': time.time(),
                        'request_count': 0
                    }
    
    async def _is_connection_healthy(self, connection) -> bool:
        """检查连接健康状态"""
        try:
            # 检查是否有致命错误
            if hasattr(connection, 'has_fatal_error') and connection.has_fatal_error:
                logger.warning(f"连接有致命错误: {connection.session_id}")
                return False
            
            # 检查连接是否还活着
            if not connection or not connection.is_connected():
                return False
            
            # 检查空闲时间
            conn_id = id(connection)
            if conn_id in self.connection_health:
                idle_time = time.time() - self.connection_health[conn_id]['last_used']
                if idle_time > self.max_idle_time:
                    logger.warning(f"连接空闲时间过长: {idle_time}s")
                    return False
            
            return True
            
        except Exception as e:
            logger.error(f"健康检查失败: {e}")
            return False
    
    async def _health_check_loop(self):
        """定期健康检查"""
        while True:
            try:
                await asyncio.sleep(settings.CONNECTION_HEALTH_CHECK_INTERVAL)
                
                idle_count = self.idle_connections.qsize()
                active_count = len(self.active_connections)
                
                logger.info(
                    f"连接池健康检查 - "
                    f"空闲: {idle_count}, "
                    f"活跃: {active_count}, "
                    f"总计: {idle_count + active_count}/{self.pool_size}"
                )
                
                # 🔧 主动清理空闲连接池中的失效连接
                unhealthy_count = 0
                temp_connections = []
                
                # 取出所有空闲连接检查
                while not self.idle_connections.empty():
                    try:
                        conn = await asyncio.wait_for(self.idle_connections.get(), timeout=1)
                        temp_connections.append(conn)
                    except asyncio.TimeoutError:
                        break
                
                # 检查并重新放回健康的连接
                for conn in temp_connections:
                    if await self._is_connection_healthy(conn):
                        await self.idle_connections.put(conn)
                    else:
                        unhealthy_count += 1
                        logger.warning(f"清理不健康的空闲连接: {conn.session_id}")
                        try:
                            await conn.close()
                        except:
                            pass
                        
                        # 创建新连接补充
                        new_conn = await self._create_connection(-1)
                        if new_conn:
                            await self.idle_connections.put(new_conn)
                            self.connection_health[id(new_conn)] = {
                                'created_at': time.time(),
                                'last_used': time.time(),
                                'request_count': 0
                            }
                
                if unhealthy_count > 0:
                    logger.info(f"🔧 清理了 {unhealthy_count} 个不健康的连接并重新创建")
                
            except Exception as e:
                logger.error(f"健康检查循环错误: {e}")
    
    def get_pool_status(self) -> dict:
        """获取连接池状态"""
        return {
            'pool_size': self.pool_size,
            'idle_connections': self.idle_connections.qsize(),
            'active_connections': len(self.active_connections),
            'total_connections': self.idle_connections.qsize() + len(self.active_connections),
            'health_check_interval': settings.CONNECTION_HEALTH_CHECK_INTERVAL,
            'max_idle_time': self.max_idle_time,
        }
    
    async def close_all(self):
        """关闭所有连接"""
        logger.info("关闭连接池...")
        
        # 关闭活跃连接
        for conn in self.active_connections:
            try:
                await conn.close()
            except Exception as e:
                logger.error(f"关闭活跃连接失败: {e}")
        
        # 关闭空闲连接
        while not self.idle_connections.empty():
            try:
                conn = await self.idle_connections.get()
                await conn.close()
            except Exception as e:
                logger.error(f"关闭空闲连接失败: {e}")
        
        self.active_connections.clear()
        self.connection_health.clear()
        self.is_initialized = False
        
        logger.info("连接池已关闭")


# 创建全局连接池实例
connection_pool = DoubaoConnectionPool()