import asyncio
import sys
from collections import defaultdict
from collections.abc import Callable, Coroutine
from typing import Any, ClassVar
try:
from zmq.asyncio import Context, Socket
except ImportError as e:
raise ImportError(
"Attempted to import zmq runner, but zmq deps are not installed. install with `noob[zmq]`",
) from e
if sys.version_info < (3, 12):
from typing_extensions import TypedDict
else:
from typing import TypedDict
from noob.logging import init_logger
from noob.network.message import Message
from noob.utils import iscoroutinefunction_partial
class _CallbackDict(TypedDict):
sync: list[Callable[[Message], Any]]
asyncio: list[Callable[[Message], Coroutine]]
[docs]
class EventloopMixin:
"""
Mixin to provide common asyncio zmq scaffolding to networked classes.
Inheriting classes should, in order
* call the ``_init_loop`` method to create the eventloop, context, and poller
* populate the private ``_sockets`` and ``_receivers`` dicts
* await the ``_poll_sockets`` method, which polls indefinitely.
Inheriting classes **must** ensure that ``_init_loop``
is called in the thread it is intended to run in,
and that thread must already have a running eventloop.
asyncio eventloops (and most of asyncio) are **not** thread safe.
To help avoid cross-threading issues, the :meth:`.context` and :meth:`.loop`
properties do *not* automatically create the objects,
raising a :class:`.RuntimeError` if they are accessed before ``_init_loop`` is called.
To handle slow subscribers in PUB/SUB, use an XPUB instead,
and each subscriber also subscribes to a topic with the :attr:`.SUBSCRIBER_PREFIX`:``id``.
The XPUB then polls for receive events with `_poll_subscriptions`,
and maintains a list of nodes that have subscribed to it.
Subscribers can then wait for the publisher to acknowledge its subscription before
marking themselves as ready.
"""
SUBSCRIBER_PREFIX: ClassVar[str] = "__subscriber__:"
def __init__(self):
self._context = None
self._loop = None
self._quitting: asyncio.Event = None # type: ignore[assignment]
self._sockets: dict[str, Socket] = {}
"""
All sockets, mapped from some common name to the socket.
The same key used here should be shared between _receivers and _callbacks
"""
self._receivers: dict[str, Socket] = {}
"""Sockets that should be polled for incoming messages"""
self._callbacks: dict[str, _CallbackDict] = defaultdict(
lambda: _CallbackDict(sync=[], asyncio=[])
)
"""Callbacks for each receiver socket"""
self._subscribers: set[str] = set()
"""The node IDs that we know to be subscribed to our XPUB socket"""
if not hasattr(self, "logger"):
self.logger = init_logger("eventloop")
@property
def context(self) -> Context:
if self._context is None:
raise RuntimeError("Loop has not been initialized with _init_loop!")
return self._context
@property
def loop(self) -> asyncio.AbstractEventLoop:
if self._loop is None:
raise RuntimeError("Loop has not been initialized with _init_loop!")
return self._loop
@property
def sockets(self) -> dict[str, Socket]:
return self._sockets
[docs]
def register_socket(self, name: str, socket: Socket, receiver: bool = False) -> None:
"""Register a socket, optionally declaring it as a receiver socket to poll"""
if name in self._sockets:
raise KeyError(f"Socket {name} already declared!")
self._sockets[name] = socket
if receiver:
self._receivers[name] = socket
[docs]
def add_callback(
self, socket: str, callback: Callable[[Message], Any] | Callable[[Message], Coroutine]
) -> None:
"""
Add a callback to be called when the socket receives a message.
Callbacks are called in the order in which they are added.
"""
if socket not in self._receivers:
raise KeyError(f"Socket {socket} does not exist or is not a receiving socket")
if iscoroutinefunction_partial(callback):
self._callbacks[socket]["asyncio"].append(callback)
else:
self._callbacks[socket]["sync"].append(callback)
[docs]
def clear_callbacks(self) -> None:
self._callbacks = defaultdict(lambda: _CallbackDict(sync=[], asyncio=[]))
def _init_loop(self) -> None:
self._loop = asyncio.get_running_loop()
self._context = Context.instance()
self._quitting = asyncio.Event()
def _stop_loop(self) -> None:
if self._quitting is None:
return
self._quitting.set()
async def _poll_receivers(self) -> None:
"""
Rather than using the zmq.asyncio.Poller which wastes a ton of time,
it turns out doing it this way is roughly 4x as fast:
just manually poll the sockets, and if you have multiple sockets,
gather multiple coroutines where you're polling the sockets.
"""
if len(self._receivers) == 1:
await self._poll_receiver(next(iter(self._receivers.keys())))
else:
await asyncio.gather(*[self._poll_receiver(name) for name in self._receivers])
async def _poll_subscriptions(
self, name: str, handler: Callable[[bool, str], Coroutine]
) -> None:
"""
Drain XPUB subscription events from an outbox socket.
XPUB delivers a single frame per (un)subscription: a leading ``\\x01`` (subscribe)
or ``\\x00`` (unsubscribe) byte followed by the subscribed topic. Subscribers
additionally subscribe to a ``__subscriber__:<node_id>`` topic so we can attribute
the event to a node id; the anonymous ``""`` data subscription is ignored here.
``handler`` is called ``(subscribed: bool, node_id: str)`` for attributable events.
"""
socket = self._sockets[name]
prefix = self.SUBSCRIBER_PREFIX.encode("utf-8")
while not self._quitting.is_set():
frame: bytes = await socket.recv()
if not frame:
continue
subscribed = frame[0] == 1
topic = frame[1:]
if not topic.startswith(prefix):
# the anonymous "" data subscription, or an unrelated topic
continue
node_id = topic[len(prefix) :].decode("utf-8")
changed = (
node_id not in self._subscribers if subscribed else node_id in self._subscribers
)
if not changed:
self.logger.debug(
"Node %s %s, but already was! Ignoring.",
node_id,
"subscribed" if subscribed else "unsubscribed",
)
continue
if subscribed:
self._subscribers.add(node_id)
else:
self._subscribers.discard(node_id)
await handler(subscribed, node_id)
await asyncio.sleep(0)
self.logger.debug("Exiting subscription polling loop")
async def _poll_receiver(self, name: str) -> None:
socket = self._receivers[name]
while not self._quitting.is_set():
msg_bytes = await socket.recv_multipart()
try:
msg = Message.from_bytes(msg_bytes)
except Exception as e:
self.logger.exception(
"Exception decoding message for socket %s: %s, %s", name, msg_bytes, e
)
continue
# purposely don't catch errors here because we want them to bubble up into the caller
# just log - otherwise they can get lost in threading tracebacks.
try:
for acb in self._callbacks[name]["asyncio"]:
await acb(msg)
for cb in self._callbacks[name]["sync"]:
self.loop.run_in_executor(None, cb, msg)
except Exception as e:
self.logger.exception("EXCEPTION IN EVENTLOOP: %s", e)
raise e
# force a context switch to avoid us getting flooded and only processing messages
await asyncio.sleep(0)
self.logger.debug("Exiting polling loop")