From c0b9ec0f0faf5447580798aa410a8ddbecf2335c Mon Sep 17 00:00:00 2001
From: Sanjay Santhanam <51058514+Sanjays2402@users.noreply.github.com>
Date: Fri, 31 Jul 2026 03:17:59 -0700
Subject: [PATCH] Keep monitoring side tasks after one returns cleanly

_monitor_side_tasks awaited asyncio.wait(..., FIRST_COMPLETED) once and
only acted when a completed task had raised. A side task that returned
cleanly therefore ended supervision of every other side task for the rest
of the worker's life.

This is the normal path, not an edge case: PeriodicDeferrer.worker()
returns immediately when no periodic task is registered, so any app
without periodic tasks loses side-task supervision at startup and a later
listener or heartbeat failure no longer stops the worker.

The monitor now loops over the remaining pending tasks until one fails or
all of them have finished, skipping cancelled tasks.

Adds a regression test in tests/unit/test_worker.py asserting the worker
is stopped by a side task that fails after another one returned cleanly.

Closes #1597
---
 procrastinate/worker.py   | 40 +++++++++++++++++++++------------------
 tests/unit/test_worker.py | 23 ++++++++++++++++++++++
 2 files changed, 45 insertions(+), 18 deletions(-)

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