from __future__ import annotations import asyncio from collections.abc import Awaitable, Callable from yuxi.utils.logging_config import logger SENTINEL = object() class EventQueue: def __init__(self, max_size: int = 10000): self._queue: asyncio.Queue = asyncio.Queue(maxsize=max_size) self._handler: Callable[[dict], Awaitable[None]] | None = None self._task: asyncio.Task | None = None self._running = False self._dropped_count = 0 async def start(self, handler: Callable[[dict], Awaitable[None]]) -> None: self._handler = handler self._running = True self._task = asyncio.create_task(self._process_loop()) logger.info("[Yuanbao] Event queue started") async def stop(self) -> None: self._running = False self._queue.put_nowait(SENTINEL) if self._task and not self._task.done(): self._task.cancel() try: await self._task except asyncio.CancelledError: pass self._task = None logger.info("[Yuanbao] Event queue stopped") async def enqueue(self, event: dict) -> None: try: self._queue.put_nowait(event) except asyncio.QueueFull: event_type = event.get("type", "unknown") msg_id = event.get("msg_id", "N/A") logger.warning( f"[Yuanbao] Event queue full, dropping event: " f"type={event_type} msg_id={msg_id}" ) self._dropped_count += 1 async def _process_loop(self) -> None: while self._running: try: event = await self._queue.get() if event is SENTINEL: break if self._handler: await self._handler(event) except asyncio.CancelledError: break except Exception as e: logger.error(f"[Yuanbao] Event queue process error: {e}", exc_info=True) @property def size(self) -> int: return self._queue.qsize()