"""
Based on :doc:`aiotools.context`, this module provides an automated lifecycle
management for multi-process servers with explicit initialization steps and
graceful shutdown steps.
.. function:: server(func)
A decorator wrapper for :class:`AsyncServerContextManager`.
Usage example:
.. code:: python
@aiotools.server
async def myserver(loop, pidx, args):
await do_init(args)
stop_sig = yield
if stop_sig == signal.SIGINT:
await do_graceful_shutdown()
else:
await do_forced_shutdown()
aiotools.start_server(myserver, ...)
"""
import asyncio
import functools
import inspect
import logging
import multiprocessing as mp
import multiprocessing.connection as mpconn
import os
import signal
import struct
import sys
import threading
from collections.abc import (
AsyncGenerator,
Callable,
Collection,
Coroutine,
Generator,
Mapping,
Sequence,
)
from contextlib import AbstractContextManager, ContextDecorator
from contextvars import ContextVar
from types import TracebackType
from typing import Any, ParamSpec, TypeVar
from .context import AbstractAsyncContextManager
from .fork import AbstractChildProcess, MPContext, afork
__all__ = (
"main_context",
"server_context",
"start_server",
"process_index",
"AsyncServerContextManager",
"ServerMainContextManager",
"InterruptedBySignal",
)
log = logging.getLogger(__spec__.name)
process_index: ContextVar[int] = ContextVar("process_index")
TYield = TypeVar("TYield")
PArgs = ParamSpec("PArgs")
[docs]
class InterruptedBySignal(BaseException):
"""
A new :class:`BaseException` that represents interruption by an arbitrary UNIX
signal.
Since this is a :class:`BaseException` instead of :class:`Exception`, it behaves
like :class:`KeyboardInterrupt` and :class:`SystemExit` exceptions (i.e.,
bypassing except clauses catching the :class:`Exception` type only)
The first argument of this exception is the signal number received.
"""
pass
[docs]
class AsyncServerContextManager(AbstractAsyncContextManager[TYield]):
"""
A modified version of :func:`contextlib.asynccontextmanager`.
The implementation detail is mostly taken from the ``contextlib`` standard
library, with a minor change to inject ``self.yield_return`` into the wrapped
async generator.
"""
yield_return: signal.Signals = signal.SIGTERM
def __init__(
self,
func: Callable[..., AsyncGenerator[TYield, signal.Signals]],
args: Sequence[Any],
kwargs: Mapping[str, Any],
) -> None:
if not inspect.isasyncgenfunction(func):
raise RuntimeError("Context manager function must be an async-generator")
self._agen: AsyncGenerator[TYield, signal.Signals] = func(*args, **kwargs)
self.func: Callable[..., AsyncGenerator[TYield, signal.Signals]] = func
self.args: Sequence[Any] = args
self.kwargs: Mapping[str, Any] = kwargs
self.yield_return: signal.Signals = signal.SIGTERM
async def __aenter__(self) -> TYield:
try:
return await self._agen.__anext__()
except StopAsyncIteration:
raise RuntimeError("async-generator didn't yield") from None
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
/,
) -> bool | None:
if exc_type is None:
try:
# Here is the modified part.
await self._agen.asend(self.yield_return)
except StopAsyncIteration:
return None
else:
raise RuntimeError("async-generator didn't stop") from None
else:
if exc_value is None:
exc_value = exc_type()
try:
await self._agen.athrow(exc_type, exc_value, traceback)
raise RuntimeError("async-generator didn't stop after athrow()")
except StopAsyncIteration as exc_new_value:
return exc_new_value is not exc_value
except RuntimeError as exc_new_value:
if exc_new_value is exc_value:
return False
if isinstance(exc_value, (StopIteration, StopAsyncIteration)):
if exc_new_value.__cause__ is exc_value:
return False
raise
except (BaseException, asyncio.CancelledError) as exc:
if exc is not exc_value:
raise
return None
[docs]
class ServerMainContextManager(AbstractContextManager[TYield], ContextDecorator):
"""
A modified version of :func:`contextlib.contextmanager`.
The implementation detail is mostly taken from the ``contextlib`` standard
library, with a minor change to inject ``self.yield_return`` into the wrapped
generator.
"""
def __init__(
self,
func: Callable[..., Generator[TYield, signal.Signals, None]],
args: Sequence[Any],
kwargs: Mapping[str, Any],
) -> None:
self.gen: Generator[TYield, signal.Signals, None] = func(*args, **kwargs)
self.func: Callable[..., Generator[TYield, signal.Signals, None]] = func
self.args: Sequence[Any] = args
self.kwargs: Mapping[str, Any] = kwargs
self.yield_return: signal.Signals = signal.SIGTERM
def __enter__(self) -> TYield:
del self.args, self.kwargs, self.func
try:
return next(self.gen)
except StopIteration:
raise RuntimeError("generator didn't yield") from None
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
/,
) -> bool | None:
if exc_type is None:
try:
self.gen.send(self.yield_return)
except StopIteration:
return False
else:
raise RuntimeError("generator didn't stop")
else:
if exc_value is None:
exc_value = exc_type()
try:
self.gen.throw(exc_type, exc_value, traceback)
except StopIteration as exc:
return exc is not exc_value
except RuntimeError as exc:
if exc is exc_value:
return False
if exc_type is StopIteration and exc.__cause__ is exc_value:
return False
raise
except: # noqa
if sys.exc_info()[1] is exc_value:
return False
raise
raise RuntimeError("generator didn't stop after throw()")
return None
# This is a dirty hack to implement "module callable".
# NOTE: only works in Python 3.5 or higher.
[docs]
def server_context(
func: Callable[
[asyncio.AbstractEventLoop, int, Sequence[Any]],
AsyncGenerator[None, signal.Signals],
],
) -> Callable[..., AsyncServerContextManager[None]]:
@functools.wraps(func)
def helper(*args: Any, **kwargs: Any) -> AsyncServerContextManager[None]:
return AsyncServerContextManager(func, args, kwargs)
return helper
class _ServerModule(sys.modules[__name__].__class__): # type: ignore
def __call__(
self,
func: Callable[
[asyncio.AbstractEventLoop, int, Sequence[Any]],
AsyncGenerator[None, signal.Signals],
],
) -> Callable[..., AsyncServerContextManager[None]]:
return server_context(func)
sys.modules[__name__].__class__ = _ServerModule
[docs]
def main_context(
func: Callable[[], Generator[TYield, signal.Signals, None]],
) -> Callable[[], ServerMainContextManager[TYield]]:
"""
A decorator wrapper for :class:`ServerMainContextManager`
Usage example:
.. code:: python
@aiotools.main_context
def mymain() -> Generator[TServerArgs, signal.Signals]:
server_args = do_init()
stop_sig = yield server_args
if stop_sig == signal.SIGINT:
do_graceful_shutdown()
else:
do_forced_shutdown()
aiotools.start_server(..., main_ctxmgr=mymain, ...)
"""
@functools.wraps(func)
def helper(*args: Any, **kwargs: Any) -> ServerMainContextManager[TYield]:
return ServerMainContextManager(func, args, kwargs)
return helper
def _get_default_stop_signal(
stop_signals: Collection[signal.Signals],
) -> signal.Signals:
if signal.SIGTERM in stop_signals: # prefer SIGTERM if included
return signal.SIGTERM
return sorted(stop_signals)[0]
def _worker_main(
worker_actxmgr: Callable[
[asyncio.AbstractEventLoop, int, Sequence[Any]],
AsyncServerContextManager[Any],
],
stop_signals: Collection[signal.Signals],
intr_write_pipe: mpconn.Connection,
proc_idx: int,
run_to_completion: bool,
args: Sequence[Any],
*,
prestart_hook: Callable[[int], None] | None = None,
runner: Callable[[Coroutine[Any, Any, None]], None] = asyncio.run,
) -> int:
process_index.set(proc_idx)
if prestart_hook:
prestart_hook(proc_idx)
async def _wrapped_worker() -> None:
loop = asyncio.get_running_loop()
def handle_stop_signal(signum: signal.Signals) -> None:
if interrupted.is_set():
pass
else:
interrupted.set()
ctx.yield_return = signum
forever_future.cancel()
for signum in stop_signals:
loop.add_signal_handler(
signum,
functools.partial(handle_stop_signal, signum),
)
interrupted = asyncio.Event()
ctx = worker_actxmgr(loop, proc_idx, args)
forever_future: asyncio.Future[None] = loop.create_future()
# Allow the worker to be interrupted during initialization
# (in case of initialization failures in other workers)
signal.pthread_sigmask(signal.SIG_UNBLOCK, stop_signals)
err_ctx = "enter"
try:
async with ctx:
err_ctx = "body"
try:
if not run_to_completion:
await forever_future
else:
ctx.yield_return = _get_default_stop_signal(stop_signals)
except asyncio.CancelledError:
pass
finally:
err_ctx = "exit"
except Exception:
if err_ctx != "body":
err_ctx_str = "initialization" if err_ctx == "enter" else "shutdown"
log.exception(
f"Worker {proc_idx}: Error during context manager {err_ctx_str}"
)
try:
intr_write_pipe.send_bytes(struct.pack("i", proc_idx))
except BrokenPipeError:
pass
raise
finally:
intr_write_pipe.close()
runner(_wrapped_worker())
return 0
def _extra_main(
extra_func: Callable[[threading.Event | None, int, Sequence[Any]], None],
stop_signals: Collection[signal.Signals],
proc_idx: int,
args: Sequence[Any],
*,
prestart_hook: Callable[[int], None] | None = None,
) -> int:
process_index.set(proc_idx)
if prestart_hook:
prestart_hook(proc_idx)
interrupted = threading.Event()
# Since signals only work for the main thread in Python,
# extra processes in use_threading=True mode should check
# the intr_event by themselves (probably in their loops).
def raise_stop(signum: int, frame: Any) -> None:
if interrupted.is_set():
pass
else:
interrupted.set()
if signum == signal.SIGINT:
raise KeyboardInterrupt
elif signum == signal.SIGTERM:
raise SystemExit
else:
raise InterruptedBySignal(signum)
# restore signal handler.
for signum in stop_signals:
signal.signal(signum, raise_stop)
signal.pthread_sigmask(signal.SIG_UNBLOCK, stop_signals)
intr_event: threading.Event | None = None
try:
if not interrupted.is_set():
extra_func(intr_event, proc_idx, args)
except (SystemExit, KeyboardInterrupt, InterruptedBySignal):
log.warning(f"extra_proc[{proc_idx}] did not handle stop signals.")
finally:
# same as in _worker_main()
signal.pthread_sigmask(signal.SIG_BLOCK, stop_signals)
return 0
[docs]
def start_server(
worker_actxmgr: Callable[
[asyncio.AbstractEventLoop, int, Sequence[Any]], AsyncServerContextManager[Any]
],
main_ctxmgr: Callable[[], ServerMainContextManager[Any]] | None = None,
extra_procs: Collection[
Callable[[threading.Event | None, int, Sequence[Any]], None]
] = tuple(),
stop_signals: Collection[signal.Signals] = (signal.SIGINT, signal.SIGTERM),
num_workers: int = 1,
args: Sequence[Any] = tuple(),
*,
wait_timeout: float | None = None,
mp_context: MPContext | None = None,
prestart_hook: Callable[[int], None] | None = None,
ignore_child_interrupts: bool = False,
run_to_completion: bool = False,
runner: Callable[[Coroutine[Any, Any, None]], None] = asyncio.run,
) -> None:
"""
Starts a multi-process server where each process has their own individual
asyncio event loop. Their lifecycles are automantically managed -- if the
main program receives one of the signals specified in ``stop_signals`` it
will initiate the shutdown routines on each worker that stops the event
loop gracefully.
Args:
worker_actxmgr:
An asynchronous context manager that dicates the
initialization and shutdown steps of each worker.
It should accept the following three arguments:
* **loop**: the asyncio event loop created and set
by aiotools
* **pidx**: the 0-based index of the worker
(use this for per-worker logging)
* **args**: a concatenated tuple of values yielded by
**main_ctxmgr** and the user-defined arguments in
**args**.
aiotools automatically installs an interruption handler
that calls ``loop.stop()`` to the given event loop,
regardless of using either threading or
multiprocessing.
main_ctxmgr:
An optional context manager that performs global
initialization and shutdown steps of the whole program.
It may yield one or more values to be passed to worker
processes along with **args** passed to this function.
There is no arguments passed to those functions since
you can directly access ``sys.argv`` to parse command
line arguments and/or read user configurations.
extra_procs:
An iterable of functions that consist of extra processes
whose lifecycles are synchronized with other workers.
They should set up their own signal handlers.
It should accept the following three arguments:
* **intr_event**: Always ``None``, kept for legacy
* **pidx**: same to **worker_actxmgr** argument
* **args**: same to **worker_actxmgr** argument
stop_signals:
A list of UNIX signals that the main program to
recognize as termination signals.
num_workers: The number of children workers.
args:
The user-defined arguments passed to workers and extra
processes. If **main_ctxmgr** yields one or more values,
they are *prepended* to this user arguments when passed to
workers and extra processes.
wait_timeout:
The timeout in seconds before forcibly killing all
remaining child processes after sending initial stop
signals.
mp_context:
The multiprocessing context to use for creating child
processes. If not specified, the default context is
used.
prestart_hook:
A function to be called once before creating the
event loop in the children. The function should
accept an int argument representing the process index.
ignore_child_interrupts:
By default, any unhandled exceptions in the
child functions are translated as active
SIGINT to the main and all other worker processes,
meaning a full shutdown.
This flag makes the main process to ignore them,
which is useful to gather all worker's results
even when some of them raises unhandled exceptions.
run_to_completion:
If True, the main/worker processes will NOT wait forever
until interrupted but *immediately exit* when the main
functions complete. This flag implies **ignore_child_interrupts**.
runner:
A function to run the root coroutine, which defaults to
``asyncio.run``. You may set it to ``uvloop.run`` or other
runner functions.
Returns:
None
.. versionchanged:: 0.3.2
The name of argument **num_proc** is changed to **num_workers**.
Even if **num_workers** is 1, a child is created instead of
doing everything at the main thread.
.. versionadded:: 0.3.2
The argument ``extra_procs`` and ``main_ctxmgr``.
.. versionadded:: 0.4.0
Now supports use of threading instead of multiprocessing via
**use_threading** option.
.. versionchanged:: 0.8.0
Now **worker_actxmgr** must be an instance of
:class:`AsyncServerContextManager` or async generators decorated by
``@aiotools.server``.
Now **main_ctxmgr** must be an instance of :class:`ServerMainContextManager`
or plain generators decorated by ``@aiotools.main``.
The usage is same to asynchronous context managers, but optionally you can
distinguish the received stop signal by retrieving the return value of the
``yield`` statement.
In **extra_procs** in non-threaded mode, stop signals are converted into
either one of :class:`KeyboardInterrupt`, :class:`SystemExit`, or
:class:`InterruptedBySignal` exception.
.. versionadded:: 0.8.4
**start_method** argument can be set to change the subprocess spawning
implementation.
.. deprecated:: 1.2.0
The **start_method** and **use_threading** arguments, in favor of our new
:func:`afork()` function which provides better synchronization and pid-fd
support.
.. versionchanged:: 1.2.0
The **extra_procs** will be always separate processes since **use_threading**
is deprecated and thus **intr_event** arguments are now always ``None``.
.. versionadded:: 1.5.5
The **wait_timeout** argument.
.. versionadded:: 1.9.0
The **mp_context**, **prestart_hook**, **ignore_child_interrupts**,
and **run_to_completion** arguments.
.. versionchanged:: 2.1.0
In Python 3.14 or higher, the "fork" mode support is REMOVED AND DISCOURAGED
as it causes silent hanging when combined with asyncio event loops.
Using custom stop signals with **extra_proc** is now STRONGLY DISCOURAGED
as it causes multiprocessing's resource tracker killed by them and
there is no way to control this behavior from our side.
.. versionadded:: 2.2.3
The **runner** argument to replace event loop implementations
(e.g., ``uvloop.run()``) in response to pending deprecation of
event loop policies in Python 3.16.
"""
@main_context
def noop_main_ctxmgr() -> Generator[None, signal.Signals, None]:
yield
assert stop_signals
try:
asyncio.get_running_loop()
except RuntimeError:
# there should be none.
pass
else:
raise RuntimeError(
"aiotools.start_server() cannot be called inside a running event loop."
)
if sys.version_info >= (3, 14):
if mp_context is not None and mp_context.get_start_method() == "fork":
raise NotImplementedError(
"In Python 3.14 or higher, the 'fork' start-method is no longer supported."
)
if main_ctxmgr is None:
main_ctxmgr = noop_main_ctxmgr
children: list[AbstractChildProcess] = []
sigblock_mask: frozenset[signal.Signals] = frozenset(stop_signals)
main_ctx: ServerMainContextManager[Any] = main_ctxmgr()
# temporarily block signals and register signal handlers to main_loop
signal.pthread_sigmask(signal.SIG_BLOCK, sigblock_mask)
async def _parent_main() -> None:
main_loop = asyncio.get_running_loop()
main_future: asyncio.Future[None] = main_loop.create_future()
# Build a main-to-worker interrupt channel using signals
def handle_stop_signal(signum: signal.Signals) -> None:
main_ctx.yield_return = signum
for child in children:
child.send_signal(signum)
# instead of main_loop.stop(), we use a future here
# NOT to interrupt the cleanup routines in
# an arbitrary timing upon async child failures.
main_future.cancel()
for signum in stop_signals:
main_loop.add_signal_handler(
signum,
functools.partial(handle_stop_signal, signum),
)
# Build a reliable worker-to-main interrupt channel using a pipe.
# This channel is used when the worker main functions raise an unhandled exception,
# so that the main program can be interrupted immediately.
def handle_child_interrupt(read_pipe: mpconn.Connection) -> None:
try:
child_idx: int = struct.unpack("i", read_pipe.recv_bytes(4))[0]
except EOFError:
# read_pipe is already closed
return
if not ignore_child_interrupts and not run_to_completion:
# self-interrupt to initiate the main-to-worker interrupts
log.debug(f"Child {child_idx} has interrupted the main program.")
signal.pthread_sigmask(signal.SIG_UNBLOCK, {signal.SIGINT})
os.kill(0, signal.SIGINT)
read_pipe, write_pipe = mp.Pipe()
main_loop.add_reader(read_pipe.fileno(), handle_child_interrupt, read_pipe)
# start
try:
with main_ctx as main_args:
# retrieve args generated by the user-defined main
main_args_tuple: tuple[Any, ...]
if main_args is None:
main_args_tuple = tuple()
elif not isinstance(main_args, tuple):
main_args_tuple = (main_args,)
else:
main_args_tuple = main_args
# spawn managed async workers
for i in range(num_workers):
try:
p = await afork(
functools.partial(
_worker_main,
worker_actxmgr,
stop_signals,
write_pipe,
i,
run_to_completion,
(*main_args_tuple, *args),
prestart_hook=prestart_hook,
runner=runner,
),
mp_context=mp_context,
)
except RuntimeError as e:
if "loop stopped" in e.args[0]:
log.warning(
"skipping spawning of child[worker]:%d due to "
"async failure(s) of other child process",
i,
)
continue
raise
children.append(p)
# spawn extra workers
for i, f in enumerate(extra_procs):
try:
p = await afork(
functools.partial(
_extra_main,
f,
stop_signals,
num_workers + i,
(*main_args_tuple, *args),
prestart_hook=prestart_hook,
),
mp_context=mp_context,
)
except RuntimeError as e:
if "loop stopped" in e.args[0]:
log.warning(
"skipping spawning of child[extra]:%d due to "
"async failure(s) of other child process",
num_workers + i,
)
continue
raise
children.append(p)
write_pipe.close()
# unblock the stop signals for user/external interrupts.
signal.pthread_sigmask(signal.SIG_UNBLOCK, sigblock_mask)
# run!
try:
if not run_to_completion:
await main_future
else:
main_ctx.yield_return = _get_default_stop_signal(stop_signals)
except asyncio.CancelledError:
pass
finally:
# If interrupted or complete, wait for workers to finish.
try:
worker_results: list[
int | BaseException
] = await asyncio.wait_for(
asyncio.gather(
*[child.wait() for child in children],
return_exceptions=True,
),
wait_timeout,
)
for child, result in zip(children, worker_results):
if isinstance(result, Exception):
log.error(
"Waiting for a child process [%d] has failed by an error.",
child.pid,
exc_info=result,
)
except asyncio.TimeoutError:
log.warning(
"Timeout during waiting for child processes; killing all",
)
for child in children:
child.send_signal(signal.SIGKILL)
finally:
main_loop.remove_reader(read_pipe.fileno())
read_pipe.close()
runner(_parent_main())