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 # Changelog
## [Unreleased] ## [Unreleased]
### Added
- Ingester queue schema v2 adds a `last_heartbeat_at` lease column on `jobs`; existing queue databases migrate in place on open.
### Fixed ### Fixed
- A failed FTS index build is logged at `WARNING` instead of `DEBUG`, so silent full-text search degradation is visible. - 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 import event
from sqlalchemy.ext.asyncio import AsyncEngine from sqlalchemy.ext.asyncio import AsyncEngine
SCHEMA_VERSION = 1 SCHEMA_VERSION = 2
metadata = sa.MetaData() metadata = sa.MetaData()
@ -24,6 +24,7 @@ jobs = sa.Table(
sa.Column("scheduled_at", sa.Text, nullable=False), sa.Column("scheduled_at", sa.Text, nullable=False),
sa.Column("claimed_at", sa.Text), sa.Column("claimed_at", sa.Text),
sa.Column("claimed_by", sa.Text), sa.Column("claimed_by", sa.Text),
sa.Column("last_heartbeat_at", sa.Text),
sa.Column("completed_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() ).scalar_one_or_none()
if current is None: if current is None:
await conn.execute(sa.insert(schema_version).values(version=SCHEMA_VERSION)) await conn.execute(sa.insert(schema_version).values(version=SCHEMA_VERSION))
elif current < SCHEMA_VERSION: # pragma: no cover - no migrations yet elif current < SCHEMA_VERSION:
# No diff migrations exist yet — future versions add ALTER/UPDATE # create_all only creates missing tables/indexes, never adds columns
# statements between create_all and the version bump. # 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)) await conn.execute(sa.update(schema_version).values(version=SCHEMA_VERSION))
return SCHEMA_VERSION return SCHEMA_VERSION

View file

@ -35,6 +35,7 @@ class Job(BaseModel):
scheduled_at: datetime scheduled_at: datetime
claimed_at: datetime | None = None claimed_at: datetime | None = None
claimed_by: str | None = None claimed_by: str | None = None
last_heartbeat_at: datetime | None = None
completed_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"]), scheduled_at=datetime.fromisoformat(row["scheduled_at"]),
claimed_at=_parse_dt(row["claimed_at"]), claimed_at=_parse_dt(row["claimed_at"]),
claimed_by=row["claimed_by"], claimed_by=row["claimed_by"],
last_heartbeat_at=_parse_dt(row["last_heartbeat_at"]),
completed_at=_parse_dt(row["completed_at"]), completed_at=_parse_dt(row["completed_at"]),
) )
@ -157,6 +158,7 @@ class JobRepo:
status="claimed", status="claimed",
claimed_at=now, claimed_at=now,
claimed_by=worker_id, claimed_by=worker_id,
last_heartbeat_at=now,
attempts=jobs.c.attempts + 1, attempts=jobs.c.attempts + 1,
) )
.returning(*jobs.c) .returning(*jobs.c)
@ -230,6 +232,7 @@ class JobRepo:
scheduled_at=scheduled, scheduled_at=scheduled,
claimed_at=None, claimed_at=None,
claimed_by=None, claimed_by=None,
last_heartbeat_at=None,
last_error=error, last_error=error,
) )
.returning(jobs.c.id) .returning(jobs.c.id)
@ -259,6 +262,7 @@ class JobRepo:
last_error=None, last_error=None,
claimed_at=None, claimed_at=None,
claimed_by=None, claimed_by=None,
last_heartbeat_at=None,
completed_at=None, completed_at=None,
scheduled_at=now, scheduled_at=now,
) )
@ -449,6 +453,7 @@ class JobRepo:
status="queued", status="queued",
claimed_at=None, claimed_at=None,
claimed_by=None, claimed_by=None,
last_heartbeat_at=None,
scheduled_at=_utcnow_iso(), scheduled_at=_utcnow_iso(),
attempts=_attempts_minus_one(), attempts=_attempts_minus_one(),
) )
@ -485,20 +490,26 @@ class JobRepo:
result = await conn.execute(stmt) result = await conn.execute(stmt)
return result.rowcount or 0 return result.rowcount or 0
async def reap_stale(self, claim_timeout_seconds: int) -> int: async def reap_stale(self, lease_ttl_seconds: int) -> int:
"""Reset claimed jobs whose claimed_at is older than the timeout """Reset claimed jobs whose lease has gone stale back to `queued`. A
back to `queued`. Decrements `attempts` to undo the increment from live worker renews its lease via `renew_claims`, so a claim is stale
`claim_next` a crashed worker isn't a consumed attempt.""" 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 = ( threshold = (
datetime.now(UTC) - timedelta(seconds=claim_timeout_seconds) datetime.now(UTC) - timedelta(seconds=lease_ttl_seconds)
).isoformat() ).isoformat()
lease = sa.func.coalesce(jobs.c.last_heartbeat_at, jobs.c.claimed_at)
stmt = ( stmt = (
sa.update(jobs) sa.update(jobs)
.where(jobs.c.status == "claimed", jobs.c.claimed_at < threshold) .where(jobs.c.status == "claimed", lease < threshold)
.values( .values(
status="queued", status="queued",
claimed_at=None, claimed_at=None,
claimed_by=None, claimed_by=None,
last_heartbeat_at=None,
attempts=_attempts_minus_one(), attempts=_attempts_minus_one(),
) )
) )
@ -506,6 +517,30 @@ class JobRepo:
result = await conn.execute(stmt) result = await conn.execute(stmt)
return result.rowcount or 0 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: class SyncStateRepo:
def __init__(self, engine: AsyncEngine): def __init__(self, engine: AsyncEngine):

