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..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) @@ -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) 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 ---