tractor/tractor/_portal.py

"""
Portal api
"""
import importlib
import inspect
import typing
from typing import Tuple, Any, Dict, Optional, Set
from functools import partial
from dataclasses import dataclass

import trio
from async_generator import asynccontextmanager

from ._state import current_actor
from ._ipc import Channel
from .log import get_logger
from ._exceptions import unpack_error, NoResult, RemoteActorError


log = get_logger('tractor')


@asynccontextmanager
async def maybe_open_nursery(
    nursery: trio.Nursery = None,
    shield: bool = False,
) -> typing.AsyncGenerator[trio.Nursery, Any]:
    """Create a new nursery if None provided.

    Blocks on exit as expected if no input nursery is provided.
    """
    if nursery is not None:
        yield nursery
    else:
        async with trio.open_nursery() as nursery:
            nursery.cancel_scope.shield = shield
            yield nursery


class StreamReceiveChannel(trio.abc.ReceiveChannel):
    """A wrapper around a ``trio._channel.MemoryReceiveChannel`` with
    special behaviour for signalling stream termination across an
    inter-actor ``Channel``. This is the type returned to a local task
    which invoked a remote streaming function using `Portal.run()`.

    Termination rules:
    - if the local task signals stop iteration a cancel signal is
      relayed to the remote task indicating to stop streaming
    - if the remote task signals the end of a stream, raise a
      ``StopAsyncIteration`` to terminate the local ``async for``

    """
    def __init__(
        self,
        cid: str,
        rx_chan: trio.abc.ReceiveChannel,
        portal: 'Portal',
    ) -> None:
        self._cid = cid
        self._rx_chan = rx_chan
        self._portal = portal

    # delegate directly to underlying mem channel
    def receive_nowait(self):
        return self._rx_chan.receive_nowait()

    async def receive(self):
        try:
            msg = await self._rx_chan.receive()
            return msg['yield']
        except trio.ClosedResourceError:
            # when the send is closed we assume the stream has
            # terminated and signal this local iterator to stop
            await self.aclose()
            raise StopAsyncIteration
        except trio.Cancelled:
            # relay cancels to the remote task
            await self.aclose()
            raise
        except KeyError:
            # internal error should never get here
            assert msg.get('cid'), (
                "Received internal error at portal?")
            raise unpack_error(msg, self._portal.channel)

    async def aclose(self):
        """Cancel associated remote actor task and local memory channel
        on close.
        """
        if self._rx_chan._closed:
            log.warning(f"{self} is already closed")
            return
        cid = self._cid
        with trio.move_on_after(0.5) as cs:
            cs.shield = True
            log.warning(
                f"Cancelling stream {cid} to "
                f"{self._portal.channel.uid}")
            # NOTE: we're telling the far end actor to cancel a task
            # corresponding to *this actor*. The far end local channel
            # instance is passed to `Actor._cancel_task()` implicitly.
            await self._portal.run('self', '_cancel_task', cid=cid)

        if cs.cancelled_caught:
            # XXX: there's no way to know if the remote task was indeed
            # cancelled in the case where the connection is broken or
            # some other network error occurred.
            if not self._portal.channel.connected():
                log.warning(
                    "May have failed to cancel remote task "
                    f"{cid} for {self._portal.channel.uid}")

        with trio.CancelScope(shield=True):
            await self._rx_chan.aclose()

    def clone(self):
        return self


