tractor/tractor/runtime/_supervise.py

# tractor: structured concurrent "actors".
# Copyright 2018-eternity Tyler Goodlet.

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

"""
``trio`` inspired apis and helpers

"""
from contextlib import asynccontextmanager as acm
from functools import partial
from typing import (
    TYPE_CHECKING,
)
import typing
import warnings

import trio


from ..devx import (
    debug,
    pformat as _pformat,
)
from ..discovery._addr import (
    UnwrappedAddress,
    mk_uuid,
)
from ._state import (
    current_actor,
    is_main_process,
)
from ..log import (
    get_logger,
    get_loglevel,
)
from ..msg import Aid
from ._runtime import Actor
from ._portal import Portal
from ..trionics import (
    is_multi_cancelled,
    collapse_eg,
)
from .._exceptions import (
    ActorTooSlowError,
    ContextCancelled,
)
from .._root import (
    open_root_actor,
)
from . import _state
from ..spawn import _spawn


if TYPE_CHECKING:
    import multiprocessing as mp
    # from ..ipc._server import IPCServer
    from ..ipc import IPCServer
    from ..spawn._spawn import ProcessType


log = get_logger()


async def _try_cancel_then_kill(
    portal: Portal,
    # `ProcessType` is `TYPE_CHECKING`-only (defined under that
    # guard in `..spawn._spawn`) so we stringify here to avoid
    # eager runtime eval of the annotation at function-def time
    # (this module has no `from __future__ import annotations`).
    proc: 'ProcessType',
    subactor: Actor,
    debug_mode_active: bool = False,
) -> None:
    '''
    Per-child cancel-then-escalate helper used by
    `ActorNursery.cancel()`.

    Sends a graceful actor-runtime cancel-RPC via
    `Portal.cancel_actor(raise_on_timeout=True)`. If the bounded-wait
    expires before the peer ack's, `ActorTooSlowError` is raised and
    we escalate via `proc.kill()` per SC-discipline:

      graceful cancel-req -> bounded wait -> hard-kill

    Without this escalation, a same-name sibling subactor whose
    cancel-RPC failed to ack within `Portal.cancel_timeout` (e.g.
    under TCP+forkserver register-RPC contention) would park the
    parent's `soft_kill()` watcher forever waiting on `proc.poll()`,
    deadlocking nursery `__aexit__`. See `ActorTooSlowError` for
    the wider write-up.

    '''
    # XXX, delay hard-kill escalation while any debugger guard
    # below is active. Killing the sub immediately would tear down
    # its tree and clobber an actor proxying a REPL session:
    #
    # - `Lock.ctx_in_debug is not None`: most precise — some
    #   actor in the tree is currently REPL-locked. Set in the
    #   root actor for the lifetime of the lock. Raceable
    #   (false negative if SIGINT arrives before lock-acquire
    #   RPC completes).
    #
    # - `_runtime_vars['_debug_mode']`: root-actor was opened
    #   with `debug_mode=True` (via `open_root_actor` /
    #   `open_nursery`). Set once at root boot, never cleared.
    #   Catches deep-descendant REPL sessions even when the
    #   intermediate nurseries didn't pass `debug_mode=` per-
    #   child.
    #
    # - `debug_mode_active`: this nursery has at least one
    #   child started with an explicit `debug_mode=True` arg
    #   (`ActorNursery._at_least_one_child_in_debug`). Catches
    #   the case where root is NOT in debug-mode but a
    #   nursery-direct child opted in.
    #
    # Independent because root may NOT be in debug-mode even
    # when a child is (only the child's `_runtime_vars` is
    # mutated by per-child `debug_mode=True`). ORing covers
    # every flavor without false-positively skipping
    # legitimate hard-kill paths in non-debug trees.
    def child_in_debug() -> bool:
        '''
        Sample child/tree debugger protection state.

        '''
        return (
            debug_mode_active
            or
            debug.Lock.ctx_in_debug is not None
        )

    debug_protected: bool = (
        child_in_debug()
        or
        _state._runtime_vars.get('_debug_mode', False)
    )

    try:
        cancelled: bool = await portal.cancel_actor(
            raise_on_timeout=not debug_protected,
        )
        if not cancelled:
            if debug_protected:
                await debug.maybe_wait_for_debugger(
                    # Re-sample after the cancel-RPC checkpoint.
                    child_in_debug=child_in_debug(),
                    header_msg=(
                        'Delaying subproc hard-reap while '
                        'debugger locked..\n'
                    ),
                )

            peer_id: str = portal.channel.aid.reprol()
            raise ActorTooSlowError(
                f'Peer {peer_id} disconnected before '
                f'acknowledging its `Actor.cancel()` RPC'
            )

    except ActorTooSlowError as too_slow:
        log.error(
            f'Cancel-ack TIMED OUT for sub-actor\n'
            f'  uid: {subactor.aid.reprol()!r}\n'
            f'  reason: {too_slow}\n'
            f'-> escalating to `proc.kill()` (hard-reap)\n'
        )
        # XXX, the `subint` backend stores an `int` interp-id in the
        # `proc` slot (not a `Process`), so it has no `.kill()`.
        # Guard here so a cancel-ack timeout doesn't `AttributeError`
        # once that backend lands; its hard-kill path is a TODO.
        if hasattr(proc, 'kill'):
            if proc.poll() is None:
                proc.kill()
        else:
            log.error(
                f'Cannot hard-kill sub-actor — backend proc-handle '
                f'{proc!r} ({type(proc).__name__!r}) has no '
                f'`.kill()`!\n'
                f'  uid: {subactor.aid.reprol()!r}\n'
                f'TODO: per-backend cancel-escalation.\n'
            )


