"""发送串行化队列。 所有 xdotool 操作经此队列串行执行,避免并发 UI 操作冲突。 内部维护最近 1 秒调用时间戳用于限流。 """ from __future__ import annotations import asyncio import logging import time from typing import Any, Awaitable, Callable, Optional from woc_bridge.models import BridgeError logger = logging.getLogger("woc-bridge") # 工厂类型:返回一个待执行的 coroutine CoroFactory = Callable[[], Awaitable[Any]] class SendQueue: """串行化发送队列 + 限流。 通过 asyncio.Queue 串行执行所有发送任务;执行间隔可配置 (默认 3000ms),单实例每秒调用上限可配置(默认 10)。 """ def __init__( self, send_delay_ms: int = 3000, max_calls_per_sec: int = 10, max_queue_size: int = 100, ) -> None: """初始化队列配置。 Args: send_delay_ms: 两次发送之间的最小间隔毫秒 max_calls_per_sec: 每秒最大调用次数 max_queue_size: 队列最大深度,满时入队抛 RATE_LIMITED """ self.send_delay_ms = send_delay_ms self.max_calls_per_sec = max_calls_per_sec self._queue: asyncio.Queue[tuple[CoroFactory, asyncio.Future, Optional[int]]] = asyncio.Queue(maxsize=max_queue_size) self._worker: asyncio.Task | None = None # 最近 1 秒内的调用时间戳 self._recent_call_times: list[float] = [] async def start(self) -> None: """启动 worker task。""" if self._worker is None or self._worker.done(): self._worker = asyncio.create_task(self._run()) async def stop(self) -> None: """取消 worker。""" if self._worker is not None and not self._worker.done(): self._worker.cancel() try: await self._worker except asyncio.CancelledError: pass self._worker = None async def enqueue( self, coro_factory: CoroFactory, delay_ms: Optional[int] = None, wait_timeout_ms: Optional[int] = None, ) -> Any: """将一个返回 coroutine 的工厂入队,等待执行结果。 Args: coro_factory: 调用后返回 coroutine 的工厂函数 delay_ms: 自定义延时毫秒,None 用默认 send_delay_ms wait_timeout_ms: 队列等待超时毫秒,None 表示无限等待; 超时抛出 BridgeError(TIMEOUT) Returns: coroutine 的执行结果 Raises: BridgeError: 限流命中、队列满或等待超时时抛错; 任务执行抛出的异常会透传给调用方 """ loop = asyncio.get_running_loop() future: asyncio.Future = loop.create_future() # 队列满时立即拒绝,避免无界增长 OOM if self._queue.maxsize > 0 and self._queue.qsize() >= self._queue.maxsize: logger.warning( "send_queue: 队列已满 (size=%d/%d),拒绝入队", self._queue.qsize(), self._queue.maxsize, ) raise BridgeError( code="RATE_LIMITED", message=f"发送队列已满({self._queue.qsize()}/{self._queue.maxsize}),请稍后重试", details={"retry_after": 3}, ) await self._queue.put((coro_factory, future, delay_ms)) pending = self._queue.qsize() logger.info("send_queue: 入队 (pending=%d)", pending) if wait_timeout_ms is not None and wait_timeout_ms > 0: try: return await asyncio.wait_for( future, timeout=wait_timeout_ms / 1000.0 ) except asyncio.TimeoutError: # 竞争窗口:worker 可能刚好在此时完成并 set_result。 # 若 future 已完成且未被取消,直接取结果;否则取消 future 并抛 TIMEOUT。 if future.done() and not future.cancelled(): logger.info( "send_queue: 等待超时但任务刚好完成,取结果 (pending=%d)", pending, ) return future.result() # 取消 future,worker 出队时发现 cancelled 会跳过执行 future.cancel() wait_sec = wait_timeout_ms / 1000.0 logger.warning( "send_queue: 等待超时 (pending=%d, wait_timeout=%.1fs)", pending, wait_sec, ) raise BridgeError( code="TIMEOUT", message=f"发送队列等待超时({pending} 个待处理,已等 {wait_sec:.1f}s)", details={"retry_after": max(1, int(self.send_delay_ms / 1000))}, ) return await future def pending_count(self) -> int: """返回当前队列中待执行任务数(供 /api/status 暴露给客户端做退避决策)。""" return self._queue.qsize() def _check_rate_limit(self) -> None: """检查限流。 清理 1 秒前的时间戳,若当前已满 max_calls_per_sec 则抛 BridgeError(RATE_LIMITED),并在 details 中携带 retry_after 秒数 供上层设置 Retry-After 响应头。 """ now = time.monotonic() # 清理 1 秒前的时间戳 self._recent_call_times = [t for t in self._recent_call_times if now - t < 1.0] if len(self._recent_call_times) >= self.max_calls_per_sec: # 计算建议等待秒数:最早一次调用距窗口边界还差多久 oldest = self._recent_call_times[0] retry_after = max(1, int(1.0 - (now - oldest)) + 1) logger.warning( "send_queue: 限流命中,拒绝执行 (recent=%d/%d, retry_after=%ds)", len(self._recent_call_times), self.max_calls_per_sec, retry_after, ) raise BridgeError( code="RATE_LIMITED", message=f"发送限流:每秒最多 {self.max_calls_per_sec} 次", details={"retry_after": retry_after}, ) logger.debug( "send_queue: 限流通过 (recent=%d/%d)", len(self._recent_call_times), self.max_calls_per_sec, ) async def _run(self) -> None: """worker 主循环。 循环取出任务执行;执行前检查限流(超限则失败该任务); 执行前记录开始时间戳(避免长任务导致 1 秒窗口内超限); 执行后 sleep delay_ms/1000(自定义延时优先于默认 send_delay_ms)。 """ while True: coro_factory, future, custom_delay_ms = await self._queue.get() # 调用方已超时取消:跳过执行与延时 if future.cancelled(): self._queue.task_done() logger.info( "send_queue: 出队任务已取消,跳过 (pending=%d)", self._queue.qsize(), ) continue logger.info("send_queue: 出队,开始处理 (pending=%d)", self._queue.qsize()) # 标记任务是否真正开始执行(用于决定 finally 是否延时) executed = False t_exec = time.perf_counter() try: # 执行前检查限流 self._check_rate_limit() # 记录开始时间戳(限流窗口基于开始时刻,避免长任务后窗口偏移) self._recent_call_times.append(time.monotonic()) executed = True # 执行任务 logger.info("send_queue: 开始执行任务") result = await coro_factory() logger.info( "send_queue: 任务执行完成 (%.0fms)", (time.perf_counter() - t_exec) * 1000, ) if not future.done(): future.set_result(result) except asyncio.CancelledError: # worker 被取消时,把取消传播给等待的调用方 if not future.done(): future.cancel() raise except Exception as e: logger.warning( "send_queue: 任务执行抛异常 %s: %s (%.0fms)", type(e).__name__, e, (time.perf_counter() - t_exec) * 1000, ) if not future.done(): future.set_exception(e) finally: self._queue.task_done() # 仅在任务真正执行过时延时,限流失败/已取消的任务不延时 if executed: delay = custom_delay_ms if custom_delay_ms is not None else self.send_delay_ms logger.info( "send_queue: 延时 %dms 后处理下一个", delay, ) await asyncio.sleep(delay / 1000.0)