Kris
ed2ee086de
主要变更： 1. 新增UIActionScheduler实现优先级UI任务调度，支持CRITICAL/HIGH/NORMAL/LOW四级优先级 2. 替换原有SendQueue为UIActionScheduler，统一所有UI操作的调度逻辑 3. 为截图、登录、重启等API添加调度器封装，支持超时、限流与fast-fail机制 4. 新增Prometheus监控指标，统计UI任务执行、等待耗时、超时、限流与快速失败情况 5. 为QrCapture添加命令执行超时保护，避免X11操作卡死 6. 扩展StatusResponse与状态接口，暴露调度器指标 7. 新增完整的单元测试覆盖调度器功能 8. 为看门狗注入调度器，实现kill前队列清空与快速失败窗口
333 lines
13 KiB
Python
333 lines
13 KiB
Python
"""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
|