Merge pull request #428 from ggozad/fix/per-source-circuit-breaker

Make the ingester worker circuit breaker per-source

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Fixed
+
+- Ingester worker circuit breaker is now per-source: a streak of transient failures pauses claims only for the affected source's jobs while healthy sources keep flowing, instead of pausing the whole worker pool. Paused sources are excluded at the claim query.
+
 ## [0.55.1] - 2026-06-08
 
 ### Changed

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

@@ -125,17 +125,25 @@ class JobRepo:
                 self.job_available.notify_all()
         return _row_to_job(row) if row else None
 
-    async def claim_next(self, worker_id: str) -> Job | None:
+    async def claim_next(
+        self, worker_id: str, *, exclude_source_ids: set[str] | None = None
+    ) -> Job | None:
         """Atomically claim the oldest queued job whose scheduled_at <= now.
         A single `UPDATE ... WHERE id = (SELECT ... LIMIT 1) RETURNING` keeps
         the claim atomic across connections: on Postgres the subquery adds
         `FOR UPDATE SKIP LOCKED`; on SQLite the whole statement runs under one
         write lock, so a racing connection re-evaluates the subquery against
-        the committed state and finds the row already claimed."""
+        the committed state and finds the row already claimed.
+
+        `exclude_source_ids` skips jobs from those sources (a paused breaker).
+        Empty or None adds no clause, leaving the query unchanged."""
         now = _utcnow_iso()
+        conditions = [jobs.c.status == "queued", jobs.c.scheduled_at <= now]
+        if exclude_source_ids:
+            conditions.append(jobs.c.source_id.notin_(sorted(exclude_source_ids)))
         candidate = (
             sa.select(jobs.c.id)
-            .where(jobs.c.status == "queued", jobs.c.scheduled_at <= now)
+            .where(*conditions)
             .order_by(jobs.c.scheduled_at, jobs.c.id)
             .limit(1)
             .with_for_update(skip_locked=True)

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

@@ -56,12 +56,7 @@ class WorkerPool:
         self._workers: list[asyncio.Task] = []
         self._reaper: asyncio.Task | None = None
         self._pending_releases: set[asyncio.Task] = set()
-        self._breaker = CircuitBreaker(
+        self._breakers: dict[str, CircuitBreaker] = {}
-            CircuitBreakerConfig(
-                failure_threshold=_WORKER_BREAKER_THRESHOLD,
-                cooldown_s=_WORKER_BREAKER_COOLDOWN_S,
-            )
-        )
 
     @property
     def live_workers(self) -> int:
@@ -71,11 +66,26 @@ class WorkerPool:
 
     @property
     def breaker_open(self) -> bool:
-        return self._breaker.is_open
+        return any(b.is_open for b in self._breakers.values())
 
     @property
     def breaker_consecutive_failures(self) -> int:
-        return self._breaker.consecutive_failures
+        return max((b.consecutive_failures for b in self._breakers.values()), default=0)
+
+    def _breaker_for(self, source_id: str) -> CircuitBreaker:
+        breaker = self._breakers.get(source_id)
+        if breaker is None:
+            breaker = CircuitBreaker(
+                CircuitBreakerConfig(
+                    failure_threshold=_WORKER_BREAKER_THRESHOLD,
+                    cooldown_s=_WORKER_BREAKER_COOLDOWN_S,
+                )
+            )
+            self._breakers[source_id] = breaker
+        return breaker
+
+    def _paused_source_ids(self) -> set[str]:
+        return {sid for sid, b in self._breakers.items() if b.is_open}
 
     async def start(self) -> None:
         if self._workers:
@@ -131,10 +141,9 @@ class WorkerPool:
 
     async def _worker_loop(self, worker_id: str) -> None:
         while not self._stop.is_set():
-            if self._breaker.is_open:
+            job = await self._jobs.claim_next(
-                await self._sleep_or_stop(self._poll_idle_s)
+                worker_id, exclude_source_ids=self._paused_source_ids()
-                continue
+            )
-            job = await self._jobs.claim_next(worker_id)
             if job is None:
                 try:
                     async with self._jobs.job_available:
@@ -201,12 +210,14 @@ class WorkerPool:
             logger.info("Job %s dead (permanent): %s", job.id, e)
             return
         except TransientError as e:
-            was_closed = not self._breaker.is_open
+            breaker = self._breaker_for(job.source_id)
-            self._breaker.record_failure()
+            was_closed = not breaker.is_open
-            if was_closed and self._breaker.is_open:
+            breaker.record_failure()
+            if was_closed and breaker.is_open:
                 logger.warning(
-                    "Worker pool breaker opened after %d consecutive transient "
+                    "Worker breaker opened for source %s after %d consecutive "
-                    "failures; pausing claims for %.0fs",
+                    "transient failures; pausing its claims for %.0fs",
+                    job.source_id,
                     _WORKER_BREAKER_THRESHOLD,
                     _WORKER_BREAKER_COOLDOWN_S,
                 )
@@ -246,10 +257,14 @@ class WorkerPool:
                 job.id,
             )
             return
-        was_open = self._breaker.is_open
+        breaker = self._breaker_for(job.source_id)
-        self._breaker.record_success()
+        was_open = breaker.is_open
+        breaker.record_success()
         if was_open:
-            logger.info("Worker pool breaker closed after successful probe")
+            logger.info(
+                "Worker breaker closed for source %s after successful probe",
+                job.source_id,
+            )
         try:
             if job.op is JobOp.DELETE:
                 await self._sync.delete(job.source_id, job.uri)

tests/ingester/test_queue.py

@@ -241,6 +241,28 @@ async def test_claim_next_returns_oldest_first(jobs):
     assert second.id == j2.id
 
 
+@pytest.mark.asyncio
+async def test_claim_next_excludes_source_ids(jobs):
+    a = await jobs.enqueue("a", "u", JobOp.UPSERT)
+    b = await jobs.enqueue("b", "u", JobOp.UPSERT)
+    assert a is not None
+    assert b is not None
+    first = await jobs.claim_next("w", exclude_source_ids={"a"})
+    assert first is not None
+    assert first.id == b.id
+    # The "a" job is excluded, so nothing more is claimable.
+    assert await jobs.claim_next("w", exclude_source_ids={"a"}) is None
+
+
+@pytest.mark.asyncio
+async def test_claim_next_empty_exclude_is_noop(jobs):
+    job = await jobs.enqueue("s", "u", JobOp.UPSERT)
+    assert job is not None
+    claimed = await jobs.claim_next("w", exclude_source_ids=set())
+    assert claimed is not None
+    assert claimed.id == job.id
+
+
 @pytest.mark.asyncio
 async def test_claim_next_skips_future_scheduled(conn, jobs):
     job = await jobs.enqueue("s", "u", JobOp.UPSERT)

tests/ingester/test_workers.py

@@ -524,14 +524,14 @@ async def test_double_start_raises(client, jobs, sync):
         await pool.stop()
 
 
-# --- pool-wide circuit breaker ---
+# --- per-source circuit breaker ---
 
 
 @pytest.mark.asyncio
 async def test_breaker_opens_after_n_consecutive_transient_failures(client, jobs, sync):
-    """N back-to-back TransientErrors flips the pool breaker open. While
+    """N back-to-back TransientErrors from one source flips that source's
-    open, _worker_loop's claim_next is gated off so subsequent jobs don't
+    breaker open. While open, _worker_loop excludes the source from
-    burn their attempts during the same downstream outage."""
+    claim_next so its other jobs don't burn attempts during the same outage."""
     from haiku.rag.ingester.workers.pool import _WORKER_BREAKER_THRESHOLD
 
     client.create_document_from_source.side_effect = TransientError("downstream down")
