AutoScraperX/scheduler/async_consumer.py

import asyncio
import json
import traceback
from typing import List
import signal

from config import settings
from core.base.async_redis_client import RedisManager
from core.utils.log.logger_manager import LoggerManager
from core.utils.trace_utils import generate_trace_id, TraceContext
from services.async_mysql_service import AsyncMysqlService
from spiders.spider_registry import get_spider_class

logger = LoggerManager.get_logger()
aliyun_logger = LoggerManager.get_aliyun_logger()


async def async_handle_topic(topic: str, stop_event: asyncio.Event):
    """
    单个 topic 的消费逻辑，运行在协程中：
    - 从 MQ 中消费消息（单条处理，处理完再拉取下一条）；
    - 根据消息内容执行对应爬虫；
    - 使用异步数据库服务查询配置；
    - 记录日志、确认消息。
    """
    # 每个 topic 创建独立的 consumer 实例（使用优化后的 AsyncRocketMQConsumer）

    from services.async_mq_consumer import AsyncRocketMQConsumer
    from services.async_redis_service import AsyncRedisService
    redis_controller = AsyncRedisService()
    consumer = AsyncRocketMQConsumer(topic_name=topic, group_id=topic)

    async def handle_single_message(message):
        """处理单条消息的业务逻辑（不含拉取和循环）"""
        message_id = message.message_id
        status = await redis_controller.get_status(message_id)
        if status == "1":
            logger.info(f"[{topic}] MessageID: {message_id} 已执行完成，直接 Ack 跳过")
            await consumer.ack_message(message.receipt_handle)
            return
        if not status:
            logger.info(f"[{topic}] MessageID: {message_id} 未执行，标记为处理中")
            await redis_controller.mark_processing(message_id)

        try:
            payload = json.loads(message.message_body)
            task_id = payload["task_id"]
            logger.info(f"[{topic}]接收到任务消息: {task_id}")

            aliyun_logger.logging(
                code="1000",
                message="任务接收成功",
                data=payload,
                account=topic
            )
            # 从数据库查询配置
            async with AsyncMysqlService() as mysql:
                user_list = await mysql.get_user_list(task_id)
                rule_dict = await mysql.get_rule_dict(task_id)

            # 执行爬虫任务
            CrawlerClass = get_spider_class(topic)
            crawler = CrawlerClass(
                rule_dict=rule_dict,
                user_list=user_list,
            )
            await crawler.run()  # 爬虫成功执行后再确认消息
            logger.info(f"[{topic}]任务 {task_id} 执行成功")
            aliyun_logger.logging(
                code="1010",
                message="任务执行成功",
                data={"task_id": task_id, "topic": topic},
                account=topic
            )
            await redis_controller.mark_done(message_id)

        except Exception as e:
            logger.error(f"[{topic}]任务处理失败 MessageID: {message_id}: {e} \n {traceback.format_exc()}")
            aliyun_logger.logging(
                code="9001",
                message=f"处理消息失败: {str(e)} \n {traceback.format_exc()}",
                data={
                    "error_type": type(e).__name__,
                    "stack_trace": traceback.format_exc(),
                    "message_body": message.message_body
                },
                account=topic
            )
            # 处理失败不 Ack，消息会被 MQ 重新投递（依赖 MQ 的重试机制）

    # 独立的消费循环：拉取消息并调用处理函数
    async def consume_loop():
        logger.info(f"[{topic}] 启动消费循环，开始拉取消息...")
        while not stop_event.is_set():  # 监听停止信号，支持优雅退出
            try:
                # 拉取单条消息（依赖优化后的 receive_message，无消息时返回 None 不报错）
                message = await consumer.receive_message()
                if message:
                    with TraceContext() as trace_id:  # 生成 trace_id 并绑定到上下文
                        # 有消息则处理，处理完成后再进入下一次循环
                        await handle_single_message(message)
                else:
                    # 无消息时短暂休眠，避免频繁空轮询
                    await asyncio.sleep(1)
            except Exception as e:
                # 非消息处理的异常（如 MQ 连接失败），记录并重试
                logger.error(f"[{topic}] 消费循环异常: {e}", exc_info=True)
                aliyun_logger.logging(
                    code="9002",
                    message=f"{topic} 消费循环异常，即将重试: {str(e)}",
                    data={"error_type": type(e).__name__, "stack_trace": traceback.format_exc()},
                    account=topic
                )
                await asyncio.sleep(5)  # 异常后延迟重试，减轻服务压力

        logger.info(f"[{topic}] 消费循环已停止（收到退出信号）")

    # 启动消费循环（这是消费逻辑的入口）
    await consume_loop()


async def run_all_topics(topics: List[str]):
    await RedisManager.init(redis_url=settings.redis_url)
    stop_event = asyncio.Event()
    loop = asyncio.get_running_loop()

    def shutdown():
        """处理停止信号（如 Ctrl+C），触发优雅退出"""
        logger.warning("[系统] 收到停止信号，准备优雅退出...")
        aliyun_logger.logging(
            code="1600",
            message="[系统] 收到停止信号，准备优雅退出...",
        )
        stop_event.set()
        asyncio.create_task(RedisManager.close())  # 关闭连接池

    # 注册信号处理（支持 Ctrl+C 和 kill 命令）
    for sig in [signal.SIGINT, signal.SIGTERM]:
        loop.add_signal_handler(sig, shutdown)

    # 为每个 topic 创建独立协程任务
    tasks = [asyncio.create_task(async_handle_topic(topic, stop_event)) for topic in topics]

    await stop_event.wait()  # 等待退出信号

    # 取消所有任务并等待结束
    logger.warning(f"[系统] 正在取消所有消费任务...")
    for task in tasks:
        task.cancel()

    # 收集任务结果，忽略取消异常
    results = await asyncio.gather(*tasks, return_exceptions=True)
    for idx, result in enumerate(results):
        if isinstance(result, Exception) and not isinstance(result, asyncio.CancelledError):
            logger.error(f"[系统] 任务 {topics[idx]} 异常退出: {result}")

    logger.warning(f"[系统] 所有任务已退出，进程已关闭...")
    aliyun_logger.logging(
        code="1602",
        message="[系统] 所有任务已退出，进程已关闭...",
        data={"task_count": len(tasks)}
    )


def handle_topic_group(topics: List[str]):
    """子进程入口函数：启动异步事件循环处理该组 topics。"""
    asyncio.run(run_all_topics(topics))