diff --git a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py index e7aaaafc..86724e65 100644 --- a/haiku_rag_slim/haiku/rag/ingester/queue/repository.py +++ b/haiku_rag_slim/haiku/rag/ingester/queue/repository.py @@ -127,34 +127,34 @@ class JobRepo: async def claim_next(self, worker_id: str) -> Job | None: """Atomically claim the oldest queued job whose scheduled_at <= now. - A `SELECT ... FOR UPDATE SKIP LOCKED` picks the row (multi-process safe - on Postgres; the clause is omitted on SQLite, where pool_size=1 keeps - the select-then-update atomic), then a guarded UPDATE claims it.""" + A single `UPDATE ... WHERE id = (SELECT ... LIMIT 1) RETURNING` keeps + the claim atomic across connections: on Postgres the subquery adds + `FOR UPDATE SKIP LOCKED`; on SQLite the whole statement runs under one + write lock, so a racing connection re-evaluates the subquery against + the committed state and finds the row already claimed.""" now = _utcnow_iso() - select_candidate = ( + candidate = ( sa.select(jobs.c.id) .where(jobs.c.status == "queued", jobs.c.scheduled_at <= now) - .order_by(jobs.c.scheduled_at) + .order_by(jobs.c.scheduled_at, jobs.c.id) .limit(1) .with_for_update(skip_locked=True) + .scalar_subquery() + ) + claim = ( + sa.update(jobs) + .where(jobs.c.id == candidate) + .values( + status="claimed", + claimed_at=now, + claimed_by=worker_id, + attempts=jobs.c.attempts + 1, + ) + .returning(*jobs.c) ) async with self._engine.begin() as conn: - job_id = (await conn.execute(select_candidate)).scalar_one_or_none() - if job_id is None: - return None - claim = ( - sa.update(jobs) - .where(jobs.c.id == job_id) - .values( - status="claimed", - claimed_at=now, - claimed_by=worker_id, - attempts=jobs.c.attempts + 1, - ) - .returning(*jobs.c) - ) - row = (await conn.execute(claim)).mappings().one() - return _row_to_job(row) + row = (await conn.execute(claim)).mappings().one_or_none() + return _row_to_job(row) if row else None async def get_job(self, job_id: str) -> Job | None: async with self._engine.connect() as conn: diff --git a/tests/ingester/test_queue.py b/tests/ingester/test_queue.py index 71575aea..c7d9529b 100644 --- a/tests/ingester/test_queue.py +++ b/tests/ingester/test_queue.py @@ -7,6 +7,7 @@ import sqlalchemy as sa from haiku.rag.config import QueueConfig from haiku.rag.ingester.queue.migrations import apply_migrations, open_queue from haiku.rag.ingester.queue.models import JobOp, JobStatus, SyncRow +from haiku.rag.ingester.queue.repository import JobRepo # --- migrations / schema --- @@ -229,6 +230,43 @@ async def test_claim_next_atomic_under_concurrency(jobs): assert {c.id for c in claimed} == {j.id for j in enqueued} +@pytest.mark.asyncio +async def test_claim_next_atomic_across_independent_engines(tmp_path): + """Two engines on one SQLite file stand in for two ingester processes. + The claim is a single UPDATE...WHERE id=(subquery), so it stays atomic + across connections — pool_size=1 only serializes within one engine.""" + cfg = QueueConfig(path=tmp_path / "shared-queue.db") + engine_a = await open_queue(cfg) + engine_b = await open_queue(cfg) + try: + repo_a = JobRepo(engine_a) + repo_b = JobRepo(engine_b) + + enqueued = [] + for i in range(10): + job = await repo_a.enqueue("s", f"u{i}", JobOp.UPSERT) + assert job is not None + enqueued.append(job) + + # Claims alternate across the two engines, contending on the same file. + results = await asyncio.gather( + *((repo_a if i % 2 == 0 else repo_b).claim_next(f"w{i}") for i in range(20)) + ) + claimed = [r for r in results if r is not None] + + assert len(claimed) == 10 + assert len({c.id for c in claimed}) == 10 + assert {c.id for c in claimed} == {j.id for j in enqueued} + for job in enqueued: + refreshed = await repo_a.get_job(job.id) + assert refreshed is not None + assert refreshed.status is JobStatus.CLAIMED + assert refreshed.attempts == 1 + finally: + await engine_a.dispose() + await engine_b.dispose() + + # --- terminal transitions ---