sycommon-python-lib 0.1.56__py3-none-any.whl → 0.1.56b2__py3-none-any.whl

sycommon/rabbitmq/rabbitmq_pool.py

@@ -5,25 +5,16 @@ from aio_pika import connect_robust, RobustChannel, Message
 from aio_pika.abc import (
     AbstractRobustConnection, AbstractQueue, AbstractExchange, AbstractMessage
 )
+from aio_pika.exceptions import ChannelClosed
+import aiormq.exceptions
+
 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
-        return self.method(obj)
-
-
 class RabbitMQConnectionPool:
-    """单连接单通道RabbitMQ客户端 (严格执行“先清理后连接”策略)"""
+    """单连接单通道RabbitMQ客户端（核心特性：依赖connect_robust原生自动重连/恢复 + 仅关闭时释放资源）"""
 
     def __init__(
         self,
@@ -32,12 +23,13 @@ class RabbitMQConnectionPool:
         username: str,
         password: str,
         virtualhost: str = "/",
-        heartbeat: int = 15,
+        heartbeat: int = 30,
         app_name: str = "",
-        connection_timeout: int = 15,
+        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主机地址")
@@ -52,347 +44,295 @@ class RabbitMQConnectionPool:
         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}")
+        logger.info(
+            f"随机选择RabbitMQ主机: {self._current_host}（依赖connect_robust原生自动重连/恢复）")
 
-        # 核心资源
-        self._connection: Optional[AbstractRobustConnection] = None
-        self._channel: Optional[RobustChannel] = None
-        self._consumer_channels: Dict[str, RobustChannel] = {}
+        # 核心资源（单连接+单通道，基于原生自动重连）
+        self._connection: Optional[AbstractRobustConnection] = None  # 原生自动重连连接
+        self._channel: Optional[RobustChannel] = None  # 单通道（原生自动恢复）
+        # 消费者通道跟踪（独立于主通道）
+        self._consumer_channels: Dict[str,
+                                      Tuple[RobustChannel, Callable, bool, dict]] = {}
 
-        # 状态控制
+        # 状态控制（并发安全+生命周期管理）
         self._lock = asyncio.Lock()
         self._initialized = False
         self._is_shutdown = False
 
-    @AsyncProperty
+    async def _is_connection_valid(self) -> bool:
+        """原子化检查连接有效性（所有状态判断均加锁，确保原子性）"""
+        async with self._lock:
+            # 优先级：先判断是否关闭，再判断是否初始化，最后判断连接状态
+            return not self._is_shutdown and self._initialized and self._connection is not None and not self._connection.is_closed
+
+    @property
     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
-            if self._channel is None or self._channel.is_closed:
-                return False
-            return True
-
-    async def _cleanup_resources(self):
-        """
-        彻底清理旧资源
-        必须在持有 self._lock 的情况下调用
-        """
-        logger.info("🧹 [CLEANUP] 开始清理旧资源...")
-
-        # 1. 清理所有消费者通道
-        if self._consumer_channels:
-            channels_to_close = list(self._consumer_channels.values())
-            self._consumer_channels.clear()
+            # 存活条件：未关闭 + 已初始化 + 连接有效 + 主通道有效
+            return self._initialized and self._connection is not None and not self._connection.is_closed and self._channel is not None and not self._channel.is_closed
 
-            for ch in channels_to_close:
-                try:
-                    if not ch.is_closed:
-                        await ch.close()
-                except Exception as e:
-                    logger.warning(f"⚠️ [CLEANUP_CH] 关闭消费者通道失败: {e}")
+    async def _create_connection(self) -> AbstractRobustConnection:
+        """创建原生自动重连连接（仅创建一次，内部自动重试）"""
+        async with self._lock:
+            if self._is_shutdown:
+                raise RuntimeError("客户端已关闭，无法创建连接")
 
-        # 2. 关闭主通道
-        if self._channel:
-            try:
-                if not self._channel.is_closed:
-                    await self._channel.close()
-            except Exception as e:
-                logger.warning(f"⚠️ [CLEANUP_MAIN_CH] 关闭主通道失败: {e}")
-            finally:
-                self._channel = None
+        conn_url = f"amqp://{self.username}:{self.password}@{self._current_host}:{self.port}/{self.virtualhost}?name={self.app_name}&heartbeat={self.heartbeat}&reconnect_interval={self.reconnect_interval}&fail_fast=1"
+        logger.info(f"尝试创建原生自动重连连接: {self._current_host}:{self.port}")
 
