From 1911b5a6ba2335fe07d82679d2f298522ece6c53 Mon Sep 17 00:00:00 2001 From: Joachim Jablon Date: Wed, 19 Aug 2026 23:25:59 +0200 Subject: [PATCH 1/4] Re-cancel side tasks that ignore their cancellation Worker._shutdown cancelled its side tasks then awaited them unbounded. A cancellation can be lost when it lands as the task takes its first step: the task keeps looping and, since a side task only wakes on its own schedule, the gather waits for it for ever. Observed with the periodic deferrer, whose sleep runs to the next cron tick, but the heartbeat and abort-polling tasks have the same shape. This is reachable from any cancellation-based shutdown, including the documented asyncio.wait_for(app.run_worker_async(), timeout=...): run() turns a cancellation into a graceful stop, so the hanging gather blocks it for good. Re-cancel while tasks are still pending; the second cancellation lands on an ordinary sleep and takes effect. Give up after CANCEL_ATTEMPTS rather than hanging, and log instead of raising: _shutdown runs in the run loop's finally, so raising would mask whatever caused the shutdown and would turn a cancelled worker into an error rather than a CancelledError. Co-Authored-By: Claude Opus 5 --- procrastinate/utils.py | 36 ++++++++++++++++++++++++ tests/unit/test_utils.py | 60 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/procrastinate/utils.py b/procrastinate/utils.py index 9ad5c42f6..cafaec029 100644 --- a/procrastinate/utils.py +++ b/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()): diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index beadb6dc6..aa66e28dc 100644 --- a/tests/unit/test_utils.py +++ b/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")] ) From 41ce190623cb1c5927f75fb61a7227d3e17375aa Mon Sep 17 00:00:00 2001 From: Joachim Jablon Date: Wed, 19 Aug 2026 23:26:00 +0200 Subject: [PATCH 2/4] Wait for the periodic defer instead of budgeting 0.1s for it test_procrastinate_periodic_defers ran the worker for a fixed 0.1 s and then indexed the resulting row, so a runner too slow to finish worker startup in that window failed with IndexError rather than a useful message. Run the worker until the deferrer has actually written its row. This needs the side task cancellation fix from the previous commit: cancelling once the row lands means cancelling during startup, which is exactly when the deferrer's cancellation could be lost. Co-Authored-By: Claude Opus 5 --- .../integration/contrib/django/test_models.py | 30 ++++++++++++++----- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/tests/integration/contrib/django/test_models.py b/tests/integration/contrib/django/test_models.py index d3b302f67..dae8881b5 100644 --- a/tests/integration/contrib/django/test_models.py +++ b/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" From 04cd2e37f76777919e5cb95fc6067dcc1e8fc586 Mon Sep 17 00:00:00 2001 From: Joachim Jablon Date: Sat, 22 Aug 2026 08:54:10 +0200 Subject: [PATCH 3/4] Cancel and wait in the same round, report errors before abandoning Two review findings on the cancellation loop: The initial cancellation sat outside the loop, so the last cancellation the loop sent was never waited on. Reaping was unaffected (that cancellation was beyond the budget either way) but it made CANCEL_ATTEMPTS mean "N waits and N+1 cancellations". Cancel at the top of each round instead, so the constant means N cancel-and-wait rounds and no cancellation goes unwaited. Giving up returned before errors were captured, so one task refusing to stop discarded another task's exception, which was then never retrieved. Capture the errors of the tasks that did stop before abandoning the rest. Co-Authored-By: Claude Opus 5 --- procrastinate/utils.py | 50 ++++++++++++++++++-------------- tests/unit/test_utils.py | 62 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 22 deletions(-) diff --git a/procrastinate/utils.py b/procrastinate/utils.py index cafaec029..398ab2ffb 100644 --- a/procrastinate/utils.py +++ b/procrastinate/utils.py @@ -230,49 +230,55 @@ def log_task_exception(task: asyncio.Task[Any], error: BaseException): }, ) + 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 - 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 + # 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). - # 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 + # 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: - logger.debug( - f"Task {task.get_name()} ignored its cancellation, cancelling again", - extra={"action": "cancel_task_again", "task_name": task.get_name()}, - ) + 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: - stuck = [task for task in tasks if not task.done()] + # 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(task.get_name() for task in stuck)}", + f"{', '.join(sorted(task.get_name() for task in pending))}", extra={ "action": "cancel_tasks_failed", - "task_names": [task.get_name() for task in stuck], + "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]): diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index aa66e28dc..21ff4694b 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -411,6 +411,68 @@ async def stubborn_task(): 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) From 14a639659a211b263efbca41e93a64164b3e4468 Mon Sep 17 00:00:00 2001 From: Joachim Jablon Date: Sat, 22 Aug 2026 09:35:32 +0200 Subject: [PATCH 4/4] Bound the wait on worker shutdown in the periodic defer test The test cancelled the worker and awaited it with no bound, so a worker that refused to shut down would hang the job rather than fail the test. That is worth avoiding in a test whose whole point is a shutdown that used to hang: a hung job holds a runner until CI's own timeout and, with the matrix cancelling siblings, reports as several versions failing for no visible reason. Co-Authored-By: Claude Opus 5 --- tests/integration/contrib/django/test_models.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/contrib/django/test_models.py b/tests/integration/contrib/django/test_models.py index dae8881b5..5eeae3a75 100644 --- a/tests/integration/contrib/django/test_models.py +++ b/tests/integration/contrib/django/test_models.py @@ -138,8 +138,10 @@ async def wait_for_periodic_defer(): 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 worker + await asyncio.wait_for(worker, timeout=10) except asyncio.CancelledError: pass