76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import heapq
|
||
|
|
import logging
|
||
|
|
from collections.abc import Awaitable, Callable
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
class OrderedEventQueue:
|
||
|
|
def __init__(self, max_size: int = 1000, max_wait_ms: int = 5000):
|
||
|
|
self._max_size = max_size
|
||
|
|
self._max_wait_ms = max_wait_ms
|
||
|
|
self._heap: list[tuple[int, int, dict]] = []
|
||
|
|
self._seq_counter = 0
|
||
|
|
self._locked = False
|
||
|
|
self._consumer_task: asyncio.Task | None = None
|
||
|
|
self._handler: Callable[[dict], Awaitable[None]] | None = None
|
||
|
|
self._event = asyncio.Event()
|
||
|
|
|
||
|
|
def set_handler(self, handler: Callable[[dict], Awaitable[None]]) -> None:
|
||
|
|
self._handler = handler
|
||
|
|
|
||
|
|
async def push(self, event: dict, timestamp: int = 0) -> None:
|
||
|
|
self._seq_counter += 1
|
||
|
|
entry = (timestamp, self._seq_counter, event)
|
||
|
|
heapq.heappush(self._heap, entry)
|
||
|
|
if len(self._heap) > self._max_size:
|
||
|
|
heapq.heappop(self._heap)
|
||
|
|
self._event.set()
|
||
|
|
|
||
|
|
async def start(self) -> None:
|
||
|
|
if self._consumer_task is not None:
|
||
|
|
return
|
||
|
|
self._consumer_task = asyncio.create_task(self._consume_loop())
|
||
|
|
|
||
|
|
async def stop(self) -> None:
|
||
|
|
if self._consumer_task:
|
||
|
|
self._consumer_task.cancel()
|
||
|
|
try:
|
||
|
|
await self._consumer_task
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
pass
|
||
|
|
self._consumer_task = None
|
||
|
|
|
||
|
|
async def _consume_loop(self) -> None:
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
if not self._heap:
|
||
|
|
await asyncio.wait_for(
|
||
|
|
self._event.wait(),
|
||
|
|
timeout=self._max_wait_ms / 1000.0,
|
||
|
|
)
|
||
|
|
self._event.clear()
|
||
|
|
if not self._heap:
|
||
|
|
continue
|
||
|
|
|
||
|
|
ts, seq, event = heapq.heappop(self._heap)
|
||
|
|
if self._handler:
|
||
|
|
try:
|
||
|
|
await self._handler(event)
|
||
|
|
except Exception:
|
||
|
|
logger.exception("OrderedEventQueue handler failed")
|
||
|
|
|
||
|
|
except TimeoutError:
|
||
|
|
continue
|
||
|
|
except asyncio.CancelledError:
|
||
|
|
break
|
||
|
|
except Exception:
|
||
|
|
logger.exception("OrderedEventQueue consume loop error")
|
||
|
|
|
||
|
|
@property
|
||
|
|
def size(self) -> int:
|
||
|
|
return len(self._heap)
|