-        # 3. 关闭连接
-        if self._connection:
-            try:
-                if not self._connection.is_closed:
-                    # close() 可能是同步的，也可能是异步的，aio_pika 中通常是异步的
-                    await self._connection.close()
-            except Exception as e:
-                logger.warning(f"⚠️ [CLEANUP_CONN] 关闭连接失败: {e}")
-            finally:
-                self._connection = None
-
-        logger.info("✅ [CLEANUP] 资源清理完成")
-
-    async def _create_connection_impl(self, host: str) -> AbstractRobustConnection:
-        conn_url = (
-            f"amqp://{self.username}:{self.password}@{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] 尝试连接节点: {host}")
         try:
-            conn = await asyncio.wait_for(
-                connect_robust(conn_url),
-                timeout=self.connection_timeout + 5
+            conn = await connect_robust(
+                conn_url,
+                timeout=self.connection_timeout,
             )
-            logger.info(f"✅ [CONNECT_OK] 节点连接成功: {host}")
+            logger.info(f"连接创建成功: {self._current_host}:{self.port}（原生自动重连已启用）")
             return conn
         except Exception as e:
-            logger.error(f"❌ [CONNECT_FAIL] 节点 {host} 连接失败: {str(e)}")
-            raise ConnectionError(f"无法连接RabbitMQ {host}") from e
-
-    async def _ensure_main_channel(self) -> RobustChannel:
-        """
-        确保主通道有效
-        逻辑：
-        1. 检查连接状态
-        2. 如果断开 -> 清理 -> 轮询重试
-        3. 如果连接在但通道断开 -> 仅重建通道
-        """
+            logger.error(f"连接创建失败: {str(e)}", exc_info=True)
+            raise ConnectionError(
+                f"无法连接RabbitMQ主机 {self._current_host}:{self.port}") from e
+
+    async def _init_single_channel(self):
+        """初始化单通道（通道自带原生自动恢复）"""
         async with self._lock:
+            # 先判断是否关闭（优先级最高）
             if self._is_shutdown:
-                raise RuntimeError("客户端已关闭")
-
-            # --- 阶段A：连接恢复逻辑 (如果连接断了) ---
-            if self._connection is None or self._connection.is_closed:
-
-                # 1. 【强制】先彻底清理所有旧资源
-                await self._cleanup_resources()
-
-                retry_hosts = self.hosts.copy()
-                random.shuffle(retry_hosts)
-                last_error = None
-                max_attempts = min(len(retry_hosts), 3)
-
-                # 2. 轮询尝试新连接
-                for _ in range(max_attempts):
-                    if not retry_hosts:
-                        break
-
-                    host = retry_hosts.pop()
-                    self._current_host = host
-                    temp_conn = None
-
-                    try:
-                        temp_conn = await self._create_connection_impl(host)
-
-                        # 3. 只有在连接成功后，才更新 self._connection
-                        self._connection = temp_conn
-                        temp_conn = None  # 转移所有权
-                        self._initialized = True
-                        last_error = None
-                        logger.info(f"🔗 [RECONNECT_OK] 切换到节点: {host}")
-                        break
-
-                    except Exception as e:
-                        logger.warning(f"⚠️ [RECONNECT_RETRY] 节点 {host} 不可用")
-                        if temp_conn is not None:
-                            # 尝试连接失败了，必须把这个“半成品”连接关掉
-                            try:
-                                await temp_conn.close()
-                            except Exception:
-                                pass
-                        last_error = e
-                        await asyncio.sleep(self.reconnect_interval)
-
-                # 4. 如果所有尝试都失败
-                if last_error:
-                    # 确保状态是干净的
-                    self._connection = None
-                    self._initialized = False
-                    logger.error("💥 [RECONNECT_FATAL] 所有节点重试失败")
-                    raise ConnectionError("所有 RabbitMQ 节点连接失败") from last_error
-
-            # --- 阶段B：通道恢复逻辑 (如果连接在但通道断了) ---
-            # 注意：这里不需要清理连接，只重置通道
-            if self._channel is None or self._channel.is_closed:
-                try:
-                    self._channel = await self._connection.channel()
-                    await self._channel.set_qos(prefetch_count=self.prefetch_count)
-                    logger.info(f"✅ [CHANNEL_OK] 主通道已恢复")
-                except Exception as e:
-                    # 如果连通道都创建不了，说明这个连接也是坏的，回滚到阶段A
-                    logger.error(f"❌ [CHANNEL_FAIL] 通道创建失败，标记连接无效: {e}")
-                    # 强制清理连接，触发下一次进入阶段A
-                    await self._cleanup_resources()
-                    raise
+                raise RuntimeError("客户端已关闭，无法初始化通道")
+            # 再判断连接是否有效
+            if not self._connection or self._connection.is_closed:
+                raise RuntimeError("无有效连接，无法初始化通道")
+
+            # 清理旧通道（如果存在）
+            if self._channel and not self._channel.is_closed:
+                await self._channel.close()
+
+            # 创建单通道并设置QOS
+            try:
+                self._channel = await self._connection.channel()
+                await self._channel.set_qos(prefetch_count=self.prefetch_count)
+                logger.info(f"单通道初始化完成（带原生自动恢复）")
+            except Exception as e:
+                logger.error(f"创建单通道失败: {str(e)}", exc_info=True)
+                raise
+
+    async def _check_and_recover_channel(self) -> RobustChannel:
+        """检查并恢复通道（确保通道有效，所有状态判断加锁）"""
+        async with self._lock:
+            # 1. 先判断是否关闭（优先级最高）
+            if self._is_shutdown:
+                raise RuntimeError("客户端已关闭，无法获取通道")
+            # 2. 检查连接状态
+            if not self._connection or self._connection.is_closed:
+                raise RuntimeError("连接已关闭（等待原生重连）")
+            # 3. 通道失效时重新创建
+            if not self._channel or self._channel.is_closed:
+                logger.warning("通道失效，重新创建（依赖原生自动恢复）")
+                await self._init_single_channel()
 
             return self._channel
 
     async def init_pools(self):
-        """初始化入口"""
+        """初始化客户端（仅执行一次）"""
         async with self._lock:
+            # 原子化判断：是否已关闭/已初始化
             if self._is_shutdown:
-                raise RuntimeError("客户端已关闭")
+                raise RuntimeError("客户端已关闭，无法初始化")
             if self._initialized:
+                logger.warning("客户端已初始化，无需重复调用")
                 return
 
-        # 在 try 之前声明变量，确保 except 块能访问
-        conn_created_in_this_try = None
-
         try:
-            # 锁外创建连接，减少锁持有时间
-            init_host = random.choice(self.hosts)
-            conn = await self._create_connection_impl(init_host)
+            # 1. 创建原生自动重连连接
+            self._connection = await self._create_connection()
 
-            # 记录本次创建的连接
-            conn_created_in_this_try = conn
+            # 2. 初始化单通道
+            await self._init_single_channel()
 
+            # 3. 标记为已初始化（加锁保护）
             async with self._lock:
-                if self._is_shutdown:
-                    raise RuntimeError("客户端已关闭")
-
-                # 提交新资源
-                self._connection = conn
-                self._channel = await self._connection.channel()
-                await self._channel.set_qos(prefetch_count=self.prefetch_count)
                 self._initialized = True
 
-                # 所有权转移成功，清空临时引用，防止 finally 重复关闭
-                conn_created_in_this_try = None
-
-                logger.info(f"🚀 [INIT_OK] 连接池初始化完成: {init_host}")
-
+            logger.info("RabbitMQ单通道客户端初始化完成（原生自动重连/恢复已启用）")
         except Exception as e:
-            logger.error(f"💥 [INIT_FAIL] 初始化异常: {str(e)}")
-
-            # 这里现在可以合法访问 conn_created_in_this_try
-            if conn_created_in_this_try is not None:
-                try:
-                    await conn_created_in_this_try.close()
-                except Exception:
-                    pass
-
-            if not self._is_shutdown:
-                await self.close()
+            logger.error(f"初始化失败: {str(e)}", exc_info=True)
+            await self.close()  # 初始化失败直接关闭
             raise
 
-    async def force_reconnect(self):
-        """
-        强制重连
-        严格执行：清理所有资源 -> 尝试建立新资源
-        """
+    async def acquire_channel(self) -> Tuple[RobustChannel, AbstractRobustConnection]:
+        """获取单通道（返回 (通道, 连接) 元组，保持API兼容）"""
         async with self._lock:
+            # 原子化状态校验
             if self._is_shutdown:
-                return
-
-            logger.warning("🔄 [FORCE_RECONNECT] 开始强制重连...")
-
-            # 1. 【关键】标记未初始化，迫使 _ensure_main_channel 走清理流程
-            self._initialized = False
-
-            # 2. 【关键】立即清理旧资源 (在锁内)
-            await self._cleanup_resources()
+                raise RuntimeError("客户端已关闭，无法获取通道")
+            if not self._initialized:
+                raise RuntimeError("客户端未初始化，请先调用init_pools()")
 
-            # 此时 self._connection 和 self._channel 均为 None
+        # 检查并恢复通道
+        channel = await self._check_and_recover_channel()
+        return channel, self._connection  # 单通道无需管理"使用中/空闲"状态
 
-        # 3. 锁外触发恢复 (避免阻塞锁太久)
-        try:
-            await self.acquire_channel()
-            logger.info("✅ [FORCE_RECONNECT_OK] 强制重连成功")
-        except Exception as e:
-            logger.error(f"❌ [FORCE_RECONNECT_FAIL] 强制重连失败: {e}")
-            raise
+    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 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
+    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)
 
     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"消息发布成功 - 交换机: {exchange.name}, 路由键: {routing_key}"
