sycommon-python-lib 0.1.56b5__py3-none-any.whl → 0.1.57b1__py3-none-any.whl

sycommon/rabbitmq/rabbitmq_pool.py

@@ -19,12 +19,11 @@ class AsyncProperty:
     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,
@@ -33,9 +32,9 @@ class RabbitMQConnectionPool:
         username: str,
         password: str,
         virtualhost: str = "/",
-        heartbeat: int = 30,
+        heartbeat: int = 15,
         app_name: str = "",
-        connection_timeout: int = 30,
+        connection_timeout: int = 15,
         reconnect_interval: int = 5,
         prefetch_count: int = 2,
     ):
@@ -68,124 +67,231 @@ class RabbitMQConnectionPool:
 
     @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
-
-            # 可选：检查主通道是否存活
             if self._channel is None or self._channel.is_closed:
-                # 如果你认为通道断了连接也算死，就保留这行；否则删除
                 return False
-
             return True
 
-    async def _create_connection_impl(self) -> AbstractRobustConnection:
+    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()
+
+            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}")
+
+        # 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
+
+        # 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}@{self._current_host}:{self.port}/"
+            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_START] 尝试创建连接 -> {self._current_host}:{self.port}")
+        logger.info(f"🔌 [CONNECT] 尝试连接节点: {host}")
         try:
-            conn = await connect_robust(conn_url, timeout=self.connection_timeout)
-            # 注意：connect_robust 返回时，底层 TCP 可能还在握手，但对象已创建
-            logger.info(f"✅ [CONNECT_OK] 连接对象创建成功: {id(conn)}")
+            conn = await asyncio.wait_for(
+                connect_robust(conn_url),
+                timeout=self.connection_timeout + 5
+            )
+            logger.info(f"✅ [CONNECT_OK] 节点连接成功: {host}")
             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
+            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. 如果连接在但通道断开 -> 仅重建通道
+        """
         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)}")
+            # --- 阶段A：连接恢复逻辑 (如果连接断了) ---
+            if self._connection is None or self._connection.is_closed:
 
-            # 2. 确保主通道存在
+                # 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:
-                logger.info("⚠️ [RECOVER_CHANNEL] 检测到主通道不存在或已关闭，开始恢复...")
-                self._channel = await self._create_channel_impl(conn)
+                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
 
             return self._channel
 
     async def init_pools(self):
-        """
-        初始化入口与异常处理 (修复泄漏的关键)
-        """
+        """初始化入口"""
         async with self._lock:
             if self._is_shutdown:
                 raise RuntimeError("客户端已关闭")
             if self._initialized:
                 return
 
+        # 在 try 之前声明变量，确保 except 块能访问
         conn_created_in_this_try = None
+
         try:
-            # 步骤 A: 创建连接 (在锁外进行，避免阻塞其他操作)
-            conn = await self._create_connection_impl()
-            conn_created_in_this_try = conn  # 记录本次创建的对象，用于失败回滚
+            # 锁外创建连接，减少锁持有时间
+            init_host = random.choice(self.hosts)
+            conn = await self._create_connection_impl(init_host)
+
+            # 记录本次创建的连接
+            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._channel = await self._connection.channel()
+                await self._channel.set_qos(prefetch_count=self.prefetch_count)
                 self._initialized = True
-                logger.info(
-                    f"🚀 [INIT_SUCCESS] 客户端初始化完成. ConnID: {id(self._connection)}")
+
+                # 所有权转移成功，清空临时引用，防止 finally 重复关闭
+                conn_created_in_this_try = None
+
+                logger.info(f"🚀 [INIT_OK] 连接池初始化完成: {init_host}")
 
         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)}")
+            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()
-                    logger.info(
-                        f"✅ [CLOSE_OK] 泄漏连接已关闭: {id(conn_created_in_this_try)}")
-                except Exception as close_err:
-                    logger.error(f"❌ [CLOSE_ERR] 关闭泄漏连接时出错: {str(close_err)}")
+                except Exception:
+                    pass
 
-            # 如果是因为中途关闭导致的错误，不需要再次调用全局 close，否则调用
             if not self._is_shutdown:
                 await self.close()
             raise
 
+    async def force_reconnect(self):
+        """
+        强制重连
+        严格执行：清理所有资源 -> 尝试建立新资源
+        """
+        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()
+
+            # 此时 self._connection 和 self._channel 均为 None
+
+        # 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 acquire_channel(self) -> Tuple[RobustChannel, AbstractRobustConnection]:
         """获取主通道"""
         if not self._initialized and not self._is_shutdown:
@@ -193,22 +299,20 @@ class RabbitMQConnectionPool:
         return await self._ensure_main_channel(), self._connection
 
     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("客户端已关闭")
@@ -218,20 +322,26 @@ class RabbitMQConnectionPool:
             if not self._connection or self._connection.is_closed:
                 raise RuntimeError("连接不可用，无法启动消费")
 
+        # 声明队列 (使用主通道)
         await self.declare_queue(queue_name, **kwargs)
 
         try:
-            # 获取原始连接对象创建新通道
-            conn = self._connection
+            # 获取最新连接
+            _, conn = await self.acquire_channel()
+
+            # 创建消费者通道
             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
+                if queue_name in self._consumer_channels:
+                    await consumer_channel.close()  # 其他协程已经创建了
+                    return
+
                 self._consumer_channels[queue_name] = consumer_channel
 
             async def consume_callback_wrapper(message: AbstractMessage):
@@ -240,77 +350,45 @@ class RabbitMQConnectionPool:
                     if not auto_ack:
                         await message.ack()
                 except Exception as e:
-                    logger.error(
-                        f"❌ [CALLBACK_ERR] 消费回调异常 {queue_name}: {str(e)}")
+                    logger.error(f"❌ [CALLBACK_ERR] {queue_name}: {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
+                queue_name, consumer_callback=consume_callback_wrapper, auto_ack=auto_ack
             )
-            logger.info(f"🎧 [CONSUME_START] 开始消费队列: {queue_name}")
+            logger.info(f"🎧 [CONSUME_START] {queue_name}")
 
         except Exception as e:
-            logger.error(f"💥 [CONSUME_ERR] 启动消费失败 {queue_name}: {str(e)}")
+            logger.error(f"💥 [CONSUME_ERR] {queue_name}: {e}")
+            # 失败时清理字典
             async with self._lock:
                 if queue_name in self._consumer_channels:
-                    del self._consumer_channels[queue_name]
+                    # 注意：这里清理的是字典里的引用，通道本身应该在 try 块里被关闭了吗？
+                    # 如果 consumer_channel 创建成功但 basic_consume 失败，需要手动关闭
+                    ch = self._consumer_channels.pop(queue_name, None)
+                    if ch:
+                        try:
+                            await ch.close()
+                        except:
+                            pass
             raise
 
     async def close(self):
-        """
-        资源销毁入口
-        """
+        """资源销毁"""
         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
+        logger.info("🛑 [CLOSE] 开始关闭连接池...")
 
-        # 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
+        # 1. 清理所有资源
+        await self._cleanup_resources()
 
-        logger.info("🏁 [CLOSE_DONE] RabbitMQ客户端已完全关闭")
+        logger.info("🏁 [CLOSE] 连接池已关闭")
 
-    # --- 辅助方法省略 (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)