diff --git a/procrastinate/utils.py b/procrastinate/utils.py index 9ad5c42f6..398ab2ffb 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,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]): diff --git a/tests/integration/contrib/django/test_models.py b/tests/integration/contrib/django/test_models.py index d3b302f67..5eeae3a75 100644 --- a/tests/integration/contrib/django/test_models.py +++ b/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" diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index beadb6dc6..21ff4694b 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,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")] )