+            )
         except Exception as e:
-            logger.error(f"❌ [PUBLISH_FAIL] 发布失败: {str(e)}")
-            raise
+            logger.error(f"发布消息失败: {str(e)}", exc_info=True)
+            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("客户端已关闭")
+                raise RuntimeError("客户端已关闭，无法启动消费")
+            if not self._initialized:
+                raise RuntimeError("客户端未初始化，请先调用init_pools()")
             if queue_name in self._consumer_channels:
-                logger.warning(f"⚠️ [CONSUMER_EXISTS] 队列 {queue_name} 已在消费中")
+                logger.warning(f"队列 {queue_name} 已在消费中，无需重复启动")
                 return
-            if not self._connection or self._connection.is_closed:
-                raise RuntimeError("连接不可用，无法启动消费")
 
-        # 声明队列 (使用主通道)
+        # 先声明队列（确保队列存在）
         await self.declare_queue(queue_name, **kwargs)
 
-        try:
-            # 获取最新连接
-            _, conn = await self.acquire_channel()
+        # 创建独立的消费者通道（不使用主单通道，避免消费阻塞发布）
+        async with self._lock:
+            if self._is_shutdown:  # 二次校验：防止创建通道前客户端被关闭
+                raise RuntimeError("客户端已关闭，无法创建消费者通道")
+            if not self._connection or self._connection.is_closed:
+                raise RuntimeError("无有效连接，无法创建消费者通道")
+            channel = await self._connection.channel()
+            await channel.set_qos(prefetch_count=self.prefetch_count)
 
