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,9 +230,41 @@ def log_task_exception(task: asyncio.Task[Any], error: BaseException):
             },
         )
 
+    if not tasks:
+        return
+
     for task in tasks:
         task.cancel()
 
+    # A cancellation is lost if it lands just as the task takes its first step: the
+    # task then keeps looping and gather() below would wait for it for ever, since a
+    # side task only wakes on its own schedule (up to a whole cron period away).
+    # Re-cancel until the tasks really stop.
+    for _attempt in range(CANCEL_ATTEMPTS):
+        _done, pending = await asyncio.wait(tasks, timeout=CANCEL_TIMEOUT)
+        if not pending:
+            break
+        for task in pending:
+            logger.debug(
+                f"Task {task.get_name()} ignored its cancellation, cancelling again",
+                extra={"action": "cancel_task_again", "task_name": task.get_name()},
+            )
+            task.cancel()
+    else:
+        stuck = [task for task in tasks if not task.done()]
+        # 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(task.get_name() for task in stuck)}",
+            extra={
+                "action": "cancel_tasks_failed",
+                "task_names": [task.get_name() for task in stuck],
+            },
+        )
+        return
+
     await asyncio.gather(*tasks, return_exceptions=True)
 
     for task in (task for task in tasks if task.done() and not task.cancelled()):

tests/integration/contrib/django/test_models.py

@@ -115,19 +115,35 @@ 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()
+                try:
+                    await worker
+                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,65 @@ 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_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")]
 )