From ba6b318ecec49384fc9d62ee50e94a30dec56e84 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 25 Jun 2026 11:44:22 +0300 Subject: [PATCH 1/5] Make ingester worker ids globally unique --- CHANGELOG.md | 1 + .../haiku/rag/ingester/workers/pool.py | 14 ++++++++- tests/ingester/test_workers.py | 30 +++++++++++++++++++ 3 files changed, 44 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 87b76ad4..a22e096b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixed - A failed FTS index build is logged at `WARNING` instead of `DEBUG`, so silent full-text search degradation is visible. +- Ingester worker ids are now globally unique (`{pid}-{uuid}-{n}`) instead of `worker-{n}`, so the `claimed_by` guards on job completion/reschedule/release distinguish workers across processes sharing one Postgres queue. ### Changed diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 464be50e..2be52df4 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -1,7 +1,9 @@ import asyncio import logging +import os import time from typing import TYPE_CHECKING +from uuid import uuid4 from haiku.rag.config import CircuitBreakerConfig from haiku.rag.ingester.exceptions import PermanentError, TransientError @@ -59,12 +61,20 @@ class WorkerPool: self._metadata_providers: dict[str, MetadataProvider] = ( dict(metadata_providers) if metadata_providers else {} ) + # Globally-unique so claimed_by distinguishes this pool's workers from + # those of any other process sharing the queue; the claimed_by guards on + # mark_succeeded/reschedule/release rely on it. The uuid guarantees + # uniqueness; the pid just makes claimed_by readable in logs/dashboard. + self._instance = f"{os.getpid()}-{uuid4().hex[:8]}" self._stop = asyncio.Event() self._workers: list[asyncio.Task] = [] self._reaper: asyncio.Task | None = None self._pending_releases: set[asyncio.Task] = set() self._breakers: dict[str, CircuitBreaker] = {} + def _worker_id(self, i: int) -> str: + return f"{self._instance}-{i}" + @property def live_workers(self) -> int: """Worker tasks that are still running. Equal to worker_count under @@ -106,7 +116,9 @@ class WorkerPool: if reset: logger.info("Boot-reaped %d stale claim(s) from previous process", reset) for i in range(self._worker_count): - self._workers.append(asyncio.create_task(self._worker_loop(f"worker-{i}"))) + self._workers.append( + asyncio.create_task(self._worker_loop(self._worker_id(i))) + ) self._reaper = asyncio.create_task(self._reaper_loop()) async def stop(self) -> None: diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index f13f826e..207e1cb5 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -31,6 +31,36 @@ def _pool(client, jobs, sync, **kwargs) -> WorkerPool: ) +# --- worker identity --- + + +@pytest.mark.asyncio +async def test_worker_ids_are_unique_across_pools(client, jobs, sync): + """Two pools built with default construction must not share worker ids; + otherwise a stale worker from one pool can satisfy the claimed_by guard of + a job re-claimed by another pool and clobber its result.""" + pool_a = _pool(client, jobs, sync) + pool_b = _pool(client, jobs, sync) + worker_a = pool_a._worker_id(0) + worker_b = pool_b._worker_id(0) + assert worker_a != worker_b + + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + + claimed = await jobs.claim_next(worker_a) + assert claimed is not None + # Reaper resets the claim; pool B re-claims and finishes first. + await jobs.reap_stale(claim_timeout_seconds=0) + reclaimed = await jobs.claim_next(worker_b) + assert reclaimed is not None and reclaimed.id == job.id + assert await jobs.mark_succeeded(reclaimed.id, worker_b) is True + + # Pool A's slow worker finishes later: with unique ids this is a no-op, so + # pool B's success is not clobbered. + assert await jobs.mark_succeeded(job.id, worker_a) is False + + # --- event-driven wakeup --- From fd73f576497a3acd163d28140d033696517e5750 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 25 Jun 2026 12:06:46 +0300 Subject: [PATCH 2/5] Add last_heartbeat_at lease column to the ingester queue --- CHANGELOG.md | 4 + haiku_rag_slim/haiku/rag/ingester/queue/db.py | 3 +- .../haiku/rag/ingester/queue/migrations.py | 16 +- .../haiku/rag/ingester/queue/models.py | 1 + .../haiku/rag/ingester/queue/repository.py | 47 ++++- .../haiku/rag/ingester/workers/pool.py | 10 +- tests/ingester/test_queue.py | 195 +++++++++++++++++- tests/ingester/test_queue_postgres.py | 25 ++- tests/ingester/test_workers.py | 49 ++++- 9 files changed, 315 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a22e096b..44fb73fb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Added + +- Ingester queue schema v2 adds a `last_heartbeat_at` lease column on `jobs`; existing queue databases migrate in place on open. + ### Fixed - A failed FTS index build is logged at `WARNING` instead of `DEBUG`, so silent full-text search degradation is visible. diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/db.py b/haiku_rag_slim/haiku/rag/ingester/queue/db.py index a0ebb872..4b6ca167 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/db.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/db.py @@ -2,7 +2,7 @@ import sqlalchemy as sa from sqlalchemy import event from sqlalchemy.ext.asyncio import AsyncEngine -SCHEMA_VERSION = 1 +SCHEMA_VERSION = 2 metadata = sa.MetaData() @@ -24,6 +24,7 @@ jobs = sa.Table( sa.Column("scheduled_at", sa.Text, nullable=False), sa.Column("claimed_at", sa.Text), sa.Column("claimed_by", sa.Text), + sa.Column("last_heartbeat_at", sa.Text), sa.Column("completed_at", sa.Text), ) diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/migrations.py b/haiku_rag_slim/haiku/rag/ingester/queue/migrations.py index 35d03e5c..ac316c1e 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/migrations.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/migrations.py @@ -51,9 +51,19 @@ async def apply_migrations(engine: AsyncEngine) -> int: ).scalar_one_or_none() if current is None: await conn.execute(sa.insert(schema_version).values(version=SCHEMA_VERSION)) - elif current < SCHEMA_VERSION: # pragma: no cover - no migrations yet - # No diff migrations exist yet — future versions add ALTER/UPDATE - # statements between create_all and the version bump. + elif current < SCHEMA_VERSION: + # create_all only creates missing tables/indexes, never adds columns + # to an existing table, so column additions need explicit ALTERs. + if current < 2: + await conn.execute( + sa.text("ALTER TABLE jobs ADD COLUMN last_heartbeat_at TEXT") + ) + await conn.execute( + sa.text( + "UPDATE jobs SET last_heartbeat_at = claimed_at " + "WHERE status = 'claimed'" + ) + ) await conn.execute(sa.update(schema_version).values(version=SCHEMA_VERSION)) return SCHEMA_VERSION diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/models.py b/haiku_rag_slim/haiku/rag/ingester/queue/models.py index ae5df4d3..2d8aa982 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/models.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/models.py @@ -35,6 +35,7 @@ class Job(BaseModel): scheduled_at: datetime claimed_at: datetime | None = None claimed_by: str | None = None + last_heartbeat_at: datetime | None = None completed_at: datetime | None = None diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index 680a63c8..85f68e41 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -58,6 +58,7 @@ def _row_to_job(row: Mapping) -> Job: scheduled_at=datetime.fromisoformat(row["scheduled_at"]), claimed_at=_parse_dt(row["claimed_at"]), claimed_by=row["claimed_by"], + last_heartbeat_at=_parse_dt(row["last_heartbeat_at"]), completed_at=_parse_dt(row["completed_at"]), ) @@ -157,6 +158,7 @@ class JobRepo: status="claimed", claimed_at=now, claimed_by=worker_id, + last_heartbeat_at=now, attempts=jobs.c.attempts + 1, ) .returning(*jobs.c) @@ -230,6 +232,7 @@ class JobRepo: scheduled_at=scheduled, claimed_at=None, claimed_by=None, + last_heartbeat_at=None, last_error=error, ) .returning(jobs.c.id) @@ -259,6 +262,7 @@ class JobRepo: last_error=None, claimed_at=None, claimed_by=None, + last_heartbeat_at=None, completed_at=None, scheduled_at=now, ) @@ -449,6 +453,7 @@ class JobRepo: status="queued", claimed_at=None, claimed_by=None, + last_heartbeat_at=None, scheduled_at=_utcnow_iso(), attempts=_attempts_minus_one(), ) @@ -485,20 +490,26 @@ class JobRepo: result = await conn.execute(stmt) return result.rowcount or 0 - async def reap_stale(self, claim_timeout_seconds: int) -> int: - """Reset claimed jobs whose claimed_at is older than the timeout - back to `queued`. Decrements `attempts` to undo the increment from - `claim_next` — a crashed worker isn't a consumed attempt.""" + async def reap_stale(self, lease_ttl_seconds: int) -> int: + """Reset claimed jobs whose lease has gone stale back to `queued`. A + live worker renews its lease via `renew_claims`, so a claim is stale + only when its owner stopped renewing (crash, wedged loop, lost DB). + Staleness is measured against the last heartbeat, falling back to + claimed_at for a claim made by a process that doesn't write the lease + (an older version sharing the queue). Decrements `attempts` to undo the + increment from `claim_next` — a reaped worker isn't a consumed attempt.""" threshold = ( - datetime.now(UTC) - timedelta(seconds=claim_timeout_seconds) + datetime.now(UTC) - timedelta(seconds=lease_ttl_seconds) ).isoformat() + lease = sa.func.coalesce(jobs.c.last_heartbeat_at, jobs.c.claimed_at) stmt = ( sa.update(jobs) - .where(jobs.c.status == "claimed", jobs.c.claimed_at < threshold) + .where(jobs.c.status == "claimed", lease < threshold) .values( status="queued", claimed_at=None, claimed_by=None, + last_heartbeat_at=None, attempts=_attempts_minus_one(), ) ) @@ -506,6 +517,30 @@ class JobRepo: result = await conn.execute(stmt) return result.rowcount or 0 + async def renew_claims(self, claims: Mapping[str, str]) -> int: + """Refresh `last_heartbeat_at` for the given `job_id -> claimed_by` + pairs, extending their lease so the reaper leaves them alone. Guarded + on `status='claimed' AND claimed_by=?` per pair, so a job already + reaped and re-claimed elsewhere is left untouched — renewal never + resurrects a lost claim. Returns the number of rows renewed.""" + if not claims: + return 0 + now = _utcnow_iso() + pair_match = sa.or_( + *( + sa.and_(jobs.c.id == job_id, jobs.c.claimed_by == worker_id) + for job_id, worker_id in claims.items() + ) + ) + stmt = ( + sa.update(jobs) + .where(jobs.c.status == "claimed", pair_match) + .values(last_heartbeat_at=now) + ) + async with self._engine.begin() as conn: + result = await conn.execute(stmt) + return result.rowcount or 0 + class SyncStateRepo: def __init__(self, engine: AsyncEngine): diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 2be52df4..85876c9f 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -108,11 +108,11 @@ class WorkerPool: if self._workers: raise RuntimeError("WorkerPool already started") self._stop.clear() - # Any rows in `claimed` at start time are owned by workers from a - # previous process that didn't get to release them (SIGKILL, OOM, - # host reboot). Reset them so fresh workers can claim immediately - # instead of waiting on the reaper's claim_timeout_s. - reset = await self._jobs.reap_stale(claim_timeout_seconds=0) + # Sweep claims left behind by a previous process (SIGKILL, OOM, host + # reboot) so fresh workers can take them over. Scoped by the lease TTL, + # not 0: a peer process sharing the queue may hold live claims, and + # those must not be wiped out from under it on our startup. + reset = await self._jobs.reap_stale(lease_ttl_seconds=self._claim_timeout_s) if reset: logger.info("Boot-reaped %d stale claim(s) from previous process", reset) for i in range(self._worker_count): diff --git a/tests/ingester/test_queue.py b/tests/ingester/test_queue.py index bf020c51..14b16f06 100644 --- a/tests/ingester/test_queue.py +++ b/tests/ingester/test_queue.py @@ -3,8 +3,11 @@ from datetime import UTC, datetime, timedelta import pytest import sqlalchemy as sa +from sqlalchemy.engine import URL +from sqlalchemy.ext.asyncio import create_async_engine from haiku.rag.config import QueueConfig +from haiku.rag.ingester.queue.db import SCHEMA_VERSION from haiku.rag.ingester.queue.db import jobs as jobs_table from haiku.rag.ingester.queue.migrations import ( apply_migrations, @@ -34,6 +37,81 @@ async def test_apply_migrations_is_idempotent(engine, conn): assert row["n"] == 1 +_V1_JOBS_DDL = """ +CREATE TABLE jobs ( + id TEXT PRIMARY KEY, + source_id TEXT NOT NULL, + uri TEXT NOT NULL, + op TEXT NOT NULL, + content_hash TEXT, + revision TEXT, + status TEXT NOT NULL, + attempts INTEGER NOT NULL DEFAULT 0, + max_attempts INTEGER NOT NULL DEFAULT 5, + last_error TEXT, + extra TEXT, + enqueued_at TEXT NOT NULL, + scheduled_at TEXT NOT NULL, + claimed_at TEXT, + claimed_by TEXT, + completed_at TEXT +) +""" + + +@pytest.mark.asyncio +async def test_migration_v1_to_v2_adds_and_backfills_heartbeat(tmp_path): + """A real v1 DB (no last_heartbeat_at, schema_version=1) gains the column + and has it backfilled from claimed_at for in-flight rows only.""" + db = tmp_path / "v1.db" + engine = create_async_engine(URL.create("sqlite+aiosqlite", database=str(db))) + try: + when = "2026-01-01T00:00:00+00:00" + claimed_when = "2026-01-01T00:05:00+00:00" + async with engine.begin() as conn: + await conn.execute(sa.text(_V1_JOBS_DDL)) + await conn.execute( + sa.text("CREATE TABLE schema_version (version INTEGER PRIMARY KEY)") + ) + await conn.execute(sa.text("INSERT INTO schema_version VALUES (1)")) + await conn.execute( + sa.text( + "INSERT INTO jobs (id, source_id, uri, op, status, attempts, " + "max_attempts, enqueued_at, scheduled_at, claimed_at, claimed_by) " + "VALUES ('claimed-1', 's', 'u1', 'upsert', 'claimed', 1, 5, " + f"'{when}', '{when}', '{claimed_when}', 'w')" + ) + ) + await conn.execute( + sa.text( + "INSERT INTO jobs (id, source_id, uri, op, status, attempts, " + "max_attempts, enqueued_at, scheduled_at) " + "VALUES ('queued-1', 's', 'u2', 'upsert', 'queued', 0, 5, " + f"'{when}', '{when}')" + ) + ) + + version = await apply_migrations(engine) + assert version == SCHEMA_VERSION + + async with engine.connect() as conn: + cols = (await conn.execute(sa.text("PRAGMA table_info(jobs)"))).fetchall() + assert "last_heartbeat_at" in {c[1] for c in cols} + rows = { + row[0]: row[1] + for row in ( + await conn.execute( + sa.text("SELECT id, last_heartbeat_at FROM jobs") + ) + ).fetchall() + } + # Backfilled from claimed_at for the in-flight row, left NULL otherwise. + assert rows["claimed-1"] == claimed_when + assert rows["queued-1"] is None + finally: + await engine.dispose() + + @pytest.mark.asyncio async def test_open_queue_creates_file_and_schema(tmp_path): path = tmp_path / "subdir" / "queue.db" @@ -389,7 +467,7 @@ async def test_mark_succeeded_with_claimed_by_guard_skips_when_resurrected(jobs) a = await jobs.claim_next("worker-A") assert a is not None # Reaper resets A's stale claim. - await jobs.reap_stale(claim_timeout_seconds=0) + await jobs.reap_stale(lease_ttl_seconds=0) # B picks the job up. b = await jobs.claim_next("worker-B") assert b is not None and b.id == job.id and b.claimed_by == "worker-B" @@ -409,7 +487,7 @@ async def test_mark_dead_with_claimed_by_guard_skips_when_resurrected(jobs): job = await jobs.enqueue("s", "u", JobOp.UPSERT) a = await jobs.claim_next("worker-A") assert a is not None - await jobs.reap_stale(claim_timeout_seconds=0) + await jobs.reap_stale(lease_ttl_seconds=0) b = await jobs.claim_next("worker-B") assert b is not None and b.claimed_by == "worker-B" assert await jobs.mark_dead(job.id, "boom", "worker-A") is False @@ -427,7 +505,7 @@ async def test_reschedule_with_claimed_by_guard_skips_when_resurrected(jobs): job = await jobs.enqueue("s", "u", JobOp.UPSERT) a = await jobs.claim_next("worker-A") assert a is not None - await jobs.reap_stale(claim_timeout_seconds=0) + await jobs.reap_stale(lease_ttl_seconds=0) b = await jobs.claim_next("worker-B") assert b is not None and b.claimed_by == "worker-B" assert await jobs.reschedule(job.id, 30.0, "transient", "worker-A") is False @@ -445,7 +523,7 @@ async def test_release_if_claimed_with_claimed_by_guard_skips_when_resurrected(j job = await jobs.enqueue("s", "u", JobOp.UPSERT) a = await jobs.claim_next("worker-A") assert a is not None - await jobs.reap_stale(claim_timeout_seconds=0) + await jobs.reap_stale(lease_ttl_seconds=0) b = await jobs.claim_next("worker-B") assert b is not None and b.claimed_by == "worker-B" assert await jobs.release_if_claimed(job.id, "worker-A") is False @@ -601,38 +679,137 @@ async def test_reap_stale_resets_old_claims(conn, jobs): assert claimed is not None assert claimed.attempts == 1 - # Backdate claimed_at to simulate a crashed worker. + # Backdate the lease to simulate a worker that stopped renewing (crash). long_ago = (datetime.now(UTC) - timedelta(hours=1)).isoformat() await conn.execute( - "UPDATE jobs SET claimed_at = ? WHERE id = ?", (long_ago, job.id) + "UPDATE jobs SET claimed_at = ?, last_heartbeat_at = ? WHERE id = ?", + (long_ago, long_ago, job.id), ) await conn.commit() - reset = await jobs.reap_stale(claim_timeout_seconds=60) + reset = await jobs.reap_stale(lease_ttl_seconds=60) assert reset == 1 refreshed = await jobs.get_job(job.id) assert refreshed is not None assert refreshed.status is JobStatus.QUEUED assert refreshed.claimed_at is None assert refreshed.claimed_by is None - # claim_next incremented to 1; reap undoes that since a crashed worker + assert refreshed.last_heartbeat_at is None + # claim_next incremented to 1; reap undoes that since a reaped worker # isn't a consumed attempt — matches release_if_claimed semantics. assert refreshed.attempts == 0 +@pytest.mark.asyncio +async def test_reap_stale_falls_back_to_claimed_at_when_no_heartbeat(conn, jobs): + """A claim made by a process that doesn't write the lease (an older version + sharing the queue) has last_heartbeat_at IS NULL. Reaping must fall back to + claimed_at instead of stranding it forever (NULL comparisons are falsy).""" + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + claimed = await jobs.claim_next("w") + assert claimed is not None + + long_ago = (datetime.now(UTC) - timedelta(hours=1)).isoformat() + await conn.execute( + "UPDATE jobs SET claimed_at = ?, last_heartbeat_at = NULL WHERE id = ?", + (long_ago, job.id), + ) + await conn.commit() + + reset = await jobs.reap_stale(lease_ttl_seconds=60) + assert reset == 1 + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.QUEUED + + @pytest.mark.asyncio async def test_reap_stale_leaves_fresh_claims_alone(jobs): job = await jobs.enqueue("s", "u", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None - reset = await jobs.reap_stale(claim_timeout_seconds=3600) + reset = await jobs.reap_stale(lease_ttl_seconds=3600) assert reset == 0 refreshed = await jobs.get_job(job.id) assert refreshed is not None assert refreshed.status is JobStatus.CLAIMED +@pytest.mark.asyncio +async def test_claim_next_sets_both_timestamps(jobs): + await jobs.enqueue("s", "u", JobOp.UPSERT) + claimed = await jobs.claim_next("w") + assert claimed is not None + assert claimed.claimed_at is not None + assert claimed.last_heartbeat_at is not None + + +@pytest.mark.asyncio +async def test_renew_claims_keeps_a_stale_claim_alive(conn, jobs): + """A claim whose claimed_at is old but whose lease was just renewed must + not be reaped — this is the slow-but-alive worker the reaper used to kill.""" + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + claimed = await jobs.claim_next("w") + assert claimed is not None + + # Job started long ago, but the worker is alive and renewing. + long_ago = (datetime.now(UTC) - timedelta(hours=1)).isoformat() + await conn.execute( + "UPDATE jobs SET claimed_at = ? WHERE id = ?", (long_ago, job.id) + ) + await conn.commit() + renewed = await jobs.renew_claims({job.id: "w"}) + assert renewed == 1 + + reset = await jobs.reap_stale(lease_ttl_seconds=60) + assert reset == 0 + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.CLAIMED + + +@pytest.mark.asyncio +async def test_renew_claims_invariant_only_touches_heartbeat(conn, jobs): + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + claimed = await jobs.claim_next("w") + assert claimed is not None and claimed.claimed_at is not None + + # Backdate the heartbeat so renewal demonstrably advances it. + old = (datetime.now(UTC) - timedelta(minutes=5)).isoformat() + await conn.execute( + "UPDATE jobs SET last_heartbeat_at = ? WHERE id = ?", (old, job.id) + ) + await conn.commit() + before = await jobs.get_job(job.id) + assert before is not None and before.last_heartbeat_at is not None + + assert await jobs.renew_claims({job.id: "w"}) == 1 + after = await jobs.get_job(job.id) + assert after is not None and after.last_heartbeat_at is not None + assert after.claimed_at == before.claimed_at + assert after.last_heartbeat_at > before.last_heartbeat_at + + +@pytest.mark.asyncio +async def test_renew_claims_skips_reclaimed_job(jobs): + """Renewal is guarded on claimed_by: a job reaped and re-claimed by another + worker must not be renewed by the original (lost) claimant.""" + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + a = await jobs.claim_next("worker-A") + assert a is not None + await jobs.reap_stale(lease_ttl_seconds=0) + b = await jobs.claim_next("worker-B") + assert b is not None and b.claimed_by == "worker-B" + + assert await jobs.renew_claims({job.id: "worker-A"}) == 0 + + +@pytest.mark.asyncio +async def test_renew_claims_empty_is_noop(jobs): + assert await jobs.renew_claims({}) == 0 + + # --- prune_dead --- diff --git a/tests/ingester/test_queue_postgres.py b/tests/ingester/test_queue_postgres.py index 859e375d..5205be78 100644 --- a/tests/ingester/test_queue_postgres.py +++ b/tests/ingester/test_queue_postgres.py @@ -106,7 +106,7 @@ async def test_reap_stale_clamps_attempts(postgres_dburi): assert job is not None claimed = await jobs.claim_next("w") assert claimed is not None and claimed.attempts == 1 - reset = await jobs.reap_stale(claim_timeout_seconds=0) + reset = await jobs.reap_stale(lease_ttl_seconds=0) assert reset == 1 refreshed = await jobs.get_job(job.id) assert refreshed is not None @@ -114,6 +114,29 @@ async def test_reap_stale_clamps_attempts(postgres_dburi): assert refreshed.attempts == 0 +@pytest.mark.asyncio +async def test_renew_claims_survives_reap_on_postgres(postgres_dburi): + """The COALESCE lease threshold and the OR-of-(id, claimed_by) renewal + predicate render and run on Postgres: a renewed claim outlives a reap even + though its claimed_at is old.""" + async with queue_engine(postgres_dburi) as engine: + jobs = JobRepo(engine) + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + assert job is not None + claimed = await jobs.claim_next("w") + assert claimed is not None + + async with engine.begin() as conn: + await conn.execute( + text("UPDATE jobs SET claimed_at = :ts WHERE id = :id"), + {"ts": "2000-01-01T00:00:00+00:00", "id": job.id}, + ) + assert await jobs.renew_claims({job.id: "w"}) == 1 + assert await jobs.reap_stale(lease_ttl_seconds=60) == 0 + refreshed = await jobs.get_job(job.id) + assert refreshed is not None and refreshed.status is JobStatus.CLAIMED + + @pytest.mark.asyncio async def test_sync_state_upsert_coalesce_preserves_revision(postgres_dburi): """ON CONFLICT DO UPDATE with COALESCE leaves an existing revision in place diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index 207e1cb5..99837052 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -51,7 +51,7 @@ async def test_worker_ids_are_unique_across_pools(client, jobs, sync): claimed = await jobs.claim_next(worker_a) assert claimed is not None # Reaper resets the claim; pool B re-claims and finishes first. - await jobs.reap_stale(claim_timeout_seconds=0) + await jobs.reap_stale(lease_ttl_seconds=0) reclaimed = await jobs.claim_next(worker_b) assert reclaimed is not None and reclaimed.id == job.id assert await jobs.mark_succeeded(reclaimed.id, worker_b) is True @@ -448,7 +448,7 @@ async def test_worker_loses_claim_to_reaper_does_not_write_sync_state( claimed_by_a = await jobs.claim_next("worker-A") assert claimed_by_a is not None # Reaper resets A's claim, worker-B re-claims. - await jobs.reap_stale(claim_timeout_seconds=0) + await jobs.reap_stale(lease_ttl_seconds=0) await jobs.claim_next("worker-B") pool = _pool(client, jobs, sync, worker_count=1) @@ -478,7 +478,7 @@ async def test_permanent_error_loses_claim_to_reaper_writes_no_marker( claimed_by_a = await jobs.claim_next("worker-A") assert claimed_by_a is not None # Reaper resets A's claim, worker-B re-claims. - await jobs.reap_stale(claim_timeout_seconds=0) + await jobs.reap_stale(lease_ttl_seconds=0) await jobs.claim_next("worker-B") pool = _pool(client, jobs, sync, worker_count=1) @@ -826,16 +826,25 @@ async def test_permanent_failure_marker_write_failure_does_not_crash_worker( @pytest.mark.asyncio -async def test_boot_reap_resets_pre_existing_claims(client, jobs, sync): - """A SIGKILL'd previous process leaves rows in `claimed` state. The new - WorkerPool.start() must reset them immediately so fresh workers can - claim them, instead of waiting on the periodic reaper's claim_timeout_s - window (default 1800s).""" +async def test_boot_reap_resets_stale_pre_existing_claims(client, jobs, sync, conn): + """A SIGKILL'd previous process leaves a stale claim (its lease stopped + being renewed). WorkerPool.start() sweeps it so fresh workers can take it + over.""" + from datetime import UTC, datetime, timedelta + await jobs.enqueue("src", "u", JobOp.UPSERT) pre_claimed = await jobs.claim_next("ghost-worker") assert pre_claimed is not None assert pre_claimed.status is JobStatus.CLAIMED + # The ghost stopped renewing an hour ago. + long_ago = (datetime.now(UTC) - timedelta(hours=1)).isoformat() + await conn.execute( + "UPDATE jobs SET claimed_at = ?, last_heartbeat_at = ? WHERE id = ?", + (long_ago, long_ago, pre_claimed.id), + ) + await conn.commit() + pool = _pool(client, jobs, sync, worker_count=0) await pool.start() try: @@ -850,6 +859,25 @@ async def test_boot_reap_resets_pre_existing_claims(client, jobs, sync): await pool.stop() +@pytest.mark.asyncio +async def test_boot_reap_leaves_a_peer_process_fresh_claim_alone(client, jobs, sync): + """A peer process sharing the queue holds a freshly-claimed job. Our + startup boot-reap must not wipe its live claim.""" + await jobs.enqueue("src", "u", JobOp.UPSERT) + peer_claim = await jobs.claim_next("peer-worker") + assert peer_claim is not None + + pool = _pool(client, jobs, sync, worker_count=0) + await pool.start() + try: + refreshed = await jobs.get_job(peer_claim.id) + assert refreshed is not None + assert refreshed.status is JobStatus.CLAIMED + assert refreshed.claimed_by == "peer-worker" + finally: + await pool.stop() + + @pytest.mark.asyncio async def test_reaper_resets_stale_claims(client, jobs, sync, conn): from datetime import UTC, datetime, timedelta @@ -858,10 +886,11 @@ async def test_reaper_resets_stale_claims(client, jobs, sync, conn): claimed = await jobs.claim_next("worker-old") assert claimed is not None - # Push claimed_at back so reap_stale picks it up. + # Push the lease back so reap_stale picks it up. long_ago = (datetime.now(UTC) - timedelta(hours=1)).isoformat() await conn.execute( - "UPDATE jobs SET claimed_at = ? WHERE id = ?", (long_ago, job.id) + "UPDATE jobs SET claimed_at = ?, last_heartbeat_at = ? WHERE id = ?", + (long_ago, long_ago, job.id), ) await conn.commit() From 4707db780f1aad8aecc78fbe62f8e6d338f73f4d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 25 Jun 2026 13:20:28 +0300 Subject: [PATCH 3/5] Renew job leases from the worker pool; lease-based reaping --- CHANGELOG.md | 12 +- haiku_rag_slim/haiku/rag/config/models.py | 43 +++-- haiku_rag_slim/haiku/rag/ingester/app.py | 13 +- .../haiku/rag/ingester/workers/pool.py | 99 ++++++++++- tests/ingester/test_config.py | 27 +++ tests/ingester/test_workers.py | 158 +++++++++++++++++- 6 files changed, 319 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 44fb73fb..7928b419 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,18 +1,16 @@ # Changelog ## [Unreleased] -### Added +### Changed -- Ingester queue schema v2 adds a `last_heartbeat_at` lease column on `jobs`; existing queue databases migrate in place on open. +- Ingester reaping is lease-based: a worker renews `last_heartbeat_at` on its in-flight jobs every `heartbeat_interval_s`, and the reaper reclaims a claim only once its lease is older than `lease_ttl_s`. A job slower than the timeout is no longer reaped and reprocessed while still running. Adds queue schema v2 (`last_heartbeat_at` on `jobs`); existing queue databases migrate in place on open. +- **Breaking:** `ingester.workers.claim_timeout_s` removed; set `lease_ttl_s` (default 120) and `heartbeat_interval_s` (default 30) instead. `WorkerConfig` rejects unknown keys. +- A bare `${VAR}` in YAML config now raises `MissingEnvVarError` when the variable is set but empty, matching the unset case. Use `${VAR:-default}` to allow an empty/absent value. ### Fixed - A failed FTS index build is logged at `WARNING` instead of `DEBUG`, so silent full-text search degradation is visible. -- Ingester worker ids are now globally unique (`{pid}-{uuid}-{n}`) instead of `worker-{n}`, so the `claimed_by` guards on job completion/reschedule/release distinguish workers across processes sharing one Postgres queue. - -### Changed - -- A bare `${VAR}` in YAML config now raises `MissingEnvVarError` when the variable is set but empty, matching the unset case. Use `${VAR:-default}` to allow an empty/absent value. +- Ingester worker ids are now globally unique (`{pid}-{uuid}-{n}`); the `claimed_by` guards on job completion/reschedule/release distinguish workers across processes sharing one Postgres queue. ## [0.61.2] - 2026-06-24 diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 08622db7..7a10e6a4 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -1,7 +1,7 @@ from pathlib import Path from typing import Annotated, Any, Literal -from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from haiku.rag.utils import get_default_data_dir @@ -309,6 +309,10 @@ class CircuitBreakerConfig(BaseModel): class WorkerConfig(BaseModel): + # Reject unknown keys so a renamed/removed setting (e.g. the former + # claim_timeout_s) fails loudly instead of being silently ignored. + model_config = ConfigDict(extra="forbid", validate_assignment=True) + worker_count: int = Field( default=4, description="Number of async worker tasks pulling from the queue. " @@ -323,14 +327,22 @@ class WorkerConfig(BaseModel): "polls. Lower = lower latency picking up new jobs, higher = less " "queue churn when the queue is usually empty.", ) - claim_timeout_s: int = Field( - default=1800, - description="A `claimed` job whose claimed_at is older than this is " - "presumed dead and reset to `queued` by the reaper. MUST exceed your " - "longest legitimate job duration — set it too short and the reaper " - "resurrects jobs still being processed, causing two workers to run " - "the same URI. Default (30min) covers typical docling conversions; " - "raise it if you ingest very large PDFs through docling-local.", + lease_ttl_s: int = Field( + default=120, + gt=0, + description="A `claimed` job whose lease has not been renewed within " + "this window is presumed dead and reset to `queued` by the reaper. A " + "live worker renews its lease every heartbeat_interval_s while " + "processing, so this need not exceed job duration — it only bounds how " + "long a crashed worker's job stays stuck before another worker takes " + "it over.", + ) + heartbeat_interval_s: int = Field( + default=30, + gt=0, + description="How often a worker renews the lease on its in-flight " + "jobs. Must be comfortably shorter than lease_ttl_s so scheduler " + "jitter or a slow DB round-trip can't let a live job's lease lapse.", ) reaper_interval_s: int = Field( default=60, @@ -341,10 +353,19 @@ class WorkerConfig(BaseModel): shutdown_grace_s: float = Field( default=60.0, description="On SIGINT/SIGTERM, how long to wait for in-flight jobs to " - "finish before forcing cancellation. Cancelled jobs stay 'claimed' in " - "the queue; the reaper resets them after claim_timeout_s.", + "finish before forcing cancellation. Cancelled jobs are released back " + "to `queued` for immediate re-claim.", ) + @model_validator(mode="after") + def _check_heartbeat_cadence(self) -> "WorkerConfig": + if self.heartbeat_interval_s > self.lease_ttl_s / 3: + raise ValueError( + "heartbeat_interval_s must be <= lease_ttl_s / 3 so a live " + "worker renews its lease several times before it could expire" + ) + return self + class APIConfig(BaseModel): """HTTP control plane settings for the ingester.""" diff --git a/haiku_rag_slim/haiku/rag/ingester/app.py b/haiku_rag_slim/haiku/rag/ingester/app.py index dd08372c..321845c0 100644 --- a/haiku_rag_slim/haiku/rag/ingester/app.py +++ b/haiku_rag_slim/haiku/rag/ingester/app.py @@ -133,7 +133,8 @@ class IngesterApp: worker_count=ingester_cfg.workers.worker_count, retry_policy=retry, poll_idle_interval_s=ingester_cfg.workers.poll_idle_interval_s, - claim_timeout_s=ingester_cfg.workers.claim_timeout_s, + lease_ttl_s=ingester_cfg.workers.lease_ttl_s, + heartbeat_interval_s=ingester_cfg.workers.heartbeat_interval_s, reaper_interval_s=ingester_cfg.workers.reaper_interval_s, retention_s=( ingester_cfg.queue.retention_days * 86400 @@ -193,12 +194,14 @@ class IngesterApp: try: await asyncio.wait_for(self._pool.stop(), timeout=grace_s) except TimeoutError: - # In-flight jobs stay 'claimed'; the reaper resets them after - # claim_timeout_s on the next start. + # Cancelling the workers triggers their cancel-cleanup, which + # releases each in-flight job back to `queued` (drained just below). + # Any release that doesn't land has its lease stop being renewed, so + # the reaper reclaims it after lease_ttl_s. logger.warning( "Shutdown grace of %.1fs elapsed with jobs still in flight; " - "cancelling — they'll be reclaimed after claim_timeout_s on " - "next start", + "cancelling — they'll be released back to the queue (or " + "reclaimed after lease_ttl_s if release doesn't land)", grace_s, ) landed = await self._pool.drain_pending_releases(timeout=2.0) diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 85876c9f..f270bcf5 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -28,9 +28,11 @@ _WORKER_BREAKER_COOLDOWN_S = 60.0 class WorkerPool: """`worker_count` async tasks each pull jobs from the queue and run them - through `run_job`. A reaper task resets claims older than - `claim_timeout_s` so a crashed worker doesn't strand its job. Lifecycle: - build it, await start(), let it run, await stop(). + through `run_job`. While a worker processes a job it renews the job's lease + on a `heartbeat_interval_s` cadence; a reaper task resets claims whose lease + has gone stale (`lease_ttl_s`) so a crashed worker doesn't strand its job, + without reaping jobs that are merely slow. Lifecycle: build it, await + start(), let it run, await stop(). """ def __init__( @@ -42,7 +44,8 @@ class WorkerPool: worker_count: int = 4, retry_policy: RetryPolicy | None = None, poll_idle_interval_s: float = 1.0, - claim_timeout_s: int = 1800, + lease_ttl_s: int = 120, + heartbeat_interval_s: int = 30, reaper_interval_s: int = 60, retention_s: int | None = None, sources: "list[Source] | None" = None, @@ -54,7 +57,8 @@ class WorkerPool: self._worker_count = worker_count self._retry = retry_policy or RetryPolicy() self._poll_idle_s = poll_idle_interval_s - self._claim_timeout_s = claim_timeout_s + self._lease_ttl_s = lease_ttl_s + self._heartbeat_interval_s = heartbeat_interval_s self._reaper_interval_s = reaper_interval_s self._retention_s = retention_s self._sources: list[Source] = list(sources) if sources else [] @@ -69,18 +73,46 @@ class WorkerPool: self._stop = asyncio.Event() self._workers: list[asyncio.Task] = [] self._reaper: asyncio.Task | None = None + self._heartbeat: asyncio.Task | None = None + # job_id -> claimed_by for jobs currently being processed by this pool; + # the heartbeat renews exactly these leases. + self._inflight: dict[str, str] = {} + # Set whenever _inflight is empty so the heartbeat can leave promptly + # once stopping instead of sleeping out a full interval. + self._idle = asyncio.Event() + self._idle.set() self._pending_releases: set[asyncio.Task] = set() self._breakers: dict[str, CircuitBreaker] = {} def _worker_id(self, i: int) -> str: return f"{self._instance}-{i}" + def _track_inflight(self, job_id: str, worker_id: str) -> None: + self._inflight[job_id] = worker_id + self._idle.clear() + + def _untrack_inflight(self, job_id: str, worker_id: str) -> None: + # Only drop our own entry: if the reaper reset this claim and a sibling + # worker re-claimed it, the entry now belongs to that worker and must + # keep being renewed — our exit must not evict it. + if self._inflight.get(job_id) == worker_id: + del self._inflight[job_id] + if not self._inflight: + self._idle.set() + @property def live_workers(self) -> int: """Worker tasks that are still running. Equal to worker_count under normal operation; less when a worker has crashed.""" return sum(1 for t in self._workers if not t.done()) + @property + def heartbeat_alive(self) -> bool: + """Whether the lease-renewal task is running. Once the pool is started + this stays True until stop(); a False here while workers are live means + in-flight leases are no longer being renewed and may be reaped.""" + return self._heartbeat is not None and not self._heartbeat.done() + @property def breaker_open(self) -> bool: return any(b.is_open for b in self._breakers.values()) @@ -112,7 +144,7 @@ class WorkerPool: # reboot) so fresh workers can take them over. Scoped by the lease TTL, # not 0: a peer process sharing the queue may hold live claims, and # those must not be wiped out from under it on our startup. - reset = await self._jobs.reap_stale(lease_ttl_seconds=self._claim_timeout_s) + reset = await self._jobs.reap_stale(lease_ttl_seconds=self._lease_ttl_s) if reset: logger.info("Boot-reaped %d stale claim(s) from previous process", reset) for i in range(self._worker_count): @@ -120,20 +152,35 @@ class WorkerPool: asyncio.create_task(self._worker_loop(self._worker_id(i))) ) self._reaper = asyncio.create_task(self._reaper_loop()) + self._heartbeat = asyncio.create_task(self._heartbeat_loop()) + self._heartbeat.add_done_callback(self._on_heartbeat_done) + + def _on_heartbeat_done(self, task: asyncio.Task) -> None: + if task.cancelled(): + return + exc = task.exception() + if exc is not None: + logger.error("Heartbeat task died: %r", exc) + elif not self._stop.is_set(): + logger.error("Heartbeat task exited while the pool was still running") 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. + # _stop also wakes the heartbeat's interval sleep. 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) + if self._heartbeat is not None: + tasks.append(self._heartbeat) if tasks: await asyncio.gather(*tasks, return_exceptions=True) self._workers.clear() self._reaper = None + self._heartbeat = None async def drain_pending_releases(self, timeout: float = 2.0) -> int: """Wait for any in-flight cancel-cleanup release Tasks to finish. @@ -180,7 +227,7 @@ class WorkerPool: await self._sleep_or_stop(self._reaper_interval_s) if self._stop.is_set(): return - reset = await self._jobs.reap_stale(self._claim_timeout_s) + reset = await self._jobs.reap_stale(self._lease_ttl_s) if reset: logger.info("Reaper reset %d stale claim(s)", reset) if self._retention_s is not None: @@ -188,6 +235,35 @@ class WorkerPool: if pruned: logger.info("Reaper pruned %d terminal job(s)", pruned) + async def _heartbeat_loop(self) -> None: + """Renew the lease on this pool's in-flight jobs so the reaper leaves + them alone while they are still being processed. Continues renewing + through graceful shutdown until the last job drains, so a peer process + doesn't reap a job we're still finishing. A forced shutdown cancels + this task along with the workers, which is correct — they are no longer + draining gracefully.""" + while True: + if self._inflight: + try: + await self._jobs.renew_claims(dict(self._inflight)) + except Exception: + logger.exception( + "Lease renewal failed; retrying on the next heartbeat" + ) + if self._stop.is_set() and not self._inflight: + return + if self._stop.is_set(): + # Draining: keep the renewal cadence, but leave as soon as the + # last in-flight job finishes instead of sleeping it out. + try: + await asyncio.wait_for( + self._idle.wait(), timeout=self._heartbeat_interval_s + ) + except TimeoutError: + pass + else: + await self._sleep_or_stop(self._heartbeat_interval_s) + async def _sleep_or_stop(self, seconds: float) -> None: try: await asyncio.wait_for(self._stop.wait(), timeout=seconds) @@ -197,6 +273,15 @@ class WorkerPool: async def _process(self, job: Job) -> None: assert job.claimed_by is not None, "_process only runs on claimed jobs" worker_id = job.claimed_by + # Track before any await so the heartbeat renews this job's lease for + # its whole lifetime; untrack on every exit (success, error, cancel). + self._track_inflight(job.id, worker_id) + try: + await self._run_job_lifecycle(job, worker_id) + finally: + self._untrack_inflight(job.id, worker_id) + + async def _run_job_lifecycle(self, job: Job, worker_id: str) -> None: started = time.monotonic() logger.info("Processing %s %s (job %s)", job.op.value, job.uri, job.id) try: diff --git a/tests/ingester/test_config.py b/tests/ingester/test_config.py index 36bcbf5b..066e68f6 100644 --- a/tests/ingester/test_config.py +++ b/tests/ingester/test_config.py @@ -12,6 +12,7 @@ from haiku.rag.config import ( IngesterConfig, RetryPolicyConfig, S3SourceConfig, + WorkerConfig, ) @@ -20,11 +21,37 @@ def test_default_ingester_config_has_sane_values(): assert cfg.sources == [] assert cfg.workers.worker_count == 4 assert cfg.workers.retry.max_attempts == 5 + assert cfg.workers.lease_ttl_s == 120 + assert cfg.workers.heartbeat_interval_s == 30 assert cfg.api.enabled is True assert cfg.api.port == 8765 assert cfg.api.root_path == "" +def test_worker_config_rejects_heartbeat_too_close_to_lease(): + with pytest.raises(ValidationError, match="heartbeat_interval_s"): + WorkerConfig(lease_ttl_s=60, heartbeat_interval_s=30) + + +def test_worker_config_accepts_heartbeat_within_a_third_of_lease(): + cfg = WorkerConfig(lease_ttl_s=60, heartbeat_interval_s=20) + assert cfg.heartbeat_interval_s == 20 + + +def test_worker_config_rejects_non_positive_timings(): + with pytest.raises(ValidationError): + WorkerConfig(lease_ttl_s=0) + with pytest.raises(ValidationError): + WorkerConfig(heartbeat_interval_s=0) + + +def test_worker_config_rejects_unknown_field(): + """The former claim_timeout_s (and any typo) must fail loudly rather than + being silently ignored.""" + with pytest.raises(ValidationError, match="claim_timeout_s"): + WorkerConfig(claim_timeout_s=1800) # ty: ignore[unknown-argument] + + @pytest.mark.parametrize( "raw, expected", [ diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index 99837052..31a3c9ed 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -24,7 +24,8 @@ def _pool(client, jobs, sync, **kwargs) -> WorkerPool: worker_count=kwargs.pop("worker_count", 2), poll_idle_interval_s=kwargs.pop("poll_idle_interval_s", 0.05), reaper_interval_s=kwargs.pop("reaper_interval_s", 60), - claim_timeout_s=kwargs.pop("claim_timeout_s", 60), + lease_ttl_s=kwargs.pop("lease_ttl_s", 60), + heartbeat_interval_s=kwargs.pop("heartbeat_interval_s", 30), retention_s=kwargs.pop("retention_s", None), retry_policy=kwargs.pop("retry_policy", RetryPolicy()), sources=kwargs.pop("sources", None), @@ -400,7 +401,7 @@ async def test_shutdown_grace_lets_inflight_job_complete(client, jobs, sync): async def test_shutdown_grace_timeout_releases_claim(client, jobs, sync): """When grace elapses and the worker is cancelled mid-job, the claim is released back to 'queued' so the next process can pick it up immediately - — no waiting on the reaper's claim_timeout_s.""" + — no waiting on the reaper's lease_ttl_s.""" cancelled = asyncio.Event() async def _hangs_forever(*args, **kwargs): @@ -432,6 +433,157 @@ async def test_shutdown_grace_timeout_releases_claim(client, jobs, sync): assert refreshed.attempts == 0 +# --- heartbeat / lease renewal --- + + +@pytest.mark.asyncio +async def test_heartbeat_keeps_long_job_from_being_reaped(client, jobs, sync): + """A job that takes longer than lease_ttl_s is kept alive by the heartbeat, + so the reaper never resets it and it completes exactly once. Without + renewal the reaper would reset the claim mid-flight and the final + mark_succeeded would be a no-op, leaving the job back in 'queued'.""" + calls = 0 + + async def _slow(*args, **kwargs): + nonlocal calls + calls += 1 + await asyncio.sleep(0.8) + return Document( + id="d", content="x", uri="u", metadata={"md5": "m", "source_revision": "r"} + ) + + client.create_document_from_source.side_effect = _slow + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + + pool = _pool( + client, + jobs, + sync, + worker_count=1, + lease_ttl_s=0.3, + heartbeat_interval_s=0.05, + reaper_interval_s=0.05, + ) + await pool.start() + try: + for _ in range(100): + refreshed = await jobs.get_job(job.id) + if refreshed is not None and refreshed.status is JobStatus.SUCCEEDED: + break + await asyncio.sleep(0.05) + finally: + await pool.stop() + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.SUCCEEDED + assert calls == 1 + + +@pytest.mark.asyncio +async def test_heartbeat_continues_during_graceful_stop(client, jobs, sync): + """While stop() waits for an in-flight job to drain, the heartbeat keeps + renewing its lease so a peer process doesn't reap it mid-drain.""" + release = asyncio.Event() + + async def _block(*args, **kwargs): + await release.wait() + return Document( + id="d", content="x", uri="u", metadata={"md5": "m", "source_revision": "r"} + ) + + client.create_document_from_source.side_effect = _block + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + + pool = _pool( + client, + jobs, + sync, + worker_count=1, + lease_ttl_s=5, + heartbeat_interval_s=0.05, + reaper_interval_s=60, + ) + await pool.start() + stop_task: asyncio.Task | None = None + try: + for _ in range(100): + refreshed = await jobs.get_job(job.id) + if refreshed is not None and refreshed.status is JobStatus.CLAIMED: + break + await asyncio.sleep(0.02) + assert refreshed is not None and refreshed.status is JobStatus.CLAIMED + + # Begin a graceful stop; it blocks until the job drains. + stop_task = asyncio.create_task(pool.stop()) + await asyncio.sleep(0.05) + first = (await jobs.get_job(job.id)).last_heartbeat_at + await asyncio.sleep(0.2) + second = (await jobs.get_job(job.id)).last_heartbeat_at + assert first is not None and second is not None + assert second > first, "heartbeat did not renew during graceful drain" + assert not stop_task.done() + + release.set() + await asyncio.wait_for(stop_task, timeout=2.0) + finally: + release.set() + if stop_task is not None and not stop_task.done(): + await stop_task + + final = await jobs.get_job(job.id) + assert final is not None and final.status is JobStatus.SUCCEEDED + + +@pytest.mark.asyncio +async def test_stale_worker_exit_keeps_a_siblings_renewal_entry(client, jobs, sync): + """Same-pool reaper race: worker A is processing J, the reaper resets it, + sibling worker B re-claims J. A's eventual exit must not evict B's renewal + entry, or B's lease would stop being renewed and J could be reaped again.""" + pool = _pool(client, jobs, sync, worker_count=2) + worker_a, worker_b = pool._worker_id(0), pool._worker_id(1) + + pool._track_inflight("J", worker_a) + # Reaper reset + sibling re-claim: B now owns J's renewal entry. + pool._track_inflight("J", worker_b) + assert pool._inflight["J"] == worker_b + + # A finally exits — must leave B's entry intact. + pool._untrack_inflight("J", worker_a) + assert pool._inflight.get("J") == worker_b + + # B's own exit clears it. + pool._untrack_inflight("J", worker_b) + assert "J" not in pool._inflight + + +@pytest.mark.asyncio +async def test_forced_shutdown_cancels_heartbeat(client, jobs, sync): + """When the shutdown grace elapses and stop() is cancelled, the heartbeat + task is cancelled along with the workers — they are no longer draining.""" + + async def _hangs_forever(*args, **kwargs): + await asyncio.sleep(60) + return Document(id="d", content="x", uri="u") + + client.create_document_from_source.side_effect = _hangs_forever + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + + pool = _pool(client, jobs, sync, worker_count=1) + await pool.start() + await asyncio.sleep(0.05) + assert pool.heartbeat_alive + + with pytest.raises(TimeoutError): + await asyncio.wait_for(pool.stop(), timeout=0.2) + # Let the cancellation settle. + await asyncio.sleep(0.05) + assert not pool.heartbeat_alive + + @pytest.mark.asyncio async def test_worker_loses_claim_to_reaper_does_not_write_sync_state( client, jobs, sync @@ -900,7 +1052,7 @@ async def test_reaper_resets_stale_claims(client, jobs, sync, conn): sync, worker_count=0, reaper_interval_s=0.05, - claim_timeout_s=1, + lease_ttl_s=1, ) await pool.start() try: From 4333ab4fe2ba476c2e7cecbe2e2561a85591e6a1 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 25 Jun 2026 13:23:53 +0300 Subject: [PATCH 4/5] Document lease-based ingester reaping --- docs/ingester.md | 50 ++++++++++++++++++++++++++++++++---------------- 1 file changed, 34 insertions(+), 16 deletions(-) diff --git a/docs/ingester.md b/docs/ingester.md index e89b6125..9cd295f8 100644 --- a/docs/ingester.md +++ b/docs/ingester.md @@ -335,7 +335,8 @@ ingester: workers: worker_count: 4 poll_idle_interval_s: 1.0 - claim_timeout_s: 1800 + lease_ttl_s: 120 + heartbeat_interval_s: 30 reaper_interval_s: 60 shutdown_grace_s: 60 # SIGTERM drains in-flight up to this long retry: @@ -352,8 +353,11 @@ rescheduled with exponential backoff plus jitter, up to `max_attempts`, then land in the dead-letter queue. `PermanentError` (unsupported extension, 4xx HTTP except 408/429, etc.) skips retry entirely. -A reaper task resets jobs whose `claimed_at` is older than -`claim_timeout_s` so a crashed worker doesn't strand its job. +While a worker processes a job it renews the job's lease every +`heartbeat_interval_s`. A reaper task resets any claim whose lease has not +been renewed within `lease_ttl_s` so a crashed worker doesn't strand its +job. Because a live worker keeps renewing, `lease_ttl_s` need not exceed +job duration — a slow job is not reaped while it is still running. **Backpressure.** Each poller skips its periodic sweep when its source already has queued or claimed jobs in the queue. The unique-index dedup @@ -363,22 +367,25 @@ still flow during a skipped sweep, so new files aren't lost. **Graceful shutdown.** On `SIGINT` / `SIGTERM`, pollers stop immediately and workers are given `shutdown_grace_s` to finish in-flight jobs. Jobs -still running after the grace window are cancelled — they stay -`claimed` in the queue and are reset by the reaper on the next start -once `claim_timeout_s` elapses. +still running after the grace window are cancelled and released back to +`queued` for immediate re-claim; any release that doesn't land has its +lease lapse and is reclaimed by the reaper after `lease_ttl_s`. **Tuning.** -- `claim_timeout_s` must exceed the longest legitimate job duration; a - shorter value lets the reaper resurrect in-flight jobs. +- `lease_ttl_s` bounds how long a crashed worker's job stays stuck before + another worker takes it over. It no longer needs to exceed job duration, + so it can be short; keep it well above `heartbeat_interval_s`. +- `heartbeat_interval_s` must be at most `lease_ttl_s / 3` so scheduler + jitter or a slow DB round-trip can't let a live job's lease lapse. - `worker_count` should match downstream capacity. docling-serve processes one task per instance, so `worker_count` above the number of `providers.docling_serve.base_url` entries over-subscribes the - fleet — extra submissions queue inside docling-serve and inflate - `claimed_at` duration toward `claim_timeout_s`. + fleet — extra submissions queue inside docling-serve. They are not + reaped while queued because the worker keeps renewing the lease. - `poll_idle_interval_s`: lower = faster pickup, more SQLite churn. - `reaper_interval_s`: worst-case post-crash reclaim is - `claim_timeout_s + reaper_interval_s`. + `lease_ttl_s + reaper_interval_s`. **Per-source override.** A source can opt out of the global retry policy: @@ -431,7 +438,10 @@ server, then pollers, then in-flight workers. LanceDB supports exactly one writer + N readers per database URI. Run exactly one `haiku-ingester serve` against a given LanceDB. Multiple -MCP servers or read-only consumers against the same DB are fine. +MCP servers or read-only consumers against the same DB are fine. Sharing +the Postgres queue across processes is safe (the claim/lease lifecycle is +cross-process-correct) but does not relax this constraint — it governs the +queue, not the LanceDB. ## HTTP control plane @@ -603,10 +613,18 @@ the same way as for SQLite: haiku-ingester queue init ``` -Workers claim jobs with `FOR UPDATE SKIP LOCKED`, so several `haiku-ingester -serve` processes can share one Postgres queue and scale out horizontally. One -caveat: idle workers wake on new work instantly only within their own process. -Workers in other processes pick up enqueued jobs on their next +Workers claim jobs with `FOR UPDATE SKIP LOCKED`, and the claim/lease lifecycle +is cross-process-safe — claims are renewed and reaped correctly no matter which +process owns them — so several `haiku-ingester serve` processes can share one +Postgres queue without double-claiming or reaping each other's live jobs. + +This does not lift the LanceDB +[single-writer constraint](#single-writer-constraint): each `serve` still owns +its own LanceDB. A shared queue therefore spans processes writing distinct +LanceDB URIs; it does not let several processes write one database. + +One caveat: idle workers wake on new work instantly only within their own +process. Workers in other processes pick up enqueued jobs on their next `poll_idle_interval_s` tick rather than immediately. ### Logs From 87f0233bcdee802df290eb8a92b2bd4acc43da6c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 25 Jun 2026 16:32:56 +0300 Subject: [PATCH 5/5] Cover heartbeat health, lease-renewal failure, and breaker recovery --- tests/ingester/test_workers.py | 138 +++++++++++++++++++++++++++++++++ 1 file changed, 138 insertions(+) diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index 31a3c9ed..c1602e40 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -584,6 +584,124 @@ async def test_forced_shutdown_cancels_heartbeat(client, jobs, sync): assert not pool.heartbeat_alive +@pytest.mark.asyncio +async def test_heartbeat_survives_a_renewal_failure( + client, jobs, sync, monkeypatch, caplog +): + """A transient DB error during lease renewal is logged and the heartbeat + keeps running, rather than dying and leaving in-flight leases unrenewed.""" + release = asyncio.Event() + + async def _block(*args, **kwargs): + await release.wait() + return Document( + id="d", content="x", uri="u", metadata={"md5": "m", "source_revision": "r"} + ) + + client.create_document_from_source.side_effect = _block + + calls = 0 + + async def _boom(_claims): + nonlocal calls + calls += 1 + raise RuntimeError("queue unavailable") + + monkeypatch.setattr(jobs, "renew_claims", _boom) + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + + pool = _pool( + client, + jobs, + sync, + worker_count=1, + lease_ttl_s=5, + heartbeat_interval_s=0.05, + reaper_interval_s=60, + ) + await pool.start() + try: + with caplog.at_level("ERROR", logger="haiku.rag.ingester.workers.pool"): + for _ in range(100): + refreshed = await jobs.get_job(job.id) + if refreshed is not None and refreshed.status is JobStatus.CLAIMED: + break + await asyncio.sleep(0.02) + await asyncio.sleep(0.2) + assert calls >= 1 + assert pool.heartbeat_alive + assert "Lease renewal failed" in caplog.text + finally: + release.set() + await pool.stop() + + +@pytest.mark.asyncio +async def test_heartbeat_done_callback_logs_unexpected_outcomes( + client, jobs, sync, caplog +): + """The done-callback surfaces a heartbeat that died or exited while the pool + was still running — both should never happen, so they log loudly.""" + pool = _pool(client, jobs, sync, worker_count=0) + + async def _raises(): + raise RuntimeError("boom") + + died = asyncio.create_task(_raises()) + with pytest.raises(RuntimeError): + await died + with caplog.at_level("ERROR", logger="haiku.rag.ingester.workers.pool"): + pool._on_heartbeat_done(died) + assert "Heartbeat task died" in caplog.text + + async def _exits(): + return None + + exited = asyncio.create_task(_exits()) + await exited + caplog.clear() + with caplog.at_level("ERROR", logger="haiku.rag.ingester.workers.pool"): + pool._on_heartbeat_done(exited) + assert "exited while the pool was still running" in caplog.text + + +@pytest.mark.asyncio +async def test_reaper_logs_when_it_resets_a_stale_claim( + client, jobs, sync, conn, caplog +): + """A claim that goes stale after startup (so boot-reap misses it) is reset + by the periodic reaper, which logs the reset count.""" + from datetime import UTC, datetime, timedelta + + pool = _pool( + client, jobs, sync, worker_count=0, reaper_interval_s=0.05, lease_ttl_s=1 + ) + await pool.start() + try: + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + claimed = await jobs.claim_next("external-worker") + assert claimed is not None + long_ago = (datetime.now(UTC) - timedelta(hours=1)).isoformat() + await conn.execute( + "UPDATE jobs SET claimed_at = ?, last_heartbeat_at = ? WHERE id = ?", + (long_ago, long_ago, job.id), + ) + await conn.commit() + + with caplog.at_level("INFO", logger="haiku.rag.ingester.workers.pool"): + for _ in range(40): + refreshed = await jobs.get_job(job.id) + if refreshed is not None and refreshed.status is JobStatus.QUEUED: + break + await asyncio.sleep(0.05) + assert refreshed is not None and refreshed.status is JobStatus.QUEUED + assert "Reaper reset" in caplog.text + finally: + await pool.stop() + + @pytest.mark.asyncio async def test_worker_loses_claim_to_reaper_does_not_write_sync_state( client, jobs, sync @@ -862,6 +980,26 @@ async def test_breaker_closes_on_successful_probe(client, jobs, sync): assert pool.breaker_consecutive_failures == 0 +@pytest.mark.asyncio +async def test_breaker_closes_when_success_lands_while_open(client, jobs, sync, caplog): + """A job that succeeds while the source's breaker is still open (drained + directly, bypassing the paused-source skip) closes it and logs recovery.""" + 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) + breaker = pool._breaker_for("src") + for _ in range(10): + breaker.record_failure() + assert breaker.is_open + + await jobs.enqueue("src", "u", JobOp.UPSERT) + with caplog.at_level("INFO", logger="haiku.rag.ingester.workers.pool"): + await pool.drain_once() + assert pool.breaker_consecutive_failures == 0 + assert "Worker breaker closed" in caplog.text + + @pytest.mark.asyncio async def test_breaker_isolates_sources(client, jobs, sync): """An open breaker pauses only the failing source. Workers keep draining