-            # 创建消费者通道
-            consumer_channel = await conn.channel()
-            await consumer_channel.set_qos(prefetch_count=self.prefetch_count)
+            # 注册消费者通道
+            self._consumer_channels[queue_name] = (
+                channel, callback, auto_ack, kwargs)
 
-            async with self._lock:
-                # 再次检查，防止并发创建
-                if self._is_shutdown:
-                    await consumer_channel.close()
-                    return
-                if queue_name in self._consumer_channels:
-                    await consumer_channel.close()  # 其他协程已经创建了
+        async def consume_callback_wrapper(message: AbstractMessage):
+            """消费回调包装（处理通道失效，依赖原生恢复）"""
+            try:
+                async with self._lock:
+                    # 原子化校验状态：客户端是否关闭 + 通道是否有效 + 连接是否有效
+                    if self._is_shutdown:
+                        logger.warning(f"客户端已关闭，拒绝处理消息（队列: {queue_name}）")
+                        if not auto_ack:
+                            await message.nack(requeue=True)
+                        return
+                    channel_valid = not channel.is_closed
+                    conn_valid = self._connection and not self._connection.is_closed
+
+                if not channel_valid or not conn_valid:
+                    logger.warning(f"消费者通道 {queue_name} 失效（等待原生自动恢复）")
+                    if not auto_ack:
+                        await message.nack(requeue=True)
                     return
 
