diff --git a/application/tests/librarian/knowledge_source_test.py b/application/tests/librarian/knowledge_source_test.py index 27ce9c901..17dc8f9b6 100644 --- a/application/tests/librarian/knowledge_source_test.py +++ b/application/tests/librarian/knowledge_source_test.py @@ -12,8 +12,10 @@ import json import os import tempfile +import threading import unittest from datetime import datetime, timezone +from typing import List, Optional from application import create_app, sqla from application.database.db import KnowledgeQueueItem as KnowledgeQueueRow @@ -167,6 +169,74 @@ def test_unmodellable_row_is_skipped_not_fatal(self) -> None: self.assertEqual([i.id for i in items], ["a"]) + def test_concurrent_readers_skip_locked_rows(self) -> None: + """Two Module C workers must never both claim the same row. + + ``queue_runner.run_librarian_queue`` reads a batch, then runs the full + retrieval/rerank pipeline on it, then finally commits. If a second + worker's read is not fenced off from the first worker's still-open + transaction, both would process (and both would persist a decision + for) the same chunk. FOR UPDATE SKIP LOCKED must make the second + worker's read exclude rows the first is holding, rather than block on + them (which would just delay the double-processing) or read them + again. Postgres-only: SKIP LOCKED is a no-op on SQLite (there is + nothing to skip — FOR UPDATE itself is ignored), so this reproduces + the real bug only against Postgres, same as the existing + ``user_model_test`` row-lock test. + """ + if "postgresql" not in str(sqla.engine.url): + self.skipTest("row-lock serialization requires Postgres (SKIP LOCKED)") + + sqla.session.add_all([_row("a"), _row("b")]) + sqla.session.commit() + + t: Optional[threading.Thread] = None + try: + # Worker 1: read (and thereby lock) both rows, then hold the + # transaction open -- exactly queue_runner.py's shape, which does + # not commit until the whole batch, LLM calls included, has + # finished. + worker1_ids = [i.id for i in DbKnowledgeSource(sqla.session).items()] + self.assertEqual(sorted(worker1_ids), ["a", "b"]) + + worker2_ids: List[str] = [] + worker2_errors: List[BaseException] = [] + + def worker2() -> None: + with self.app.app_context(): + try: + items = list(DbKnowledgeSource(sqla.session).items()) + worker2_ids.extend(i.id for i in items) + except BaseException as exc: # noqa: BLE001 - surface below + worker2_errors.append(exc) + finally: + sqla.session.remove() + + t = threading.Thread(target=worker2) + t.start() + t.join(timeout=5) + + # A still-running thread means SKIP LOCKED failed to exclude the + # locked rows and worker2 is blocked waiting on them instead -- + # that is a failure, not a pass, so confirm it actually finished. + self.assertFalse(t.is_alive(), "worker2 did not finish -- it is blocked") + self.assertEqual(worker2_errors, []) + + # Worker 2 must see neither row: both are still locked by worker + # 1's open transaction, so SKIP LOCKED excludes them instead of + # blocking or (worse) reading and reprocessing them a second time. + self.assertEqual(worker2_ids, []) + finally: + # Release worker 1's row locks regardless of outcome, so a failed + # assertion above cannot leak a held lock into the next test. + sqla.session.rollback() + # If worker2 was still blocked when an assertion above failed, + # the rollback just now unblocks it -- join again so it is fully + # finished before tearDown() drops the tables out from under it. + if t is not None: + t.join(timeout=5) + if __name__ == "__main__": unittest.main() + diff --git a/application/utils/librarian/knowledge_source.py b/application/utils/librarian/knowledge_source.py index 87939dc9d..2b06324fc 100644 --- a/application/utils/librarian/knowledge_source.py +++ b/application/utils/librarian/knowledge_source.py @@ -84,6 +84,18 @@ class DbKnowledgeSource(KnowledgeSource): Rows are ordered by ``created_at`` then ``id``: the timestamp alone is not unique (B inserts a batch inside one transaction), and an unstable order would make a ``limit``ed run non-reproducible. + + **Concurrency.** ``items()`` claims every row it yields with + ``SELECT ... FOR UPDATE SKIP LOCKED`` (Postgres only — a no-op on SQLite, + matching ``db.py``'s existing ``with_for_update()`` use). The lock is held + for the life of the caller's transaction, i.e. until + ``queue_runner.run_librarian_queue`` commits after ``mark_consumed`` — which + is the whole point: without it, two concurrent runs (a retry overlapping a + scheduled pass, two orchestrator workers) would both read the same + unconsumed rows, both pay for the retrieval/rerank work, and both persist a + decision envelope for the same chunk before either reaches + ``mark_consumed``. SKIP LOCKED means a second reader simply excludes rows + the first is holding instead of blocking on them or re-reading them. """ def __init__( @@ -119,6 +131,8 @@ def _query(self) -> object: query = query.order_by(KnowledgeQueueRow.created_at, KnowledgeQueueRow.id) if self._limit is not None: query = query.limit(self._limit) + # Claim the batch: see the concurrency note on the class docstring. + query = query.with_for_update(skip_locked=True) return query def items(self) -> Iterator[KnowledgeQueueItem]: