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._agen = self._iter_packets()
|
||||||
self._send_lock = trio.StrictFIFOLock()
|
self._send_lock = trio.StrictFIFOLock()
|
||||||
|
|
||||||
|
# public i guess?
|
||||||
|
self.drained = []
|
||||||
|
|
||||||
async def _iter_packets(self) -> typing.AsyncGenerator[dict, None]:
|
async def _iter_packets(self) -> typing.AsyncGenerator[dict, None]:
|
||||||
"""Yield packets from the underlying stream.
|
"""Yield packets from the underlying stream.
|
||||||
"""
|
"""
|
||||||
|
@ -132,7 +135,7 @@ class MsgpackTCPStream:
|
||||||
|
|
||||||
if data == b'':
|
if data == b'':
|
||||||
raise TransportClosed(
|
raise TransportClosed(
|
||||||
f'transport {self} was already closed prior ro read'
|
f'transport {self} was already closed prior to read'
|
||||||
)
|
)
|
||||||
|
|
||||||
unpacker.feed(data)
|
unpacker.feed(data)
|
||||||
|
@ -156,6 +159,14 @@ class MsgpackTCPStream:
|
||||||
async def recv(self) -> Any:
|
async def recv(self) -> Any:
|
||||||
return await self._agen.asend(None)
|
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):
|
def __aiter__(self):
|
||||||
return self._agen
|
return self._agen
|
||||||
|
|
||||||
|
@ -164,7 +175,8 @@ class MsgpackTCPStream:
|
||||||
|
|
||||||
|
|
||||||
class MsgspecTCPStream(MsgpackTCPStream):
|
class MsgspecTCPStream(MsgpackTCPStream):
|
||||||
'''A ``trio.SocketStream`` delivering ``msgpack`` formatted data
|
'''
|
||||||
|
A ``trio.SocketStream`` delivering ``msgpack`` formatted data
|
||||||
using ``msgspec``.
|
using ``msgspec``.
|
||||||
|
|
||||||
'''
|
'''
|
||||||
|
@ -259,9 +271,12 @@ def get_msg_transport(
|
||||||
|
|
||||||
|
|
||||||
class Channel:
|
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__(
|
def __init__(
|
||||||
|
@ -299,10 +314,12 @@ class Channel:
|
||||||
# set after handshake - always uid of far end
|
# set after handshake - always uid of far end
|
||||||
self.uid: Optional[Tuple[str, str]] = None
|
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._agen = self._aiter_recv()
|
||||||
|
self._exc: Optional[Exception] = None # set if far end actor errors
|
||||||
self._closed: bool = False
|
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
|
@classmethod
|
||||||
def from_stream(
|
def from_stream(
|
||||||
|
|
Loading…
Reference in New Issue