diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index 76a5cef7..4156ac9a 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -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() diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 7709cd0a..9e8bdacb 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -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, diff --git a/tests/ingester/test_queue.py b/tests/ingester/test_queue.py index fc6066a0..e7b50b9e 100644 --- a/tests/ingester/test_queue.py +++ b/tests/ingester/test_queue.py @@ -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 diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index 9644ab5b..6c1c3d97 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -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