Tolan/backend/http_pool.py

"""
HTTP连接池管理器 - 类似SQLAlchemy的连接池模式
统一管理DashScope API的HTTP客户端连接，避免重复创建连接池
"""
import httpx
import logging
import threading
import asyncio
import time
from typing import Optional
from config import settings, get_dashscope_http_base_url

logger = logging.getLogger(__name__)


class HTTPConnectionPool:
    """
    HTTP连接池管理器（单例模式）
    类似SQLAlchemy的连接池，统一管理所有HTTP客户端连接
    """
    
    _instance: Optional['HTTPConnectionPool'] = None
    _lock = threading.Lock()
    
    def __new__(cls):
        """单例模式实现"""
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        """初始化连接池（只执行一次）"""
        if self._initialized:
            return
        
        self._client: Optional[httpx.Client] = None
        self._keepalive_task: Optional[asyncio.Task] = None
        self._keepalive_interval = 240  # 4分钟发送一次心跳（小于keepalive_expiry）
        self._last_keepalive_time = 0.0
        self._initialized = True
        logger.info("🔧 HTTP连接池管理器初始化")
    
    def get_client(self) -> httpx.Client:
        """
        获取HTTP客户端（连接池复用）
        类似SQLAlchemy的session.get_bind()，返回共享的连接池客户端
        
        Returns:
            httpx.Client: 配置好的HTTP客户端实例
        """
        if self._client is None:
            with self._lock:
                if self._client is None:
                    self._client = self._create_client()
                    logger.info("✅ HTTP连接池已创建")
        
        return self._client
    
    def _create_client(self) -> httpx.Client:
        """
        创建HTTP客户端（带连接池配置）
        类似SQLAlchemy的Engine创建，配置连接池参数
        """
        import os
        
        # 临时清除代理环境变量（只在这个进程中）
        env_backup = {}
        proxy_vars = ['HTTP_PROXY', 'HTTPS_PROXY', 'http_proxy', 'https_proxy', 'ALL_PROXY', 'all_proxy']
        for var in proxy_vars:
            if var in os.environ:
                env_backup[var] = os.environ[var]
                del os.environ[var]
        
        try:
            # 🔥 创建高性能 HTTP 客户端 - 连接池配置
            # 类似SQLAlchemy的pool_size、max_overflow等参数
            http_client = httpx.Client(
                proxy=None,  # 禁用代理
                timeout=60.0,
                trust_env=False,  # 不信任环境变量中的代理设置
                # 连接池优化：复用 TCP 连接
                limits=httpx.Limits(
                    max_keepalive_connections=settings.HTTPX_MAX_KEEPALIVE_CONNECTIONS,
                    max_connections=settings.HTTPX_MAX_CONNECTIONS,
                    keepalive_expiry=settings.HTTPX_KEEPALIVE_EXPIRY,
                ),
                # HTTP/2支持 - 多路复用，减少连接开销
                http2=True,
            )
            
            logger.info(
                f"📊 HTTP连接池配置: "
                f"keepalive={settings.HTTPX_MAX_KEEPALIVE_CONNECTIONS}, "
                f"max_connections={settings.HTTPX_MAX_CONNECTIONS}, "
                f"keepalive_expiry={settings.HTTPX_KEEPALIVE_EXPIRY}s"
            )
            
            return http_client
        finally:
            # 恢复环境变量（避免影响其他模块）
            for var, value in env_backup.items():
                os.environ[var] = value
    
    def recreate_client(self):
        """
        重建HTTP客户端（连接失效时）
        类似SQLAlchemy的Engine.dispose()，关闭旧连接并创建新连接
        """
        with self._lock:
            if self._client is not None:
                try:
                    self._client.close()
                    logger.info("🔄 关闭旧HTTP连接池")
                except Exception as e:
                    logger.warning(f"⚠️ 关闭HTTP连接池时出错: {e}")
            
            self._client = self._create_client()
            logger.info("✅ HTTP连接池已重建")
    
    def start_keepalive(self, ai_service=None):
        """
        启动连接保活任务（后台定期发送心跳）
        在应用启动时调用，保持连接活跃
        
        Args:
            ai_service: AIService实例，用于发送保活请求（可选）
        """
        if self._keepalive_task is not None:
            return  # 已经启动
        
        if ai_service is None:
            logger.warning("⚠️ 未提供AIService实例，连接保活功能将使用轻量级方式")
        
        async def keepalive_loop():
            """连接保活循环"""
            while True:
                try:
                    await asyncio.sleep(self._keepalive_interval)
                    
                    # 检查连接是否存在
                    if self._client is None:
                        continue
                    
                    # 发送心跳请求（使用极小的测试请求）
                    try:
                        # 方法1：如果有AIService，发送一个极小的LLM请求
                        if ai_service is not None:
                            def send_heartbeat():
                                # 发送一个极小的LLM请求来保持连接活跃
                                from openai import OpenAI
                                from config import settings, get_dashscope_http_base_url
                                
                                client = OpenAI(
                                    api_key=settings.active_dashscope_api_key,
                                    base_url=get_dashscope_http_base_url(),
                                    http_client=self._client,  # 使用连接池
                                )
                                
                                # 发送一个极小的请求（只请求1个token，快速返回）
                                try:
                                    client.chat.completions.create(
                                        model=settings.llm_model,
                                        messages=[{"role": "user", "content": "Hi"}],
                                        max_tokens=1,
                                        timeout=5.0,  # 5秒超时
                                        extra_body=settings.llm_openai_extra_body,
                                    )
                                except Exception:
                                    pass  # 忽略错误，只为了保持连接活跃
                            
                            # 在后台线程中执行（避免阻塞事件循环）
                            await asyncio.to_thread(send_heartbeat)
                        
                        # 记录保活时间
                        self._last_keepalive_time = time.time()
                        logger.debug("💓 HTTP连接保活心跳")
                    except Exception as e:
                        logger.warning(f"⚠️ 连接保活心跳失败: {e}")
                        
                except asyncio.CancelledError:
                    logger.info("🛑 连接保活任务已停止")
                    break
                except Exception as e:
                    logger.error(f"❌ 连接保活任务异常: {e}")
                    await asyncio.sleep(60)  # 出错后等待1分钟再继续
        
        # 启动保活任务
        try:
            loop = asyncio.get_event_loop()
            self._keepalive_task = loop.create_task(keepalive_loop())
            logger.info(f"✅ HTTP连接保活任务已启动（间隔{self._keepalive_interval}秒）")
        except RuntimeError:
            # 如果没有事件循环，在后台线程中创建
            def start_in_thread():
                new_loop = asyncio.new_event_loop()
                asyncio.set_event_loop(new_loop)
                self._keepalive_task = new_loop.create_task(keepalive_loop())
                new_loop.run_forever()
            
            thread = threading.Thread(target=start_in_thread, daemon=True)
            thread.start()
            logger.info(f"✅ HTTP连接保活任务已启动（后台线程，间隔{self._keepalive_interval}秒）")
    
    def stop_keepalive(self):
        """停止连接保活任务"""
        if self._keepalive_task is not None:
            self._keepalive_task.cancel()
            self._keepalive_task = None
            logger.info("🛑 HTTP连接保活任务已停止")
    
    def close(self):
        """
        关闭连接池（应用关闭时）
        类似SQLAlchemy的Engine.dispose()，清理所有连接
        """
        # 先停止保活任务
        self.stop_keepalive()
        
        with self._lock:
            if self._client is not None:
                try:
                    self._client.close()
                    logger.info("🔒 HTTP连接池已关闭")
                except Exception as e:
                    logger.warning(f"⚠️ 关闭HTTP连接池时出错: {e}")
                finally:
                    self._client = None


# 全局连接池实例（单例）
http_pool = HTTPConnectionPool()