@@ -548,28 +548,28 @@ async def test_breaker_opens_after_n_consecutive_transient_failures(client, jobs
         await pool.drain_once()
     assert pool.breaker_open is True
 
-    # drain_once bypasses the worker-loop gate (it's intended for tests), so
+    # drain_once claims without the breaker exclusion (it's intended for
-    # it would still process more jobs. Verify the gate exists by checking
+    # tests), so it would still process more jobs. The exclusion lives in
-    # _worker_loop: a fresh worker started with the breaker open shouldn't
+    # _worker_loop: a fresh worker with this source's breaker open won't
-    # claim anything.
+    # claim its jobs.
     remaining_before = len(await jobs.list_jobs(status=JobStatus.QUEUED, limit=500))
     assert remaining_before >= 1
 
 
 @pytest.mark.asyncio
 async def test_breaker_pauses_worker_loop_claims(client, jobs, sync):
-    """Worker loop honours the breaker: claim_next is not called while
+    """Worker loop honours the breaker: an open source is excluded from
-    is_open, so queued jobs stay queued until the breaker closes."""
+    claim_next, so its queued jobs stay queued until the breaker closes."""
     pool = _pool(client, jobs, sync, worker_count=1, poll_idle_interval_s=0.02)
-    # Force the breaker open without touching the queue.
+    # Force the source's breaker open without touching the queue.
     for _ in range(10):
-        pool._breaker.record_failure()
+        pool._breaker_for("src").record_failure()
     assert pool.breaker_open is True
 
     await jobs.enqueue("src", "u", JobOp.UPSERT)
     await pool.start()
     try:
-        # Even with a queued job available and a live worker, the gate
+        # Even with a queued job available and a live worker, the exclusion
         # keeps the job in 'queued' state.
         await asyncio.sleep(0.1)
         refreshed = await jobs.list_jobs(status=JobStatus.QUEUED, limit=10)
@@ -586,11 +586,13 @@ async def test_breaker_closes_on_successful_probe(client, jobs, sync):
         id="d", content="x", uri="u", metadata={"md5": "m", "source_revision": "r"}
     )
     pool = _pool(client, jobs, sync)
