sycommon-python-lib 0.1.32__py3-none-any.whl → 0.1.56b5__py3-none-any.whl

sycommon/rabbitmq/rabbitmq_pool.py

@@ -1,15 +1,30 @@
-from typing import Optional, List
-from aio_pika import connect_robust, Channel
-from aio_pika.abc import AbstractRobustConnection
-from aio_pika.pool import Pool
-
+import asyncio
+import random
+from typing import Optional, List, Dict, Callable, Tuple
+from aio_pika import connect_robust, RobustChannel, Message
+from aio_pika.abc import (
+    AbstractRobustConnection, AbstractQueue, AbstractExchange, AbstractMessage
+)
 from sycommon.logging.kafka_log import SYLogger
 
 logger = SYLogger
 
 
+class AsyncProperty:
+    """实现 await obj.attr 的支持"""
+
+    def __init__(self, method):
+        self.method = method
+
+    def __get__(self, obj, objtype=None):
+        if obj is None:
+            return self
+        # 关键：当访问 obj.attr 时，直接返回协程对象，而不是方法本身
+        return self.method(obj)
+
+
 class RabbitMQConnectionPool:
-    """RabbitMQ连接池管理，负责创建和管理连接池与通道池"""
+    """单连接单通道RabbitMQ客户端 (增强版日志)"""
 
     def __init__(
         self,
@@ -18,87 +33,288 @@ class RabbitMQConnectionPool:
         username: str,
         password: str,
         virtualhost: str = "/",
-        connection_pool_size: int = 2,
-        channel_pool_size: int = 10,
-        heartbeat: int = 10,
-        app_name: str = ""
+        heartbeat: int = 30,
+        app_name: str = "",
+        connection_timeout: int = 30,
+        reconnect_interval: int = 5,
+        prefetch_count: int = 2,
     ):
         self.hosts = [host.strip() for host in hosts if host.strip()]
         if not self.hosts:
             raise ValueError("至少需要提供一个RabbitMQ主机地址")
+
         self.port = port
         self.username = username
         self.password = password
         self.virtualhost = virtualhost
         self.app_name = app_name or "rabbitmq-client"
         self.heartbeat = heartbeat
+        self.connection_timeout = connection_timeout
+        self.reconnect_interval = reconnect_interval
+        self.prefetch_count = prefetch_count
+
+        self._current_host: str = random.choice(self.hosts)
+        logger.info(f"[INIT] 随机选择RabbitMQ主机: {self._current_host}")
+
+        # 核心资源
+        self._connection: Optional[AbstractRobustConnection] = None
+        self._channel: Optional[RobustChannel] = None
+        self._consumer_channels: Dict[str, RobustChannel] = {}
+
+        # 状态控制
+        self._lock = asyncio.Lock()
+        self._initialized = False
+        self._is_shutdown = False
+
+    @AsyncProperty
+    async def is_alive(self) -> bool:
+        """对外暴露的连接存活状态（原子化判断）"""
+        async with self._lock:
+            if self._is_shutdown:
+                return False
+
+            if not self._initialized:
+                return False
+
+            if self._connection is None or self._connection.is_closed:
+                return False
 