-                self._consumer_channels[queue_name] = consumer_channel
+                # 执行业务回调
+                await callback(message)
+                if not auto_ack:
+                    await message.ack()
+            except ChannelClosed as e:
+                logger.error(f"消费者通道 {queue_name} 关闭: {str(e)}", exc_info=True)
+                if not auto_ack:
+                    await message.nack(requeue=True)
+            except aiormq.exceptions.ChannelInvalidStateError as e:
+                logger.error(
+                    f"消费者通道 {queue_name} 状态异常: {str(e)}", exc_info=True)
+                if not auto_ack:
+                    await message.nack(requeue=True)
+            except Exception as e:
+                logger.error(
+                    f"消费消息失败（队列: {queue_name}）: {str(e)}", exc_info=True)
+                if not auto_ack:
+                    await message.nack(requeue=True)
 
-            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}: {e}")
-                    if not auto_ack:
-                        await message.nack(requeue=True)
+        logger.info(f"开始消费队列: {queue_name}（通道带原生自动恢复）")
 
-            await consumer_channel.basic_consume(
-                queue_name, consumer_callback=consume_callback_wrapper, auto_ack=auto_ack
+        try:
+            await 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}: {e}")
-            # 失败时清理字典
-            async with self._lock:
-                if queue_name in self._consumer_channels:
-                    # 注意：这里清理的是字典里的引用，通道本身应该在 try 块里被关闭了吗？
-                    # 如果 consumer_channel 创建成功但 basic_consume 失败，需要手动关闭
-                    ch = self._consumer_channels.pop(queue_name, None)
-                    if ch:
-                        try:
-                            await ch.close()
-                        except:
-                            pass
+            logger.error(f"启动消费失败（队列: {queue_name}）: {str(e)}", exc_info=True)
+            # 清理异常资源
+            try:
+                async with self._lock:
+                    if not channel.is_closed:
+                        await channel.close()
+                    # 移除无效的消费者通道注册
+                    if queue_name in self._consumer_channels:
+                        del self._consumer_channels[queue_name]
+            except Exception as close_e:
+                logger.warning(f"关闭消费者通道失败: {str(close_e)}")
             raise
 
     async def close(self):
-        """资源销毁"""
+        """关闭客户端（释放所有资源，原子化状态管理）"""
         async with self._lock:
             if self._is_shutdown:
+                logger.warning("客户端已关闭，无需重复操作")
                 return
+            # 先标记为关闭，阻止后续所有操作（原子化修改）
             self._is_shutdown = True
             self._initialized = False
 
-        logger.info("🛑 [CLOSE] 开始关闭连接池...")
+        logger.info("开始关闭RabbitMQ单通道客户端（释放所有资源）...")
 
-        # 1. 清理所有资源
-        await self._cleanup_resources()
+        # 1. 关闭所有消费者通道
+        consumer_channels = []
+        async with self._lock:
+            consumer_channels = list(self._consumer_channels.values())
+            self._consumer_channels.clear()
+        for channel, _, _, _ in consumer_channels:
+            try:
+                if not channel.is_closed:
+                    await channel.close()
+            except Exception as e:
+                logger.warning(f"关闭消费者通道失败: {str(e)}")
 
-        logger.info("🏁 [CLOSE] 连接池已关闭")
+        # 2. 关闭主单通道
+        if self._channel:
+            try:
+                async with self._lock:
+                    if not self._channel.is_closed:
+                        await self._channel.close()
+            except Exception as e:
+                logger.warning(f"关闭主通道失败: {str(e)}")
+            self._channel = None
 
-    async def declare_queue(self, queue_name: str, **kwargs) -> AbstractQueue:
-        channel, _ = await self.acquire_channel()
-        return await channel.declare_queue(queue_name, **kwargs)
+        # 3. 关闭连接（终止原生自动重连）
+        if self._connection:
+            try:
+                async with self._lock:
+                    if not self._connection.is_closed:
+                        await self._connection.close()
+                logger.info(
+                    f"已关闭连接: {self._current_host}:{self.port}（终止原生自动重连）")
+            except Exception as e:
+                logger.warning(f"关闭连接失败: {str(e)}")
+            self._connection = None
 
-    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)
+        logger.info("RabbitMQ单通道客户端已完全关闭")