Keep monitoring side tasks after one returns cleanly

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}",

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()