RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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.
This commit is contained in:
Aakash Kattelu 2026-08-19 09:36:29 -04:00
parent 8e0c3afd75
commit f34f3e7f55
2 changed files with 146 additions and 254 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

230
src/crud/document.py
View File

@ -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(

170
tests/crud/test_document.py
View File

@ -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