126 lines
3.9 KiB
Python
126 lines
3.9 KiB
Python
"""A tiny in-process publish/subscribe bus.
|
|
|
|
The robot layer publishes; the WebSocket layer subscribes and forwards to the
|
|
browser. Keeping this in the middle means the robot adapters never know that a
|
|
browser exists, and the API layer never polls the robot.
|
|
|
|
Delivery is best-effort and non-blocking: a slow or wedged browser tab can never
|
|
stall the robot loop. If a subscriber's queue overflows we drop its oldest event
|
|
and mark the stream as lossy rather than applying backpressure upstream.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
MAX_QUEUE = 256
|
|
|
|
|
|
class EventType(str):
|
|
"""String constants - kept as plain strings so they serialise for free."""
|
|
|
|
ROBOT_STATUS = "robot.status"
|
|
SPEECH_PROGRESS = "speech.progress"
|
|
SPEECH_RESULT = "speech.result"
|
|
HISTORY_UPDATED = "history.updated"
|
|
CONFIG_UPDATED = "config.updated"
|
|
NOTICE = "notice"
|
|
|
|
|
|
@dataclass
|
|
class Event:
|
|
type: str
|
|
data: Dict[str, Any] = field(default_factory=dict)
|
|
at: float = field(default_factory=time.time)
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
return {"type": self.type, "data": self.data, "at": self.at}
|
|
|
|
|
|
class Subscription:
|
|
"""An async iterator over bus events, scoped to one subscriber."""
|
|
|
|
def __init__(self, bus: "EventBus", name: str) -> None:
|
|
self._bus = bus
|
|
self._queue: asyncio.Queue = asyncio.Queue(maxsize=MAX_QUEUE)
|
|
self.name = name
|
|
self.dropped = 0
|
|
|
|
def _offer(self, event: Event) -> None:
|
|
try:
|
|
self._queue.put_nowait(event)
|
|
except asyncio.QueueFull:
|
|
self.dropped += 1
|
|
try: # make room, keep the newest - stale status is worse than none
|
|
self._queue.get_nowait()
|
|
self._queue.put_nowait(event)
|
|
except (asyncio.QueueEmpty, asyncio.QueueFull): # pragma: no cover
|
|
pass
|
|
|
|
async def get(self, timeout: Optional[float] = None) -> Optional[Event]:
|
|
if timeout is None:
|
|
return await self._queue.get()
|
|
try:
|
|
return await asyncio.wait_for(self._queue.get(), timeout=timeout)
|
|
except asyncio.TimeoutError:
|
|
return None
|
|
|
|
def __enter__(self) -> "Subscription":
|
|
return self
|
|
|
|
def __exit__(self, *exc: Any) -> None:
|
|
self.close()
|
|
|
|
def close(self) -> None:
|
|
self._bus.unsubscribe(self)
|
|
|
|
|
|
class EventBus:
|
|
def __init__(self) -> None:
|
|
self._subscribers: List[Subscription] = []
|
|
self._last: Dict[str, Event] = {}
|
|
|
|
# -- subscription -------------------------------------------------------- #
|
|
def subscribe(self, name: str = "anonymous") -> Subscription:
|
|
sub = Subscription(self, name)
|
|
self._subscribers.append(sub)
|
|
logger.debug("event subscriber added: %s (total=%d)", name, len(self._subscribers))
|
|
return sub
|
|
|
|
def unsubscribe(self, sub: Subscription) -> None:
|
|
if sub in self._subscribers:
|
|
self._subscribers.remove(sub)
|
|
logger.debug("event subscriber removed: %s (total=%d)", sub.name, len(self._subscribers))
|
|
|
|
@property
|
|
def subscriber_count(self) -> int:
|
|
return len(self._subscribers)
|
|
|
|
# -- publishing ---------------------------------------------------------- #
|
|
def publish(self, event_type: str, data: Optional[Dict[str, Any]] = None) -> Event:
|
|
event = Event(type=event_type, data=data or {})
|
|
self._last[event_type] = event
|
|
for sub in list(self._subscribers):
|
|
sub._offer(event)
|
|
return event
|
|
|
|
def last(self, event_type: str) -> Optional[Event]:
|
|
"""Most recent event of a type - used to prime a newly-opened socket."""
|
|
return self._last.get(event_type)
|
|
|
|
|
|
_bus: Optional[EventBus] = None
|
|
|
|
|
|
def get_event_bus() -> EventBus:
|
|
global _bus
|
|
if _bus is None:
|
|
_bus = EventBus()
|
|
return _bus
|