Drain orphan cancel-cleanup releases before closing the queue

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

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)

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)