first commit

2026-05-26 14:16:16 +08:00
commit 535d7f4ac5
40 changed files with 4571 additions and 0 deletions
@@ -0,0 +1 @@
+# 消息队列模块包初始化
@@ -0,0 +1,137 @@
+import asyncio
+import json
+import time
+from typing import Any, Dict
+
+import redis.asyncio as aioredis
+from loguru import logger
+
+from src.config.settings import get_settings
+from src.task_queue.dispatcher import TaskDispatcher
+
+_CLAIM_DUE_MEMBERS_LUA = """
+local key = KEYS[1]
+local now = tonumber(ARGV[1])
+local count = tonumber(ARGV[2])
+local vals = redis.call('ZRANGEBYSCORE', key, '-inf', now, 'LIMIT', 0, count)
+if #vals > 0 then
+  redis.call('ZREM', key, unpack(vals))
+end
+return vals
+"""
+
+
+class RedisZSetDelayListener:
+    def __init__(self, dispatcher: TaskDispatcher) -> None:
+        self.settings = get_settings()
+        self.dispatcher = dispatcher
+        self.redis: aioredis.Redis | None = None
+        self.running = False
+
+        self.node_id = self.settings.node_id
+        self.redis_delay_zset_key = self.settings.redis_delay_zset_key
+        self.zset_key = f"{self.redis_delay_zset_key}:{self.node_id}"
+        self.batch_size = self.settings.redis_delay_batch_size
+
+    async def initialize(self) -> None:
+        logger.debug(f"正在连接 Redis(延时队列): {self.settings.redis_host}:{self.settings.redis_port}...")
+        self.redis = aioredis.Redis(
+            host=self.settings.redis_host,
+            port=self.settings.redis_port,
+            password=self.settings.redis_password,
+            db=self.settings.redis_db,
+        )
+        await self.redis.ping()
+        logger.debug("Redis(延时队列) 连接成功。")
+
+    async def listen(self) -> None:
+        if not self.redis:
+            await self.initialize()
+
+        self.running = True
+        logger.info(f"Redis ZSET 延时队列监听循环已启动，key='{self.zset_key}'。")
+
+        while self.running:
+            try:
+                due_members = await self._claim_due_members()
+                if not due_members:
+                    await asyncio.sleep(3)
+                    continue
+
+                for member in due_members:
+                    member_str = self._decode_member(member)
+                    message_data = self._parse_member_json(member_str)
+                    logger.debug(f"解析到到期任务: {message_data}")
+                    await self.dispatcher.dispatch(message_data)
+
+            except aioredis.ConnectionError as ce:
+                logger.error(f"Redis(延时队列) 连接断开: {ce}，5 秒后尝试重连...")
+                await self._handle_reconnect()
+            except Exception as e:
+                logger.exception(f"延时队列监听循环发生未预期错误: {e}，5 秒后重试...")
+                await asyncio.sleep(5)
+
+    async def _claim_due_members(self) -> list[Any]:
+        if not self.redis:
+            return []
+
+        now_seconds = int(time.time())
+        try:
+            members = await self.redis.eval(
+                _CLAIM_DUE_MEMBERS_LUA,
+                1,
+                self.zset_key,
+                now_seconds,
+                self.batch_size,
+            )
+            return members or []
+        except Exception as e:
+            logger.error(f"延时队列认领到期任务失败: {e}")
+            return []
+
+    def _decode_member(self, member: Any) -> str:
+        if isinstance(member, bytes):
+            return member.decode("utf-8")
+        return str(member)
+
+    def _parse_member_json(self, member_str: str) -> Dict[str, Any]:
+        try:
+            data = json.loads(member_str)
+            if not isinstance(data, dict):
+                raise ValueError("member JSON 不是对象结构")
+            now = int(time.time())
+            task_data = {
+                "task_id": f"id_{now}",
+                "topic": "surugaya_trade_monitor",
+                "purchase_account": "trdian023@tokugen.jp",
+                "site_id": "7",
+                "record_video": "",
+                "payload": {
+                    "trade_code": data.get("trade_code"),
+                },
+            }
+
+            return task_data
+        except Exception as e:
+            raise ValueError(f"延时队列 member 解析失败: {e} | member={member_str}")
+
+    async def _handle_reconnect(self) -> None:
+        self.redis = None
+        while self.running and not self.redis:
+            try:
+                await asyncio.sleep(5)
+                await self.initialize()
+            except Exception as e:
+                logger.error(f"重连失败: {e}，继续重试...")
+
+    async def stop(self) -> None:
+        logger.info("正在停止 Redis ZSET 延时队列监听器...")
+        self.running = False
+        if self.redis:
+            try:
+                await self.redis.close()
+                logger.info("Redis(延时队列) 连接已关闭。")
+            except Exception as e:
+                logger.error(f"关闭 Redis(延时队列) 连接时出错: {e}")
+            finally:
+                self.redis = None
@@ -0,0 +1,119 @@
+import asyncio
+from collections import defaultdict
+from typing import Callable, Coroutine, Dict, Any, Type
+from loguru import logger
+
+from src.tasks.surugaya import SurugayaPurchaseTask
+from src.tasks.surugaya_trade_monitor import SurugayaTradeMonitorTask
+
+# topic 到采购任务类的映射注册表
+TOPIC_TO_TASK_MAP: Dict[str, Type[Any]] = {
+    "surugaya": SurugayaPurchaseTask,  # 骏河屋采购任务
+    "surugaya_trade_monitor": SurugayaTradeMonitorTask,  # 骏河屋交易订单获取任务
+}
+
+
+class TaskDispatcher:
+    """
+    任务分发器：处理来自消息队列的采购任务。
+    - 在后台协程中执行任务，与 Redis 监听器解耦。
+    - 按账号维度使用 asyncio.Lock 串行化浏览器操作，避免冲突。
+    - 支持按 topic 模块化路由到不同工作流。
+    """
+
+    def __init__(self) -> None:
+        # 按 account_id 维护锁字典，新账号会自动创建锁
+        # 单线程事件循环下完全异步安全
+        self._account_locks: Dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
+        logger.info("任务分发器已初始化（支持按账号串行锁）。")
+
+    async def dispatch(
+            self,
+            message_data: Dict[str, Any],
+            on_complete_callback: Callable[[str, bool, Any], Coroutine[Any, Any, None]] = None
+    ) -> None:
+        """
+        启动非阻塞后台协程处理任务，使监听器无需等待浏览器自动化完成即可继续轮询 Redis。
+        
+        Args:
+            message_data: 已解析的任务消息字典。
+            on_complete_callback: 任务完成（成功或失败）时执行的异步回调。
+        """
+        task_id = message_data.get("task_id", "N/A")
+        # logger.debug(f"正在为任务 '{task_id}' 启动后台协程...")
+
+        # 投递到事件循环，fire-and-forget
+        asyncio.create_task(self._execute_task_safe(message_data, on_complete_callback))
+
+    async def _execute_task_safe(
+            self,
+            message_data: Dict[str, Any],
+            on_complete_callback: Callable[[str, bool, Any], Coroutine[Any, Any, None]] = None
+    ) -> None:
+        """
+        在账号锁约束下执行具体采购流程，并统一记录异常。
+        """
+        task_id = message_data.get("task_id", "N/A")
+        topic = str(message_data.get("topic"))
+        account_id = str(message_data.get("purchase_account"))
+        callback_id = message_data.get("_stream_msg_id") or message_data.get("_callback_id") or task_id
+
+        logger.info(f"[任务 {task_id}] 开始分发，topic='{topic}'，账号='{account_id}'。")
+
+        # 1. 校验必填字段
+        if not topic or not account_id:
+            logger.error(f"[任务 {task_id}] 缺少必填字段 topic 或 purchase_account，任务已中止。")
+            if on_complete_callback:
+                await on_complete_callback(str(callback_id), False, {"error": "缺少关键路由字段。", "result": None})
+            return
+
+        # 2. 根据 topic 匹配已注册的工作流
+        task_class = TOPIC_TO_TASK_MAP.get(topic.lower())
+        if not task_class:
+            logger.error(f"[任务 {task_id}] 未找到 topic '{topic}' 对应的采购任务类，任务已中止。")
+            if on_complete_callback:
+                await on_complete_callback(
+                    str(callback_id),
+                    False,
+                    {"error": f"未知的 topic: {topic}", "result": None},
+                )
+            return
+
+        # 3. 获取该账号的串行锁
+        lock = self._account_locks[account_id]
+        logger.info(f"[任务 {task_id}] 账号 '{account_id}' 忙碌中，已进入串行锁队列等待...")
+
+        # 异步锁保证同一 account_id 同一时刻只有一个协程在执行
+        async with lock:
+            logger.success(f"[任务 {task_id}] 已获取账号 '{account_id}' 的锁，正在启动浏览器工作流...")
+
+            success = False
+            error_msg = None
+            result_data: Any = None
+            task_instance: Any = None
+
+            try:
+                # 实例化任务并执行
+                task_instance = task_class(message_data)
+                await task_instance.run()
+                success = True
+                result_data = getattr(task_instance, "task_result_data", None)
+                logger.success(f"[任务 {task_id}] 工作流执行成功。")
+            except Exception as e:
+                error_msg = str(e)
+                if task_instance:
+                    result_data = getattr(task_instance, "task_result_data", None)
+                logger.error(f"[任务 {task_id}] 工作流因未处理异常而终止: {e}")
+            finally:
+                # 4. 触发完成回调（例如 Redis Stream ACK）
+                if on_complete_callback:
+                    try:
+                        await on_complete_callback(
+                            str(callback_id),
+                            success,
+                            {"error": error_msg, "result": result_data},
+                        )
+                    except Exception as cb_err:
+                        logger.error(f"[任务 {task_id}] 完成回调执行失败: {cb_err}")
+
+                logger.debug(f"[任务 {task_id}] 已释放账号 '{account_id}' 的锁。")
@@ -0,0 +1,205 @@
+import asyncio
+import json
+from typing import Dict, Any
+import redis.asyncio as aioredis
+from loguru import logger
+from src.config.settings import get_settings
+from src.task_queue.dispatcher import TaskDispatcher
+
+
+class RedisStreamListener:
+    """
+    Redis Stream 持续消息监听器。
+    - 使用 Redis 消费者组在多消费者间分担负载。
+    - 启动时读取 Pending 任务，恢复崩溃前未完成的任务。
+    - 非阻塞地将任务交给 TaskDispatcher，执行完毕后再 ACK。
+    - 连接异常时安全重连。
+    """
+
+    def __init__(self, dispatcher: TaskDispatcher) -> None:
+        self.settings = get_settings()
+        self.dispatcher = dispatcher
+        self.redis: aioredis.Redis | None = None
+        self.running = False
+
+        # 从全局配置读取 Stream 相关参数
+        self.stream_key = self.settings.redis_stream_key
+        self.group_name = self.settings.redis_consumer_group
+        self.consumer_name = self.settings.redis_consumer_name
+
+    async def initialize(self) -> None:
+        """
+        连接 Redis，若消费者组不存在则创建。
+        """
+        logger.debug(f"正在连接 Redis: {self.settings.redis_host}:{self.settings.redis_port}...")
+        self.redis = aioredis.Redis(
+            host=self.settings.redis_host,
+            port=self.settings.redis_port,
+            password=self.settings.redis_password,
+            db=self.settings.redis_db
+        )
+
+        # 测试连接
+        await self.redis.ping()
+        logger.debug("Redis 连接成功。")
+
+        # 确保 Stream 消费者组存在
+        try:
+            # mkstream=True：若 Stream 不存在则自动创建
+            await self.redis.xgroup_create(
+                name=self.stream_key,
+                groupname=self.group_name,
+                id="0",
+                mkstream=True
+            )
+            logger.debug(f"已创建消费者组 '{self.group_name}'，Stream: '{self.stream_key}'。")
+        except Exception as e:
+            # BUSYGROUP 表示消费者组已存在，属于正常情况
+            if "BUSYGROUP" in str(e):
+                logger.debug(f"消费者组 '{self.group_name}' 已存在。")
+            else:
+                logger.error(f"创建消费者组失败: {e}")
+                raise e
+
+    async def listen(self) -> None:
+        """
+        启动监听循环：先处理 Pending 任务，再持续读取新消息。
+        """
+        if not self.redis:
+            await self.initialize()
+
+        self.running = True
+        logger.info("Redis Stream 监听循环已启动。")
+
+        # 1. 先恢复分配给本消费者但未 ACK 的 Pending 消息（崩溃重启清理）
+        await self._consume_pending_messages()
+
+        # 2. 主循环：持续消费新任务
+        while self.running:
+            try:
+                # 读取新消息（">" 表示从未投递给其他消费者的消息）
+                streams = await self.redis.xreadgroup(
+                    groupname=self.group_name,
+                    consumername=self.consumer_name,
+                    streams={self.stream_key: ">"},
+                    count=1,
+                    block=3000  # 阻塞 3 秒，保持事件循环响应
+                )
+
+                if not streams:
+                    # 超时无新消息，继续下一轮
+                    continue
+
+                for stream, messages in streams:
+                    for msg_id, raw_data in messages:
+                        parsed_msg = self._parse_message(msg_id, raw_data)
+                        # 交给分发器处理，并注册 ACK 回调
+                        await self.dispatcher.dispatch(parsed_msg, on_complete_callback=self.ack_message)
+
+            except aioredis.ConnectionError as ce:
+                logger.error(f"Redis 连接断开: {ce}，5 秒后尝试重连...")
+                await self._handle_reconnect()
+            except Exception as e:
+                logger.exception(f"监听循环发生未预期错误: {e}，5 秒后重试...")
+                await asyncio.sleep(5)
+
+    async def _consume_pending_messages(self) -> None:
+        """
+        检查并处理本消费者在 PEL（Pending Entries List）中的未确认任务。
+        保证崩溃后可恢复中断的作业。
+        """
+        logger.info("正在检查 PEL 中的待恢复任务...")
+        try:
+            # 读取已投递给本消费者但未 ACK 的消息（id="0"）
+            streams = await self.redis.xreadgroup(
+                groupname=self.group_name,
+                consumername=self.consumer_name,
+                streams={self.stream_key: "0"},
+                count=10,
+                block=500
+            )
+
+            if not streams:
+                logger.info("没有需要恢复的待处理任务。")
+                return
+
+            for stream, messages in streams:
+                if messages:
+                    logger.warning(f"发现 {len(messages)} 条待恢复任务，正在重新处理...")
+                    for msg_id, raw_data in messages:
+                        parsed_msg = self._parse_message(msg_id, raw_data)
+                        await self.dispatcher.dispatch(parsed_msg, on_complete_callback=self.ack_message)
+        except Exception as e:
+            logger.error(f"恢复待处理任务失败: {e}")
+
+    def _parse_message(self, msg_id: bytes, raw_data: Dict[bytes, bytes]) -> Dict[str, Any]:
+        """
+        将 Redis 二进制数据解码为字符串，并组装为任务字典。
+        """
+        decoded_data: Dict[str, Any] = {}
+        for k, v in raw_data.items():
+            key_str = k.decode("utf-8")
+            val_str = v.decode("utf-8")
+            decoded_data[key_str] = val_str
+
+        # 若 payload 为 JSON 字符串，则反序列化
+        if "payload" in decoded_data:
+            try:
+                payload_value = decoded_data["payload"]
+                if isinstance(payload_value, str):
+                    decoded_data = json.loads(payload_value)
+            except Exception:
+                # 非 JSON 格式时保持原字符串
+                pass
+
+        # 注入 Stream 消息 ID，供后续 ACK 使用
+        decoded_data["_stream_msg_id"] = msg_id.decode("utf-8")
+        logger.debug(f"解析消息 '{msg_id.decode('utf-8')}'，内容: {decoded_data}")
+        return decoded_data
+
+    async def ack_message(self, msg_id: str, success: bool, error_msg: Any = None) -> None:
+        """
+        任务结束后的回调：执行 Redis Stream XACK，从 PEL 中移除已处理消息。
+        """
+        if not self.redis:
+            logger.error(f"Redis 未连接，无法 ACK 消息 '{msg_id}'。")
+            return
+
+        try:
+            await self.redis.xack(self.stream_key, self.group_name, msg_id)
+            if success:
+                logger.debug(f"消息 '{msg_id}' 已 ACK，任务执行成功。")
+            else:
+                detail = error_msg
+                if isinstance(error_msg, dict) and "error" in error_msg:
+                    detail = error_msg.get("error")
+                logger.warning(f"消息 '{msg_id}' 已 ACK，任务执行失败: {detail}")
+        except Exception as e:
+            logger.error(f"ACK 消息 '{msg_id}' 失败: {e}")
+
+    async def _handle_reconnect(self) -> None:
+        """
+        周期性尝试重连 Redis，直至成功。
+        """
+        self.redis = None
+        while not self.redis:
+            try:
+                await asyncio.sleep(5)
+                await self.initialize()
+            except Exception as e:
+                logger.error(f"重连失败: {e}，继续重试...")
+
+    async def stop(self) -> None:
+        """
+        停止监听循环并安全关闭 Redis 连接。
+        """
+        logger.info("正在停止 Redis Stream 监听器...")
+        self.running = False
+        if self.redis:
+            try:
+                await self.redis.close()
+                logger.info("Redis 连接已关闭。")
+            except Exception as e:
+                logger.error(f"关闭 Redis 连接时出错: {e}")
+            finally:
+                self.redis = None
@@ -0,0 +1,59 @@
+from functools import lru_cache
+
+import redis.asyncio as aioredis
+from loguru import logger
+
+from src.config.settings import get_settings
+
+
+class RedisManager:
+    """
+    管理任务侧共享 Redis 客户端的生命周期。
+    """
+
+    def __init__(self) -> None:
+        self.settings = get_settings()
+        self._redis: aioredis.Redis | None = None
+
+    async def initialize(self) -> None:
+        if self._redis is not None:
+            return
+
+        logger.debug(
+            f"正在初始化任务侧 RedisManager: {self.settings.redis_host}:{self.settings.redis_port}..."
+        )
+        redis_client = aioredis.Redis(
+            host=self.settings.redis_host,
+            port=self.settings.redis_port,
+            password=self.settings.redis_password,
+            db=self.settings.redis_db,
+            health_check_interval=30,  # 定期心跳保活，防止长连接被防火墙静默掐断
+            max_connections=50,        # 限制最大连接数，防止突发并发耗尽 Redis 资源
+        )
+        await redis_client.ping()
+        self._redis = redis_client
+        logger.info("RedisManager 初始化成功。")
+
+    def get_client(self) -> aioredis.Redis:
+        if self._redis is None:
+            raise RuntimeError("RedisManager 尚未初始化，请先在应用启动阶段调用 initialize()。")
+        return self._redis
+
+    async def close(self) -> None:
+        if self._redis is None:
+            return
+
+        try:
+            await self._redis.close()
+            logger.info("任务侧 RedisManager 已关闭。")
+        finally:
+            self._redis = None
+
+
+@lru_cache
+def get_redis_manager() -> RedisManager:
+    return RedisManager()
+
+
+def get_redis_client() -> aioredis.Redis:
+    return get_redis_manager().get_client()