diff --git a/src/acp/task/sender.py b/src/acp/task/sender.py index 613535bf..f3670506 100644 --- a/src/acp/task/sender.py +++ b/src/acp/task/sender.py @@ -23,9 +23,27 @@ def __init__(self, writer: asyncio.StreamWriter, supervisor: TaskSupervisor) -> self._writer = writer self._queue: asyncio.Queue[_PendingSend | None] = asyncio.Queue() self._closed = False + self._active: _PendingSend | None = None + self._terminal_error: BaseException | None = None + self._terminal_cancelled = False self._task = supervisor.create(self._loop(), name="acp.Sender.loop", on_error=self._on_error) + self._task.add_done_callback(self._on_task_done) async def send(self, payload: dict[str, Any]) -> None: + if self._terminal_error is not None: + if self._terminal_cancelled: + raise asyncio.CancelledError + error = self._terminal_error + raise error + if self._closed: + raise asyncio.CancelledError + if self._task.done(): + if self._task.cancelled(): + self._set_terminal(ConnectionError("Sender closed"), cancel_pending=True) + raise asyncio.CancelledError + error = self._task.exception() or ConnectionError("Sender closed") + self._set_terminal(error) + raise error data = (json.dumps(payload, separators=(",", ":")) + "\n").encode("utf-8") future: asyncio.Future[None] = asyncio.get_running_loop().create_future() await self._queue.put(_PendingSend(data, future)) @@ -46,18 +64,61 @@ async def _loop(self) -> None: item = await self._queue.get() if item is None: return + self._active = item try: self._writer.write(item.payload) await self._writer.drain() + except asyncio.CancelledError: + self._set_terminal(ConnectionError("Sender closed"), cancel_pending=True) + raise except Exception as exc: - if not item.future.done(): - item.future.set_exception(exc) + self._set_terminal(exc) raise else: if not item.future.done(): item.future.set_result(None) + finally: + self._active = None except asyncio.CancelledError: + self._set_terminal(ConnectionError("Sender closed"), cancel_pending=True) + raise + + def _set_terminal(self, error: BaseException, *, cancel_pending: bool = False) -> None: + if self._terminal_error is not None: return + self._terminal_error = error + self._terminal_cancelled = cancel_pending + if cancel_pending: + self._cancel_pending() + return + if self._active is not None and not self._active.future.done(): + self._active.future.set_exception(error) + while True: + try: + item = self._queue.get_nowait() + except asyncio.QueueEmpty: + break + if item is not None and not item.future.done(): + item.future.set_exception(error) + + def _cancel_pending(self) -> None: + if self._active is not None and not self._active.future.done(): + self._active.future.cancel() + while True: + try: + item = self._queue.get_nowait() + except asyncio.QueueEmpty: + break + if item is not None and not item.future.done(): + item.future.cancel() + + def _on_task_done(self, task: asyncio.Task[Any]) -> None: + if task.cancelled(): + self._set_terminal(ConnectionError("Sender closed"), cancel_pending=True) + else: + error = task.exception() + if error is not None: + self._set_terminal(error) def _on_error(self, task: asyncio.Task[Any], exc: BaseException) -> None: logging.exception("Send loop failed", exc_info=exc) diff --git a/tests/test_connection_recovery.py b/tests/test_connection_recovery.py index 67d58d30..5505adde 100644 --- a/tests/test_connection_recovery.py +++ b/tests/test_connection_recovery.py @@ -1,8 +1,9 @@ from __future__ import annotations import asyncio +import contextlib import json -from typing import Any +from typing import Any, cast from unittest.mock import AsyncMock, MagicMock import pytest @@ -29,6 +30,20 @@ def _make_connection( return conn, reader +class _FailingWriter: + def __init__(self) -> None: + self.drain_started = asyncio.Event() + self.release = asyncio.Event() + + def write(self, _payload: bytes) -> None: + pass + + async def drain(self) -> None: + self.drain_started.set() + await self.release.wait() + raise BrokenPipeError("writer failed") + + @pytest.mark.asyncio async def test_receive_loop_handles_oversized_frame(caplog: pytest.LogCaptureFixture) -> None: conn, reader = _make_connection(limit=128) @@ -149,3 +164,44 @@ def tracking_process(message: dict[str, Any]) -> None: await conn.close() assert processed == ["survivor"] + + +@pytest.mark.asyncio +async def test_send_request_removes_pending_after_writer_failure() -> None: + reader = asyncio.StreamReader() + writer = _FailingWriter() + conn = Connection(_noop_handler, cast(asyncio.StreamWriter, writer), reader, listening=False) + requests = [ + asyncio.create_task(conn.send_request("first"), name="request-first"), + asyncio.create_task(conn.send_request("second"), name="request-second"), + ] + + try: + await asyncio.wait_for(writer.drain_started.wait(), timeout=1) + for _ in range(100): + if conn._transport._sender._queue.qsize() == 1: # type: ignore[attr-defined] + break + await asyncio.sleep(0) + else: + raise AssertionError("second request was not queued") + + writer.release.set() + done, pending = await asyncio.wait(requests, timeout=1) + try: + assert not pending, f"requests still pending: {pending}" + for request in done: + with pytest.raises(BrokenPipeError, match="writer failed"): + request.result() + assert conn._pending == {} + finally: + for request in pending: + request.cancel() + await asyncio.gather(*requests, return_exceptions=True) + finally: + writer.release.set() + for request in requests: + if not request.done(): + request.cancel() + await asyncio.gather(*requests, return_exceptions=True) + with contextlib.suppress(BrokenPipeError): + await conn.close() diff --git a/tests/test_message_sender.py b/tests/test_message_sender.py new file mode 100644 index 00000000..4808ae52 --- /dev/null +++ b/tests/test_message_sender.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import asyncio +import json +from typing import Any, cast + +import pytest + +from acp.task import MessageSender, TaskSupervisor + + +class _ControlledWriter: + def __init__(self, *, error: BaseException | None = None) -> None: + self.error = error + self.drain_started = asyncio.Event() + self.release = asyncio.Event() + self.writes: list[bytes] = [] + + def write(self, payload: bytes) -> None: + self.writes.append(payload) + + async def drain(self) -> None: + self.drain_started.set() + await self.release.wait() + if self.error is not None: + raise self.error + + +async def _wait_for_queue_size(sender: MessageSender, size: int) -> None: + async def queue_reached_size() -> None: + while sender._queue.qsize() < size: + await asyncio.sleep(0) + + await asyncio.wait_for(queue_reached_size(), timeout=1) + + +async def _assert_tasks_finish(tasks: list[asyncio.Task[Any]]) -> None: + _done, pending = await asyncio.wait(tasks, timeout=1) + try: + assert not pending, f"send tasks still pending: {pending}" + finally: + for task in pending: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +async def _close_supervisor(supervisor: TaskSupervisor, tasks: list[asyncio.Task[Any]]) -> None: + await supervisor.shutdown() + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_writer_failure_settles_active_queued_and_late_sends() -> None: + error = BrokenPipeError("writer failed") + writer = _ControlledWriter(error=error) + supervisor = TaskSupervisor(source="test") + sender = MessageSender(cast(asyncio.StreamWriter, writer), supervisor) + tasks: list[asyncio.Task[Any]] = [] + + try: + first = asyncio.create_task(sender.send({"id": 1}), name="send-1") + tasks.append(first) + await asyncio.wait_for(writer.drain_started.wait(), timeout=1) + + second = asyncio.create_task(sender.send({"id": 2}), name="send-2") + tasks.append(second) + await _wait_for_queue_size(sender, 1) + + writer.release.set() + with pytest.raises(BrokenPipeError, match="writer failed"): + await first + + late = asyncio.create_task(sender.send({"id": 3}), name="send-3") + tasks.append(late) + await _assert_tasks_finish(tasks) + + for task in (second, late): + assert isinstance(task.exception(), BrokenPipeError) + assert str(task.exception()) == "writer failed" + finally: + await _close_supervisor(supervisor, tasks) + + +@pytest.mark.asyncio +async def test_supervisor_shutdown_settles_active_queued_and_late_sends() -> None: + writer = _ControlledWriter() + supervisor = TaskSupervisor(source="test") + sender = MessageSender(cast(asyncio.StreamWriter, writer), supervisor) + tasks: list[asyncio.Task[Any]] = [] + + try: + first = asyncio.create_task(sender.send({"id": 1}), name="send-1") + tasks.append(first) + await asyncio.wait_for(writer.drain_started.wait(), timeout=1) + + second = asyncio.create_task(sender.send({"id": 2}), name="send-2") + tasks.append(second) + await _wait_for_queue_size(sender, 1) + + await supervisor.shutdown() + late = asyncio.create_task(sender.send({"id": 3}), name="send-3") + tasks.append(late) + await _assert_tasks_finish(tasks) + + for task in tasks[:2]: + assert task.cancelled() + assert late.cancelled() + finally: + await _close_supervisor(supervisor, tasks) + + +@pytest.mark.asyncio +async def test_shutdown_before_sender_loop_starts_rejects_later_send() -> None: + writer = _ControlledWriter() + supervisor = TaskSupervisor(source="test") + sender = MessageSender(cast(asyncio.StreamWriter, writer), supervisor) + tasks: list[asyncio.Task[Any]] = [] + + try: + await supervisor.shutdown() + late = asyncio.create_task(sender.send({"id": 1}), name="send-after-shutdown") + tasks.append(late) + await _assert_tasks_finish(tasks) + assert late.cancelled() + finally: + await _close_supervisor(supervisor, tasks) + + +@pytest.mark.asyncio +async def test_healthy_sends_preserve_order_and_graceful_close() -> None: + writer = _ControlledWriter() + supervisor = TaskSupervisor(source="test") + sender = MessageSender(cast(asyncio.StreamWriter, writer), supervisor) + tasks: list[asyncio.Task[Any]] = [] + closing: asyncio.Task[None] | None = None + + try: + tasks.extend([ + asyncio.create_task(sender.send({"id": 1}), name="send-1"), + asyncio.create_task(sender.send({"id": 2}), name="send-2"), + ]) + await asyncio.wait_for(writer.drain_started.wait(), timeout=1) + await _wait_for_queue_size(sender, 1) + + closing = asyncio.create_task(sender.close(), name="sender-close") + await asyncio.sleep(0) + assert not closing.done() + assert [json.loads(payload)["id"] for payload in writer.writes] == [1] + + writer.release.set() + await asyncio.gather(*tasks) + await asyncio.wait_for(closing, timeout=1) + assert [json.loads(payload)["id"] for payload in writer.writes] == [1, 2] + + late = asyncio.create_task(sender.send({"id": 3}), name="send-after-close") + tasks.append(late) + await _assert_tasks_finish([late]) + assert late.cancelled() + finally: + if closing is not None and not closing.done(): + closing.cancel() + if closing is not None: + await asyncio.gather(closing, return_exceptions=True) + await _close_supervisor(supervisor, tasks)