Fix worker shutdown hanging on side tasks that ignore cancellation

procrastinate/utils.py

@@ -212,6 +212,10 @@ async def _inner_coro() -> U:
         return _inner_coro().__await__()
 
 
+CANCEL_ATTEMPTS = 5
+CANCEL_TIMEOUT = 1
+
+
 async def cancel_and_capture_errors(tasks: list[asyncio.Task[Any]]):
     """
     Cancel all tasks and capture any error returned by any of those tasks (except the CancellationError itself)
@@ -226,17 +230,55 @@ def log_task_exception(task: asyncio.Task[Any], error: BaseException):
             },
         )
 
-    for task in tasks:
-        task.cancel()
+    def capture_errors(tasks: Iterable[asyncio.Task[Any]]):
+        for task in (task for task in tasks if task.done() and not task.cancelled()):
+            error = task.exception()
+            if error:
+                log_task_exception(task, error=error)
+            else:
+                logger.debug(f"Cancelled task {task.get_name()}")
+
+    if not tasks:
+        return
+
+    # A cancellation is lost if it lands just as the task takes its first step: the
+    # task then keeps looping and waiting for it below would never return, since a
+    # side task only wakes on its own schedule (up to a whole cron period away).
+    # Cancel again on every round, so each cancellation gets its own chance to be
+    # acted upon.
+    pending: set[asyncio.Task[Any]] = set(tasks)
+    for attempt in range(CANCEL_ATTEMPTS):
+        for task in pending:
+            if attempt:
+                logger.debug(
+                    f"Task {task.get_name()} ignored its cancellation, cancelling again",
+                    extra={"action": "cancel_task_again", "task_name": task.get_name()},
+                )
+            task.cancel()
+
+        _done, pending = await asyncio.wait(pending, timeout=CANCEL_TIMEOUT)
+        if not pending:
+            break
+    else:
+        # The tasks that did stop may have something to report, and their exceptions
+        # would otherwise never be retrieved.
+        capture_errors(task for task in tasks if task not in pending)
+        # Not raising: this runs in the run loop's finally, so an exception here
+        # would mask whatever caused the shutdown, and would turn a cancelled
+        # worker into an error rather than a CancelledError.
+        logger.error(
+            f"Abandoning tasks that did not stop when cancelled: "
+            f"{', '.join(sorted(task.get_name() for task in pending))}",
+            extra={
+                "action": "cancel_tasks_failed",
+                "task_names": sorted(task.get_name() for task in pending),
+            },
+        )
+        return
 
     await asyncio.gather(*tasks, return_exceptions=True)
 
-    for task in (task for task in tasks if task.done() and not task.cancelled()):
-        error = task.exception()
-        if error:
-            log_task_exception(task, error=error)
-        else:
-            logger.debug(f"Cancelled task {task.get_name()}")
+    capture_errors(tasks)
 
 
 async def wait_any(*coros_or_futures: Coroutine[Any, Any, Any] | asyncio.Future[Any]):

tests/integration/contrib/django/test_models.py

@@ -115,19 +115,37 @@ async def test_procrastinate_periodic_defers(db):
     def my_task(timestamp):
         pass
 
+    async def list_periodic_defers():
+        return [
+            element
+            async for element in models.ProcrastinatePeriodicDefer.objects.values().all()
+        ]
+
+    async def wait_for_periodic_defer():
+        while not await list_periodic_defers():
+            await asyncio.sleep(0.01)
+
     django_app = procrastinate.contrib.django.app
     with django_app.replace_connector(
         django_app.connector.get_worker_connector()
     ) as app:
         async with app.open_async():
+            worker = asyncio.create_task(app.run_worker_async())
             try:
-                await asyncio.wait_for(app.run_worker_async(), timeout=0.1)
-            except asyncio.TimeoutError:
-                pass
-
-    periodic_defers = []
-    async for element in models.ProcrastinatePeriodicDefer.objects.values().all():
-        periodic_defers.append(element)
+                # Run the worker until the deferrer has actually written its row.
+                # Giving it a fixed budget instead makes the test depend on worker
+                # startup fitting in that window, which it doesn't on a loaded runner.
+                await asyncio.wait_for(wait_for_periodic_defer(), timeout=5)
+            finally:
+                worker.cancel()
+                # Bounded, so that a worker refusing to shut down fails the test
+                # instead of hanging the job until CI's own timeout.
+                try:
+                    await asyncio.wait_for(worker, timeout=10)
+                except asyncio.CancelledError:
+                    pass
+
+    periodic_defers = await list_periodic_defers()
 
     assert periodic_defers[-1]["periodic_id"] == "bar"
     assert periodic_defers[-1]["task_name"] == "foo"

tests/unit/test_utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import datetime
 import functools
 import logging
@@ -385,6 +386,127 @@ async def task_2():
     assert len(caplog.records) == expected_error_count
 
 
+async def test_cancel_and_capture_errors__task_swallows_first_cancellation(mocker):
+    # A side task can miss its cancellation when it lands as the task takes its
+    # first step. It then goes back to sleep until its next tick, which can be a
+    # whole cron period away, and an unbounded wait would hang the shutdown.
+    mocker.patch.object(utils, "CANCEL_TIMEOUT", 0.01)
+    swallowed = asyncio.Event()
+
+    async def stubborn_task():
+        while True:
+            try:
+                await asyncio.sleep(30)
+            except asyncio.CancelledError:
+                if swallowed.is_set():
+                    raise
+                swallowed.set()
+
+    tasks = [asyncio.create_task(stubborn_task(), name="stubborn")]
+    await asyncio.sleep(0.01)
+
+    await asyncio.wait_for(utils.cancel_and_capture_errors(tasks), timeout=5)
+
+    assert swallowed.is_set()
+    assert tasks[0].cancelled()
+
+
+async def test_cancel_and_capture_errors__task_stops_on_last_attempt(mocker):
+    mocker.patch.object(utils, "CANCEL_TIMEOUT", 0.01)
+    mocker.patch.object(utils, "CANCEL_ATTEMPTS", 3)
+    cancellations = 0
+
+    async def late_task():
+        nonlocal cancellations
+        while True:
+            try:
+                await asyncio.sleep(30)
+            except asyncio.CancelledError:
+                cancellations += 1
+                # Stops only on the last cancellation we're willing to send
+                if cancellations == 3:
+                    raise
+
+    tasks = [asyncio.create_task(late_task(), name="late")]
+    await asyncio.sleep(0.01)
+
+    await asyncio.wait_for(utils.cancel_and_capture_errors(tasks), timeout=5)
+
+    assert cancellations == 3
+    assert tasks[0].cancelled()
+
+
+async def test_cancel_and_capture_errors__captures_errors_despite_stuck_task(
+    caplog, mocker
+):
+    caplog.set_level(logging.ERROR)
+    mocker.patch.object(utils, "CANCEL_TIMEOUT", 0.01)
+    mocker.patch.object(utils, "CANCEL_ATTEMPTS", 2)
+    release = asyncio.Event()
+
+    async def failing_task():
+        raise ValueError("Nope from failing_task")
+
+    async def unkillable_task():
+        while True:
+            try:
+                await asyncio.sleep(30)
+            except asyncio.CancelledError:
+                if release.is_set():
+                    raise
+
+    tasks = [
+        asyncio.create_task(failing_task(), name="failing"),
+        asyncio.create_task(unkillable_task(), name="unkillable"),
+    ]
+    await asyncio.sleep(0.01)
+
+    await asyncio.wait_for(utils.cancel_and_capture_errors(tasks), timeout=5)
+
+    # A task refusing to stop must not hide another task's error
+    assert "Nope from failing_task" in caplog.text
+    assert "unkillable" in caplog.text
+
+    release.set()
+    tasks[1].cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await tasks[1]
+
+
+async def test_cancel_and_capture_errors__task_never_stops(caplog, mocker):
+    caplog.set_level(logging.ERROR)
+    mocker.patch.object(utils, "CANCEL_TIMEOUT", 0.01)
+    mocker.patch.object(utils, "CANCEL_ATTEMPTS", 2)
+
+    # Lets the test reap the task once it has made its point
+    release = asyncio.Event()
+
+    async def unkillable_task():
+        while True:
+            try:
+                await asyncio.sleep(30)
+            except asyncio.CancelledError:
+                if release.is_set():
+                    raise
+
+    tasks = [asyncio.create_task(unkillable_task(), name="unkillable")]
+    await asyncio.sleep(0.01)
+
+    # Gives up rather than hanging for ever
+    await asyncio.wait_for(utils.cancel_and_capture_errors(tasks), timeout=5)
+
+    assert "unkillable" in caplog.text
+
+    release.set()
+    tasks[0].cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await tasks[0]
+
+
+async def test_cancel_and_capture_errors__no_task():
+    await utils.cancel_and_capture_errors([])
+
+
 @pytest.mark.parametrize(
     "queues, result", [(None, "all queues"), (["foo", "bar"], "queues foo, bar")]
 )