From c59638767db03865f1ada531c9ad45506f81ec90 Mon Sep 17 00:00:00 2001 From: "Owen T. Heisler" Date: Mon, 27 Jul 2026 00:30:48 +0000 Subject: [PATCH] Fix job abortion during worker graceful shutdown During worker graceful shutdown, stop the "poll_jobs_to_abort" side task *after* running jobs have ended rather than before. Without this change, it is impossible to abort a running job via abort polling after a worker has been asked to stop. Also add a test that fails without this change. Closes . --- procrastinate/worker.py | 10 +++++++- tests/unit/test_worker.py | 53 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/procrastinate/worker.py b/procrastinate/worker.py index 7dbfa93d3..858f962a0 100644 --- a/procrastinate/worker.py +++ b/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") diff --git a/tests/unit/test_worker.py b/tests/unit/test_worker.py index bb949b35d..ef0404b4f 100644 --- a/tests/unit/test_worker.py +++ b/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