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

Extend reaper-resurrection guard to reschedule and release_if_claimed

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-05-26 16:04:24 +03:00
parent 6799a4a6d4
commit 761956eb70
No known key found for this signature in database
4 changed files with 72 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -174,7 +174,13 @@ class JobRepo:
await self._conn.commit()
return row is not None
async def reschedule(self, job_id: str, delay_seconds: float, error: str) -> None:
async def reschedule(
self, job_id: str, delay_seconds: float, error: str, claimed_by: str
) -> bool:
"""Reset a still-claimed job back to `queued` with a future
scheduled_at. Guarded on `status='claimed' AND claimed_by=?` so a
slow worker can't clobber a re-claim that happened after the reaper
reset its claim. Returns True when the row was updated."""
scheduled = (datetime.now(UTC) + timedelta(seconds=delay_seconds)).isoformat()
async with self._lock:
async with self._conn.execute(
@ -185,12 +191,14 @@ class JobRepo:
claimed_at=NULL,
claimed_by=NULL,
last_error=?
WHERE id=?
WHERE id=? AND status='claimed' AND claimed_by=?
RETURNING id
""",
(scheduled, error, job_id),
):
pass
(scheduled, error, job_id, claimed_by),
) as cursor:
row = await cursor.fetchone()
await self._conn.commit()
return row is not None
async def retry(self, job_id: str) -> Job:
"""Reset a `dead` or `queued` job: status='queued', attempts=0,
@ -328,12 +336,13 @@ class JobRepo:
rows = await cursor.fetchall()
return {row["source_id"]: row["n"] for row in rows}
async def release_if_claimed(self, job_id: str) -> bool:
async def release_if_claimed(self, job_id: str, claimed_by: str) -> bool:
"""Reset a still-claimed job back to queued, immediately reclaimable.
Idempotent — a no-op if the job already transitioned to
succeeded/dead/rescheduled. Decrements attempts to undo the increment
from `claim_next`, since a cancellation isn't a failed attempt.
Returns True if the row was released."""
Guarded on `status='claimed' AND claimed_by=?` so the cancel-cleanup
of a slow worker doesn't strip the claim of a different worker that
re-claimed after a reaper reset. Decrements attempts to undo the
increment from `claim_next`, since a cancellation isn't a failed
attempt. Returns True if the row was released."""
now = _utcnow_iso()
async with self._lock:
async with self._conn.execute(
@ -344,10 +353,10 @@ class JobRepo:
claimed_by=NULL,
scheduled_at=?,
attempts=MAX(0, attempts - 1)
WHERE id=? AND status='claimed'
WHERE id=? AND status='claimed' AND claimed_by=?
RETURNING id
""",
(now, job_id),
(now, job_id, claimed_by),
) as cursor:
row = await cursor.fetchone()
await self._conn.commit()

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

@ -137,7 +137,7 @@ class WorkerPool:
# runs. The reaper is still the backstop, but releasing eagerly
# lets a restart re-pick the job immediately.
try:
await asyncio.shield(self._jobs.release_if_claimed(job.id))
await asyncio.shield(self._jobs.release_if_claimed(job.id, worker_id))
except asyncio.CancelledError:
logger.info(
"Job %s cancel-cleanup interrupted; reaper will reclaim", job.id
@ -157,7 +157,13 @@ class WorkerPool:
)
return
delay = compute_backoff(job.attempts, self._retry)
await self._jobs.reschedule(job.id, delay, str(e))
if not await self._jobs.reschedule(job.id, delay, str(e), worker_id):
logger.warning(
"Job %s lost claim before reschedule (likely reaper race); "
"letting the re-claiming worker drive retry instead",
job.id,
)
return
logger.info(
"Job %s rescheduled in %.1fs (attempt %d/%d): %s",
job.id,

46
tests/ingester/test_queue.py
View file

@ -307,6 +307,42 @@ async def test_mark_dead_with_claimed_by_guard_skips_when_resurrected(jobs):
assert refreshed.last_error is None
@pytest.mark.asyncio
async def test_reschedule_with_claimed_by_guard_skips_when_resurrected(jobs):
"""Worker A times out → reaper resets → worker B re-claims → A surfaces
with a TransientError and calls reschedule. The guard turns A's call
into a no-op so B's claim isn't stripped back to `queued`."""
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)
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
refreshed = await jobs.get_job(job.id)
assert refreshed is not None
assert refreshed.status is JobStatus.CLAIMED
assert refreshed.claimed_by == "worker-B"
assert refreshed.last_error is None
@pytest.mark.asyncio
async def test_release_if_claimed_with_claimed_by_guard_skips_when_resurrected(jobs):
"""Cancel-cleanup path: worker A is cancelled while reaper-resurrected.
release_if_claimed must not strip worker B's fresh claim."""
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)
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
refreshed = await jobs.get_job(job.id)
assert refreshed is not None
assert refreshed.status is JobStatus.CLAIMED
assert refreshed.claimed_by == "worker-B"
# --- reschedule + retry ---
@ -315,7 +351,7 @@ async def test_reschedule_pushes_scheduled_at_into_future(jobs):
job = await jobs.enqueue("s", "u", JobOp.UPSERT)
claimed = await jobs.claim_next("w")
assert claimed is not None
await jobs.reschedule(claimed.id, delay_seconds=30.0, error="transient")
await jobs.reschedule(claimed.id, 30.0, "transient", "w")
refreshed = await jobs.get_job(job.id)
assert refreshed is not None
assert refreshed.status is JobStatus.QUEUED
@ -330,7 +366,7 @@ async def test_reschedule_then_claim_skips_until_due(conn, jobs):
job = await jobs.enqueue("s", "u", JobOp.UPSERT)
claimed = await jobs.claim_next("w")
assert claimed is not None
await jobs.reschedule(claimed.id, delay_seconds=60.0, error="transient")
await jobs.reschedule(claimed.id, 60.0, "transient", "w")
assert await jobs.claim_next("w") is None
# Backdate scheduled_at to simulate the delay elapsing.
@ -471,7 +507,7 @@ async def test_release_if_claimed_resets_claimed_job_and_decrements_attempts(job
assert claimed is not None
assert claimed.attempts == 1
released = await jobs.release_if_claimed(job.id)
released = await jobs.release_if_claimed(job.id, "w")
assert released is True
refreshed = await jobs.get_job(job.id)
@ -487,7 +523,7 @@ async def test_release_if_claimed_noop_on_already_queued(jobs):
job = await jobs.enqueue("s", "u", JobOp.UPSERT)
assert job is not None
released = await jobs.release_if_claimed(job.id)
released = await jobs.release_if_claimed(job.id, "w")
assert released is False
refreshed = await jobs.get_job(job.id)
assert refreshed is not None
@ -503,7 +539,7 @@ async def test_release_if_claimed_noop_on_succeeded(jobs):
assert claimed is not None
await jobs.mark_succeeded(job.id, "w")
released = await jobs.release_if_claimed(job.id)
released = await jobs.release_if_claimed(job.id, "w")
assert released is False
refreshed = await jobs.get_job(job.id)
assert refreshed is not None

4
tests/ingester/test_workers.py
View file

@ -357,11 +357,11 @@ async def test_cancel_cleanup_survives_second_cancel(client, jobs, sync, monkeyp
real_release = jobs.release_if_claimed
async def _slow_release(job_id):
async def _slow_release(job_id, claimed_by):
release_entered.set()
# Long enough for the second cancel to arrive mid-update.
await asyncio.sleep(0.2)
result = await real_release(job_id)
result = await real_release(job_id, claimed_by)
release_done.set()
return result