-        # 连接池和通道池
-        self.connection_pool: Optional[Pool] = None
-        self.channel_pool: Optional[Pool] = None
-        self.connection_pool_size = connection_pool_size
-        self.channel_pool_size = channel_pool_size
+            # 可选：检查主通道是否存活
+            if self._channel is None or self._channel.is_closed:
+                # 如果你认为通道断了连接也算死，就保留这行；否则删除
+                return False
+
+            return True
+
+    async def _create_connection_impl(self) -> AbstractRobustConnection:
+        """
+        连接创建入口
+        """
+        conn_url = (
+            f"amqp://{self.username}:{self.password}@{self._current_host}:{self.port}/"
+            f"{self.virtualhost}?name={self.app_name}&heartbeat={self.heartbeat}"
+            f"&reconnect_interval={self.reconnect_interval}&fail_fast=1"
+        )
+        logger.info(
+            f"🔌 [CONNECT_START] 尝试创建连接 -> {self._current_host}:{self.port}")
+        try:
+            conn = await connect_robust(conn_url, timeout=self.connection_timeout)
+            # 注意：connect_robust 返回时，底层 TCP 可能还在握手，但对象已创建
+            logger.info(f"✅ [CONNECT_OK] 连接对象创建成功: {id(conn)}")
+            return conn
+        except Exception as e:
+            logger.error(f"❌ [CONNECT_FAIL] 连接创建失败: {str(e)}")
+            raise ConnectionError(f"无法连接RabbitMQ {self._current_host}") from e
+
+    async def _create_channel_impl(self, connection: AbstractRobustConnection) -> RobustChannel:
+        """创建通道"""
+        try:
+            channel = await connection.channel()
+            await channel.set_qos(prefetch_count=self.prefetch_count)
+            logger.debug(f"✅ [CHANNEL_OK] 通道创建成功: {id(channel)}")
+            return channel
+        except Exception as e:
+            logger.error(f"❌ [CHANNEL_FAIL] 通道创建失败: {str(e)}")
+            raise
+
+    async def _ensure_main_channel(self) -> RobustChannel:
+        """确保主通道有效 (原子操作)"""
+        async with self._lock:
+            if self._is_shutdown:
+                raise RuntimeError("客户端已关闭")
+
+            # 1. 确保底层连接存在
+            conn = self._connection
+            if conn is None or conn.is_closed:
+                logger.info("⚠️ [RECONNECT] 检测到连接不存在或已关闭，开始重建...")
+                conn = await self._create_connection_impl()
+                self._connection = conn
+                logger.info(f"🔗 [UPDATE] 连接对象已更新: {id(conn)}")
+
+            # 2. 确保主通道存在
+            if self._channel is None or self._channel.is_closed:
+                logger.info("⚠️ [RECOVER_CHANNEL] 检测到主通道不存在或已关闭，开始恢复...")
+                self._channel = await self._create_channel_impl(conn)
+
+            return self._channel
 
     async def init_pools(self):
-        """初始化连接池和通道池"""
-        # 连接创建函数（支持集群节点轮询）
-        async def create_connection() -> AbstractRobustConnection:
-            # 轮询选择主机（简单负载均衡）
-            hosts = self.hosts.copy()
-            while hosts:
-                host = hosts.pop(0)
+        """
+        初始化入口与异常处理 (修复泄漏的关键)
+        """
+        async with self._lock:
+            if self._is_shutdown:
+                raise RuntimeError("客户端已关闭")
+            if self._initialized:
+                return
+
+        conn_created_in_this_try = None
+        try:
+            # 步骤 A: 创建连接 (在锁外进行，避免阻塞其他操作)
+            conn = await self._create_connection_impl()
+            conn_created_in_this_try = conn  # 记录本次创建的对象，用于失败回滚
+
+            # 步骤 B: 更新状态和初始化通道 (在锁内进行，保证原子性)
+            async with self._lock:
+                if self._is_shutdown:
+                    # 如果在创建连接期间，外部调用了 close，则必须立即清理刚创建的连接
+                    logger.warning("⚠️ [ABORT] 检测到关闭信号，放弃初始化并清理资源")
+                    raise RuntimeError("客户端已关闭")
+
+                self._connection = conn
+                self._channel = await self._create_channel_impl(conn)
+                self._initialized = True
+                logger.info(
+                    f"🚀 [INIT_SUCCESS] 客户端初始化完成. ConnID: {id(self._connection)}")
+
+        except Exception as e:
+            logger.error(f"💥 [INIT_ERROR] 初始化流程异常: {str(e)}", exc_info=True)
+            # 如果步骤A成功但步骤B失败（例如通道创建失败），或者步骤B中出错，
+            # 必须显式关闭在步骤A中创建的连接，否则它会变成“游离连接”。
+            if conn_created_in_this_try:
+                logger.warning(
+                    f"🧹 [LEAK_PREVENTION] 检测到初始化失败，正在显式关闭刚创建的连接: {id(conn_created_in_this_try)}")
                 try:
