From 9cb3ce40ad4032309ce53c63092355c64f116764 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 26 May 2026 13:57:35 +0300 Subject: [PATCH] Shield cancel-cleanup release in worker pool --- .../haiku/rag/ingester/workers/pool.py | 19 +++++-- tests/ingester/test_workers.py | 52 +++++++++++++++++++ 2 files changed, 66 insertions(+), 5 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 709c7898..418a2771 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -128,11 +128,20 @@ class WorkerPool: try: result = await run_job(self._client, job, sources=self._sources) except asyncio.CancelledError: - # Graceful shutdown cancelled us mid-flight. Release the claim so - # the next process can pick the job up immediately instead of - # waiting on the reaper's claim_timeout_s. - await self._jobs.release_if_claimed(job.id) - logger.info("Job %s released back to queue on cancel", job.id) + # Graceful shutdown cancelled us mid-flight. Release the claim + # under shield so a second cancel (e.g. shutdown_grace_s elapses + # and IngesterApp's wait_for cancels the gather again) can't + # interrupt the SQL update and strand the claim until the reaper + # runs. The reaper is still the backstop, but releasing eagerly + # lets a restart re-pick the job immediately. + try: + await asyncio.shield(self._jobs.release_if_claimed(job.id)) + except asyncio.CancelledError: + logger.info( + "Job %s cancel-cleanup interrupted; reaper will reclaim", job.id + ) + else: + logger.info("Job %s released back to queue on cancel", job.id) raise except PermanentError as e: await self._jobs.mark_dead(job.id, str(e)) diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index 74fa65b2..b115e305 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -314,6 +314,58 @@ async def test_shutdown_grace_timeout_releases_claim(client, jobs, sync): assert refreshed.attempts == 0 +@pytest.mark.asyncio +async def test_cancel_cleanup_survives_second_cancel(client, jobs, sync, monkeypatch): + """A second cancel arriving while the cancel-handler is awaiting + release_if_claimed must not strand the claim. The shielded await may + raise CancelledError, but the underlying SQL update keeps running and + completes the release as an orphan task.""" + release_entered = asyncio.Event() + release_done = asyncio.Event() + + real_release = jobs.release_if_claimed + + async def _slow_release(job_id): + release_entered.set() + # Long enough for the second cancel to arrive mid-update. + await asyncio.sleep(0.2) + result = await real_release(job_id) + release_done.set() + return result + + monkeypatch.setattr(jobs, "release_if_claimed", _slow_release) + + async def _hangs_forever(*args, **kwargs): + await asyncio.sleep(60) + return Document(id="doc", content="x", uri="u") + + client.create_document_from_source.side_effect = _hangs_forever + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + + pool = _pool(client, jobs, sync, worker_count=1, max_concurrent=1) + await pool.start() + try: + await asyncio.sleep(0.05) + worker_task = pool._workers[0] + worker_task.cancel() + # Wait until the worker is inside the shielded release call. + await asyncio.wait_for(release_entered.wait(), timeout=1.0) + # Second cancel mid-cleanup. Shield holds the SQL update upright. + worker_task.cancel() + with pytest.raises(asyncio.CancelledError): + await worker_task + # Background release Task still alive; let it finish. + await asyncio.wait_for(release_done.wait(), timeout=1.0) + finally: + await pool.stop() + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.QUEUED + assert refreshed.claimed_by is None + + @pytest.mark.asyncio async def test_double_start_raises(client, jobs, sync): pool = _pool(client, jobs, sync, worker_count=1)