Extend reaper-resurrection guard to reschedule and release_if_claimed

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),
+                (scheduled, error, job_id, claimed_by),
-            ):
+            ) as cursor:
-                pass
+                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
+        Guarded on `status='claimed' AND claimed_by=?` so the cancel-cleanup
-        succeeded/dead/rescheduled. Decrements attempts to undo the increment
+        of a slow worker doesn't strip the claim of a different worker that
-        from `claim_next`, since a cancellation isn't a failed attempt.
+        re-claimed after a reaper reset. Decrements attempts to undo the
-        Returns True if the row was released."""
+        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()

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,

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

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