Renew job leases from the worker pool; lease-based reaping

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.
+- 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.
-
-### 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.
 
 ## [0.61.2] - 2026-06-24
 

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(
+    lease_ttl_s: int = Field(
-        default=1800,
+        default=120,
-        description="A `claimed` job whose claimed_at is older than this is "
+        gt=0,
-        "presumed dead and reset to `queued` by the reaper. MUST exceed your "
+        description="A `claimed` job whose lease has not been renewed within "
-        "longest legitimate job duration — set it too short and the reaper "
+        "this window is presumed dead and reset to `queued` by the reaper. A "
-        "resurrects jobs still being processed, causing two workers to run "
+        "live worker renews its lease every heartbeat_interval_s while "
-        "the same URI. Default (30min) covers typical docling conversions; "
+        "processing, so this need not exceed job duration — it only bounds how "
-        "raise it if you ingest very large PDFs through docling-local.",
+        "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 "
+        "finish before forcing cancellation. Cancelled jobs are released back "
-        "the queue; the reaper resets them after claim_timeout_s.",
+        "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."""

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
+            # Cancelling the workers triggers their cancel-cleanup, which
-            # claim_timeout_s on the next start.
+            # 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 "
+                "cancelling — they'll be released back to the queue (or "
-                "next start",
+                "reclaimed after lease_ttl_s if release doesn't land)",
                 grace_s,
             )
         landed = await self._pool.drain_pending_releases(timeout=2.0)

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
+    through `run_job`. While a worker processes a job it renews the job's lease
-    `claim_timeout_s` so a crashed worker doesn't strand its job. Lifecycle:
+    on a `heartbeat_interval_s` cadence; a reaper task resets claims whose lease
-    build it, await start(), let it run, await stop().
+    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:

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",
     [

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: