diff --git a/src/ezmsg/core/frameproto.py b/src/ezmsg/core/frameproto.py new file mode 100644 index 00000000..6c56aa9d --- /dev/null +++ b/src/ezmsg/core/frameproto.py @@ -0,0 +1,198 @@ +""" +Framed :class:`asyncio.Protocol` support for ezmsg's per-message hot path. + +The high-level streams API (:func:`asyncio.open_connection`) delivers incoming +bytes by resolving a future and scheduling the reading task, so every message +costs a full event-loop iteration before any work happens. A +:class:`asyncio.Protocol` has ``data_received`` called synchronously from the +transport's read callback, so the message can be handled in that callback +instead. On stock asyncio this is worth roughly 25us per hop. + +Only public asyncio API is used here, which keeps uvloop a drop-in replacement +on POSIX: uvloop implements the same ``AbstractEventLoop.create_connection`` / +``create_server`` contract for :class:`asyncio.Protocol`, and under uvloop the +streams path is already as fast, so the same code is optimal on both loops. + +A connection runs in two phases: + +* **handshake** -- sequential request/response, driven by ``await read_exactly()`` + and friends. One task wakeup per read, which is fine: it happens once. +* **dispatch** -- entered via :meth:`FramedProtocol.start_dispatch`. From then on + :meth:`FramedProtocol.frames_available` is called synchronously from + ``data_received`` and consumes whole frames out of :attr:`buffer`. +""" + +import asyncio +import logging +import socket + +from .netprotocol import UINT64_SIZE, BYTEORDER + +logger = logging.getLogger("ezmsg") + + +class FramedProtocol(asyncio.Protocol): + """ + Base protocol that buffers incoming bytes and supports a handshake phase + followed by an inline dispatch phase. + + Subclasses implement :meth:`frames_available` to consume complete frames + from :attr:`buffer` during the dispatch phase. + """ + + def __init__(self) -> None: + self._buffer = bytearray() + self._transport: asyncio.Transport | None = None + self._read_waiter: "asyncio.Future[None] | None" = None + self._need = 0 + self._dispatching = False + self._closed: "asyncio.Future[None] | None" = None + self._close_exc: BaseException | None = None + + # ------------------------------------------------------------------ # + # asyncio.Protocol interface + # ------------------------------------------------------------------ # + + def connection_made(self, transport: asyncio.BaseTransport) -> None: + self._transport = transport # type: ignore[assignment] + self._closed = asyncio.get_running_loop().create_future() + sock = transport.get_extra_info("socket") + if sock is not None: + # Notifications are tiny and latency-critical; never coalesce them. + try: + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + except (OSError, AttributeError): + pass + + def data_received(self, data: bytes) -> None: + self._buffer.extend(data) + + if self._dispatching: + self._drain_frames() + else: + self._wake_reader() + + def eof_received(self) -> bool: + return False # let the transport close us + + def connection_lost(self, exc: BaseException | None) -> None: + self._close_exc = exc + waiter, self._read_waiter = self._read_waiter, None + if waiter is not None and not waiter.done(): + waiter.set_exception( + exc if exc is not None else asyncio.IncompleteReadError(b"", None) + ) + if self._closed is not None and not self._closed.done(): + self._closed.set_result(None) + + # ------------------------------------------------------------------ # + # dispatch phase + # ------------------------------------------------------------------ # + + def start_dispatch(self) -> None: + """ + Leave the handshake phase. :meth:`frames_available` will be called + inline from ``data_received`` from now on, starting with whatever is + already buffered. + """ + self._dispatching = True + if self._buffer: + self._drain_frames() + + def frames_available(self) -> None: + """ + Consume as many whole frames as :attr:`buffer` holds. + + Called synchronously from the transport's read callback, so it must not + block and must leave any partial trailing frame in the buffer. + """ + raise NotImplementedError + + def _drain_frames(self) -> None: + try: + self.frames_available() + except Exception as exc: + # A raise here would be swallowed by the transport, silently + # wedging the connection, so surface it and tear down instead. + logger.exception("%s: error dispatching frame", type(self).__name__) + self._close_exc = exc + self.abort() + + # ------------------------------------------------------------------ # + # handshake phase + # ------------------------------------------------------------------ # + + def _wake_reader(self) -> None: + waiter = self._read_waiter + if waiter is not None and not waiter.done() and len(self._buffer) >= self._need: + self._read_waiter = None + waiter.set_result(None) + + async def read_exactly(self, n: int) -> bytes: + """Await exactly ``n`` bytes. Only valid during the handshake phase.""" + while len(self._buffer) < n: + if self._closed is not None and self._closed.done(): + raise asyncio.IncompleteReadError(bytes(self._buffer), n) + self._need = n + self._read_waiter = asyncio.get_running_loop().create_future() + await self._read_waiter + out = bytes(self._buffer[:n]) + del self._buffer[:n] + return out + + async def read_uint64(self) -> int: + return int.from_bytes(await self.read_exactly(UINT64_SIZE), BYTEORDER) + + async def read_str(self) -> str: + return (await self.read_exactly(await self.read_uint64())).decode("utf-8") + + # ------------------------------------------------------------------ # + # misc + # ------------------------------------------------------------------ # + + @property + def buffer(self) -> bytearray: + """Unconsumed bytes. Subclasses consume from the front of this.""" + return self._buffer + + @property + def transport(self) -> asyncio.Transport: + assert self._transport is not None, "connection_made has not run" + return self._transport + + def write(self, data: bytes) -> None: + if self._transport is not None and not self._transport.is_closing(): + self._transport.write(data) + + def pause_reading(self) -> None: + if self._transport is not None and not self._transport.is_closing(): + self._transport.pause_reading() + + def resume_reading(self) -> None: + if self._transport is not None and not self._transport.is_closing(): + self._transport.resume_reading() + + def close(self) -> None: + """ + Close gracefully, flushing anything still queued for write. + + Note that ``connection_lost`` -- and therefore :meth:`wait_closed` -- + is deferred until that queue drains. Use :meth:`abort` for teardown, + where a peer that has stopped reading must not be able to wedge us. + """ + if self._transport is not None and not self._transport.is_closing(): + self._transport.close() + + def abort(self) -> None: + """ + Close immediately, discarding anything still queued for write. + + ``connection_lost`` fires without waiting for the peer, so + :meth:`wait_closed` always completes. + """ + if self._transport is not None: + self._transport.abort() + + async def wait_closed(self) -> None: + if self._closed is not None: + await self._closed diff --git a/src/ezmsg/core/messagechannel.py b/src/ezmsg/core/messagechannel.py index e6479486..6fab584b 100644 --- a/src/ezmsg/core/messagechannel.py +++ b/src/ezmsg/core/messagechannel.py @@ -11,18 +11,24 @@ from .backpressure import Backpressure from .messagecache import MessageCache, CacheMiss from .graphserver import GraphService +from .frameproto import FramedProtocol from .netprotocol import ( Command, Address, AddressType, + BYTEORDER, read_str, - read_int, uint64_to_bytes, encode_str, close_stream_writer, ) from .graphmeta import ProfileChannelType +# TX_SHM and TX_TCP share a fixed prefix: command byte, msg_id, then a length +# that covers the rest of the frame (the SHM segment name, or the serialized +# message respectively). +_FRAME_PREFIX = 1 + 8 + 8 + logger = logging.getLogger("ezmsg") @@ -70,6 +76,118 @@ def put_nowait(self, item: typing.Tuple[UUID, int]) -> None: NotificationQueue = asyncio.Queue[typing.Tuple[UUID, int]] | LeakyQueue +class ChannelProtocol(FramedProtocol): + """ + Receives message notifications from a Publisher. + + During dispatch, :meth:`frames_available` runs synchronously in the + transport's read callback, so an incoming message is cached and its + subscribers notified without the extra task wakeup the streams API costs. + """ + + def __init__(self) -> None: + super().__init__() + self._channel: "Channel | None" = None + + def bind(self, channel: "Channel") -> None: + self._channel = channel + + def connection_lost(self, exc: BaseException | None) -> None: + super().connection_lost(exc) + if self._channel is not None: + self._channel._on_disconnected() + + def frames_available(self) -> None: + chan = self._channel + assert chan is not None, "protocol dispatching before bind()" + buf = self._buffer + + while True: + if len(buf) < _FRAME_PREFIX: + return + cmd = bytes(buf[0:1]) + msg_id = int.from_bytes(buf[1:9], BYTEORDER) + tail = int.from_bytes(buf[9:_FRAME_PREFIX], BYTEORDER) + end = _FRAME_PREFIX + tail + if len(buf) < end: + return + + if cmd == Command.TX_SHM.value: + shm_name = bytes(buf[_FRAME_PREFIX:end]).decode("utf-8") + if chan.shm is None or chan.shm.name != shm_name: + # Attaching is async and cannot happen in a read callback. + # Stop reading, leave this frame buffered, and hand off. + self.pause_reading() + asyncio.get_running_loop().create_task( + self._reattach_shm(shm_name, end, msg_id), + name=f"chan-{chan.id}: reattach_shm", + ) + return + del buf[:end] + chan._deliver_from_shm(msg_id) + + elif cmd == Command.TX_TCP.value: + payload = bytes(buf[_FRAME_PREFIX:end]) + del buf[:end] + chan._deliver_from_tcp(msg_id, payload) + + else: + raise ValueError(f"unimplemented data telemetry: {cmd!r}") + + async def _reattach_shm(self, shm_name: str, frame_end: int, msg_id: int) -> None: + """ + Swap to a new SHM generation, then resume dispatch. + + The triggering frame is still buffered: on success we re-dispatch it + against the new segment, on failure we drop it and release its + backpressure so the publisher is not stalled by it. + """ + chan = self._channel + assert chan is not None + + try: + preserved = chan._snapshot_cached_messages() + chan.cache.clear() + for preserved_msg in preserved: + chan.cache.put_from_mem(preserved_msg) + + if chan.shm is not None: + old_shm = chan.shm + chan.shm = None + old_shm.close() + await old_shm.wait_closed() + + try: + chan.shm = await GraphService(chan._graph_address).attach_shm(shm_name) + except ValueError: + logger.warning( + "Channel %s received stale SHM %s for publisher %s; waiting for next valid SHM", + chan.id, + shm_name, + chan.pub_id, + ) + chan.shm = None + + if chan.shm is None: + # Drop the frame we parked, otherwise dispatch would retry the + # attach against the same name forever. + logger.warning( + "Channel %s dropping message %s from publisher %s because its SHM generation is stale", + chan.id, + msg_id, + chan.pub_id, + ) + del self._buffer[:frame_end] + chan._release_backpressure(msg_id, chan.id) + except Exception: + logger.exception("Channel %s failed to reattach SHM", chan.id) + self.abort() + return + + self.resume_reading() + self._drain_frames() + + class Channel: """ Channel is a "middle-man" that receives messages from a particular Publisher, @@ -96,8 +214,7 @@ class Channel: backpressure: Backpressure _graph_task: asyncio.Task[None] - _pub_task: asyncio.Task[None] - _pub_writer: asyncio.StreamWriter + _proto: ChannelProtocol _graph_address: AddressType | None _local_backpressure: Backpressure | None _channel_kind: ProfileChannelType @@ -163,30 +280,37 @@ async def create( id_str = await read_str(graph_reader) pub_address = await Address.from_stream(graph_reader) - reader, writer = await asyncio.open_connection(*pub_address) - writer.write(Command.CHANNEL.value) - writer.write(encode_str(id_str)) + # The per-message path uses a Protocol rather than the streams API so + # that an incoming notification is handled in the transport's read + # callback instead of costing an extra task wakeup. The handshake below + # is sequential and runs once, so it reads in the ordinary awaiting way. + loop = asyncio.get_running_loop() + _, proto = await loop.create_connection(ChannelProtocol, *pub_address) + proto.write(Command.CHANNEL.value + encode_str(id_str)) - topic = await read_str(reader) + topic = await proto.read_str() shm = None - shm_name = await read_str(reader) + shm_name = await proto.read_str() try: shm = await graph_service.attach_shm(shm_name) - writer.write(Command.SHM_OK.value) + proto.write(Command.SHM_OK.value) except (ValueError, OSError): shm = None - writer.write(Command.SHM_ATTACH_FAILED.value) - writer.write(uint64_to_bytes(os.getpid())) + proto.write(Command.SHM_ATTACH_FAILED.value) + proto.write(uint64_to_bytes(os.getpid())) - result = await reader.read(1) + result = await proto.read_exactly(1) if result != Command.COMPLETE.value: # NOTE: The only reason this would happen is if the # publisher's writer is closed due to a crash or shutdown + proto.close() raise ValueError(f"failed to create channel {pub_id=}") - num_buffers = await read_int(reader) - assert num_buffers > 0, "publisher reports invalid num_buffers" + num_buffers = await proto.read_uint64() + if num_buffers <= 0: + proto.close() + raise ValueError("publisher reports invalid num_buffers") chan = cls(UUID(id_str), pub_id, num_buffers, shm, graph_address, _guard=cls._SENTINEL) chan.topic = topic @@ -196,11 +320,11 @@ async def create( name=f"chan-{chan.id}: _graph_connection", ) - chan._pub_writer = writer - chan._pub_task = asyncio.create_task( - chan._publisher_connection(reader), - name=f"chan-{chan.id}: _publisher_connection", - ) + chan._proto = proto + proto.bind(chan) + # Anything the publisher sent between the handshake and here is already + # buffered; start_dispatch drains it before returning. + proto.start_dispatch() logger.debug(f"created channel {chan.id=} {pub_id=} {pub_address=}") @@ -210,15 +334,18 @@ def close(self) -> None: """ Mark the Channel for shutdown and resource deallocation """ - self._pub_task.cancel() + # abort() rather than close(): a graceful close defers connection_lost + # until queued acks flush, so a publisher that has already stopped + # reading them would leave wait_closed() hanging. The publisher frees + # this channel's backpressure on disconnect, so dropping them is safe. + self._proto.abort() self._graph_task.cancel() async def wait_closed(self) -> None: """ Wait until the Channel has properly shutdown and its resources have been deallocated. """ - with suppress(asyncio.CancelledError): - await self._pub_task + await self._proto.wait_closed() with suppress(asyncio.CancelledError): await self._graph_task if self.shm is not None: @@ -247,115 +374,63 @@ async def _graph_connection( finally: await close_stream_writer(writer) - async def _publisher_connection(self, reader: asyncio.StreamReader) -> None: + def _deliver_from_shm(self, msg_id: int) -> None: """ - The task that handles communication between the Channel and the Publisher it receives messages from. + Cache the message sitting in our SHM slot and notify clients. + + Called inline from :meth:`ChannelProtocol.frames_available`, so the + caller has already established that ``self.shm`` matches the segment the + publisher named. """ + assert self.shm is not None + shm_buf = self.shm[msg_id % self.num_buffers] + # The slot for this msg_id may be uninitialized after a mid-stream + # resize; msg_id() raises UninitializedMemory in that case. Treat it as + # a mismatch (drop + release) rather than letting it kill the channel. try: - while True: - msg = await reader.read(1) - - if not msg: - break - - msg_id = await read_int(reader) - buf_idx = msg_id % self.num_buffers - channel_kind = ProfileChannelType.UNKNOWN - - if msg == Command.TX_SHM.value: - channel_kind = ProfileChannelType.SHM - shm_name = await read_str(reader) - - if self.shm is None or self.shm.name != shm_name: - preserved_cache = self._snapshot_cached_messages() - self.cache.clear() - for preserved_msg in preserved_cache: - self.cache.put_from_mem(preserved_msg) - - if self.shm is not None: - old_shm = self.shm - old_shm.close() - await old_shm.wait_closed() - - try: - self.shm = await GraphService( - self._graph_address - ).attach_shm(shm_name) - except ValueError: - logger.warning( - "Channel %s received stale SHM %s for publisher %s; waiting for next valid SHM", - self.id, - shm_name, - self.pub_id, - ) - self.shm = None - - if self.shm is None: - logger.warning( - "Channel %s dropping message %s from publisher %s because its SHM generation is stale", - self.id, - msg_id, - self.pub_id, - ) - self._release_backpressure(msg_id, self.id) - continue - - shm_buf = self.shm[buf_idx] - # The slot for this msg_id may be uninitialized after a - # mid-stream resize; msg_id() raises UninitializedMemory in - # that case. Treat it as a mismatch (drop + release) instead - # of letting it escape and silently kill the channel task. - try: - slot_msg_id = MessageMarshal.msg_id(shm_buf) - except UninitializedMemory: - slot_msg_id = None - if slot_msg_id != msg_id: - logger.warning( - "Channel %s skipping stale SHM contents for message %s from publisher %s; will use next valid SHM generation", - self.id, - msg_id, - self.pub_id, - ) - self._release_backpressure(msg_id, self.id) - continue - - self.cache.put_from_mem(shm_buf) - - elif msg == Command.TX_TCP.value: - channel_kind = ProfileChannelType.TCP - buf_size = await read_int(reader) - obj_bytes = await reader.readexactly(buf_size) - assert MessageMarshal.msg_id(obj_bytes) == msg_id - self.cache.put_from_mem(memoryview(obj_bytes).toreadonly()) - - else: - raise ValueError(f"unimplemented data telemetry: {msg}") - - self._set_channel_kind(channel_kind) - - if not self._notify_clients(msg_id): - # Nobody is listening; need to ack! - self.cache.release(msg_id) - self._acknowledge(msg_id) - - except (ConnectionResetError, BrokenPipeError, asyncio.IncompleteReadError): - logger.debug(f"connection fail: channel:{self.id} - pub:{self.pub_id}") - except Exception: - logger.exception( - "Channel %s publisher connection crashed for pub %s", + slot_msg_id = MessageMarshal.msg_id(shm_buf) + except UninitializedMemory: + slot_msg_id = None + if slot_msg_id != msg_id: + logger.warning( + "Channel %s skipping stale SHM contents for message %s from publisher %s; will use next valid SHM generation", self.id, + msg_id, self.pub_id, ) - raise + self._release_backpressure(msg_id, self.id) + return - finally: - self.cache.clear() - if self.shm is not None: - self.shm.close() + self.cache.put_from_mem(shm_buf) + self._set_channel_kind(ProfileChannelType.SHM) + self._finish_delivery(msg_id) - await close_stream_writer(self._pub_writer) + def _deliver_from_tcp(self, msg_id: int, obj_bytes: bytes) -> None: + """ + Cache a message that arrived inline over TCP and notify clients. - logger.debug(f"disconnected: channel:{self.id} -> pub:{self.pub_id}") + Called inline from :meth:`ChannelProtocol.frames_available`. + """ + assert MessageMarshal.msg_id(obj_bytes) == msg_id + self.cache.put_from_mem(memoryview(obj_bytes).toreadonly()) + self._set_channel_kind(ProfileChannelType.TCP) + self._finish_delivery(msg_id) + + def _finish_delivery(self, msg_id: int) -> None: + if not self._notify_clients(msg_id): + # Nobody is listening; need to ack! + self.cache.release(msg_id) + self._acknowledge(msg_id) + + def _on_disconnected(self) -> None: + """ + Release per-connection resources. Invoked from the protocol's + ``connection_lost``, which replaces the old task's ``finally`` block. + """ + self.cache.clear() + if self.shm is not None: + self.shm.close() + logger.debug(f"disconnected: channel:{self.id} -> pub:{self.pub_id}") def _set_channel_kind(self, kind: ProfileChannelType) -> None: if self._channel_kind == ProfileChannelType.UNKNOWN: @@ -480,7 +555,7 @@ def _release_backpressure(self, msg_id: int, client_id: UUID) -> None: def _acknowledge(self, msg_id: int) -> None: try: ack = Command.RX_ACK.value + uint64_to_bytes(msg_id) - self._pub_writer.write(ack) + self._proto.write(ack) except (BrokenPipeError, ConnectionResetError): logger.info(f"ack fail: channel:{self.id} -> pub:{self.pub_id}") diff --git a/tests/test_channel.py b/tests/test_channel.py index 9db30ddb..ffcffdc0 100644 --- a/tests/test_channel.py +++ b/tests/test_channel.py @@ -3,7 +3,8 @@ import pytest -from ezmsg.core.messagechannel import Channel +from ezmsg.core.frameproto import FramedProtocol +from ezmsg.core.messagechannel import Channel, ChannelProtocol from ezmsg.core.messagecache import CacheMiss from ezmsg.core.messagemarshal import MessageMarshal from ezmsg.core.netprotocol import Command, uint64_to_bytes @@ -11,6 +12,8 @@ class DummyWriter: + """Stands in for Channel._proto where only outbound acks matter.""" + def __init__(self): self.buffer: list[bytes] = [] @@ -24,6 +27,38 @@ async def wait_closed(self) -> None: return None +class FakeTransport: + """Minimal asyncio.Transport stand-in for driving a ChannelProtocol.""" + + def __init__(self): + self.buffer: list[bytes] = [] + self.reading = True + + def write(self, data: bytes) -> None: + self.buffer.append(data) + + def is_closing(self) -> bool: + return False + + def close(self) -> None: + return None + + def pause_reading(self) -> None: + self.reading = False + + def resume_reading(self) -> None: + self.reading = True + + def get_extra_info(self, name, default=None): + return default + + +async def _drain_tasks(turns: int = 20): + """Let tasks spawned from within a read callback (e.g. SHM reattach) run.""" + for _ in range(turns): + await asyncio.sleep(0) + + def _resolved_task(): loop = asyncio.get_running_loop() fut = loop.create_future() @@ -56,8 +91,7 @@ def _raw_message(msg_id: int, payload) -> memoryview: @pytest.mark.asyncio async def test_channel_acknowledges_remote_messages(): channel = Channel(uuid4(), uuid4(), 2, None, None, Channel._SENTINEL) - channel._pub_writer = DummyWriter() - channel._pub_task = _resolved_task() + channel._proto = DummyWriter() channel._graph_task = _resolved_task() client_id = uuid4() @@ -84,14 +118,13 @@ async def test_channel_acknowledges_remote_messages(): assert channel.backpressure.buffers[buf_idx].is_empty expected_ack = Command.RX_ACK.value + uint64_to_bytes(msg_id) - assert channel._pub_writer.buffer[-1] == expected_ack + assert channel._proto.buffer[-1] == expected_ack @pytest.mark.asyncio async def test_channel_releases_local_backpressure(monkeypatch): channel = Channel(uuid4(), uuid4(), 2, None, None, Channel._SENTINEL) - channel._pub_writer = DummyWriter() - channel._pub_task = _resolved_task() + channel._proto = DummyWriter() channel._graph_task = _resolved_task() local_bp = Backpressure(channel.num_buffers) @@ -113,7 +146,7 @@ async def test_channel_releases_local_backpressure(monkeypatch): buf_idx = msg_id % channel.num_buffers assert local_bp.buffers[buf_idx].is_empty - assert channel._pub_writer.buffer == [] + assert channel._proto.buffer == [] def test_channel_put_local_requires_local_backpressure(): @@ -128,8 +161,6 @@ async def test_channel_preserves_cached_message_during_shm_reattach(monkeypatch) new_slots = [_raw_message(4, {"value": 4}), _raw_message(1, {"value": 1}), _raw_message(2, {"value": 2})] channel = Channel(uuid4(), uuid4(), 3, None, None, Channel._SENTINEL) - channel._pub_writer = DummyWriter() - channel._pub_task = _resolved_task() channel._graph_task = _resolved_task() preserved_during_reattach = False @@ -147,10 +178,69 @@ async def fake_attach_shm(self, shm_name): monkeypatch.setattr("ezmsg.core.messagechannel.GraphService.attach_shm", fake_attach_shm) - reader = asyncio.StreamReader() - reader.feed_data(Command.TX_SHM.value + uint64_to_bytes(2) + uint64_to_bytes(3) + b"new") - reader.feed_eof() + client_id = uuid4() + queue: asyncio.Queue = asyncio.Queue() + channel.register_client(client_id, queue) - await channel._publisher_connection(reader) + proto = ChannelProtocol() + proto.connection_made(FakeTransport()) + proto.bind(channel) + channel._proto = proto + proto.start_dispatch() + + # Naming an unattached segment parks the frame and hands off to a task; the + # frame is then re-dispatched against the new segment. + proto.data_received( + Command.TX_SHM.value + uint64_to_bytes(2) + uint64_to_bytes(3) + b"new" + ) + await _drain_tasks() assert preserved_during_reattach + assert channel.shm.name == "new" + assert queue.get_nowait() == (channel.pub_id, 2) + assert channel.cache[2] == {"value": 2} + assert not proto.buffer, "frame should be fully consumed" + + +class _EchoProtocol(FramedProtocol): + def frames_available(self) -> None: + self._buffer.clear() + + +@pytest.mark.asyncio +async def test_teardown_completes_when_peer_stops_reading(): + """ + Channel teardown must not wait on a peer that has stopped reading. + + A graceful transport.close() defers connection_lost until queued writes + flush, so a publisher that has already torn down its ack reader could + otherwise leave wait_closed() hanging forever. + """ + + # Server.wait_closed() waits on connection handlers, so the peer needs a + # way out or the test's own teardown would hang and mask the result. + release_peer = asyncio.Event() + + async def deaf_peer(reader, writer): + await release_peer.wait() # accepts, never reads + writer.close() + + server = await asyncio.start_server(deaf_peer, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + try: + _, proto = await asyncio.get_running_loop().create_connection( + _EchoProtocol, "127.0.0.1", port + ) + proto.start_dispatch() + + # Fill the socket and the transport's own write buffer. + for _ in range(200): + proto.write(b"x" * 65536) + assert proto.transport.get_write_buffer_size() > 0 + + proto.abort() + await asyncio.wait_for(proto.wait_closed(), timeout=5.0) + finally: + release_peer.set() + server.close() + await server.wait_closed()