From f34f3e7f5540727d66994ccb173ac96ad3daf3a4 Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Wed, 19 Aug 2026 09:36:29 -0400 Subject: [PATCH] fix(deriver): replace create_documents advisory lock with id-ordered row locks Advisory locks are database-scoped and would serialize every writer to a collection, including across Groudon tenants that share names. Collect reinforcement and replace ops during the loop, lock target rows with SELECT ... ORDER BY id FOR UPDATE, then apply. populate_existing reloads times_derived so a prefetched identity-map row cannot lose a concurrent increment. --- src/crud/document.py | 230 +++++++++++++++++++----------------- tests/crud/test_document.py | 170 ++++---------------------- 2 files changed, 146 insertions(+), 254 deletions(-) diff --git a/src/crud/document.py b/src/crud/document.py index 677a6b5d..2261479f 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -4,9 +4,9 @@ from collections.abc import Sequence from dataclasses import dataclass, field from enum import Enum from logging import getLogger -from typing import Any, cast +from typing import Any, Literal, cast -from sqlalchemy import delete, select, text, update +from sqlalchemy import delete, select, update from sqlalchemy.engine import CursorResult from sqlalchemy.exc import DBAPIError, IntegrityError, SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession @@ -211,8 +211,7 @@ def _uses_pgvector() -> bool: ) -# Semantic-dup candidate search parameters, shared by is_rejected_duplicate -# and the pre-lock candidate resolution in create_documents. +# Shared by is_rejected_duplicate and create_documents candidate resolution. _SEMANTIC_DUP_MAX_DISTANCE = 0.05 _SEMANTIC_DUP_TOP_K = 1 @@ -489,6 +488,13 @@ def _dedup_key( ) +@dataclass(frozen=True, slots=True) +class _DocumentRowOp: + kind: Literal["reinforce", "replace"] + document_id: str + incoming_times_derived: int = 1 + + @dataclass class CreateDocumentsResult: created_documents: list[schemas.DocumentCreate] = field(default_factory=list) @@ -531,11 +537,8 @@ async def create_documents( # Store (document_model, embedding) pairs - IDs aren't available until after commit docs_with_embeddings: list[tuple[models.Document, list[float]]] = [] - # Resolve external-vector-store dup candidates up front, before the first - # DB statement: the advisory lock's critical section below must contain no - # network round trips, and the lazy session has no open transaction yet. - # A None entry falls back to in-place resolution in is_rejected_duplicate - # (the pgvector path, or a document with no valid merge partner). + # Resolve external-store dup candidates before the first DB statement. + # None falls back to in-place resolution (pgvector, or no merge partner). semantic_candidates: list[list[str] | None] = [None] * len(documents) if deduplicate and not _uses_pgvector(): @@ -602,27 +605,11 @@ async def create_documents( existing_doc, ) - # Serialize writers per (workspace, observer, observed) collection: two - # batches reinforcing the same rows in different orders deadlock otherwise - # (DEV-1975), and work-unit keys don't prevent concurrent writers to one - # collection. Skipped when the only writes are INSERTs of new rows, which - # can't deadlock on document rows. This advisory lock must stay the - # OUTERMOST lock taken in this transaction — don't add row-locking reads - # earlier in this function. SET LOCAL persists for the transaction, so the - # timeout also bounds the FOR KEY SHARE wait at the final INSERT; - # lock_timeout raises 55P03, which the queue layer classifies retryable. - # hashtextextended needs PG 11+; \x1f because names can contain ':'. - if documents and (deduplicate or existing_by_key): - lock_key = f"honcho:documents:{workspace_name}\x1f{observer}\x1f{observed}" - await db.execute(text("SET LOCAL lock_timeout = '30s'")) - await db.execute( - text("SELECT pg_advisory_xact_lock(hashtextextended(:key, 0))"), - {"key": lock_key}, - ) - # Tracks dedup keys already accepted from this batch so exact # duplicates within a single inference call collapse to one document. seen_in_batch: set[tuple[str, str, str | None]] = set() + row_ops: list[_DocumentRowOp] = [] + pending_times_derived: dict[str, int] = {} exact_dup_existing_count = 0 exact_dup_in_batch_count = 0 @@ -658,26 +645,20 @@ async def create_documents( # the re-derivation as reinforcement on the existing row. existing_match = existing_by_key.get(dedup_key) if existing_match is not None: - # Reinforce the existing row. greatest(...) keeps the bump atomic - # server-side (concurrent workers can't lose an increment) while - # still honoring an incoming doc that already carries accumulated - # reinforcement (times_derived > 1, e.g. a future re-ingestion or - # collection-merge path). Mirrors the superior-replacement branch - # in is_rejected_duplicate. - # Row lock held until commit; ordered by the advisory lock above. - existing_match.times_derived = func.greatest( - models.Document.times_derived + 1, - doc.times_derived, + current_td = pending_times_derived.get( + existing_match.id, existing_match.times_derived + ) + pending_times_derived[existing_match.id] = max( + current_td + 1, doc.times_derived + ) + row_ops.append( + _DocumentRowOp("reinforce", existing_match.id, doc.times_derived) ) - await db.flush() exact_dup_existing_count += 1 continue - # for each document, if deduplicate is True, perform a process - # that checks against existing documents and either rejects this document - # as a duplicate OR deletes an existing document that is a duplicate. if deduplicate: - duplicate_result = await is_rejected_duplicate( + duplicate_result, existing_dup = await _semantic_dup_decision( db, doc, workspace_name, @@ -685,11 +666,29 @@ async def create_documents( observed=observed, candidate_document_ids=semantic_candidates[index], ) - if duplicate_result is SemanticRejectionResult.REPLACED_EXISTING: - # Existing doc was soft-deleted in favor of this one; the - # new doc still gets inserted below. + if ( + duplicate_result is SemanticRejectionResult.REPLACED_EXISTING + and existing_dup is not None + ): + current_td = pending_times_derived.get( + existing_dup.id, existing_dup.times_derived + ) + doc.times_derived = max(doc.times_derived, current_td + 1) + row_ops.append(_DocumentRowOp("replace", existing_dup.id)) semantic_dup_replaced_count += 1 - elif duplicate_result is SemanticRejectionResult.REJECTED: + elif ( + duplicate_result is SemanticRejectionResult.REJECTED + and existing_dup is not None + ): + current_td = pending_times_derived.get( + existing_dup.id, existing_dup.times_derived + ) + pending_times_derived[existing_dup.id] = max( + current_td + 1, doc.times_derived + ) + row_ops.append( + _DocumentRowOp("reinforce", existing_dup.id, doc.times_derived) + ) semantic_dup_rejected_count += 1 continue @@ -753,6 +752,13 @@ async def create_documents( continue try: + await _apply_document_row_updates( + db, + row_ops, + workspace_name=workspace_name, + observer=observer, + observed=observed, + ) db.add_all(honcho_documents) # NOTE # If the process crashes after this commit but before vector upsert completes, @@ -1228,13 +1234,52 @@ async def create_observations( return honcho_documents +async def _apply_document_row_updates( + db: AsyncSession, + ops: list[_DocumentRowOp], + *, + workspace_name: str, + observer: str, + observed: str, +) -> None: + """Lock target rows by id, then apply reinforcements and replacements.""" + if not ops: + return + ids = sorted({op.document_id for op in ops}) + result = await db.execute( + select(models.Document) + .where( + models.Document.id.in_(ids), + models.Document.workspace_name == workspace_name, + models.Document.observer == observer, + models.Document.observed == observed, + ) + .order_by(models.Document.id) + .with_for_update() + .execution_options(populate_existing=True) + ) + locked = {doc.id: doc for doc in result.scalars()} + now = datetime.datetime.now(datetime.timezone.utc) + for op in ops: + row = locked.get(op.document_id) + if row is None: + continue + if op.kind == "replace": + row.deleted_at = now + continue + if row.deleted_at is not None: + continue + row.times_derived = max(row.times_derived + 1, op.incoming_times_derived) + await db.flush() + + class SemanticRejectionResult(Enum): NOT_DUPLICATE = 0 REPLACED_EXISTING = 1 REJECTED = 2 -async def is_rejected_duplicate( +async def _semantic_dup_decision( db: AsyncSession, doc: schemas.DocumentCreate, workspace_name: str, @@ -1242,45 +1287,12 @@ async def is_rejected_duplicate( observer: str, observed: str, candidate_document_ids: list[str] | None = None, -) -> SemanticRejectionResult: - """ - Check if a document is a duplicate of an existing document. - - Uses: 1) Cosine similarity (>=0.95), 2) Token diff for retention. - - Returns True if both: - - the document is deemed a duplicate of an existing document - - the existing document is deemed a superior duplicate - - If the document is not a duplicate, returns False. - - If the document is a duplicate AND the new document is superior, - deletes the existing document and returns False. In this case - ``doc.times_derived`` is updated in place to carry the replaced - document's reinforcement count forward. - - If the document is a duplicate AND the existing document is superior, - increments the existing document's ``times_derived`` to record the - reinforcement, then returns True. - - Merges are scoped so they never cross document levels, and never cross - sessions for explicit-level documents (session-purity invariant: an - explicit document records what was derived from exactly one session, so - a near-duplicate from another session must not reinforce or replace it). - - ``candidate_document_ids`` carries pre-resolved external-vector-store - candidates (see create_documents), keeping vector-store round trips out - of the caller's transaction; the ``deleted_at IS NULL`` filter in - ``fetch_documents_by_ids`` re-validates candidates that went stale since - resolution. None means resolve in place. - """ +) -> tuple[SemanticRejectionResult, models.Document | None]: + """Classify a semantic duplicate without writing.""" filters = _semantic_dup_filters(doc) if filters is None: - # create_documents refuses session-less explicit documents; if one - # reaches here anyway it has no valid merge partner. - return SemanticRejectionResult.NOT_DUPLICATE + return SemanticRejectionResult.NOT_DUPLICATE, None - # Step 1: Find potential duplicates using cosine similarity if candidate_document_ids is not None: similar_docs: Sequence[models.Document] = await fetch_documents_by_ids( db=db, @@ -1304,46 +1316,50 @@ async def is_rejected_duplicate( ) if not similar_docs: - return SemanticRejectionResult.NOT_DUPLICATE + return SemanticRejectionResult.NOT_DUPLICATE, None existing_doc = similar_docs[0] - - # Step 2: Determine which has more information using token set difference tokens_new = set(embedding_client.encoding.encode(doc.content)) tokens_existing = set(embedding_client.encoding.encode(existing_doc.content)) - unique_new = len(tokens_new - tokens_existing) unique_existing = len(tokens_existing - tokens_new) - score_new = len(tokens_new) + (unique_new * 10) score_existing = len(tokens_existing) + (unique_existing * 10) - - # If new document has more or equal information, keep it and delete existing if score_new >= score_existing: + return SemanticRejectionResult.REPLACED_EXISTING, existing_doc + return SemanticRejectionResult.REJECTED, existing_doc + + +async def is_rejected_duplicate( + db: AsyncSession, + doc: schemas.DocumentCreate, + workspace_name: str, + *, + observer: str, + observed: str, + candidate_document_ids: list[str] | None = None, +) -> SemanticRejectionResult: + """Classify a semantic duplicate and apply the corresponding row write.""" + result, existing_doc = await _semantic_dup_decision( + db, + doc, + workspace_name, + observer=observer, + observed=observed, + candidate_document_ids=candidate_document_ids, + ) + if existing_doc is None: + return result + if result is SemanticRejectionResult.REPLACED_EXISTING: logger.debug( "[DUPLICATE DETECTION] Deleting existing in favor of new. new=%r, existing=%r.", doc.content, existing_doc.content, ) - # Carry the reinforcement count forward so replacing a duplicate counts as - # another derivation rather than resetting times_derived to 1. doc.times_derived = max(doc.times_derived, existing_doc.times_derived + 1) - # Soft-delete the existing document - reconciliation will clean up vectors and hard-delete - # Row lock held until commit; concurrent create_documents writers are - # serialized by the advisory lock taken there. existing_doc.deleted_at = datetime.datetime.now(datetime.timezone.utc) await db.flush() - return ( - SemanticRejectionResult.REPLACED_EXISTING - ) # Don't reject the new document - - # Existing document has more information, reject the new one but record the - # reinforcement: a semantic duplicate was derived again. greatest(...) keeps - # the increment atomic server-side -- concurrent workers reinforcing the same - # document must not lose updates -- while still honoring an incoming doc that - # already carries accumulated reinforcement (times_derived > 1). - # Row lock held until commit; concurrent create_documents writers are - # serialized by the advisory lock taken there. + return result existing_doc.times_derived = func.greatest( models.Document.times_derived + 1, doc.times_derived, @@ -1354,7 +1370,7 @@ async def is_rejected_duplicate( doc.content, existing_doc.content, ) - return SemanticRejectionResult.REJECTED + return result async def cleanup_soft_deleted_documents( diff --git a/tests/crud/test_document.py b/tests/crud/test_document.py index f216f2a1..65059791 100644 --- a/tests/crud/test_document.py +++ b/tests/crud/test_document.py @@ -1340,16 +1340,7 @@ class TestSessionPurityInvariant: class TestCreateDocumentsConcurrency: - """Concurrency regression tests for create_documents (DEV-1975). - - Two deriver work units can write the same (workspace, observer, observed) - collection concurrently: work-unit keys include session but not observer, - so one peer active in two sessions produces two writers, and the dreamer - adds a third through a different key. Reinforcement UPDATEs issued in - batch order used to deadlock when two writers touched the same rows in - opposite orders; create_documents now serializes writers per collection - with a transaction-scoped advisory lock. - """ + """Concurrent same-collection reinforcements lock rows in id order.""" N_DOCS: int = 20 N_ROUNDS: int = 5 @@ -1411,8 +1402,7 @@ class TestCreateDocumentsConcurrency: db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ): - """Concurrent same-collection batches locking rows in opposite orders - must serialize, not deadlock (DEV-1975).""" + """Opposing-order batches on one collection must not deadlock.""" test_workspace, test_peer = sample_data test_peer2, test_session = await self._setup( db_session, test_workspace, test_peer @@ -1472,8 +1462,7 @@ class TestCreateDocumentsConcurrency: class TestCreateDocumentsErrorHandling: - """A dead transaction aborts the batch loudly; per-document failures - still skip just that document (DEV-1975).""" + """A dead transaction aborts the batch; per-document failures skip one document.""" async def _setup( self, @@ -1515,8 +1504,7 @@ class TestCreateDocumentsErrorHandling: db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], ): - """A DB error on an in-loop flush must raise out of create_documents - instead of cascading through a dead session, and commit nothing.""" + """A DB error while applying row updates raises and commits nothing.""" test_workspace, test_peer = sample_data test_peer2, test_session = await self._setup( db_session, test_workspace, test_peer @@ -1539,8 +1527,6 @@ class TestCreateDocumentsErrorHandling: sqlstate: str = "40P01" deadlock = OperationalError("UPDATE documents", {}, FakePGError()) - # Batch: an exact dup (triggers the reinforcement flush that "deadlocks") - # followed by a new document that must NOT be committed. with ( patch.object(db_session, "flush", AsyncMock(side_effect=deadlock)), pytest.raises(OperationalError), @@ -1614,37 +1600,32 @@ class TestCreateDocumentsErrorHandling: ] -class TestAdvisoryLockBehavior: - """The per-collection advisory lock in create_documents: held for the - transaction, keyed by (workspace, observer, observed), and skipped when - the batch can only INSERT.""" +class TestExternalCandidateHoist: + """External-store dup candidates resolve before the first DB statement.""" async def _setup( self, db_session: AsyncSession, test_workspace: models.Workspace, test_peer: models.Peer, - n_observed: int = 1, - ) -> tuple[list[models.Peer], models.Session]: - observed_peers = [ - models.Peer(name=str(generate_nanoid()), workspace_name=test_workspace.name) - for _ in range(n_observed) - ] + ) -> tuple[models.Peer, models.Session]: + observed_peer = models.Peer( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) test_session = models.Session( name=str(generate_nanoid()), workspace_name=test_workspace.name ) - db_session.add_all([*observed_peers, test_session]) + db_session.add_all([observed_peer, test_session]) await db_session.flush() - for observed_peer in observed_peers: - db_session.add( - models.Collection( - workspace_name=test_workspace.name, - observer=test_peer.name, - observed=observed_peer.name, - ) + db_session.add( + models.Collection( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=observed_peer.name, ) + ) await db_session.commit() - return observed_peers, test_session + return observed_peer, test_session def _doc(self, content: str, session_name: str) -> schemas.DocumentCreate: return schemas.DocumentCreate( @@ -1658,121 +1639,16 @@ class TestAdvisoryLockBehavior: ) @pytest.mark.asyncio - async def test_lock_blocks_same_key_not_other_keys( - self, - db_engine: AsyncEngine, - db_session: AsyncSession, - sample_data: tuple[models.Workspace, models.Peer], - ): - """While one writer holds a collection's lock, a second writer on the - same collection blocks and a writer on a different collection does not.""" - test_workspace, test_peer = sample_data - (observed_a, observed_b), test_session = await self._setup( - db_session, test_workspace, test_peer, n_observed=2 - ) - session_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False) - - holder_in_lock = asyncio.Event() - release_holder = asyncio.Event() - - real_is_rejected_duplicate = crud.document.is_rejected_duplicate - - async def pausing_is_rejected_duplicate(*args: Any, **kwargs: Any) -> Any: - holder_in_lock.set() - await release_holder.wait() - return await real_is_rejected_duplicate(*args, **kwargs) - - async def _create(observed: str, content: str) -> None: - async with session_factory() as db: - await crud.create_documents( - db, - [self._doc(content, test_session.name)], - workspace_name=test_workspace.name, - observer=test_peer.name, - observed=observed, - deduplicate=True, - ) - - with patch( - "src.crud.document.is_rejected_duplicate", - side_effect=pausing_is_rejected_duplicate, - ): - holder = asyncio.create_task(_create(observed_a.name, "holder fact")) - await asyncio.wait_for(holder_in_lock.wait(), timeout=10) - - # Holder is parked inside its critical section; the patch is no longer - # active so contender/other run the real dedup path. - contender = asyncio.create_task(_create(observed_a.name, "contender fact")) - other_key = asyncio.create_task(_create(observed_b.name, "other fact")) - - # A different collection's writer completes while the lock is held... - await asyncio.wait_for(other_key, timeout=10) - # ...but the same collection's writer stays blocked. - await asyncio.sleep(0.3) - assert not contender.done() - - release_holder.set() - await asyncio.wait_for(asyncio.gather(holder, contender), timeout=10) - - @pytest.mark.asyncio - async def test_lock_skipped_for_insert_only_batch( - self, - db_session: AsyncSession, - sample_data: tuple[models.Workspace, models.Peer], - ): - """deduplicate=False with no exact-dup matches only INSERTs, so no - advisory lock is taken; a batch with an existing exact dup takes it.""" - test_workspace, test_peer = sample_data - (observed_peer,), test_session = await self._setup( - db_session, test_workspace, test_peer - ) - - executed_sql: list[str] = [] - real_execute = db_session.execute - - async def spying_execute(statement: Any, *args: Any, **kwargs: Any) -> Any: - executed_sql.append(str(statement)) - return await real_execute(statement, *args, **kwargs) - - def lock_statements() -> list[str]: - return [s for s in executed_sql if "pg_advisory_xact_lock" in s] - - with patch.object(db_session, "execute", side_effect=spying_execute): - await crud.create_documents( - db_session, - [self._doc("all new fact", test_session.name)], - workspace_name=test_workspace.name, - observer=test_peer.name, - observed=observed_peer.name, - deduplicate=False, - ) - assert not lock_statements() - - # Same content again: the exact-dup prefetch now matches, so the - # reinforcement UPDATE must run under the lock. - await crud.create_documents( - db_session, - [self._doc("all new fact", test_session.name)], - workspace_name=test_workspace.name, - observer=test_peer.name, - observed=observed_peer.name, - deduplicate=False, - ) - assert len(lock_statements()) == 1 - - @pytest.mark.asyncio - async def test_external_candidates_resolved_before_lock( + async def test_external_candidates_resolved_before_db( self, db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer], monkeypatch: pytest.MonkeyPatch, ): - """On the external-vector-store path, every candidate resolution - happens before the advisory lock is acquired.""" from src.config import settings test_workspace, test_peer = sample_data - (observed_peer,), test_session = await self._setup( + observed_peer, test_session = await self._setup( db_session, test_workspace, test_peer ) monkeypatch.setattr(settings.VECTOR_STORE, "TYPE", "turbopuffer") @@ -1782,8 +1658,7 @@ class TestAdvisoryLockBehavior: real_execute = db_session.execute async def spying_execute(statement: Any, *args: Any, **kwargs: Any) -> Any: - if "pg_advisory_xact_lock" in str(statement): - events.append("lock") + events.append("execute") return await real_execute(statement, *args, **kwargs) async def fake_resolve(*_args: Any, **_kwargs: Any) -> list[str]: @@ -1810,4 +1685,5 @@ class TestAdvisoryLockBehavior: ) assert len(result.created_documents) == 2 - assert events == ["resolve", "resolve", "lock"] + assert events[:2] == ["resolve", "resolve"] + assert "execute" in events