-                    return await connect_robust(
-                        host=host,
-                        port=self.port,
-                        login=self.username,
-                        password=self.password,
-                        virtualhost=self.virtualhost,
-                        heartbeat=self.heartbeat,
-                        client_properties={
-                            "connection_name": f"{self.app_name}@{host}"
-                        }
-                    )
-                except Exception as e:
-                    logger.warning(
-                        f"连接主机 {host}:{self.port} 失败，尝试下一个节点: {str(e)}")
-                    if not hosts:
-                        raise  # 所有节点都失败时抛出异常
-
-        # 初始化连接池
-        self.connection_pool = Pool(
-            create_connection,
-            max_size=self.connection_pool_size
-        )
+                    await conn_created_in_this_try.close()
+                    logger.info(
+                        f"✅ [CLOSE_OK] 泄漏连接已关闭: {id(conn_created_in_this_try)}")
+                except Exception as close_err:
+                    logger.error(f"❌ [CLOSE_ERR] 关闭泄漏连接时出错: {str(close_err)}")
 
-        # 通道创建函数
-        async def create_channel() -> Channel:
-            async with self.connection_pool.acquire() as connection:
-                channel = await connection.channel()
-                return channel
+            # 如果是因为中途关闭导致的错误，不需要再次调用全局 close，否则调用
+            if not self._is_shutdown:
+                await self.close()
+            raise
 
-        # 初始化通道池
-        self.channel_pool = Pool(
-            create_channel,
-            max_size=self.channel_pool_size
-        )
+    async def acquire_channel(self) -> Tuple[RobustChannel, AbstractRobustConnection]:
+        """获取主通道"""
+        if not self._initialized and not self._is_shutdown:
+            await self.init_pools()
+        return await self._ensure_main_channel(), self._connection
 
-        logger.info(
-            f"RabbitMQ连接池初始化完成 - 连接池大小: {self.connection_pool_size}, "
-            f"通道池大小: {self.channel_pool_size}, 集群节点: {self.hosts}"
-        )
+    async def publish_message(self, routing_key: str, message_body: bytes, exchange_name: str = "", **kwargs):
+        """发布消息"""
+        channel, _ = await self.acquire_channel()
+        try:
+            exchange = channel.default_exchange if not exchange_name else await channel.get_exchange(exchange_name)
+            message = Message(body=message_body, **kwargs)
+            await exchange.publish(message, routing_key=routing_key)
+            logger.debug(f"📤 [PUBLISH] 消息发布成功 - RK: {routing_key}")
+        except Exception as e:
+            logger.error(f"❌ [PUBLISH_FAIL] 发布失败: {str(e)}")
+            raise
+
+    async def consume_queue(self, queue_name: str, callback: Callable[[AbstractMessage], asyncio.Future], auto_ack: bool = False, **kwargs):
+        """消费队列"""
+        if not self._initialized:
+            await self.init_pools()
+
+        async with self._lock:
+            if self._is_shutdown:
+                raise RuntimeError("客户端已关闭")
+            if queue_name in self._consumer_channels:
+                logger.warning(f"⚠️ [CONSUMER_EXISTS] 队列 {queue_name} 已在消费中")
+                return
+            if not self._connection or self._connection.is_closed:
+                raise RuntimeError("连接不可用，无法启动消费")
+
+        await self.declare_queue(queue_name, **kwargs)
+
+        try:
+            # 获取原始连接对象创建新通道
+            conn = self._connection
+            consumer_channel = await conn.channel()
+            await consumer_channel.set_qos(prefetch_count=self.prefetch_count)
+            logger.info(
+                f"✅ [CONSUMER_CHANNEL_OK] 消费者通道创建: {id(consumer_channel)}")
+
+            async with self._lock:
+                if self._is_shutdown:
+                    await consumer_channel.close()
+                    return
+                self._consumer_channels[queue_name] = consumer_channel
+
+            async def consume_callback_wrapper(message: AbstractMessage):
+                try:
+                    await callback(message)
+                    if not auto_ack:
+                        await message.ack()
+                except Exception as e:
+                    logger.error(
+                        f"❌ [CALLBACK_ERR] 消费回调异常 {queue_name}: {str(e)}")
+                    if not auto_ack:
+                        await message.nack(requeue=True)
+
+            await consumer_channel.basic_consume(
+                queue_name, consumer_callback=consume_callback_wrapper, auto_ack=auto_ack, **kwargs
+            )
+            logger.info(f"🎧 [CONSUME_START] 开始消费队列: {queue_name}")
+
+        except Exception as e:
+            logger.error(f"💥 [CONSUME_ERR] 启动消费失败 {queue_name}: {str(e)}")
+            async with self._lock:
+                if queue_name in self._consumer_channels:
+                    del self._consumer_channels[queue_name]
+            raise
 
     async def close(self):
