From a822b754b78b16a7c7504ed7eb618bdb923f8a5a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 26 May 2026 14:07:44 +0300 Subject: [PATCH] Guard mark_succeeded/mark_dead against reaper resurrection --- .../haiku/rag/ingester/queue/repository.py | 30 ++++--- .../haiku/rag/ingester/workers/pool.py | 20 ++++- tests/ingester/test_api.py | 43 +++++++--- tests/ingester/test_pollers.py | 6 +- tests/ingester/test_queue.py | 86 +++++++++++++++---- tests/ingester/test_workers.py | 32 +++++++ 6 files changed, 168 insertions(+), 49 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index 3890b49e..76a5cef7 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -146,23 +146,33 @@ class JobRepo: row = await cursor.fetchone() return _row_to_job(row) if row else None - async def mark_succeeded(self, job_id: str) -> None: + async def mark_succeeded(self, job_id: str, claimed_by: str) -> bool: + """Transition a still-claimed job to `succeeded`. Guarded on + `status='claimed' AND claimed_by=?` so a reaper-resurrected job + picked up by a different worker isn't clobbered by the original + slow worker. Returns True when the row was updated.""" async with self._lock: async with self._conn.execute( - "UPDATE jobs SET status='succeeded', completed_at=? WHERE id=?", - (_utcnow_iso(), job_id), - ): - pass + "UPDATE jobs SET status='succeeded', completed_at=? " + "WHERE id=? AND status='claimed' AND claimed_by=? RETURNING id", + (_utcnow_iso(), job_id, claimed_by), + ) as cursor: + row = await cursor.fetchone() await self._conn.commit() + return row is not None - async def mark_dead(self, job_id: str, error: str) -> None: + async def mark_dead(self, job_id: str, error: str, claimed_by: str) -> bool: + """Transition a still-claimed job to `dead`. See `mark_succeeded` + for the guard semantics.""" async with self._lock: async with self._conn.execute( - "UPDATE jobs SET status='dead', completed_at=?, last_error=? WHERE id=?", - (_utcnow_iso(), error, job_id), - ): - pass + "UPDATE jobs SET status='dead', completed_at=?, last_error=? " + "WHERE id=? AND status='claimed' AND claimed_by=? RETURNING id", + (_utcnow_iso(), error, job_id, claimed_by), + ) as cursor: + row = await cursor.fetchone() await self._conn.commit() + return row is not None async def reschedule(self, job_id: str, delay_seconds: float, error: str) -> None: scheduled = (datetime.now(UTC) + timedelta(seconds=delay_seconds)).isoformat() diff --git a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py index 418a2771..7709cd0a 100644 --- a/haiku_rag_slim/haiku/rag/ingester/workers/pool.py +++ b/haiku_rag_slim/haiku/rag/ingester/workers/pool.py @@ -123,6 +123,8 @@ class WorkerPool: pass async def _process(self, job: Job) -> None: + assert job.claimed_by is not None, "_process only runs on claimed jobs" + worker_id = job.claimed_by started = time.monotonic() logger.info("Processing %s %s (job %s)", job.op.value, job.uri, job.id) try: @@ -144,12 +146,12 @@ class WorkerPool: logger.info("Job %s released back to queue on cancel", job.id) raise except PermanentError as e: - await self._jobs.mark_dead(job.id, str(e)) + await self._jobs.mark_dead(job.id, str(e), worker_id) logger.info("Job %s dead (permanent): %s", job.id, e) return except TransientError as e: if job.attempts >= job.max_attempts: - await self._jobs.mark_dead(job.id, str(e)) + await self._jobs.mark_dead(job.id, str(e), worker_id) logger.info( "Job %s dead (max attempts %d): %s", job.id, job.max_attempts, e ) @@ -167,11 +169,21 @@ class WorkerPool: return except Exception as e: # pragma: no cover - pipeline classifier net # Defensive: pipeline classifier should have caught everything. - await self._jobs.mark_dead(job.id, f"unclassified: {e!r}") + await self._jobs.mark_dead(job.id, f"unclassified: {e!r}", worker_id) logger.exception("Unclassified error in job %s", job.id) return - await self._jobs.mark_succeeded(job.id) + # Guard against the reaper race: if our claim was reset and another + # worker re-claimed the job, mark_succeeded is a no-op. Don't write + # sync_state in that case — the new worker will write it when it + # finishes. + if not await self._jobs.mark_succeeded(job.id, worker_id): + logger.warning( + "Job %s lost claim before mark_succeeded (likely reaper race); " + "skipping sync_state write", + job.id, + ) + return if job.op is JobOp.DELETE: await self._sync.delete(job.source_id, job.uri) else: diff --git a/tests/ingester/test_api.py b/tests/ingester/test_api.py index 8d80ea83..949030d0 100644 --- a/tests/ingester/test_api.py +++ b/tests/ingester/test_api.py @@ -64,10 +64,14 @@ def _client(state, *, auth_token: str | None = None) -> httpx.AsyncClient: @pytest.mark.asyncio async def test_health_ok_with_counts(state, jobs): - await jobs.enqueue("src", "u1", JobOp.UPSERT) + # Enqueue j2 first so claim_next reaches it before u1; then transition + # via claim → mark_dead matches the production path. j2 = await jobs.enqueue("src", "u2", JobOp.UPSERT) assert j2 is not None - await jobs.mark_dead(j2.id, "boom") + claimed = await jobs.claim_next("w") + assert claimed is not None and claimed.id == j2.id + await jobs.mark_dead(j2.id, "boom", "w") + await jobs.enqueue("src", "u1", JobOp.UPSERT) async with _client(state) as client: resp = await client.get("/health") @@ -150,7 +154,9 @@ async def test_mutation_endpoints_require_auth(state, jobs): cancel jobs or reset the DLQ.""" j = await jobs.enqueue("src", "u", JobOp.UPSERT) assert j is not None - await jobs.mark_dead(j.id, "boom") + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_dead(j.id, "boom", "w") async with _client(state, auth_token="secret") as client: # Cancel: blocked without token @@ -191,10 +197,13 @@ async def test_list_jobs_returns_recent_first(state, jobs): @pytest.mark.asyncio async def test_list_jobs_filters_by_source_and_status(state, jobs): - await jobs.enqueue("a", "u", JobOp.UPSERT) + # Enqueue b first so claim_next reaches it before the a row. j = await jobs.enqueue("b", "u", JobOp.UPSERT) assert j is not None - await jobs.mark_dead(j.id, "err") + claimed = await jobs.claim_next("w") + assert claimed is not None and claimed.id == j.id + await jobs.mark_dead(j.id, "err", "w") + await jobs.enqueue("a", "u", JobOp.UPSERT) async with _client(state) as client: resp = await client.get("/jobs?source_id=b&status=dead") @@ -224,7 +233,9 @@ async def test_get_job_404(state): async def test_retry_revives_dead_job(state, jobs): job = await jobs.enqueue("src", "u", JobOp.UPSERT) assert job is not None - await jobs.mark_dead(job.id, "err") + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_dead(job.id, "err", "w") async with _client(state) as client: resp = await client.post(f"/jobs/{job.id}/retry") @@ -259,7 +270,7 @@ async def test_cancel_succeeded_returns_404(state, jobs): assert job is not None claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_succeeded(claimed.id) + await jobs.mark_succeeded(claimed.id, "w") async with _client(state) as client: resp = await client.delete(f"/jobs/{job.id}") @@ -271,10 +282,14 @@ async def test_cancel_succeeded_returns_404(state, jobs): @pytest.mark.asyncio async def test_dlq_lists_dead_jobs_only(state, jobs): - j1 = await jobs.enqueue("src", "u1", JobOp.UPSERT) + # Enqueue j2 first so claim_next picks it before j1. j2 = await jobs.enqueue("src", "u2", JobOp.UPSERT) - assert j1 is not None and j2 is not None - await jobs.mark_dead(j2.id, "err") + assert j2 is not None + claimed = await jobs.claim_next("w") + assert claimed is not None and claimed.id == j2.id + await jobs.mark_dead(j2.id, "err", "w") + j1 = await jobs.enqueue("src", "u1", JobOp.UPSERT) + assert j1 is not None async with _client(state) as client: resp = await client.get("/dlq") @@ -287,7 +302,9 @@ async def test_dlq_lists_dead_jobs_only(state, jobs): async def test_dlq_retry_resurrects(state, jobs): job = await jobs.enqueue("src", "u", JobOp.UPSERT) assert job is not None - await jobs.mark_dead(job.id, "err") + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_dead(job.id, "err", "w") async with _client(state) as client: resp = await client.post(f"/dlq/{job.id}/retry") @@ -429,10 +446,10 @@ async def test_stats_aggregates_real_queue(state, jobs): claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_succeeded(claimed.id) + await jobs.mark_succeeded(claimed.id, "w") dead = await jobs.claim_next("w") assert dead is not None - await jobs.mark_dead(dead.id, "boom") + await jobs.mark_dead(dead.id, "boom", "w") async with _client(state) as client: resp = await client.get("/stats") diff --git a/tests/ingester/test_pollers.py b/tests/ingester/test_pollers.py index d1e06e23..34e62d38 100644 --- a/tests/ingester/test_pollers.py +++ b/tests/ingester/test_pollers.py @@ -202,7 +202,7 @@ async def test_skipped_sweep_records_pending_work_reason(fs_config, jobs, sync): # Drain the queue, sweep again, reason clears. claimed = await jobs.claim_next("worker") assert claimed is not None - await jobs.mark_succeeded(claimed.id) + await jobs.mark_succeeded(claimed.id, "worker") await poller._sweep_once() assert poller.last_skip_reason is None @@ -224,7 +224,7 @@ async def test_dead_job_does_not_clear_sync_state_revision(fs_config, jobs, sync claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_dead(claimed.id, "transient blew up") + await jobs.mark_dead(claimed.id, "transient blew up", "w") row = await sync.get_row("src", "file:///a.md") assert row is not None @@ -244,7 +244,7 @@ async def test_sweep_resumes_after_queue_drains(fs_config, jobs, sync): claimed = await jobs.claim_next("worker") assert claimed is not None - await jobs.mark_succeeded(claimed.id) + await jobs.mark_succeeded(claimed.id, "worker") assert await poller._sweep_once() is True assert source.discover_calls == 2 diff --git a/tests/ingester/test_queue.py b/tests/ingester/test_queue.py index 762c5db0..fc6066a0 100644 --- a/tests/ingester/test_queue.py +++ b/tests/ingester/test_queue.py @@ -70,7 +70,9 @@ async def test_open_queue_creates_file_and_schema(tmp_path): @pytest.mark.asyncio async def test_dlq_view_exposes_dead_jobs(conn, jobs): job = await jobs.enqueue("s", "u", JobOp.UPSERT) - await jobs.mark_dead(job.id, "permanent") + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_dead(job.id, "permanent", "w") cursor = await conn.execute("SELECT id FROM dlq") rows = await cursor.fetchall() @@ -113,7 +115,9 @@ async def test_enqueue_returns_none_on_live_conflict(jobs): @pytest.mark.asyncio async def test_enqueue_after_dead_succeeds(jobs): first = await jobs.enqueue("s", "u", JobOp.UPSERT) - await jobs.mark_dead(first.id, "error") + claimed = await jobs.claim_next("w") + assert claimed is not None + await jobs.mark_dead(first.id, "error", "w") second = await jobs.enqueue("s", "u", JobOp.UPSERT) assert second is not None assert second.id != first.id @@ -124,7 +128,7 @@ async def test_enqueue_after_succeeded_succeeds(jobs): await jobs.enqueue("s", "u", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_succeeded(claimed.id) + await jobs.mark_succeeded(claimed.id, "w") second = await jobs.enqueue("s", "u", JobOp.UPSERT) assert second is not None @@ -152,12 +156,12 @@ async def test_has_pending_false_for_terminal_states(jobs): await jobs.enqueue("src", "u-ok", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_succeeded(claimed.id) + await jobs.mark_succeeded(claimed.id, "w") await jobs.enqueue("src", "u-bad", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_dead(claimed.id, "permanent") + await jobs.mark_dead(claimed.id, "permanent", "w") assert await jobs.has_pending("src") is False @@ -246,7 +250,7 @@ async def test_mark_succeeded(jobs): job = await jobs.enqueue("s", "u", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_succeeded(claimed.id) + await jobs.mark_succeeded(claimed.id, "w") refreshed = await jobs.get_job(job.id) assert refreshed is not None assert refreshed.status is JobStatus.SUCCEEDED @@ -258,13 +262,51 @@ async def test_mark_dead_records_error(jobs): job = await jobs.enqueue("s", "u", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_dead(claimed.id, "permanent failure") + await jobs.mark_dead(claimed.id, "permanent failure", "w") refreshed = await jobs.get_job(job.id) assert refreshed is not None assert refreshed.status is JobStatus.DEAD assert refreshed.last_error == "permanent failure" +@pytest.mark.asyncio +async def test_mark_succeeded_with_claimed_by_guard_skips_when_resurrected(jobs): + """Reaper race: A claims, reaper resets, B re-claims, A finishes. A's + mark_succeeded must be a no-op so B's in-flight work isn't clobbered.""" + job = await jobs.enqueue("s", "u", JobOp.UPSERT) + 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) + # 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" + # A finally finishes and tries to mark succeeded — guarded, no-op. + assert await jobs.mark_succeeded(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" + # B's own mark_succeeded does land. + assert await jobs.mark_succeeded(job.id, "worker-B") is True + + +@pytest.mark.asyncio +async def test_mark_dead_with_claimed_by_guard_skips_when_resurrected(jobs): + """Same guard semantics as mark_succeeded for the dead transition.""" + 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.mark_dead(job.id, "boom", "worker-A") is False + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.CLAIMED + assert refreshed.last_error is None + + # --- reschedule + retry --- @@ -303,7 +345,7 @@ async def test_retry_revives_dead_job(jobs): job = await jobs.enqueue("s", "u", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_dead(claimed.id, "boom") + await jobs.mark_dead(claimed.id, "boom", "w") revived = await jobs.retry(job.id) assert revived.status is JobStatus.QUEUED @@ -347,7 +389,7 @@ async def test_retry_refuses_succeeded_job(jobs): job = await jobs.enqueue("s", "u", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_succeeded(claimed.id) + await jobs.mark_succeeded(claimed.id, "w") with pytest.raises(KeyError): await jobs.retry(job.id) @@ -372,7 +414,7 @@ async def test_cancel_succeeded_returns_false(jobs): job = await jobs.enqueue("s", "u", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_succeeded(claimed.id) + await jobs.mark_succeeded(claimed.id, "w") assert await jobs.cancel(job.id) is False @@ -459,7 +501,7 @@ async def test_release_if_claimed_noop_on_succeeded(jobs): assert job is not None claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_succeeded(job.id) + await jobs.mark_succeeded(job.id, "w") released = await jobs.release_if_claimed(job.id) assert released is False @@ -478,7 +520,7 @@ async def test_list_jobs_with_filters(jobs): j3 = await jobs.enqueue("s1", "u3", JobOp.UPSERT) claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_succeeded(claimed.id) + await jobs.mark_succeeded(claimed.id, "w") all_jobs = await jobs.list_jobs() assert len(all_jobs) == 3 @@ -498,8 +540,10 @@ async def test_counts_by_status(jobs): claimed = await jobs.claim_next("w") assert claimed is not None - await jobs.mark_succeeded(claimed.id) - await jobs.mark_dead(j2.id, "err") + await jobs.mark_succeeded(claimed.id, "w") + claimed_j2 = await jobs.claim_next("w") + assert claimed_j2 is not None and claimed_j2.id == j2.id + await jobs.mark_dead(j2.id, "err", "w") counts = await jobs.counts_by_status() assert counts == {"queued": 1, "succeeded": 1, "dead": 1} @@ -516,7 +560,7 @@ async def test_count_succeeded_since_only_includes_recent(jobs, conn): old_claim = await jobs.claim_next("w") assert old_claim is not None - await jobs.mark_succeeded(old_claim.id) + await jobs.mark_succeeded(old_claim.id, "w") long_ago = (datetime.now(UTC) - timedelta(hours=2)).isoformat() await conn.execute( "UPDATE jobs SET completed_at = ? WHERE id = ?", (long_ago, old_claim.id) @@ -525,7 +569,7 @@ async def test_count_succeeded_since_only_includes_recent(jobs, conn): new_claim = await jobs.claim_next("w") assert new_claim is not None - await jobs.mark_succeeded(new_claim.id) + await jobs.mark_succeeded(new_claim.id, "w") assert await jobs.count_succeeded_since(60) == 1 assert await jobs.count_succeeded_since(86400) == 2 @@ -566,11 +610,15 @@ async def test_oldest_queued_age_seconds_ignores_future_scheduled(jobs, conn): @pytest.mark.asyncio async def test_counts_by_source_groups_correctly(jobs): - await jobs.enqueue("s1", "u1", JobOp.UPSERT) - await jobs.enqueue("s1", "u2", JobOp.UPSERT) + # Enqueue s2 first so claim_next picks it up before the s1 rows; then + # mark_dead routes through the production claim→terminal transition. j3 = await jobs.enqueue("s2", "u3", JobOp.UPSERT) assert j3 is not None - await jobs.mark_dead(j3.id, "boom") + claimed = await jobs.claim_next("w") + assert claimed is not None and claimed.id == j3.id + await jobs.mark_dead(j3.id, "boom", "w") + await jobs.enqueue("s1", "u1", JobOp.UPSERT) + await jobs.enqueue("s1", "u2", JobOp.UPSERT) assert await jobs.counts_by_source("queued") == {"s1": 2} assert await jobs.counts_by_source("dead") == {"s2": 1} diff --git a/tests/ingester/test_workers.py b/tests/ingester/test_workers.py index b115e305..9644ab5b 100644 --- a/tests/ingester/test_workers.py +++ b/tests/ingester/test_workers.py @@ -314,6 +314,38 @@ async def test_shutdown_grace_timeout_releases_claim(client, jobs, sync): assert refreshed.attempts == 0 +@pytest.mark.asyncio +async def test_worker_loses_claim_to_reaper_does_not_write_sync_state( + client, jobs, sync +): + """If the reaper resets a slow worker's claim and another worker re-claims + the job, the original worker's mark_succeeded must be a no-op and its + sync_state.upsert must not run — otherwise we'd overwrite freshly-written + state from the re-claiming worker.""" + client.create_document_from_source.return_value = Document( + id="doc-A", content="x", uri="u", metadata={"md5": "A", "source_revision": "A"} + ) + job = await jobs.enqueue("src", "u", JobOp.UPSERT) + assert job is not None + 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.claim_next("worker-B") + + pool = _pool(client, jobs, sync, worker_count=1, max_concurrent=1) + # Drive A's _process directly with A's (now stale) Job snapshot. + await pool._process(claimed_by_a) + + refreshed = await jobs.get_job(job.id) + assert refreshed is not None + # B still owns the claim — A's mark_succeeded was a no-op. + assert refreshed.status is JobStatus.CLAIMED + assert refreshed.claimed_by == "worker-B" + # And sync_state must be untouched. + assert await sync.get_snapshot("src") == {} + + @pytest.mark.asyncio async def test_cancel_cleanup_survives_second_cancel(client, jobs, sync, monkeypatch): """A second cancel arriving while the cancel-handler is awaiting