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

Merge pull request #470 from ggozad/fix/dangling-workers

Browse source 
Lease-renewal reaping for ingester workers
This commit is contained in:
Yiorgis Gozadinos 2026-06-26 09:19:43 +03:00 committed by GitHub
commit 9f45addbe9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 845 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
CHANGELOG.md
View file

@ -1,13 +1,16 @@
# Changelog
## [Unreleased]
### Changed
- 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.
### 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

50
docs/ingester.md
View file

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

43
haiku_rag_slim/haiku/rag/config/models.py
View file

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

13
haiku_rag_slim/haiku/rag/ingester/app.py
View file

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

3
haiku_rag_slim/haiku/rag/ingester/queue/db.py
View file

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

16
haiku_rag_slim/haiku/rag/ingester/queue/migrations.py
View file

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

1
haiku_rag_slim/haiku/rag/ingester/queue/models.py
View file

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

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

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

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

@ -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
@ -26,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__(
@ -40,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,
@ -52,25 +57,62 @@ 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 []
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._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())
@ -98,30 +140,47 @@ 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._lease_ttl_s)
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())
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.
@ -168,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:
@ -176,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)
@ -185,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:

27
tests/ingester/test_config.py
View file

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

195
tests/ingester/test_queue.py
View file

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

25
tests/ingester/test_queue_postgres.py
View file

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

373
tests/ingester/test_workers.py
View file

@ -24,13 +24,44 @@ 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),
)
# --- 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(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
# 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 ---
@ -370,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):
@ -402,6 +433,275 @@ 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_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
@ -418,7 +718,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)
@ -448,7 +748,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)
@ -680,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
@ -796,16 +1116,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:
@ -820,6 +1149,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
@ -828,10 +1176,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()
@ -841,7 +1190,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: