From 4fa8a012b4cc9e52591c405896da6f0b79bd76d8 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 1 Jun 2026 06:28:16 -0400 Subject: [PATCH 1/2] Use asyncio.Condition for event-driven worker wakeup Workers previously polled the queue with a fixed 1s sleep between claim attempts, adding ~500ms average latency to job pickup. Now JobRepo.job_available (an asyncio.Condition) is notified on every successful enqueue, waking idle workers immediately. The poll interval remains as a timeout fallback for stop signals and breaker state changes. --- haiku_rag_slim/haiku/rag/ingester/queue/repository.py | 6 ++++++ haiku_rag_slim/haiku/rag/ingester/workers/pool.py | 9 ++++++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index c2801e53..799c0529 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/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: diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 122866b8..32bb6056 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -130,7 +130,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) From 720ff233577d7efab4d409ef259c22a05df7d069 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 1 Jun 2026 09:21:24 -0400 Subject: [PATCH 2/2] Fix shutdown regression: notify idle workers on stop() Workers parked on job_available.wait() were not woken by stop(), causing them to sleep out the full poll_idle_interval_s before noticing _stop. With the default 1.0s interval, stop() took ~0.8s instead of ~0.007s. Notify all waiters on the condition in stop() so idle workers exit immediately. Add tests for fast job pickup via notification and fast shutdown with idle workers. --- .../haiku/rag/ingester/workers/pool.py | 4 ++ tests/ingester/test_workers.py | 41 +++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 32bb6056..64896ac4 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/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) diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index 0df4cdca..1ea0ea3a 100644 --- a/tests/ingester/test_workers.py +++ b/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 ---