From dcef2ea40ec8c66f4401a49662392d0e226ad003 Mon Sep 17 00:00:00 2001 From: ShivaniSharma Date: Sat, 31 Oct 2026 00:30:44 +0530 Subject: [PATCH 1/2] Fix watch subscription timeout handling --- python/src/docket/cli/_support.py | 38 +++++++++++++++---- python/tests/cli/test_iterate_with_timeout.py | 20 ++++++++++ 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/python/src/docket/cli/_support.py b/python/src/docket/cli/_support.py index 582dd00a..c2f86ac9 100644 --- a/python/src/docket/cli/_support.py +++ b/python/src/docket/cli/_support.py @@ -24,27 +24,49 @@ async def iterate_with_timeout( """Iterate over an async iterator with timeout, ensuring proper cleanup. Wraps an async iterator to add timeout support and guaranteed cleanup. - On timeout, yields None to allow the caller to handle polling fallback. + On timeout, yields None without cancelling the pending read. Args: - iterator: An async generator (must have __anext__ and aclose methods) - timeout: Timeout in seconds for each iteration + iterator: An async iterator with __anext__ and aclose methods. + timeout: Timeout in seconds for each iteration. Yields: - Items from the iterator, or None if timeout expires + Items from the iterator, or None if the timeout expires. """ import asyncio + pending: asyncio.Task[dict[str, Any]] | None = None + try: while True: - try: - yield await asyncio.wait_for(iterator.__anext__(), timeout=timeout) - except asyncio.TimeoutError: - # Yield None to signal timeout, allowing caller to handle polling + if pending is None: + pending = asyncio.create_task(iterator.__anext__()) + + done, _ = await asyncio.wait({pending}, timeout=timeout) + + if not done: yield None + continue + + task = pending + pending = None + + try: + item = task.result() except StopAsyncIteration: break + + yield item finally: + if pending is not None: + if not pending.done(): + pending.cancel() + + try: + await pending + except (asyncio.CancelledError, StopAsyncIteration): + pass + await iterator.aclose() diff --git a/python/tests/cli/test_iterate_with_timeout.py b/python/tests/cli/test_iterate_with_timeout.py index 2ef326ae..10656691 100644 --- a/python/tests/cli/test_iterate_with_timeout.py +++ b/python/tests/cli/test_iterate_with_timeout.py @@ -99,3 +99,23 @@ async def aclose(self) -> None: pass # pragma: no cover - exception always raised before this assert close_called, "aclose() should be called even on exception" + + +async def test_iterate_with_timeout_preserves_iterator_after_timeout(): + """A timeout must not cancel the pending read from the iterator.""" + + async def delayed_iterator(): + await asyncio.sleep(0.3) + yield {"value": "completed"} + + results: list[dict[str, Any] | None] = [] + + async for item in iterate_with_timeout( + delayed_iterator(), timeout=0.1 + ): + results.append(item) + if item is not None: + break + + assert results[-1] == {"value": "completed"} + assert None in results[:-1] From f8ba2854288c4460e838278d47ed1faef192714b Mon Sep 17 00:00:00 2001 From: ShivaniSharma Date: Sat, 31 Oct 2026 01:01:06 +0530 Subject: [PATCH 2/2] Ensure iterator cleanup after failed pending read --- python/src/docket/cli/_support.py | 21 ++++++----- python/tests/cli/test_iterate_with_timeout.py | 37 +++++++++++++++++++ 2 files changed, 48 insertions(+), 10 deletions(-) diff --git a/python/src/docket/cli/_support.py b/python/src/docket/cli/_support.py index c2f86ac9..3595b3b6 100644 --- a/python/src/docket/cli/_support.py +++ b/python/src/docket/cli/_support.py @@ -58,16 +58,17 @@ async def iterate_with_timeout( yield item finally: - if pending is not None: - if not pending.done(): - pending.cancel() - - try: - await pending - except (asyncio.CancelledError, StopAsyncIteration): - pass - - await iterator.aclose() + try: + if pending is not None: + if not pending.done(): + pending.cancel() + + try: + await pending + except (asyncio.CancelledError, StopAsyncIteration): + pass + finally: + await iterator.aclose() class LogLevel(str, enum.Enum): diff --git a/python/tests/cli/test_iterate_with_timeout.py b/python/tests/cli/test_iterate_with_timeout.py index 10656691..a85c3db9 100644 --- a/python/tests/cli/test_iterate_with_timeout.py +++ b/python/tests/cli/test_iterate_with_timeout.py @@ -119,3 +119,40 @@ async def delayed_iterator(): assert results[-1] == {"value": "completed"} assert None in results[:-1] + + +async def test_iterate_with_timeout_closes_iterator_when_pending_read_fails(): + """Close the iterator even if the pending read raises an exception.""" + + class FailingIterator: + def __init__(self): + self.fail_read = asyncio.Event() + self.read_finished = asyncio.Event() + self.close_called = False + + def __aiter__(self): + return self + + async def __anext__(self): + await self.fail_read.wait() + try: + raise RuntimeError("read failed") + finally: + self.read_finished.set() + + async def aclose(self): + self.close_called = True + + iterator = FailingIterator() + wrapped = iterate_with_timeout(iterator, timeout=0.01) + + assert await wrapped.__anext__() is None + + iterator.fail_read.set() + await iterator.read_finished.wait() + await asyncio.sleep(0) + + with pytest.raises(RuntimeError, match="read failed"): + await wrapped.aclose() + + assert iterator.close_called \ No newline at end of file