""" Message stream types and APIs. """ import inspect from contextlib import contextmanager, asynccontextmanager from dataclasses import dataclass from typing import ( Any, Iterator, Optional, Callable, AsyncGenerator, Dict, ) import warnings import trio from ._ipc import Channel from ._exceptions import unpack_error, ContextCancelled from ._state import current_actor from .log import get_logger log = get_logger(__name__) # TODO: generic typing like trio's receive channel # but with msgspec messages? # class ReceiveChannel(AsyncResource, Generic[ReceiveType]): class ReceiveMsgStream(trio.abc.ReceiveChannel): '''A IPC message stream for receiving logically sequenced values over an inter-actor ``Channel``. This is the type returned to a local task which entered either ``Portal.open_stream_from()`` or ``Context.open_stream()``. Termination rules: - on cancellation the stream is **not** implicitly closed and the surrounding ``Context`` is expected to handle how that cancel is relayed to any task on the remote side. - if the remote task signals the end of a stream the ``ReceiveChannel`` semantics dictate that a ``StopAsyncIteration`` to terminate the local ``async for``. ''' def __init__( self, ctx: 'Context', # typing: ignore # noqa rx_chan: trio.abc.ReceiveChannel, ) -> None: self._ctx = ctx self._rx_chan = rx_chan # flag to denote end of stream self._eoc: bool = False # delegate directly to underlying mem channel def receive_nowait(self): msg = self._rx_chan.receive_nowait() return msg['yield'] async def receive(self): '''Async receive a single msg from the IPC transport, the next in sequence for this stream. ''' # see ``.aclose()`` for notes on the old behaviour prior to # introducing this if self._eoc: raise trio.EndOfChannel try: msg = await self._rx_chan.receive() return msg['yield'] except KeyError: # internal error should never get here assert msg.get('cid'), ("Received internal error at portal?") # TODO: handle 2 cases with 3.10 match syntax # - 'stop' # - 'error' # possibly just handle msg['stop'] here! if msg.get('stop'): log.debug(f"{self} was stopped at remote end") # # when the send is closed we assume the stream has # # terminated and signal this local iterator to stop # await self.aclose() # XXX: this causes ``ReceiveChannel.__anext__()`` to # raise a ``StopAsyncIteration`` **and** in our catch # block below it will trigger ``.aclose()``. raise trio.EndOfChannel # TODO: test that shows stream raising an expected error!!! elif msg.get('error'): # raise the error message raise unpack_error(msg, self._ctx.chan) else: raise except ( trio.ClosedResourceError, # by self._rx_chan trio.EndOfChannel, # by self._rx_chan or `stop` msg from far end ): # XXX: we close the stream on any of these error conditions: # a ``ClosedResourceError`` indicates that the internal # feeder memory receive channel was closed likely by the # runtime after the associated transport-channel # disconnected or broke. # an ``EndOfChannel`` indicates either the internal recv # memchan exhausted **or** we raisesd it just above after # receiving a `stop` message from the far end of the stream. # Previously this was triggered by calling ``.aclose()`` on # the send side of the channel inside # ``Actor._push_result()`` (should still be commented code # there - which should eventually get removed), but now the # 'stop' message handling has been put just above. # TODO: Locally, we want to close this stream gracefully, by # terminating any local consumers tasks deterministically. # One we have broadcast support, we **don't** want to be # closing this stream and not flushing a final value to # remaining (clone) consumers who may not have been # scheduled to receive it yet. # when the send is closed we assume the stream has # terminated and signal this local iterator to stop await self.aclose() raise # propagate async def aclose(self): """Cancel associated remote actor task and local memory channel on close. """ # XXX: keep proper adherance to trio's `.aclose()` semantics: # https://trio.readthedocs.io/en/stable/reference-io.html#trio.abc.AsyncResource.aclose rx_chan = self._rx_chan if rx_chan._closed: log.warning(f"{self} is already closed") # this stream has already been closed so silently succeed as # per ``trio.AsyncResource`` semantics. # https://trio.readthedocs.io/en/stable/reference-io.html#trio.abc.AsyncResource.aclose return self._eoc = True # NOTE: this is super subtle IPC messaging stuff: # Relay stop iteration to far end **iff** we're # in bidirectional mode. If we're only streaming # *from* one side then that side **won't** have an # entry in `Actor._cids2qs` (maybe it should though?). # So any `yield` or `stop` msgs sent from the caller side # will cause key errors on the callee side since there is # no entry for a local feeder mem chan since the callee task # isn't expecting messages to be sent by the caller. # Thus, we must check that this context DOES NOT # have a portal reference to ensure this is indeed the callee # side and can relay a 'stop'. # In the bidirectional case, `Context.open_stream()` will create # the `Actor._cids2qs` entry from a call to # `Actor.get_memchans()` and will send the stop message in # ``__aexit__()`` on teardown so it **does not** need to be # called here. if not self._ctx._portal: # Only for 2 way streams can we can send stop from the # caller side. try: # NOTE: if this call is cancelled we expect this end to # handle as though the stop was never sent (though if it # was it shouldn't matter since it's unlikely a user # will try to re-use a stream after attemping to close # it). await self._ctx.send_stop() except ( trio.BrokenResourceError, trio.ClosedResourceError ): # the underlying channel may already have been pulled # in which case our stop message is meaningless since # it can't traverse the transport. log.debug(f'Channel for {self} was already closed') # Do we close the local mem chan ``self._rx_chan`` ??!? # NO, DEFINITELY NOT if we're a bi-dir ``MsgStream``! # BECAUSE this same core-msg-loop mem recv-chan is used to deliver # the potential final result from the surrounding inter-actor # `Context` so we don't want to close it until that context has # run to completion. # XXX: Notes on old behaviour: # await rx_chan.aclose() # In the receive-only case, ``Portal.open_stream_from()`` used # to rely on this call explicitly on teardown such that a new # call to ``.receive()`` after ``rx_chan`` had been closed, would # result in us raising a ``trio.EndOfChannel`` (since we # remapped the ``trio.ClosedResourceError`). However, now if for some # reason the stream's consumer code tries to manually receive a new # value before ``.aclose()`` is called **but** the far end has # stopped `.receive()` **must** raise ``trio.EndofChannel`` in # order to avoid an infinite hang on ``.__anext__()``; this is # why we added ``self._eoc`` to denote stream closure indepedent # of ``rx_chan``. # In theory we could still use this old method and close the # underlying msg-loop mem chan as above and then **not** check # for ``self._eoc`` in ``.receive()`` (if for some reason we # think that check is a bottle neck - not likely) **but** then # we would need to map the resulting # ``trio.ClosedResourceError`` to a ``trio.EndOfChannel`` in # ``.receive()`` (as it originally was before bi-dir streaming # support) in order to trigger stream closure. The old behaviour # is arguably more confusing since we lose detection of the # runtime's closure of ``rx_chan`` in the case where we may # still need to consume msgs that are "in transit" from the far # end (eg. for ``Context.result()``). class MsgStream(ReceiveMsgStream, trio.abc.Channel): """ Bidirectional message stream for use within an inter-actor actor ``Context```. """ async def send( self, data: Any ) -> None: '''Send a message over this stream to the far end. ''' await self._ctx.chan.send({'yield': data, 'cid': self._ctx.cid}) # TODO: but make it broadcasting to consumers def clone(self): """Clone this receive channel allowing for multi-task consumption from the same channel. """ return MsgStream( self._ctx, self._rx_chan.clone(), ) @dataclass class Context: '''An inter-actor task communication context. Allows maintaining task or protocol specific state between 2 communicating actor tasks. A unique context is created on the callee side/end for every request to a remote actor from a portal. A context can be cancelled and (possibly eventually restarted) from either side of the underlying IPC channel. A context can be used to open task oriented message streams and can be thought of as an IPC aware inter-actor cancel scope. ''' chan: Channel cid: str # only set on the caller side _portal: Optional['Portal'] = None # type: ignore # noqa _recv_chan: Optional[trio.MemoryReceiveChannel] = None _result: Optional[Any] = False _cancel_called: bool = False # only set on the callee side _scope_nursery: Optional[trio.Nursery] = None async def send_yield(self, data: Any) -> None: warnings.warn( "`Context.send_yield()` is now deprecated. " "Use ``MessageStream.send()``. ", DeprecationWarning, stacklevel=2, ) await self.chan.send({'yield': data, 'cid': self.cid}) async def send_stop(self) -> None: await self.chan.send({'stop': True, 'cid': self.cid}) def _error_from_remote_msg( self, msg: Dict[str, Any], ) -> None: '''Unpack and raise a msg error into the local scope nursery for this context. Acts as a form of "relay" for a remote error raised in the corresponding remote callee task. ''' assert self._scope_nursery async def raiser(): raise unpack_error(msg, self.chan) self._scope_nursery.start_soon(raiser) async def cancel(self) -> None: '''Cancel this inter-actor-task context. Request that the far side cancel it's current linked context, Timeout quickly in an attempt to sidestep 2-generals... ''' side = 'caller' if self._portal else 'callee' log.warning(f'Cancelling {side} side of context to {self.chan}') self._cancel_called = True if side == 'caller': if not self._portal: raise RuntimeError( "No portal found, this is likely a callee side context" ) 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_from_ns('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(): if not self.chan.connected(): log.warning( "May have failed to cancel remote task " f"{cid} for {self._portal.channel.uid}") else: # callee side remote task # TODO: should we have an explicit cancel message # or is relaying the local `trio.Cancelled` as an # {'error': trio.Cancelled, cid: "blah"} enough? # This probably gets into the discussion in # https://github.com/goodboy/tractor/issues/36 assert self._scope_nursery self._scope_nursery.cancel_scope.cancel() if self._recv_chan: await self._recv_chan.aclose() @asynccontextmanager async def open_stream( self, ) -> AsyncGenerator[MsgStream, None]: '''Open a ``MsgStream``, a bi-directional stream connected to the cross-actor (far end) task for this ``Context``. This context manager must be entered on both the caller and callee for the stream to logically be considered "connected". A ``MsgStream`` is currently "one-shot" use, meaning if you close it you can not "re-open" it for streaming and instead you must re-establish a new surrounding ``Context`` using ``Portal.open_context()``. In the future this may change but currently there seems to be no obvious reason to support "re-opening": - pausing a stream can be done with a message. - task errors will normally require a restart of the entire scope of the inter-actor task context due to the nature of ``trio``'s cancellation system. ''' actor = current_actor() # here we create a mem chan that corresponds to the # far end caller / callee. # NOTE: in one way streaming this only happens on the # caller side inside `Actor.send_cmd()` so if you try # to send a stop from the caller to the callee in the # single-direction-stream case you'll get a lookup error # currently. _, recv_chan = actor.get_memchans( self.chan.uid, self.cid ) # Likewise if the surrounding context has been cancelled we error here # since it likely means the surrounding block was exited or # killed if self._cancel_called: task = trio.lowlevel.current_task().name raise ContextCancelled( f'Context around {actor.uid[0]}:{task} was already cancelled!' ) # XXX: If the underlying channel feeder receive mem chan has # been closed then likely client code has already exited # a ``.open_stream()`` block prior or there was some other # unanticipated error or cancellation from ``trio``. if recv_chan._closed: raise trio.ClosedResourceError( 'The underlying channel for this stream was already closed!?') async with MsgStream( ctx=self, rx_chan=recv_chan, ) as rchan: if self._portal: self._portal._streams.add(rchan) try: # ensure we aren't cancelled before delivering # the stream # await trio.lowlevel.checkpoint() yield rchan except trio.EndOfChannel: # likely the far end sent us a 'stop' message to # terminate the stream. raise else: # XXX: Make the stream "one-shot use". On exit, signal # ``trio.EndOfChannel``/``StopAsyncIteration`` to the # far end. await self.send_stop() finally: if self._portal: self._portal._streams.remove(rchan) async def result(self) -> Any: '''From a caller side, wait for and return the final result from the callee side task. ''' assert self._portal, "Context.result() can not be called from callee!" assert self._recv_chan if self._result is False: if not self._recv_chan._closed: # type: ignore # wait for a final context result consuming # and discarding any bi dir stream msgs still # in transit from the far end. while True: msg = await self._recv_chan.receive() try: self._result = msg['return'] break except KeyError: if 'yield' in msg: # far end task is still streaming to us.. log.warning(f'Remote stream deliverd {msg}') # do disard continue elif 'stop' in msg: log.debug('Remote stream terminated') continue # internal error should never get here assert msg.get('cid'), ( "Received internal error at portal?") raise unpack_error(msg, self._portal.channel) return self._result async def started(self, value: Optional[Any] = None) -> None: if self._portal: raise RuntimeError( f"Caller side context {self} can not call started!") await self.chan.send({'started': value, 'cid': self.cid}) # TODO: do we need a restart api? # async def restart(self) -> None: # pass def stream(func: Callable) -> Callable: """Mark an async function as a streaming routine with ``@stream``. """ # annotate # TODO: apply whatever solution ``mypy`` ends up picking for this: # https://github.com/python/mypy/issues/2087#issuecomment-769266912 func._tractor_stream_function = True # type: ignore sig = inspect.signature(func) params = sig.parameters if 'stream' not in params and 'ctx' in params: warnings.warn( "`@tractor.stream decorated funcs should now declare a `stream` " " arg, `ctx` is now designated for use with @tractor.context", DeprecationWarning, stacklevel=2, ) if ( 'ctx' not in params and 'to_trio' not in params and 'stream' not in params ): raise TypeError( "The first argument to the stream function " f"{func.__name__} must be `ctx: tractor.Context` " "(Or ``to_trio`` if using ``asyncio`` in guest mode)." ) return func def context(func: Callable) -> Callable: """Mark an async function as a streaming routine with ``@context``. """ # annotate # TODO: apply whatever solution ``mypy`` ends up picking for this: # https://github.com/python/mypy/issues/2087#issuecomment-769266912 func._tractor_context_function = True # type: ignore sig = inspect.signature(func) params = sig.parameters if 'ctx' not in params: raise TypeError( "The first argument to the context function " f"{func.__name__} must be `ctx: tractor.Context`" ) return func