Fix job abortion during worker graceful shutdown

procrastinate/worker.py

@@ -554,7 +554,12 @@ async def _shutdown(self, side_tasks: list[asyncio.Task[Any]]):
         Gracefully shutdown the worker by cancelling side tasks
         and waiting for all pending jobs.
         """
-        await utils.cancel_and_capture_errors(side_tasks)
+        # Stop side tasks that can be stopped early, but exclude those that
+        # should be stopped *after* running jobs have stopped.
+        side_tasks_to_stop_late = ["poll_jobs_to_abort"]
+        await utils.cancel_and_capture_errors(
+            [t for t in side_tasks if t.get_name() not in side_tasks_to_stop_late]
+        )
 
         now = time.time()
         for context in self._running_jobs.values():
@@ -586,6 +591,9 @@ async def _shutdown(self, side_tasks: list[asyncio.Task[Any]]):
             )
             await self._abort_running_jobs()
 
+        # Now stop *all* side tasks.
+        await utils.cancel_and_capture_errors(side_tasks)
+
         assert self.worker_id is not None
         await self.app.job_manager.unregister_worker(self.worker_id)
         logger.debug(f"Unregistered finished worker {self.worker_id} from the database")

tests/unit/test_worker.py

@@ -583,6 +583,59 @@ async def task_func():
     assert status == Status.ABORTED
 
 
+async def test_abort_async_job_during_graceful_shutdown(app: App, caplog):
+    """
+    Tests that a running job can be successfully aborted after the worker that is running
+    it has received a graceful shutdown request.
+    """
+    caplog.set_level("INFO")
+
+    async def wait_for_msg(msg):
+        """Poll until specified status is seen in log."""
+        while True:
+            for record in caplog.records:
+                if msg in record.msg:
+                    return
+            await asyncio.sleep(0.01)
+
+    async def wait_for_status(job_manager, job_id, status):
+        """Poll until specified job status is seen."""
+        while True:
+            if status == await job_manager.get_job_status_async(job_id):
+                return
+            await asyncio.sleep(0.01)
+
+    @app.task()
+    async def task_func():
+        await asyncio.Event().wait()
+
+    # Defer the task, start a worker to process it, and wait for the job have
+    # DOING status.
+    job_id = await task_func.defer_async()
+    worker = Worker(
+        app,
+        abort_job_polling_interval=0.01,
+        fetch_job_polling_interval=0.01,
+        # Disable listen_notify to test abort polling.
+        listen_notify=False,
+        shutdown_graceful_timeout=None,
+    )
+    run_task = await start_worker(worker)
+    await asyncio.wait_for(wait_for_status(app.job_manager, job_id, Status.DOING), 1)
+
+    # Ask the worker to shutdown, verify it has begun to shut down, and verify
+    # the job still has DOING status.
+    worker.stop()
+    await asyncio.wait_for(wait_for_msg("Waiting for job to finish"), 1)
+    assert await app.job_manager.get_job_status_async(job_id) == Status.DOING
+
+    # Abort the job and verify the abortion was successful.
+    await app.job_manager.cancel_job_by_id_async(job_id, abort=True)
+    await asyncio.wait_for(wait_for_status(app.job_manager, job_id, Status.ABORTED), 1)
+
+    await run_task
+
+
 async def test_abort_async_job_while_finishing(app: App, worker, mocker: MockerFixture):
     """
     Tests that aborting a job after that job completes but before the job status is updated