Add last_heartbeat_at lease column to the ingester queue

This commit is contained in:
Yiorgis Gozadinos 2026-06-25 12:06:46 +03:00
parent ba6b318ece
commit fd73f57649
No known key found for this signature in database
9 changed files with 315 additions and 35 deletions

View file

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

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

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

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

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

View file

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

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

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

View file

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