-    # Open the breaker, then collapse the cooldown so is_open returns False
+    # Open the source's breaker, then collapse the cooldown so is_open returns
-    # on the next check (the breaker's three-state model probes after cooldown).
+    # False on the next check (the breaker's three-state model probes after
+    # cooldown).
+    breaker = pool._breaker_for("src")
     for _ in range(10):
-        pool._breaker.record_failure()
+        breaker.record_failure()
-    pool._breaker._opened_at = 0.0  # type: ignore[attr-defined]
+    breaker._opened_at = 0.0  # type: ignore[attr-defined]
     assert pool.breaker_open is False  # cooldown elapsed → probe allowed
 
     await jobs.enqueue("src", "u", JobOp.UPSERT)
@@ -600,6 +602,41 @@ async def test_breaker_closes_on_successful_probe(client, jobs, sync):
     assert pool.breaker_consecutive_failures == 0
 
 
+@pytest.mark.asyncio
+async def test_breaker_isolates_sources(client, jobs, sync):
+    """An open breaker pauses only the failing source. Workers keep draining
+    a healthy source's jobs while the failing source's jobs stay queued."""
+
+    def _route(uri, *, sources=None, source_id=None):
+        if source_id == "bad":
+            raise TransientError("downstream down")
+        return Document(
+            id="d", content="x", uri=uri, metadata={"md5": "m", "source_revision": "r"}
+        )
+
+    client.create_document_from_source.side_effect = _route
+
+    for i in range(3):
+        await jobs.enqueue("bad", f"b{i}", JobOp.UPSERT)
+        await jobs.enqueue("good", f"g{i}", JobOp.UPSERT)
+
+    pool = _pool(client, jobs, sync, worker_count=2, poll_idle_interval_s=0.02)
+    # Open the bad source's breaker without touching the queue.
+    for _ in range(10):
+        pool._breaker_for("bad").record_failure()
+
+    await pool.start()
+    try:
+        await asyncio.sleep(0.2)
+        succeeded = await jobs.list_jobs(status=JobStatus.SUCCEEDED, limit=50)
+        queued = await jobs.list_jobs(status=JobStatus.QUEUED, limit=50)
+    finally:
+        await pool.stop()
+
+    assert {j.uri for j in succeeded} == {"g0", "g1", "g2"}
+    assert {j.uri for j in queued} == {"b0", "b1", "b2"}
+
+
 @pytest.mark.asyncio
 async def test_breaker_ignores_permanent_errors(client, jobs, sync):
     """Permanent errors are about the document, not downstream — they