"""UIActionScheduler —— 基于 SendQueue 的优先级调度器。 在父类 SendQueue 串行化队列之上扩展: 1. 优先级调度(PriorityQueue + IntEnum 优先级),CRITICAL 任务优先出队且不受 默认延时约束; 2. fast-fail 机制:微信进程死亡期间,非 CRITICAL 任务立即失败,避免堆积; 3. Prometheus metrics 暴露:执行计数 / 等待与执行耗时 / 超时 / 限流 / fast-fail。 元素类型为 5 元组 `(priority_int, seq, coro_factory, future, delay_ms)`, 其中 seq 为单调递增序列号,保证同优先级 FIFO。 """ from __future__ import annotations import asyncio import enum import logging import time from typing import Any, Awaitable, Callable, Optional from woc_bridge.messaging.send_queue import SendQueue, CoroFactory from woc_bridge.models import BridgeError from woc_bridge.ui.metrics import ( ui_action_executed, # Counter, label=priority ui_action_wait_duration, # Histogram ui_action_exec_duration, # Histogram, label=priority ui_action_timeout, # Counter ui_action_rate_limited, # Counter ui_action_fast_fail, # Counter ) logger = logging.getLogger("woc-bridge") def _priority_name(priority: int) -> str: """将 priority int 转为名称,非法值返回 "?"。""" try: return Priority(priority).name except ValueError: return "?" class Priority(enum.IntEnum): """UI 操作优先级。 数值越小优先级越高;CRITICAL 不受 fast-fail 与默认延时影响。 """ CRITICAL = 0 # wechat_restart / diagnostic_autofix HIGH = 1 # logout / friend_accept / login_qr_capture NORMAL = 2 # send_text / send_file / set_remark / add_friend LOW = 3 # screenshot / moments_like / moments_comment class UIActionScheduler(SendQueue): """优先级 UI 调度器。 覆盖父类的 ``_queue`` 为 :class:`asyncio.PriorityQueue`,元素为 5 元组 ``(priority, seq, coro_factory, future, delay_ms)``。提供 :meth:`execute` 入口携带优先级与 trace_id,并暴露 metrics / drain 等运维接口。 """ def __init__( self, send_delay_ms: int = 3000, max_calls_per_sec: int = 10, max_queue_size: int = 100, ) -> None: super().__init__( send_delay_ms=send_delay_ms, max_calls_per_sec=max_calls_per_sec, max_queue_size=max_queue_size, ) # 覆盖父类的 Queue 为 PriorityQueue;元素 5 元组 self._queue: asyncio.PriorityQueue[ tuple[int, int, CoroFactory, asyncio.Future, Optional[int]] ] = asyncio.PriorityQueue(maxsize=max_queue_size) # 单调递增序列号,保证同优先级 FIFO self._seq: int = 0 # fast-fail 截止时间戳(monotonic);>0 表示微信死亡期间 self._wechat_dead_until: float = 0.0 # metrics 计数器 self._metrics: dict[str, Any] = { "executed_total": 0, "executed_by_priority": {p.name: 0 for p in Priority}, "wait_duration_ms_sum": 0.0, "exec_duration_ms_sum": 0.0, "timeout_total": 0, "rate_limited_total": 0, "fast_fail_total": 0, } def mark_wechat_dead(self, duration_sec: float = 10.0) -> None: """标记微信进程死亡,期间非 CRITICAL 任务 fast-fail。 Args: duration_sec: 死亡窗口秒数,期间入队/出队的非 CRITICAL 任务 会被立即以 WECHAT_NOT_READY 失败 """ self._wechat_dead_until = time.monotonic() + duration_sec logger.warning( "[scheduler] mark_wechat_dead %.1fs, %d pending tasks will fast-fail", duration_sec, self._queue.qsize(), ) async def execute( self, action: CoroFactory, priority: Priority = Priority.NORMAL, delay_ms: Optional[int] = None, wait_timeout_ms: Optional[int] = None, trace_id: str = "", ) -> Any: """按优先级入队并等待执行结果。 CRITICAL 任务若未指定 delay_ms,自动设为 0(立即执行)。 Args: action: 返回 coroutine 的工厂函数 priority: 优先级,默认 NORMAL delay_ms: 自定义执行后延时毫秒,None 用默认 send_delay_ms wait_timeout_ms: 等待结果超时毫秒,None 表示无限等待 trace_id: 追踪 ID,用于日志关联 Returns: coroutine 的执行结果 Raises: BridgeError: 队列满 / 等待超时 / 任务执行异常透传 """ if priority == Priority.CRITICAL and delay_ms is None: delay_ms = 0 return await self._enqueue_with_priority( action, priority, delay_ms, wait_timeout_ms, trace_id, ) async def _enqueue_with_priority( self, coro_factory: CoroFactory, priority: Priority, delay_ms: Optional[int], wait_timeout_ms: Optional[int], trace_id: str, ) -> Any: """实际入队逻辑:构造 5 元组入 PriorityQueue,等待 future 结果。""" loop = asyncio.get_running_loop() future: asyncio.Future = loop.create_future() self._seq += 1 item = (int(priority), self._seq, coro_factory, future, delay_ms) try: self._queue.put_nowait(item) except asyncio.QueueFull: self._metrics["rate_limited_total"] += 1 ui_action_rate_limited.inc() logger.warning( "[scheduler] 队列已满,拒绝入队 (maxsize=%d, priority=%s, trace_id=%s)", self._queue.maxsize, priority.name, trace_id, ) raise BridgeError( code="RATE_LIMITED", message=f"UI 调度器队列已满({self._queue.maxsize}),请稍后重试", details={"retry_after": 3}, ) pending = self._queue.qsize() logger.info( "[scheduler] 入队 priority=%s pending=%d trace_id=%s", priority.name, pending, trace_id, ) wait_start = time.monotonic() if wait_timeout_ms is not None and wait_timeout_ms > 0: try: result = await asyncio.wait_for( future, timeout=wait_timeout_ms / 1000.0, ) except asyncio.TimeoutError: # 竞争窗口:worker 可能刚好在此时完成并 set_result。 if future.done() and not future.cancelled(): logger.info( "[scheduler] 等待超时但任务刚好完成,取结果 (pending=%d)", pending, ) return future.result() # 取消 future,worker 出队时发现 cancelled 会跳过执行 future.cancel() wait_sec = wait_timeout_ms / 1000.0 self._metrics["timeout_total"] += 1 ui_action_timeout.inc() logger.warning( "[scheduler] 等待超时 (pending=%d, wait_timeout=%.1fs, trace_id=%s)", pending, wait_sec, trace_id, ) raise BridgeError( code="TIMEOUT", message=f"UI 调度器等待超时({pending} 个待处理,已等 {wait_sec:.1f}s)", details={"retry_after": max(1, int(self.send_delay_ms / 1000))}, ) else: result = await future wait_ms = (time.monotonic() - wait_start) * 1000 self._metrics["wait_duration_ms_sum"] += wait_ms ui_action_wait_duration.observe(wait_ms / 1000.0) return result async def enqueue( self, coro_factory: CoroFactory, delay_ms: Optional[int] = None, wait_timeout_ms: Optional[int] = None, ) -> Any: """向后兼容包装:以 NORMAL 优先级入队。""" return await self._enqueue_with_priority( coro_factory, Priority.NORMAL, delay_ms, wait_timeout_ms, "", ) async def _run(self) -> None: """worker 主循环(覆盖父类)。 解包 5 元组 ``(priority, seq, coro_factory, future, delay_ms)``, 按优先级出队执行;含 fast-fail 与 metrics 埋点。 """ while True: try: priority, seq, coro_factory, future, custom_delay_ms = await self._queue.get() except asyncio.CancelledError: raise # 调用方已超时取消:跳过执行与延时 if future.cancelled(): self._queue.task_done() logger.info( "[scheduler] 出队任务已取消,跳过 (pending=%d, priority=%s)", self._queue.qsize(), _priority_name(priority), ) continue # fast-fail 检查:微信已死期间,非 CRITICAL 任务直接失败 if (self._wechat_dead_until > time.monotonic() and priority > Priority.CRITICAL): # CRITICAL (=0) 不受影响 logger.info( "[scheduler] fast-fail (wechat dead, %.1fs remaining, priority=%s)", self._wechat_dead_until - time.monotonic(), _priority_name(priority), ) if not future.done(): future.set_exception(BridgeError( code="WECHAT_NOT_READY", message="微信进程未就绪(重启中),请稍后重试", details={"retry_after": 5}, )) self._metrics["fast_fail_total"] += 1 ui_action_fast_fail.inc() # Prometheus self._queue.task_done() continue # 不延时,立即处理下一个 logger.info( "[scheduler] 出队,开始处理 (pending=%d, priority=%s)", self._queue.qsize(), _priority_name(priority), ) executed = False t_exec = time.perf_counter() try: self._check_rate_limit() # 继承父类 self._recent_call_times.append(time.monotonic()) executed = True pname = _priority_name(priority) logger.info("[scheduler] 开始执行任务 (priority=%s)", pname) result = await coro_factory() exec_ms = (time.perf_counter() - t_exec) * 1000 logger.info( "[scheduler] 任务执行完成 (%.0fms, priority=%s)", exec_ms, pname, ) if not future.done(): future.set_result(result) except asyncio.CancelledError: if not future.done(): future.cancel() raise except Exception as e: exec_ms = (time.perf_counter() - t_exec) * 1000 logger.warning( "[scheduler] 任务执行抛异常 %s: %s (%.0fms, priority=%s)", type(e).__name__, e, exec_ms, _priority_name(priority), ) if not future.done(): future.set_exception(e) finally: self._queue.task_done() if executed: exec_ms = (time.perf_counter() - t_exec) * 1000 self._metrics["executed_total"] += 1 self._metrics["exec_duration_ms_sum"] += exec_ms try: p = Priority(priority) self._metrics["executed_by_priority"][p.name] += 1 ui_action_executed.labels(priority=p.name).inc() # Prometheus ui_action_exec_duration.labels(priority=p.name).observe(exec_ms / 1000.0) # Prometheus except ValueError: pass delay = custom_delay_ms if custom_delay_ms is not None else self.send_delay_ms logger.info( "[scheduler] 延时 %dms 后处理下一个 (priority=%s)", delay, _priority_name(priority), ) await asyncio.sleep(delay / 1000.0) def get_metrics(self) -> dict: """返回调度器 metrics 快照(含 pending_count)。""" return {**self._metrics, "pending_count": self._queue.qsize()} async def drain(self, timeout: float = 5.0) -> bool: """等待队列排空。 Args: timeout: 最长等待秒数 Returns: True 表示队列已排空;False 表示超时仍有任务待处理 """ try: await asyncio.wait_for(self._queue.join(), timeout=timeout) return True except asyncio.TimeoutError: logger.warning( "[scheduler] drain timeout %.1fs, %d tasks pending", timeout, self._queue.qsize(), ) return False