diff --git a/python/src/docket/cli/_support.py b/python/src/docket/cli/_support.py index 582dd00a..3595b3b6 100644 --- a/python/src/docket/cli/_support.py +++ b/python/src/docket/cli/_support.py @@ -24,28 +24,51 @@ 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: - 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 2ef326ae..a85c3db9 100644 --- a/python/tests/cli/test_iterate_with_timeout.py +++ b/python/tests/cli/test_iterate_with_timeout.py @@ -99,3 +99,60 @@ 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] + + +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