fix: lock knowledge_queue rows with SELECT ... FOR UPDATE SKIP LOCKED to prevent duplicate consumption (#1025)

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()
+

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]: