You've already forked pruchase_jhw
zmq
120 lines
5.2 KiB
Python
120 lines
5.2 KiB
Python
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}' 的锁。")
|