import asyncio
from inspect import iscoroutinefunction
from multiprocessing import Queue
from multiprocessing.connection import Connection
from os import set_blocking
from typing import Any, Callable, Optional
from edri.abstract.manager.manager_base import ManagerBase
from edri.abstract.manager.worker import Worker
from edri.dataclass.event import Event
from edri.events.edri.manager import WorkerQuit
__all__ = ["AsyncManagerBase"]
[docs]
class AsyncManagerBase(ManagerBase):
"""
A :class:`~edri.abstract.ManagerBase` whose event loop is asyncio-driven
instead of a blocking ``multiprocessing.connection.wait()`` loop.
Everything else — the metaclass's ``@request``/``@response`` registration,
``_collect_vitals``/``health_check``, ``start_worker``/``WorkerProcess``,
``run()``'s startup sequence — is inherited from :class:`ManagerBase`
unchanged; this class only overrides how incoming events are waited for
and dispatched.
A ``@request``/``@response`` handler may now be ``async def``. Each async
handler is dispatched as its own ``asyncio.Task``: the loop does not wait
for it to finish before reading the next pipe, so a handler that needs to
do I/O can simply ``await`` it instead of calling ``start_worker`` with a
``WorkerThread`` just to avoid blocking the loop. Plain ``def`` handlers
keep working exactly as they do on ``ManagerBase`` today — called inline,
synchronously, in the order their pipe becomes readable — so existing
synchronous handlers need no code change.
``WorkerProcess`` (CPU-bound work needing real parallelism) is untouched:
asyncio does not grant extra cores, so CPU-bound work still belongs in a
``WorkerProcess``, offloaded via ``start_worker`` exactly as before.
Async handlers reuse ``ManagerBase``'s cache, bookkeeping, response
preparation, and router-output stages. This keeps their behavior aligned
with synchronous handlers, including calls to ``prepare_outgoing_event``.
Pipe reads deliberately remain synchronous once the event loop reports a
descriptor readable. That mirrors ``ManagerBase.wait() -> Connection.recv()``
and, critically, preserves ``multiprocessing.Connection``'s blocking framed
protocol. ``uvloop``/libuv changes fds passed to ``add_reader()`` to
``O_NONBLOCK``; this class restores blocking mode immediately after every
reader registration.
"""
# ``router_queue`` is annotated Any: the multiprocessing.Queue factory is
# not subscriptable at runtime, and PEP 649 lazy annotation evaluation
# (Python 3.14) evaluates __init__ signature annotations in
# inspect.signature(); see manager_base.py for details.
def __init__(self, router_queue: Optional[Any] = None, *args: Any, **kwargs: Any) -> None:
super().__init__(router_queue, *args, **kwargs)
self._async_tasks: set[asyncio.Task[None]] = set()
self._registered_readers: set[Connection] = set()
self._async_resolvers: frozenset[Callable[..., Any]] = frozenset()
self._worker_wakeup: asyncio.Event | None = None
def _prepare_resolvers(self) -> None:
super()._prepare_resolvers()
# _requests/_responses are fixed after this point (built once here,
# never mutated at runtime), so classify each resolver's sync/async-ness
# exactly once instead of calling iscoroutinefunction() on every event.
self._async_resolvers = frozenset(
resolver
for resolver in (*self._requests.values(), *self._responses.values())
if iscoroutinefunction(resolver)
)
[docs]
def resolve(self, event: Event) -> None:
resolver = self.get_event_resolver(event)
if resolver in self._async_resolvers:
task = asyncio.get_running_loop().create_task(
self._run_async_handler(resolver, event)
)
self._async_tasks.add(task)
task.add_done_callback(self._async_tasks.discard)
# "Resolved" is stamped inside _run_async_handler once the task
# actually finishes, not here — resolve() only schedules it.
return
super().resolve(event)
event._timing.stamp(self.__class__.__name__, "Resolved")
async def _run_async_handler(self, resolver: Callable[..., Any], event: Event) -> None:
"""Run one async resolver through the same stages as ``ManagerBase``.
``ManagerBase.resolve_router_event()`` catches failures from the whole
synchronous resolve pipeline, not only from the user resolver itself.
Async dispatch returns before that pipeline runs, so this task must own
the equivalent exception boundary around cache lookup, the handler,
bookkeeping, response preparation, and router output.
"""
try:
had_response = event.has_response()
cache = self.get_resolver_cache_state(event, resolver)
if self.resolve_not_modified(event, cache):
return
await resolver(event)
self.record_resolved_event(event)
if not self.should_send_resolved_event(event, had_response):
return
self.prepare_resolved_event(event, cache)
self.send_to_router(event)
except Exception as error:
self.handle_resolver_exception(event, error)
finally:
# This also covers the cache/304 short-circuit. The previous
# implementation returned before entering its try/finally and left
# those events without a terminal timing stamp.
event._timing.stamp(self.__class__.__name__, "Resolved")
def _register_pipe(self, pipe: Connection) -> None:
"""Register one pipe reader while preserving Connection semantics.
``uvloop.add_reader()`` is backed by libuv polling, which changes the
watched descriptor to non-blocking mode. ``Connection.recv()`` cannot
safely recover from EAGAIN after partially consuming its length-prefixed
frame, and sends on the same full-duplex descriptor have the same issue.
Restore blocking mode synchronously before control returns to the event
loop, so no other asyncio task can observe the temporary state.
"""
if pipe in self._registered_readers:
return
loop = asyncio.get_running_loop()
try:
fd = pipe.fileno()
loop.add_reader(fd, self._handle_ready_pipe, pipe)
try:
set_blocking(fd, True)
except Exception:
loop.remove_reader(fd)
raise
except OSError as error:
# Registration failure means this pipe is already unusable. Reuse
# ManagerBase's owner-aware cleanup policy rather than leaving a dead
# descriptor to be retried once per reconciliation cycle forever.
self.handle_pipe_receive_error(pipe, error)
return
self._registered_readers.add(pipe)
def _handle_ready_pipe(self, pipe: Connection) -> None:
"""Receive and dispatch one frame from an fd reported readable.
The callback intentionally performs ``recv()`` inline. Once blocking mode
has been restored this has the same semantics as ManagerBase's
``wait() -> recv()`` path, avoids corruptible non-blocking framing, and
avoids spawning executor reads whose lifetime could outlive resolver
shutdown. Async *handlers* still run independently as asyncio tasks.
"""
try:
event: Event = pipe.recv()
except (EOFError, OSError) as error:
# Remove the event-loop watcher while fileno() is still usable;
# ManagerBase's error policy may close the pipe afterwards.
self._unregister_pipe(pipe)
self.handle_pipe_receive_error(pipe, error)
return
if self.is_router_pipe(pipe):
self.logger.debug("Received event from router: %s", event)
event._timing.stamp(self.__class__.__name__, "Received")
try:
self.resolve(event)
except Exception as error:
self.handle_resolver_exception(event, error)
return
self.logger.debug("Received event: %s", event)
# ManagerBase.resolve_non_router_event() removes a worker when it sees
# WorkerQuit, and _remove_worker() closes that worker's pipe. Remove the
# asyncio reader first so a closed/reused fd can never retain a stale
# callback in the loop.
if isinstance(event, WorkerQuit):
try:
self._find_worker(pipe)
except ValueError:
pass
else:
self._unregister_pipe(pipe)
self.resolve_non_router_event(event, pipe)
self.workers._process_timeouts()
def _unregister_pipe(self, pipe: Connection) -> None:
"""Remove a registered pipe reader, safely and idempotently."""
if pipe not in self._registered_readers:
return
self._registered_readers.discard(pipe)
try:
fd = pipe.fileno()
except OSError:
return
try:
asyncio.get_running_loop().remove_reader(fd)
except OSError:
pass
def _worker_runtime_finishing(self, worker_id: str, worker: Worker) -> None:
if worker.pipe is not None:
self._unregister_pipe(worker.pipe)
def _worker_runtime_started(self, worker_id: str, worker: Worker) -> None:
"""Expose a new runtime immediately and interrupt deadline sleep."""
if not self._resolver_running:
return
if worker.pipe is not None:
self._register_pipe(worker.pipe)
if self._worker_wakeup is not None:
self._worker_wakeup.set()
def _worker_deadline_changed(self) -> None:
if self._worker_wakeup is not None:
self._worker_wakeup.set()
async def _wait_for_reconciliation(self, delay: float) -> None:
"""Sleep until reconciliation is due or a runtime changes the wait set."""
wakeup = self._worker_wakeup
if wakeup is None:
await asyncio.sleep(delay)
return
# No await occurs between clear() and task creation by the caller, so a
# runtime-start hook cannot be lost in this transition.
wakeup.clear()
sleep_task = asyncio.create_task(asyncio.sleep(delay))
wake_task = asyncio.create_task(wakeup.wait())
_, pending = await asyncio.wait(
(sleep_task, wake_task), return_when=asyncio.FIRST_COMPLETED
)
for task in pending:
task.cancel()
if pending:
await asyncio.gather(*pending, return_exceptions=True)
async def _async_run_resolver(self) -> None:
self._resolver_running = True
self._worker_wakeup = asyncio.Event()
try:
while self._resolver_running:
self.workers._process_timeouts()
current = self.get_pipes()
for pipe in current - self._registered_readers:
self._register_pipe(pipe)
if not self._resolver_running:
break
if not self._resolver_running:
break
# Pipes closed/removed elsewhere are reconciled here. WorkerQuit
# and read failures are normally unregistered proactively above.
for pipe in self._registered_readers - current:
self._unregister_pipe(pipe)
# Reconciliation cadence only — actual events are dispatched
# immediately by add_reader callbacks. Mirror ManagerBase's two
# configurable delays instead of hard-coding one second: an empty
# manager uses resolver_idle_delay, otherwise the regular wait
# timeout bounds how quickly newly-created pipes are discovered.
maximum = self.resolver_wait_timeout if current else self.resolver_idle_delay
delay = self.workers._next_timeout(maximum) or 0.0
await self._wait_for_reconciliation(delay)
finally:
self._resolver_running = False
self._worker_wakeup = None
for pipe in list(self._registered_readers):
self._unregister_pipe(pipe)
# Preserve the existing graceful semantics: already-started async
# resolvers are allowed to finish before the manager process exits.
if self._async_tasks:
await asyncio.gather(*self._async_tasks, return_exceptions=True)
[docs]
def run_resolver(self) -> None:
try:
asyncio.run(self._async_run_resolver())
except KeyboardInterrupt:
pass
def _start_and_resolve(self) -> None:
# Start-up runs inside the event loop, so events queued while
# subscribing can schedule async handlers and _prepare()/after_start()
# create loop-bound resources on the loop the handlers use.
try:
asyncio.run(self._async_start_and_resolve())
except KeyboardInterrupt:
pass
async def _async_start_and_resolve(self) -> None:
self._start_up()
await self._async_run_resolver()