""" 豆包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()