-        """关闭连接池和通道池"""
-        if self.channel_pool:
-            await self.channel_pool.close()
-        if self.connection_pool:
-            await self.connection_pool.close()
-        logger.info("RabbitMQ连接池已关闭")
-
-    async def __aenter__(self):
-        await self.init_pools()
-        return self
-
-    async def __aexit__(self, exc_type, exc, tb):
-        await self.close()
+        """
+        资源销毁入口
+        """
+        async with self._lock:
+            if self._is_shutdown:
+                return
+            self._is_shutdown = True
+            self._initialized = False
+            # 记录即将关闭的连接ID
+            conn_to_close_id = id(
+                self._connection) if self._connection else None
+
+        logger.info(f"🛑 [CLOSE_START] 开始关闭客户端... (准备关闭连接: {conn_to_close_id})")
+
+        # 1. 关闭消费者通道
+        channels_to_close = []
+        async with self._lock:
+            channels_to_close = list(self._consumer_channels.values())
+            self._consumer_channels.clear()
+
+        for ch in channels_to_close:
+            try:
+                if not ch.is_closed:
+                    await ch.close()
+                    logger.debug(f"✅ [CLOSE_CHANNEL] 消费者通道已关闭")
+            except Exception as e:
+                logger.warning(f"❌ [CLOSE_CHANNEL_ERR] 关闭消费者通道失败: {str(e)}")
+
+        # 2. 关闭主通道
+        if self._channel:
+            try:
+                if not self._channel.is_closed:
+                    await self._channel.close()
+                    logger.info(f"✅ [CLOSE_CHANNEL] 主通道已关闭")
+            except Exception:
+                pass
+            self._channel = None
+
+        # 3. 关闭连接
+        # 确保在 finally 或显式 close 中调用 connection.close()
+        if self._connection:
+            try:
+                # 打印关闭操作
+                logger.info(f"🔌 [CLOSE_CONN] 正在关闭连接: {id(self._connection)}")
+                await self._connection.close()
+                logger.info(f"✅ [CLOSE_OK] 连接已成功关闭: {id(self._connection)}")
+            except Exception as e:
+                logger.warning(f"❌ [CLOSE_ERR] 关闭连接失败: {str(e)}")
+            self._connection = None
+
+        logger.info("🏁 [CLOSE_DONE] RabbitMQ客户端已完全关闭")
+
+    # --- 辅助方法省略 (declare_queue 等) ---
+    async def declare_queue(self, queue_name: str, **kwargs) -> AbstractQueue:
+        channel, _ = await self.acquire_channel()
+        return await channel.declare_queue(queue_name, **kwargs)
+
+    async def declare_exchange(self, exchange_name: str, exchange_type: str = "direct", **kwargs) -> AbstractExchange:
+        channel, _ = await self.acquire_channel()
+        return await channel.declare_exchange(exchange_name, exchange_type, **kwargs)