Add `.drain()` support to msg streams
Enables "draining" the last set of messages after a channel/stream has been terminated mostly for the purposes of receiving a final ACK to a remote cancel command. Also, add an internal `Channel._cancel_called` flag which can be set by `Portal.cancel_actor()`.acked_backup
parent
0ac3397dbb
commit
1976e61d1a
|
@ -93,6 +93,9 @@ class MsgpackTCPStream:
|
|||
self._agen = self._iter_packets()
|
||||
self._send_lock = trio.StrictFIFOLock()
|
||||
|
||||
# public i guess?
|
||||
self.drained = []
|
||||
|
||||
async def _iter_packets(self) -> typing.AsyncGenerator[dict, None]:
|
||||
"""Yield packets from the underlying stream.
|
||||
"""
|
||||
|
@ -132,7 +135,7 @@ class MsgpackTCPStream:
|
|||
|
||||
if data == b'':
|
||||
raise TransportClosed(
|
||||
f'transport {self} was already closed prior ro read'
|
||||
f'transport {self} was already closed prior to read'
|
||||
)
|
||||
|
||||
unpacker.feed(data)
|
||||
|
@ -156,6 +159,14 @@ class MsgpackTCPStream:
|
|||
async def recv(self) -> Any:
|
||||
return await self._agen.asend(None)
|
||||
|
||||
async def drain(self):
|
||||
try:
|
||||
async for msg in self._iter_packets():
|
||||
self.drained.append(msg)
|
||||
except TransportClosed:
|
||||
for msg in self.drained:
|
||||
yield msg
|
||||
|
||||
def __aiter__(self):
|
||||
return self._agen
|
||||
|
||||
|
@ -164,7 +175,8 @@ class MsgpackTCPStream:
|
|||
|
||||
|
||||
class MsgspecTCPStream(MsgpackTCPStream):
|
||||
'''A ``trio.SocketStream`` delivering ``msgpack`` formatted data
|
||||
'''
|
||||
A ``trio.SocketStream`` delivering ``msgpack`` formatted data
|
||||
using ``msgspec``.
|
||||
|
||||
'''
|
||||
|
@ -259,9 +271,12 @@ def get_msg_transport(
|
|||
|
||||
|
||||
class Channel:
|
||||
'''An inter-process channel for communication between (remote) actors.
|
||||
'''
|
||||
An inter-process channel for communication between (remote) actors.
|
||||
|
||||
Currently the only supported transport is a ``trio.SocketStream``.
|
||||
Wraps a ``MsgStream``: transport + encoding IPC connection.
|
||||
Currently we only support ``trio.SocketStream`` for transport
|
||||
(aka TCP).
|
||||
|
||||
'''
|
||||
def __init__(
|
||||
|
@ -299,10 +314,12 @@ class Channel:
|
|||
# set after handshake - always uid of far end
|
||||
self.uid: Optional[Tuple[str, str]] = None
|
||||
|
||||
# set if far end actor errors internally
|
||||
self._exc: Optional[Exception] = None
|
||||
self._agen = self._aiter_recv()
|
||||
self._exc: Optional[Exception] = None # set if far end actor errors
|
||||
self._closed: bool = False
|
||||
# flag set on ``Portal.cancel_actor()`` indicating
|
||||
# remote (peer) cancellation of the far end actor runtime.
|
||||
self._cancel_called: bool = False # set on ``Portal.cancel_actor()``
|
||||
|
||||
@classmethod
|
||||
def from_stream(
|
||||
|
|
Loading…
Reference in New Issue