danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Make the ingester worker circuit breaker per-source

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-06-09 09:41:54 +03:00
parent 46b8fa3fef
commit bd548837e5
No known key found for this signature in database
5 changed files with 126 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

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

14
haiku_rag_slim/haiku/rag/ingester/queue/repository.py
View file

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

55
haiku_rag_slim/haiku/rag/ingester/workers/pool.py
View file

@ -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(
CircuitBreakerConfig(
failure_threshold=_WORKER_BREAKER_THRESHOLD,
cooldown_s=_WORKER_BREAKER_COOLDOWN_S,
)
)
self._breakers: dict[str, CircuitBreaker] = {}
@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:
await self._sleep_or_stop(self._poll_idle_s)
continue
job = await self._jobs.claim_next(worker_id)
job = await self._jobs.claim_next(
worker_id, exclude_source_ids=self._paused_source_ids()
)
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
self._breaker.record_failure()
if was_closed and self._breaker.is_open:
breaker = self._breaker_for(job.source_id)
was_closed = not breaker.is_open
breaker.record_failure()
if was_closed and breaker.is_open:
logger.warning(
"Worker pool breaker opened after %d consecutive transient "
"failures; pausing claims for %.0fs",
"Worker breaker opened for source %s after %d consecutive "
"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
self._breaker.record_success()
breaker = self._breaker_for(job.source_id)
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)

22
tests/ingester/test_queue.py
View file

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

71
tests/ingester/test_workers.py
View file

@ -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
open, _worker_loop's claim_next is gated off so subsequent jobs don't
burn their attempts during the same downstream outage."""
"""N back-to-back TransientErrors from one source flips that source's
breaker open. While open, _worker_loop excludes the source from
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
# it would still process more jobs. Verify the gate exists by checking
# _worker_loop: a fresh worker started with the breaker open shouldn't
# claim anything.
# drain_once claims without the breaker exclusion (it's intended for
# tests), so it would still process more jobs. The exclusion lives in
# _worker_loop: a fresh worker with this source's breaker open won't
# 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
is_open, so queued jobs stay queued until the breaker closes."""
"""Worker loop honours the breaker: an open source is excluded from
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
# on the next check (the breaker's three-state model probes after cooldown).
# Open the source's breaker, then collapse the cooldown so is_open returns
# 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()
pool._breaker._opened_at = 0.0 # type: ignore[attr-defined]
breaker.record_failure()
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