class ActorNursery:
    '''
    The fundamental actor supervision construct: spawn and manage
    explicit lifetime and capability restricted, bootstrapped,
    ``trio.run()`` scheduled sub-processes.

    Though the concept of a "process nursery" is different in complexity
    and slightly different in semantics then a tradtional single
    threaded task nursery, much of the interface is the same. New
    processes each require a top level "parent" or "root" task which is
    itself no different then any task started by a tradtional
    ``trio.Nursery``. The main difference is that each "actor" (a
    process + ``trio.run()``) contains a full, paralell executing
    ``trio``-task-tree. The following super powers ensue:

    - starting tasks in a child actor are completely independent of
      tasks started in the current process. They execute in *parallel*
      relative to tasks in the current process and are scheduled by their
      own actor's ``trio`` run loop.
    - tasks scheduled in a remote process still maintain an SC protocol
      across memory boundaries using a so called "structured concurrency
      dialogue protocol" which ensures task-hierarchy-lifetimes are linked.
    - remote tasks (in another actor) can fail and relay failure back to
      the caller task (in some other actor) via a seralized
      ``RemoteActorError`` which means no zombie process or RPC
      initiated task can ever go off on its own.

    '''
    def __init__(
        self,
        # TODO: maybe def these as fields of a struct looking type?
        actor: Actor,
        da_nursery: trio.Nursery,
        errors: dict[tuple[str, str], BaseException],

    ) -> None:
        # self.supervisor = supervisor  # TODO
        self._actor: Actor = actor

        # TODO: rename to `._tn` for our conventional "task-nursery"
        self._da_nursery = da_nursery

        self._children: dict[
            tuple[str, str],
            tuple[
                Actor,
                trio.Process | mp.Process,
                Portal | None,
            ]
        ] = {}

        self._join_procs = trio.Event()
        self._child_reap_requests: dict[
            Aid,
            trio.Event,
        ] = {}
        self._child_reaped: dict[
            Aid,
            trio.Event,
        ] = {}
        self._at_least_one_child_in_debug: bool = False
        self.errors = errors
        self._scope_error: BaseException|None = None
        self.exited = trio.Event()

        # NOTE: when no explicit call is made to
        # `.open_root_actor()` by application code,
        # `.open_nursery()` will implicitly call it to start the
        # actor-tree runtime. In this case we mark ourselves as
        # such so that runtime components can be aware for logging
        # and syncing purposes to any actor opened nurseries.
        self._implicit_runtime_started: bool = False

        # trio.Nursery-like cancel (request) statuses
        self._cancelled_caught: bool = False
        self._cancel_called: bool = False

    @property
    def cancel_called(self) -> bool:
        '''
        Records whether cancellation has been requested for this
        actor-nursery by a call to  `.cancel()` either due to,

        - an explicit call by some actor-local-task,
        - an implicit call due to an error/cancel emited inside
          the `tractor.open_nursery()` block.

        '''
        return self._cancel_called

    @property
    def cancelled_caught(self) -> bool:
        '''
        Set when this nursery was able to cance all spawned subactors
        gracefully via an (implicit) call to `.cancel()`.

        '''
        return self._cancelled_caught

    # TODO! remove internal/test-suite usage!
    @property
    def cancelled(self) -> bool:
        warnings.warn(
            "`ActorNursery.cancelled` is now deprecated, use "
            " `.cancel_called` instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return (
            self._cancel_called
            # and
            # self._cancelled_caught
        )

    def _register_child_reap(
        self,
        aid: Aid,
    ) -> tuple[trio.Event, trio.Event]:
        '''
        Register a child monitor's process-reap events.

        '''
        reap_request = trio.Event()
        reaped = trio.Event()
        self._child_reap_requests[aid] = reap_request
        self._child_reaped[aid] = reaped
        if self._join_procs.is_set():
            reap_request.set()
        return reap_request, reaped

    def _register_child(
        self,
        subactor: Actor,
        proc: 'ProcessType',
        portal: Portal|None,
    ) -> tuple[trio.Event, trio.Event, bool]:
        '''
        Atomically publish one child and its reap coordination.

        '''
        aid: Aid = subactor.aid
        uid: tuple[str, str] = aid.uid
        self._children[uid] = (
            subactor,
            proc,
            portal,
        )
        reap_request, reaped = self._register_child_reap(aid)
        return (
            reap_request,
            reaped,
            self._cancel_called,
        )

    def _request_reap_all(self) -> None:
        '''
        Release every child monitor into its process-join phase.

        '''
        self._join_procs.set()
        for reap_request in tuple(
            self._child_reap_requests.values()
        ):
            reap_request.set()

    def _mark_child_reaped(
        self,
        aid: Aid,
    ) -> None:
        '''
        Publish completed child-process teardown to its waiter.

        '''
        uid: tuple[str, str] = aid.uid
        self._children.pop(uid, None)
        reap_request: trio.Event|None = (
            self._child_reap_requests.pop(aid, None)
        )
        reaped: trio.Event|None = self._child_reaped.pop(
            aid,
            None,
        )
        assert (
            (reap_request is None)
            ==
            (reaped is None)
        )
        if reaped is not None:
            reaped.set()

    async def _cancel_and_reap_child(
        self,
        portal: Portal,
    ) -> None:
        '''
        Cancel, join and unregister one nursery-owned child.

        '''
        aid: Aid = portal.channel.aid
        uid: tuple[str, str] = aid.uid
        child_entry = self._children.get(uid)
        if child_entry is None:
            return

        subactor, proc, _ = child_entry
        reap_request: trio.Event = self._child_reap_requests[aid]
        reaped: trio.Event = self._child_reaped[aid]

        with trio.CancelScope(shield=True):
            try:
                await _try_cancel_then_kill(
                    portal,
                    proc,
                    subactor,
                    self._at_least_one_child_in_debug,
                )
            finally:
                reap_request.set()
                await reaped.wait()

    async def start_actor(
        self,
        name: str,

        *,

        bind_addrs: list[UnwrappedAddress]|None = None,
        rpc_module_paths: list[str]|None = None,
        enable_transports: list[str] = [_state._def_tpt_proto],
        enable_modules: list[str]|None = None,
        loglevel: str|None = None,  # set log level per subactor
        debug_mode: bool|None = None,
        infect_asyncio: bool = False,
        inherit_parent_main: bool = True,
        proc_kwargs: dict[str, typing.Any] | None = None,

    ) -> Portal:
        '''
        Start a (daemon) actor: an process that has no designated
        "main task" besides the runtime.

        Pass ``inherit_parent_main=False`` to keep this child on its
        own bootstrap module for the trio spawn backend instead of
        applying the parent ``__main__`` re-exec fixup during startup.
        This does not affect ``multiprocessing`` ``spawn`` or
        ``forkserver`` which reconstruct the parent's ``__main__`` as
        part of their normal stdlib bootstrap.

        '''
        __runtimeframe__: int = 1  # noqa
        if self._cancel_called:
            raise RuntimeError(
                'Cannot start an actor in a cancelling '
                '`ActorNursery`'
            )

        loglevel: str = (
            loglevel
            or self._actor.loglevel
            or get_loglevel()
        )

        # configure and pass runtime state
        _rtv = _state._runtime_vars.copy()
        _rtv['_is_root'] = False
        _rtv['_is_infected_aio'] = infect_asyncio

        # allow setting debug policy per actor
        if debug_mode is not None:
            _rtv['_debug_mode'] = debug_mode
            self._at_least_one_child_in_debug |= debug_mode

        enable_modules = list(enable_modules or [])
        proc_kwargs = dict(proc_kwargs or {})

        if rpc_module_paths:
            warnings.warn(
                "`rpc_module_paths` is now deprecated, use "
                " `enable_modules` instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            enable_modules.extend(rpc_module_paths)

        subactor = Actor(
            name=name,
            uuid=mk_uuid(),

            # modules allowed to invoked funcs from
            enable_modules=enable_modules,
            loglevel=loglevel,
            inherit_parent_main=inherit_parent_main,

            # verbatim relay this actor's registrar addresses
            registry_addrs=current_actor().registry_addrs,
        )
        parent_addr: UnwrappedAddress = self._actor.accept_addr
        assert parent_addr

        # start a task to spawn a process
        # blocks until process has been started and a portal setup
        # XXX: the type ignore is actually due to a `mypy` bug
        return await self._da_nursery.start(  # type: ignore
            partial(
                _spawn.new_proc,
                name,
                self,
                subactor,
                self.errors,
                bind_addrs,
                parent_addr,
                _rtv,  # run time vars
                infect_asyncio=infect_asyncio,
                proc_kwargs=proc_kwargs
            )
        )

    # @api_frame
    async def cancel(
        self,
        hard_kill: bool = False,

    ) -> None:
        '''
        Cancel this actor-nursery by instructing each subactor's
        runtime to cancel and wait for all underlying sub-processes
        to terminate.

        If `hard_kill` is set then kill the processes directly using
        the spawning-backend's API/OS-machinery without any attempt
        at (graceful) `trio`-style cancellation using our
        `Actor.cancel()`.

        '''
        __runtimeframe__: int = 1  # noqa
        self._cancel_called = True

        # TODO: impl a repr for spawn more compact
        # then `._children`..
        children: tuple = tuple(self._children.values())
        child_count: int = len(children)
        msg: str = f'Cancelling actor nursery with {child_count} children\n'

        server: IPCServer = self._actor.ipc_server

        with trio.move_on_after(3) as cs:
            async with (
                collapse_eg(),
                trio.open_nursery() as tn,
            ):

                subactor: Actor
                proc: trio.Process
                portal: Portal
                for (
                    subactor,
                    proc,
                    portal,
                ) in children:

                    # TODO: are we ever even going to use this or
                    # is the spawning backend responsible for such
                    # things? I'm thinking latter.
                    if hard_kill:
                        proc.terminate()

                    else:
                        if portal is None:  # actor hasn't fully spawned yet
                            event: trio.Event = server._peer_connected[
                                subactor.aid.uid
                            ]
                            log.warning(
                                f"{subactor.aid.uid} never 't finished spawning?"
                            )

                            await event.wait()

                            # channel/portal should now be up
                            _, _, portal = self._children[
                                subactor.aid.uid
                            ]

                            # XXX should be impossible to get here
                            # unless method was called from within
                            # shielded cancel scope.
                            if portal is None:
                                # cancelled while waiting on the event
                                # to arrive
                                chan = server._peers[subactor.aid.uid][-1]
                                if chan:
                                    portal = Portal(chan)
                                else:  # there's no other choice left
                                    proc.terminate()

                        # spawn per-child cancel tasks; the helper
                        # escalates to hard-kill on
                        # `ActorTooSlowError` rather than silently
                        # swallowing the cancel-ack timeout, EXCEPT
                        # when this nursery has any debug-eligible
                        # child (in which case we keep legacy
                        # fire-and-forget semantics to avoid
                        # clobbering an active REPL).
                        assert portal
                        if portal.channel.connected():
                            tn.start_soon(
                                _try_cancel_then_kill,
                                portal,
                                proc,
                                subactor,
                                self._at_least_one_child_in_debug,
                            )

                log.cancel(msg)
        # if we cancelled the cancel (we hung cancelling remote actors)
        # then hard kill all sub-processes
        if cs.cancelled_caught:
            log.error(
                f'Failed to cancel {self}?\n'
                'Hard killing underlying subprocess tree!\n'
            )
            subactor: Actor
            proc: trio.Process
            portal: Portal
            for (
                subactor,
                proc,
                portal,
                ) in children:
                log.warning(f"Hard killing process {proc}")
                proc.terminate()
        else:
            self._cancelled_caught

        # mark ourselves as having (tried to have) cancelled all subactors
        self._request_reap_all()


@acm
async def _open_and_supervise_one_cancels_all_nursery(
    actor: Actor,
    hide_tb: bool = True,

) -> typing.AsyncGenerator[ActorNursery, None]:

    # normally don't need to show user by default
    __tracebackhide__: bool = hide_tb

    # the collection of errors retreived from spawned sub-actors
    errors: dict[tuple[str, str], BaseException] = {}

    # The single "daemon actor" nursery into which ALL subactors
    # are spawned; one-shot (`to_actor.run()`) subactors are
    # result-waited and reaped in their caller's own task-scope
    # (see the #477 `.run_in_actor()`/`._ria_nursery` removal);
    # errors from this nursery bubble up to the caller.
    async with (
        collapse_eg(),
        trio.open_nursery() as da_nursery,
    ):
        an = ActorNursery(
            actor,
            da_nursery,
            errors
        )
        try:
            # spawning of actors happens in the caller's scope
            # after we yield upwards
            yield an

            # When we didn't error in the caller's scope,
            # signal all process-monitor-tasks to conduct
            # the "hard join phase".
            log.runtime(
                'Waiting on subactors to complete:\n'
                f'>}} {len(an._children)}\n'
            )
            an._request_reap_all()

        # Single one-cancels-all handler for the (now single)
        # daemon nursery. Pre-#477 a 2ndary `._ria_nursery`
        # required a separate *outer* handler to catch errors
        # bubbling from its task-reaping `__aexit__`; with that
        # nursery gone this lone handler covers every scope
        # error. NB: we deliberately do NOT re-raise here — the
        # `finally` below raises the collected `errors` (as a
        # single exc or `BaseExceptionGroup`), which already
        # superseded the old outer handler's `raise` anyway
        # since `errors` is populated (below) before any await.
        except BaseException as _scope_err:
            an._scope_error = _scope_err
            errors[actor.aid.uid] = _scope_err

            # If we error in the root but the debugger is
            # engaged we don't want to prematurely kill (and
            # thus clobber access to) the local tty since it
            # will make the pdb repl unusable.
            # Instead try to wait for pdb to be released before
            # tearing down.
            await debug.maybe_wait_for_debugger(
                child_in_debug=an._at_least_one_child_in_debug
            )

            # if the caller's scope errored then we activate our
            # one-cancels-all supervisor strategy (don't
            # worry more are coming).
            an._request_reap_all()

            # XXX NOTE XXX: hypothetically an error could
            # be raised and then a cancel signal shows up
            # slightly after in which case the `else:`
            # block here might not complete?  For now,
            # shield both.
            with trio.CancelScope(shield=True):
                etype: type = type(_scope_err)
                if etype in (
                    trio.Cancelled,
                    KeyboardInterrupt,
                ) or (
                    is_multi_cancelled(_scope_err)
                ):
                    log.cancel(
                        f'Actor-nursery cancelled by {etype}\n\n'

                        f'{current_actor().aid.uid}\n'
                        f' |_{an}\n\n'

                        # TODO: show tb str?
                        # f'{tb_str}'
                    )
                elif etype in {
                    ContextCancelled,
                }:
                    log.cancel(
                        'Actor-nursery caught remote cancellation\n'
                        '\n'
                        f'{_scope_err.tb_str}'
                    )
                else:
                    log.exception(
                        'Nursery errored with:\n'

                        # TODO: same thing as in
                        # `._invoke()` to compute how to
                        # place this div-line in the
                        # middle of the above msg
                        # content..
                        # -[ ] prolly helper-func it too
                        #   in our `.log` module..
                        # '------ - ------'
                    )

                # cancel all subactors
                await an.cancel()

        finally:
            # an error was stashed by the handler above (or by
            # a spawn task via the shared `errors` dict) so
            # cancel any remaining subactors, summarize and
            # re-raise.
            if errors:
                if an._children:
                    with trio.CancelScope(shield=True):
                        await an.cancel()

                # use `BaseExceptionGroup` as needed
                if len(errors) > 1:
                    raise BaseExceptionGroup(
                        'tractor.ActorNursery errored with',
                        tuple(errors.values()),
                    )
                else:
                    raise list(errors.values())[0]

            # show frame on any (likely) internal error
            if (
                not an.cancel_called
                and an._scope_error
            ):
                __tracebackhide__: bool = False

        # da_nursery scope end - nursery checkpoint
    # final exit


_shutdown_msg: str = (
    'Actor-runtime-shutdown'
)


@acm
# @api_frame
async def open_nursery(
    *,  # named params only!
    hide_tb: bool = True,
    **kwargs,
    # ^TODO, paramspec for `open_root_actor()`

) -> typing.AsyncGenerator[ActorNursery, None]:
    '''
    Create and yield a new ``ActorNursery`` to be used for spawning
    structured concurrent subactors.

    When an actor is spawned a new trio task invokes one of the
    process spawning backends to create and start a new subprocess.
    These tasks are started in the supervisor's process nursery.
    Spawning from a task is required because ``trio_run_in_process``
    creates an internal nursery which the opening task **must** close;
    this also makes each task's cancellation scope correspond to its
    spawned subactor.

    '''
    __tracebackhide__: bool = hide_tb
    implicit_runtime: bool = False
    actor: Actor = current_actor(err_on_no_runtime=False)
    an: ActorNursery|None = None
    try:
        if (
            actor is None
            and is_main_process()
        ):
            # if we are the parent process start the
            # actor runtime implicitly
            log.info("Starting actor runtime!")

            # mark us for teardown on exit
            implicit_runtime: bool = True

            async with open_root_actor(
                hide_tb=hide_tb,
                **kwargs,
            ) as actor:
                assert actor is current_actor()

                try:
                    async with _open_and_supervise_one_cancels_all_nursery(
                        actor
                    ) as an:

                        # NOTE: mark this nursery as having
                        # implicitly started the root actor so
                        # that `._runtime` machinery can avoid
                        # certain teardown synchronization
                        # blocking/waits and any associated (warn)
                        # logging when it's known that this
                        # nursery shouldn't be exited before the
                        # root actor is.
                        an._implicit_runtime_started = True
                        yield an
                finally:
                    # XXX: this event will be set after the root actor
                    # runtime is already torn down, so we want to
                    # avoid any blocking on it.
                    an.exited.set()

        else:  # sub-nursery case

            try:
                async with _open_and_supervise_one_cancels_all_nursery(
                    actor
                ) as an:
                    yield an
            finally:
                an.exited.set()

    finally:
        # show frame on any internal runtime-scope error
        if (
            an
            and
            not an.cancel_called
            and
            an._scope_error
        ):
            __tracebackhide__: bool = False


        op_nested_an_repr: str = _pformat.nest_from_op(
            input_op=')>',
            text=f'{an}',
            # nest_prefix='|_',
            nest_indent=1,  # under >
        )
        an_msg: str = (
            f'Actor-nursery exited\n'
            f'{op_nested_an_repr}\n'
        )
        # keep noise low during std operation.
        log.runtime(an_msg)

        if implicit_runtime:
            # shutdown runtime if it was started and report noisly
            # that we're did so.
            msg: str = (
                '\n'
                '\n'
                f'{_shutdown_msg} )>\n'
            )
            log.info(msg)