Merge pull request #391 from mcdonc/perf/event-driven-worker-wakeup

perf: event-driven worker wakeup via asyncio.Condition

haiku_rag_slim/haiku/rag/ingester/queue/repository.py

@@ -66,6 +66,9 @@ class JobRepo:
         # JobRepo and SyncStateRepo share the same connection, callers must
         # pass the same lock instance so cross-repo calls also serialize.
         self._lock = lock or asyncio.Lock()
+        # Notified after a successful enqueue so workers can wake
+        # immediately instead of polling on a fixed sleep interval.
+        self.job_available = asyncio.Condition()
 
     async def enqueue(
         self,
@@ -110,6 +113,9 @@ class JobRepo:
             ) as cursor:
                 row = await cursor.fetchone()
             await self._conn.commit()
+        if row is not None:
+            async with self.job_available:
+                self.job_available.notify_all()
         return _row_to_job(row) if row else None
 
     async def claim_next(self, worker_id: str) -> Job | None:

haiku_rag_slim/haiku/rag/ingester/workers/pool.py

@@ -92,6 +92,10 @@ class WorkerPool:
 
     async def stop(self) -> None:
         self._stop.set()
+        # Wake workers parked on job_available.wait() so they notice _stop
+        # immediately instead of sleeping out the full poll_idle interval.
+        async with self._jobs.job_available:
+            self._jobs.job_available.notify_all()
         tasks = list(self._workers)
         if self._reaper is not None:
             tasks.append(self._reaper)
@@ -130,7 +134,14 @@ class WorkerPool:
                 continue
             job = await self._jobs.claim_next(worker_id)
             if job is None:
-                await self._sleep_or_stop(self._poll_idle_s)
+                try:
+                    async with self._jobs.job_available:
+                        await asyncio.wait_for(
+                            self._jobs.job_available.wait(),
+                            timeout=self._poll_idle_s,
+                        )
+                except TimeoutError:
+                    pass
                 continue
             await self._process(job)
 

tests/ingester/test_workers.py

@@ -58,6 +58,47 @@ def _pool(client, jobs, sync, **kwargs) -> WorkerPool:
     )
 
 
+# --- event-driven wakeup ---
+
+
+@pytest.mark.asyncio
+async def test_idle_worker_picks_up_job_quickly(client, jobs, sync):
+    """An idle worker should wake up well under poll_idle_s when a job is
+    enqueued, thanks to the job_available condition notification."""
+    client.create_document_from_source.return_value = Document(
+        id="d", content="x", uri="u", metadata={"md5": "m", "source_revision": "r"}
+    )
+    pool = _pool(client, jobs, sync, worker_count=1, poll_idle_interval_s=5.0)
+    await pool.start()
+    try:
+        await jobs.enqueue("src", "u", JobOp.UPSERT)
+        for _ in range(50):
+            listed = await jobs.list_jobs(status=JobStatus.SUCCEEDED)
+            if listed:
+                break
+            await asyncio.sleep(0.05)
+        assert len(listed) == 1, "job was not picked up within 2.5s"
+    finally:
+        await pool.stop()
+
+
+@pytest.mark.asyncio
+async def test_stop_completes_with_idle_workers(client, jobs, sync):
+    """stop() must notify workers parked on job_available.wait() so they
+    exit promptly. Without the notify, workers sleep for the full
+    poll_idle_interval_s before noticing _stop."""
+    pool = _pool(client, jobs, sync, worker_count=2, poll_idle_interval_s=10.0)
+    await pool.start()
+    await asyncio.sleep(0.1)
+    try:
+        await asyncio.wait_for(pool.stop(), timeout=2.0)
+    except TimeoutError:
+        pytest.fail(
+            "stop() did not complete within 2s — idle workers were not woken"
+        )
+    assert pool.live_workers == 0
+
+
 # --- drain_once: covers _process logic deterministically ---