class Portal:
    """A 'portal' to a(n) (remote) ``Actor``.

    Allows for invoking remote routines and receiving results through an
    underlying ``tractor.Channel`` as though the remote (async)
    function / generator was invoked locally.

    Think of this like a native async IPC API.
    """
    def __init__(self, channel: Channel) -> None:
        self.channel = channel
        # when this is set to a tuple returned from ``_submit()`` then
        # it is expected that ``result()`` will be awaited at some point
        # during the portal's lifetime
        self._result: Optional[Any] = None
        # set when _submit_for_result is called
        self._expect_result: Optional[
            Tuple[str, Any, str, Dict[str, Any]]
        ] = None
        self._streams: Set[StreamReceiveChannel] = set()
        self.actor = current_actor()

    async def _submit(
        self,
        ns: str,
        func: str,
        kwargs,
    ) -> Tuple[str, trio.abc.ReceiveChannel, str, Dict[str, Any]]:
        """Submit a function to be scheduled and run by actor, return the
        associated caller id, response queue, response type str,
        first message packet as a tuple.

        This is an async call.
        """
        # ship a function call request to the remote actor
        cid, recv_chan = await self.actor.send_cmd(
            self.channel, ns, func, kwargs)

        # wait on first response msg and handle (this should be
        # in an immediate response)

        first_msg = await recv_chan.receive()
        functype = first_msg.get('functype')

        if functype == 'function' or functype == 'asyncfunction':
            resp_type = 'return'
        elif functype == 'asyncgen':
            resp_type = 'yield'
        elif 'error' in first_msg:
            raise unpack_error(first_msg, self.channel)
        else:
            raise ValueError(f"{first_msg} is an invalid response packet?")

        return cid, recv_chan, resp_type, first_msg

    async def _submit_for_result(self, ns: str, func: str, **kwargs) -> None:
        assert self._expect_result is None, \
                "A pending main result has already been submitted"
        self._expect_result = await self._submit(ns, func, kwargs)

    async def run(self, ns: str, func: str, **kwargs) -> Any:
        """Submit a remote function to be scheduled and run by actor,
        wrap and return its (stream of) result(s).

        This is a blocking call and returns either a value from the
        remote rpc task or a local async generator instance.
        """
        return await self._return_from_resptype(
            *(await self._submit(ns, func, kwargs))
        )

    async def _return_from_resptype(
        self,
        cid: str,
        recv_chan: trio.abc.ReceiveChannel,
        resptype: str,
        first_msg: dict
    ) -> Any:
        # TODO: not this needs some serious work and thinking about how
        # to make async-generators the fundamental IPC API over channels!
        # (think `yield from`, `gen.send()`, and functional reactive stuff)
        if resptype == 'yield':  # stream response
            rchan = StreamReceiveChannel(cid, recv_chan, self)
            self._streams.add(rchan)
            return rchan

        elif resptype == 'return':  # single response
            msg = await recv_chan.receive()
            try:
                return msg['return']
            except KeyError:
                # internal error should never get here
                assert msg.get('cid'), "Received internal error at portal?"
                raise unpack_error(msg, self.channel)
        else:
            raise ValueError(f"Unknown msg response type: {first_msg}")

    async def result(self) -> Any:
        """Return the result(s) from the remote actor's "main" task.
        """
        # Check for non-rpc errors slapped on the
        # channel for which we always raise
        exc = self.channel._exc
        if exc:
            raise exc

        # not expecting a "main" result
        if self._expect_result is None:
            log.warning(
                f"Portal for {self.channel.uid} not expecting a final"
                " result?\nresult() should only be called if subactor"
                " was spawned with `ActorNursery.run_in_actor()`")
            return NoResult

        # expecting a "main" result
        assert self._expect_result
        if self._result is None:
            try:
                self._result = await self._return_from_resptype(
                    *self._expect_result
                )
            except RemoteActorError as err:
                self._result = err

        # re-raise error on every call
        if isinstance(self._result, RemoteActorError):
            raise self._result

        return self._result

    async def _cancel_streams(self):
        # terminate all locally running async generator
        # IPC calls
        if self._streams:
            log.warning(
                f"Cancelling all streams with {self.channel.uid}")
            for stream in self._streams.copy():
                await stream.aclose()

    async def aclose(self):
        log.debug(f"Closing {self}")
        # TODO: once we move to implementing our own `ReceiveChannel`
        # (including remote task cancellation inside its `.aclose()`)
        # we'll need to .aclose all those channels here
        await self._cancel_streams()

    async def cancel_actor(self):
        """Cancel the actor on the other end of this portal.
        """
        if not self.channel.connected():
            log.warning("This portal is already closed can't cancel")
            return False

        await self._cancel_streams()

        log.warning(
            f"Sending actor cancel request to {self.channel.uid} on "
            f"{self.channel}")
        try:
            # send cancel cmd - might not get response
            # XXX: sure would be nice to make this work with a proper shield
            # with trio.CancelScope() as cancel_scope:
            # with trio.CancelScope(shield=True) as cancel_scope:
            with trio.move_on_after(0.5) as cancel_scope:
                cancel_scope.shield = True
                await self.run('self', 'cancel')
                return True
            if cancel_scope.cancelled_caught:
                log.warning(f"May have failed to cancel {self.channel.uid}")

            # if we get here some weird cancellation case happened
            return False
        except trio.ClosedResourceError:
            log.warning(
                f"{self.channel} for {self.channel.uid} was already closed?")
            return False


@dataclass
class LocalPortal:
    """A 'portal' to a local ``Actor``.

    A compatibility shim for normal portals but for invoking functions
    using an in process actor instance.
    """
    actor: 'Actor'  # type: ignore
    channel: Channel

    async def run(self, ns: str, func_name: str, **kwargs) -> Any:
        """Run a requested function locally and return it's result.
        """
        obj = self.actor if ns == 'self' else importlib.import_module(ns)
        func = getattr(obj, func_name)
        if inspect.iscoroutinefunction(func):
            return await func(**kwargs)
        else:
            return func(**kwargs)


@asynccontextmanager
async def open_portal(
    channel: Channel,
    nursery: Optional[trio.Nursery] = None,
    start_msg_loop: bool = True,
    shield: bool = False,
) -> typing.AsyncGenerator[Portal, None]:
    """Open a ``Portal`` through the provided ``channel``.

    Spawns a background task to handle message processing.
    """
    actor = current_actor()
    assert actor
    was_connected = False

    async with maybe_open_nursery(nursery, shield=shield) as nursery:
        if not channel.connected():
            await channel.connect()
            was_connected = True

        if channel.uid is None:
            await actor._do_handshake(channel)

        msg_loop_cs: Optional[trio.CancelScope] = None
        if start_msg_loop:
            msg_loop_cs = await nursery.start(
                partial(
                    actor._process_messages,
                    channel,
                    # if the local task is cancelled we want to keep
                    # the msg loop running until our block ends
                    shield=True,
                )
            )
        portal = Portal(channel)
        try:
            yield portal
        finally:
            await portal.aclose()

            if was_connected:
                # cancel remote channel-msg loop
                await channel.send(None)

            # cancel background msg loop task
            if msg_loop_cs:
                msg_loop_cs.cancel()

            nursery.cancel_scope.cancel()