Propagate external cancellation out of Job.close() (#612)

Author: rodrigobnogueira

CHANGES/606.bugfix

@@ -0,0 +1 @@
+Fixed ``Job.close()`` and ``Job.wait()`` swallowing an external cancellation of the calling task (exact on Python 3.11+ via ``Task.cancelling()``; on older versions the cancellation propagates whenever the job is still running).

aiojobs/_job.py

@@ -132,9 +132,28 @@ async def _close(self, timeout: Optional[float]) -> None:
         scheduler = self._scheduler
         try:
             async with asyncio_timeout(timeout):
-                await self._task
+                if sys.version_info >= (3, 11):
+                    await self._task
+                else:
+                    # Cancelling the current task would be forwarded to
+                    # self._task through _fut_waiter, making the two
+                    # cancellation sources indistinguishable; shield the
+                    # job so an external cancellation leaves it running
+                    # and detectable below.
+                    await asyncio.shield(self._task)
         except asyncio.CancelledError:
-            pass
+            # Either the cancelled job finished or the task running
+            # close() was itself cancelled; re-raise in the second case
+            # so callers stay cancellable.
+            if sys.version_info >= (3, 11):
+                ctask = asyncio.current_task()
+                assert ctask is not None
+                if ctask.cancelling() > 0:
+                    raise
+            elif not self._task.done():
+                # The job is still running, so the CancelledError cannot
+                # have come from awaiting it.
+                raise
         except asyncio.TimeoutError as exc:
             if self._explicit:
                 raise

tests/test_job.py

@@ -1,4 +1,5 @@
 import asyncio
+import sys
 from collections.abc import Awaitable
 from contextlib import suppress
 from typing import Callable, NoReturn
@@ -239,6 +240,100 @@ async def coro2() -> None:
     await job.close()
 
 
+async def test_job_close_cancelled_from_outside(scheduler: Scheduler) -> None:
+    """An external cancellation of the task running close() must propagate.
+
+    The CancelledError raised by awaiting the job's cancelled task is
+    expected and swallowed, but a cancellation aimed at the caller of
+    close() itself is not ours to swallow.
+    """
+    cancel_seen = asyncio.Event()
+    unblock = asyncio.Event()
+
+    async def coro() -> None:
+        try:
+            await asyncio.sleep(60)
+        except asyncio.CancelledError:
+            cancel_seen.set()
+            await unblock.wait()
+            raise
+
+    job = await scheduler.spawn(coro())
+    closer = asyncio.ensure_future(job.close())
+    # Once the job saw the cancellation, close() is suspended awaiting
+    # the job task, which is blocked until unblock is set.
+    await cancel_seen.wait()
+    closer.cancel()
+    await asyncio.wait({closer})
+    assert closer.cancelled()
+
+    unblock.set()
+    with suppress(asyncio.CancelledError):
+        await job.wait()
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason="distinguishing this race needs Task.cancelling()",
+)
+async def test_job_close_cancelled_from_outside_completion_race(
+    scheduler: Scheduler,
+) -> None:
+    """The job task may finish in the same loop iteration the caller of
+    close() is cancelled in; the external cancellation must still
+    propagate."""
+    cancel_seen = asyncio.Event()
+    unblock = asyncio.Event()
+
+    async def coro() -> None:
+        try:
+            await asyncio.sleep(60)
+        except asyncio.CancelledError:
+            cancel_seen.set()
+            await unblock.wait()
+            raise
+
+    job = await scheduler.spawn(coro())
+    closer = asyncio.ensure_future(job.close())
+    await cancel_seen.wait()
+    # Unblock the job before cancelling the closer: the job task then
+    # completes and schedules its done callbacks before the closer wakes
+    # up and detaches them.
+    unblock.set()
+    closer.cancel()
+    await asyncio.wait({closer})
+    assert closer.cancelled()
+    with suppress(asyncio.CancelledError):
+        await job.wait()
+
+
+async def test_job_close_timeout_source_traceback(
+    make_scheduler: _MakeScheduler,
+) -> None:
+    """In debug mode the close timeout report carries the source traceback."""
+    loop = asyncio.get_running_loop()
+    loop.set_debug(True)
+    try:
+        handler = mock.Mock()
+        scheduler = await make_scheduler(close_timeout=0.01, exception_handler=handler)
+
+        async def coro() -> None:
+            try:
+                await asyncio.sleep(60)
+            except asyncio.CancelledError:
+                await asyncio.sleep(60)
+
+        job = await scheduler.spawn(coro())
+        await scheduler.close()
+        assert job.closed
+        assert handler.called
+        context = handler.call_args[0][1]
+        assert context["message"] == "Job closing timed out"
+        assert "source_traceback" in context
+    finally:
+        loop.set_debug(False)
+
+
 async def test_job_await_closed(scheduler: Scheduler) -> None:
     async def coro() -> int:
         return 5