diff --git a/procrastinate/worker.py b/procrastinate/worker.py index 7dbfa93d3..316d4c91a 100644 --- a/procrastinate/worker.py +++ b/procrastinate/worker.py @@ -622,24 +622,28 @@ def _start_side_tasks(self) -> list[asyncio.Task[Any]]: async def _monitor_side_tasks(self, side_tasks: list[asyncio.Task[Any]]): """Monitor side tasks and stop the worker if any task fails""" try: - done, _pending = await asyncio.wait( - side_tasks, return_when=asyncio.FIRST_COMPLETED - ) - for task in done: - if exc := task.exception(): - self.logger.error( - f"Side task {task.get_name()} failed with exception: {exc}, stopping worker", - extra=self._log_extra( - action="side_task_failed", - context=None, - job_result=None, - task_name=task.get_name(), - exception=str(exc), - ), - exc_info=exc, - ) - self.stop() - return + pending = set(side_tasks) + while pending: + done, pending = await asyncio.wait( + pending, return_when=asyncio.FIRST_COMPLETED + ) + for task in done: + if task.cancelled(): + continue + if exc := task.exception(): + self.logger.error( + f"Side task {task.get_name()} failed with exception: {exc}, stopping worker", + extra=self._log_extra( + action="side_task_failed", + context=None, + job_result=None, + task_name=task.get_name(), + exception=str(exc), + ), + exc_info=exc, + ) + self.stop() + return except Exception as exc: self.logger.exception( f"Side task monitor failed: {exc}", diff --git a/tests/unit/test_worker.py b/tests/unit/test_worker.py index bb949b35d..57e99dce2 100644 --- a/tests/unit/test_worker.py +++ b/tests/unit/test_worker.py @@ -905,3 +905,26 @@ async def failing_update_heartbeat(self): assert "Simulated heartbeat failure" in error_record.message assert "stopping worker" in error_record.message assert error_record.task_name == "update_heartbeats" + + +async def test_worker_stops_when_side_task_fails_after_another_returned(app: App): + async def fail_later(): + await asyncio.sleep(0.01) + raise ValueError("Simulated side task failure") + + worker = Worker(app, install_signal_handlers=False) + + # The periodic deferrer returns immediately when no periodic task is + # registered, which must not end supervision of the other side tasks. + returns_immediately = asyncio.create_task(asyncio.sleep(0), name="deferrer_like") + fails_later = asyncio.create_task(fail_later(), name="fails_later") + + monitor = asyncio.create_task( + worker._monitor_side_tasks([returns_immediately, fails_later]) + ) + await asyncio.sleep(0.05) + + assert fails_later.done() + assert worker._stop_event.is_set() + + monitor.cancel()