diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index 4e1ba821..6747d5bf 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -144,6 +144,16 @@ class IngesterApp: "claim_timeout_s on next start", grace_s, ) + # Drain any cancel-cleanup release Tasks the worker pool + # spawned but didn't get to await (timeout path). They + # need the queue connection that the outer finally is + # about to close. + landed = await self._pool.drain_pending_releases(timeout=2.0) + if landed: + logger.info( + "Drained %d cancel-cleanup release(s) before close", + landed, + ) finally: # Close the queue connection unconditionally. aiosqlite runs the # underlying sqlite3 in a background thread; leaving it open holds diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 9e8bdacb..6bfab434 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -52,6 +52,12 @@ class WorkerPool: self._stop = asyncio.Event() self._workers: list[asyncio.Task] = [] self._reaper: asyncio.Task | None = None + # Tracks release_if_claimed Tasks spawned by the cancel-cleanup path. + # The worker task may exit (a second cancel during shutdown_grace + # timeout) before its release Task finishes; the lifecycle owner + # drains these via drain_pending_releases() so the SQL update lands + # before the queue connection closes. + self._pending_releases: set[asyncio.Task] = set() @property def live_workers(self) -> int: @@ -77,6 +83,18 @@ class WorkerPool: self._workers.clear() self._reaper = None + async def drain_pending_releases(self, timeout: float = 2.0) -> int: + """Wait for any in-flight cancel-cleanup release Tasks to finish. + Returns how many completed. Called by the lifecycle owner after + stop() (success or timeout) so orphans land their SQL update before + the queue connection closes. Tasks left running after `timeout` + will be reclaimed by the reaper on the next start instead.""" + pending = list(self._pending_releases) + if not pending: + return 0 + done, _ = await asyncio.wait(pending, timeout=timeout) + return len(done) + async def drain_once(self, worker_id: str = "drain") -> int: """Drain every currently-claimable job to completion on the calling coroutine. Used by run-once and by tests; not by `start()`.""" @@ -130,17 +148,24 @@ 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 - # 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. + # Graceful shutdown cancelled us mid-flight. Spawn the release + # as an independent Task tracked in _pending_releases — that + # way a second cancel (e.g. shutdown_grace_s elapses and + # wait_for cancels stop() again) can interrupt our await + # without interrupting the SQL update, and the lifecycle owner + # can drain the orphans before closing the queue connection. + release_task = asyncio.create_task( + self._jobs.release_if_claimed(job.id, worker_id) + ) + self._pending_releases.add(release_task) + release_task.add_done_callback(self._pending_releases.discard) try: - await asyncio.shield(self._jobs.release_if_claimed(job.id, worker_id)) + await asyncio.shield(release_task) except asyncio.CancelledError: logger.info( - "Job %s cancel-cleanup interrupted; reaper will reclaim", job.id + "Job %s cancel-cleanup interrupted; orphan release Task " + "will be drained by the lifecycle owner", + job.id, ) else: logger.info("Job %s released back to queue on cancel", job.id) diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index 6c1c3d97..bb59f365 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -398,6 +398,65 @@ async def test_cancel_cleanup_survives_second_cancel(client, jobs, sync, monkeyp assert refreshed.claimed_by is None +@pytest.mark.asyncio +async def test_drain_pending_releases_waits_for_orphan_releases( + client, jobs, sync, monkeypatch +): + """When the worker is cancelled twice (shutdown_grace_s timeout path) it + exits before its release_if_claimed Task completes, leaving an orphan. + drain_pending_releases waits for that orphan so the SQL update lands + before the lifecycle owner closes the queue connection.""" + real_release = jobs.release_if_claimed + release_entered = asyncio.Event() + + async def _slow_release(job_id, claimed_by): + release_entered.set() + await asyncio.sleep(0.15) + return await real_release(job_id, claimed_by) + + 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() + await asyncio.wait_for(release_entered.wait(), timeout=1.0) + worker_task.cancel() + with pytest.raises(asyncio.CancelledError): + await worker_task + + # At this point the orphan release is still running. drain returns + # the count that landed within the timeout. + assert len(pool._pending_releases) == 1 + landed = await pool.drain_pending_releases(timeout=1.0) + assert landed == 1 + assert pool._pending_releases == set() + finally: + await pool.stop() + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.QUEUED + + +@pytest.mark.asyncio +async def test_drain_pending_releases_with_no_orphans_is_noop(client, jobs, sync): + """Common case: nothing to drain — drain returns 0 immediately, no + asyncio.wait against an empty set.""" + pool = _pool(client, jobs, sync, worker_count=1, max_concurrent=1) + assert await pool.drain_pending_releases() == 0 + + @pytest.mark.asyncio async def test_double_start_raises(client, jobs, sync): pool = _pool(client, jobs, sync, worker_count=1)