View file

@ -108,11 +108,11 @@ class WorkerPool:
if self._workers: if self._workers:
raise RuntimeError("WorkerPool already started") raise RuntimeError("WorkerPool already started")
self._stop.clear() self._stop.clear()
# Any rows in `claimed` at start time are owned by workers from a # Sweep claims left behind by a previous process (SIGKILL, OOM, host
# previous process that didn't get to release them (SIGKILL, OOM, # reboot) so fresh workers can take them over. Scoped by the lease TTL,
# host reboot). Reset them so fresh workers can claim immediately # not 0: a peer process sharing the queue may hold live claims, and
# instead of waiting on the reaper's claim_timeout_s. # those must not be wiped out from under it on our startup.
reset = await self._jobs.reap_stale(claim_timeout_seconds=0) reset = await self._jobs.reap_stale(lease_ttl_seconds=self._claim_timeout_s)
if reset: if reset:
logger.info("Boot-reaped %d stale claim(s) from previous process", reset) logger.info("Boot-reaped %d stale claim(s) from previous process", reset)
for i in range(self._worker_count): for i in range(self._worker_count):

View file

@ -3,8 +3,11 @@ from datetime import UTC, datetime, timedelta
import pytest import pytest
import sqlalchemy as sa 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.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.db import jobs as jobs_table
from haiku.rag.ingester.queue.migrations import ( from haiku.rag.ingester.queue.migrations import (
apply_migrations, apply_migrations,
@ -34,6 +37,81 @@ async def test_apply_migrations_is_idempotent(engine, conn):
assert row["n"] == 1 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 @pytest.mark.asyncio
async def test_open_queue_creates_file_and_schema(tmp_path): async def test_open_queue_creates_file_and_schema(tmp_path):
path = tmp_path / "subdir" / "queue.db" 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") a = await jobs.claim_next("worker-A")
assert a is not None assert a is not None
# Reaper resets A's stale claim. # 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 picks the job up.
b = await jobs.claim_next("worker-B") b = await jobs.claim_next("worker-B")
assert b is not None and b.id == job.id and b.claimed_by == "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) job = await jobs.enqueue("s", "u", JobOp.UPSERT)
a = await jobs.claim_next("worker-A") a = await jobs.claim_next("worker-A")
assert a is not None 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") b = await jobs.claim_next("worker-B")
assert b is not None and b.claimed_by == "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 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) job = await jobs.enqueue("s", "u", JobOp.UPSERT)
a = await jobs.claim_next("worker-A") a = await jobs.claim_next("worker-A")
assert a is not None 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") b = await jobs.claim_next("worker-B")
assert b is not None and b.claimed_by == "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 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) job = await jobs.enqueue("s", "u", JobOp.UPSERT)
a = await jobs.claim_next("worker-A") a = await jobs.claim_next("worker-A")
assert a is not None 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") b = await jobs.claim_next("worker-B")
assert b is not None and b.claimed_by == "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 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 is not None
assert claimed.attempts == 1 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() long_ago = (datetime.now(UTC) - timedelta(hours=1)).isoformat()
await conn.execute( 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() await conn.commit()
reset = await jobs.reap_stale(claim_timeout_seconds=60) reset = await jobs.reap_stale(lease_ttl_seconds=60)
assert reset == 1 assert reset == 1
refreshed = await jobs.get_job(job.id) refreshed = await jobs.get_job(job.id)
assert refreshed is not None assert refreshed is not None
assert refreshed.status is JobStatus.QUEUED assert refreshed.status is JobStatus.QUEUED
assert refreshed.claimed_at is None assert refreshed.claimed_at is None
assert refreshed.claimed_by 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. # isn't a consumed attempt — matches release_if_claimed semantics.
assert refreshed.attempts == 0 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 @pytest.mark.asyncio
async def test_reap_stale_leaves_fresh_claims_alone(jobs): async def test_reap_stale_leaves_fresh_claims_alone(jobs):
job = await jobs.enqueue("s", "u", JobOp.UPSERT) job = await jobs.enqueue("s", "u", JobOp.UPSERT)
claimed = await jobs.claim_next("w") claimed = await jobs.claim_next("w")
assert claimed is not None 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 assert reset == 0
refreshed = await jobs.get_job(job.id) refreshed = await jobs.get_job(job.id)
assert refreshed is not None assert refreshed is not None
assert refreshed.status is JobStatus.CLAIMED 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 --- # --- prune_dead ---

View file

@ -106,7 +106,7 @@ async def test_reap_stale_clamps_attempts(postgres_dburi):
assert job is not None assert job is not None
claimed = await jobs.claim_next("w") claimed = await jobs.claim_next("w")
assert claimed is not None and claimed.attempts == 1 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 assert reset == 1
refreshed = await jobs.get_job(job.id) refreshed = await jobs.get_job(job.id)
assert refreshed is not None assert refreshed is not None
@ -114,6 +114,29 @@ async def test_reap_stale_clamps_attempts(postgres_dburi):
assert refreshed.attempts == 0 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 @pytest.mark.asyncio
async def test_sync_state_upsert_coalesce_preserves_revision(postgres_dburi): async def test_sync_state_upsert_coalesce_preserves_revision(postgres_dburi):
"""ON CONFLICT DO UPDATE with COALESCE leaves an existing revision in place """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) claimed = await jobs.claim_next(worker_a)
assert claimed is not None assert claimed is not None
# Reaper resets the claim; pool B re-claims and finishes first. # 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) reclaimed = await jobs.claim_next(worker_b)
assert reclaimed is not None and reclaimed.id == job.id assert reclaimed is not None and reclaimed.id == job.id
assert await jobs.mark_succeeded(reclaimed.id, worker_b) is True 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") claimed_by_a = await jobs.claim_next("worker-A")
assert claimed_by_a is not None assert claimed_by_a is not None
# Reaper resets A's claim, worker-B re-claims. # 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") await jobs.claim_next("worker-B")
pool = _pool(client, jobs, sync, worker_count=1) 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") claimed_by_a = await jobs.claim_next("worker-A")
assert claimed_by_a is not None assert claimed_by_a is not None
# Reaper resets A's claim, worker-B re-claims. # 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") await jobs.claim_next("worker-B")
pool = _pool(client, jobs, sync, worker_count=1) 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 @pytest.mark.asyncio
async def test_boot_reap_resets_pre_existing_claims(client, jobs, sync): async def test_boot_reap_resets_stale_pre_existing_claims(client, jobs, sync, conn):
"""A SIGKILL'd previous process leaves rows in `claimed` state. The new """A SIGKILL'd previous process leaves a stale claim (its lease stopped
WorkerPool.start() must reset them immediately so fresh workers can being renewed). WorkerPool.start() sweeps it so fresh workers can take it
claim them, instead of waiting on the periodic reaper's claim_timeout_s over."""
window (default 1800s).""" from datetime import UTC, datetime, timedelta
await jobs.enqueue("src", "u", JobOp.UPSERT) await jobs.enqueue("src", "u", JobOp.UPSERT)
pre_claimed = await jobs.claim_next("ghost-worker") pre_claimed = await jobs.claim_next("ghost-worker")
assert pre_claimed is not None assert pre_claimed is not None
assert pre_claimed.status is JobStatus.CLAIMED 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) pool = _pool(client, jobs, sync, worker_count=0)
await pool.start() await pool.start()
try: try:
@ -850,6 +859,25 @@ async def test_boot_reap_resets_pre_existing_claims(client, jobs, sync):
await pool.stop() 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 @pytest.mark.asyncio
async def test_reaper_resets_stale_claims(client, jobs, sync, conn): async def test_reaper_resets_stale_claims(client, jobs, sync, conn):
from datetime import UTC, datetime, timedelta 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") claimed = await jobs.claim_next("worker-old")
assert claimed is not None 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() long_ago = (datetime.now(UTC) - timedelta(hours=1)).isoformat()
await conn.execute( 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() await conn.commit()