From 8e0c3afd757c76d6927767ab7949a37fe735547f Mon Sep 17 00:00:00 2001 From: Aakash Kattelu Date: Tue, 18 Aug 2026 11:04:55 -0400 Subject: [PATCH] fix(deriver): eliminate create_documents deadlock and stop silently burning batches on transient errors Two concurrent work units writing the same (workspace, observer, observed) collection deadlocked on times_derived reinforcement UPDATEs issued in batch order (DEV-1975, 682 events in 90 days). The deadlock was swallowed per-document, the loop cascaded PendingRollbackErrors against the dead session, the whole batch was lost, and the queue item was marked processed. - serialize writers per collection with a transaction-scoped advisory lock (pg_advisory_xact_lock + SET LOCAL lock_timeout), skipped for insert-only batches; covers all three row-lock sites in one move - hoist external-vector-store dup-candidate resolution ahead of the first DB statement so the lock's critical section contains no network calls - abort the batch on SQLAlchemyError instead of continuing through an aborted transaction; per-document skip semantics kept for non-DB errors - classify transient errors (new src/utils/retryable_errors.py) and retry them via a bounded in-process counter instead of marking items errored --- src/crud/document.py | 144 ++++++-- src/deriver/queue_manager.py | 62 +++- src/utils/retryable_errors.py | 86 +++++ tests/crud/test_document.py | 479 ++++++++++++++++++++++++- tests/deriver/test_queue_processing.py | 233 ++++++++++++ tests/utils/test_retryable_errors.py | 117 ++++++ 6 files changed, 1090 insertions(+), 31 deletions(-) create mode 100644 src/utils/retryable_errors.py create mode 100644 tests/utils/test_retryable_errors.py diff --git a/src/crud/document.py b/src/crud/document.py index ba80712d..677a6b5d 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -1,3 +1,4 @@ +import asyncio import datetime from collections.abc import Sequence from dataclasses import dataclass, field @@ -5,9 +6,9 @@ from enum import Enum from logging import getLogger from typing import Any, cast -from sqlalchemy import delete, select, update +from sqlalchemy import delete, select, text, update from sqlalchemy.engine import CursorResult -from sqlalchemy.exc import IntegrityError +from sqlalchemy.exc import DBAPIError, IntegrityError, SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.sql import Select from sqlalchemy.sql.functions import func @@ -210,6 +211,24 @@ def _uses_pgvector() -> bool: ) +# Semantic-dup candidate search parameters, shared by is_rejected_duplicate +# and the pre-lock candidate resolution in create_documents. +_SEMANTIC_DUP_MAX_DISTANCE = 0.05 +_SEMANTIC_DUP_TOP_K = 1 + + +def _semantic_dup_filters(doc: schemas.DocumentCreate) -> dict[str, Any] | None: + """Merge scope for semantic dedup: never across levels, never across + sessions for explicit documents. None when the document has no valid + merge partner (session-less explicit).""" + filters: dict[str, Any] = {"level": doc.level} + if doc.level == "explicit": + if doc.session_name is None: + return None + filters["session_name"] = doc.session_name + return filters + + async def query_external_vector_document_ids( workspace_name: str, observer: str, @@ -512,6 +531,32 @@ 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). + semantic_candidates: list[list[str] | None] = [None] * len(documents) + if deduplicate and not _uses_pgvector(): + + async def _resolve_candidates(index: int, doc: schemas.DocumentCreate) -> None: + filters = _semantic_dup_filters(doc) + if filters is None: + return + semantic_candidates[index] = await query_external_vector_document_ids( + workspace_name=workspace_name, + observer=observer, + observed=observed, + embedding=doc.embedding, + top_k=_SEMANTIC_DUP_TOP_K, + max_distance=_SEMANTIC_DUP_MAX_DISTANCE, + filters=filters, + ) + + await asyncio.gather( + *(_resolve_candidates(i, doc) for i, doc in enumerate(documents)) + ) + # exact-content dedup (independent of `deduplicate`): pre-fetch # existing live documents whose normalized content matches anything in this # batch, scoped to (workspace, observer, observed). The SQL normalization must @@ -557,6 +602,24 @@ 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() @@ -565,7 +628,7 @@ async def create_documents( exact_dup_in_batch_count = 0 semantic_dup_rejected_count = 0 semantic_dup_replaced_count = 0 - for doc in documents: + for index, doc in enumerate(documents): try: # Session-purity invariant: an explicit document must always carry # the session it was derived from. Refuse to write session-less @@ -601,6 +664,7 @@ async def create_documents( # 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, @@ -614,7 +678,12 @@ async def create_documents( # as a duplicate OR deletes an existing document that is a duplicate. if deduplicate: duplicate_result = await is_rejected_duplicate( - db, doc, workspace_name, observer=observer, observed=observed + db, + doc, + workspace_name, + observer=observer, + 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 @@ -670,7 +739,14 @@ async def create_documents( if doc.embedding: docs_with_embeddings.append((new_doc, doc.embedding)) + except SQLAlchemyError: + # The session/transaction is dead; continuing the loop would only + # cascade PendingRollbackErrors and lose the whole batch silently. + await db.rollback() + raise except Exception as e: + # Genuinely per-document failures (bad content, metadata, token + # overflow) skip the document without poisoning the batch. logger.error( f"Error adding new document to {workspace_name}/{doc.session_name}/{observer}/{observed}: {e}" ) @@ -772,6 +848,11 @@ async def create_documents( raise ValidationException( "Failed to create documents due to integrity constraint violation" ) from e + except DBAPIError: + # Leave the session clean for callers that own it (e.g. a deadlock + # at the final commit); the queue layer classifies and retries. + await db.rollback() + raise return CreateDocumentsResult( created_documents=accepted_documents, @@ -1160,6 +1241,7 @@ 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. @@ -1185,27 +1267,41 @@ async def is_rejected_duplicate( 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. """ - filters: dict[str, Any] = {"level": doc.level} - if doc.level == "explicit": - if doc.session_name is None: - # create_documents refuses session-less explicit documents; if one - # reaches here anyway it has no valid merge partner. - return SemanticRejectionResult.NOT_DUPLICATE - filters["session_name"] = doc.session_name + 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 # Step 1: Find potential duplicates using cosine similarity - similar_docs = await query_documents( - db=db, - workspace_name=workspace_name, - query=doc.content, - observer=observer, - observed=observed, - filters=filters, - max_distance=0.05, - top_k=1, - embedding=doc.embedding, - ) + if candidate_document_ids is not None: + similar_docs: Sequence[models.Document] = await fetch_documents_by_ids( + db=db, + workspace_name=workspace_name, + observer=observer, + observed=observed, + document_ids=candidate_document_ids, + filters=filters, + ) + else: + similar_docs = await query_documents( + db=db, + workspace_name=workspace_name, + query=doc.content, + observer=observer, + observed=observed, + filters=filters, + max_distance=_SEMANTIC_DUP_MAX_DISTANCE, + top_k=_SEMANTIC_DUP_TOP_K, + embedding=doc.embedding, + ) if not similar_docs: return SemanticRejectionResult.NOT_DUPLICATE @@ -1233,6 +1329,8 @@ async def is_rejected_duplicate( # 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 ( @@ -1244,6 +1342,8 @@ async def is_rejected_duplicate( # 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. existing_doc.times_derived = func.greatest( models.Document.times_derived + 1, doc.times_derived, diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 498372ed..cc19461a 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -43,6 +43,7 @@ from src.reconciler import ( from src.schemas import ResolvedConfiguration from src.telemetry import prometheus_metrics from src.telemetry.sentry import initialize_sentry +from src.utils.retryable_errors import is_retryable_error from src.utils.work_unit import parse_work_unit_key from src.webhooks.events import ( QueueEmptyEvent, @@ -53,6 +54,11 @@ logger = getLogger(__name__) load_dotenv(override=True) +# Total processing attempts per work unit for transient errors, counted +# per instance (N deriver instances give N x this many attempts). +MAX_RETRYABLE_ATTEMPTS = 3 +RETRY_BACKOFF_SECONDS = 1.0 + class WorkerOwnership(NamedTuple): """Represents the instance of a work unit that a worker is processing.""" @@ -129,6 +135,10 @@ class QueueManager: self.worker_ownership: dict[str, WorkerOwnership] = {} self.queue_empty_flag: asyncio.Event = asyncio.Event() + # Transient-failure attempts per work-unit key. Entries are removed + # on success or terminal failure, never on the retry path itself. + self._retry_attempts: dict[str, int] = {} + # Current adaptive polling interval; grows while idle/erroring and # resets to the base interval as soon as work is claimed. self._current_poll_interval: float = ( @@ -579,11 +589,19 @@ class QueueManager: items: list[QueueItem], work_unit_key: str, context: str, - ) -> None: + ) -> bool: """ - Handle processing errors by marking queue items as errored, logging, and forwarding to Sentry. - We only mark the first queue item as errored so we don't potentially throw away a batch. This allows us - to incrementally attempt to process the batch while still maintaining progress in a work unit. + Handle a processing error. Returns True when the caller should stop + processing and release the work unit for a later re-claim. + + Transient errors (is_retryable_error) get up to MAX_RETRYABLE_ATTEMPTS + attempts per work unit: items stay unprocessed with no error recorded. + Reprocessing is safe because a retried batch re-derives the same + observations and exact dedup collapses them into reinforcement. + + Terminal errors mark only the first queue item as errored so we don't + potentially throw away a batch. This allows us to incrementally attempt + to process the batch while still maintaining progress in a work unit. Args: error: The exception that occurred @@ -591,6 +609,21 @@ class QueueManager: work_unit_key: The work unit key for the queue items context: Context string describing what was being processed (e.g., "processing representation batch") """ + if is_retryable_error(error): + attempts = self._retry_attempts.get(work_unit_key, 0) + 1 + if attempts < MAX_RETRYABLE_ATTEMPTS: + self._retry_attempts[work_unit_key] = attempts + logger.warning( + "Transient error %s for work unit %s (attempt %d/%d); leaving items unprocessed for retry", + context, + work_unit_key, + attempts, + MAX_RETRYABLE_ATTEMPTS, + exc_info=error, + ) + return True + + self._retry_attempts.pop(work_unit_key, None) error_msg = f"{error.__class__.__name__}: {str(error)}" try: if items: @@ -609,6 +642,7 @@ class QueueManager: ) if settings.SENTRY.ENABLED: sentry_sdk.capture_exception(error) + return False async def process_work_unit(self, work_unit_key: str, worker_id: str) -> None: """Process all queue items for a specific work unit by routing to the correct handler.""" @@ -672,14 +706,21 @@ class QueueManager: await self.mark_queue_items_as_processed( items_to_process, work_unit_key ) + self._retry_attempts.pop(work_unit_key, None) queue_item_count += len(items_to_process) except Exception as e: - await self._handle_processing_error( + if await self._handle_processing_error( e, items_to_process, work_unit_key, f"processing {work_unit.task_type} batch", - ) + ): + # Release the work unit (via the finally + # below) and let a later poll re-claim it. + await asyncio.sleep( + self._jitter(RETRY_BACKOFF_SECONDS) + ) + break else: queue_item = await self.get_next_queue_item( @@ -696,14 +737,19 @@ class QueueManager: await self.mark_queue_items_as_processed( [queue_item], work_unit_key ) + self._retry_attempts.pop(work_unit_key, None) queue_item_count += 1 except Exception as e: - await self._handle_processing_error( + if await self._handle_processing_error( e, [queue_item], work_unit_key, "processing queue item", - ) + ): + await asyncio.sleep( + self._jitter(RETRY_BACKOFF_SECONDS) + ) + break except Exception as e: logger.error( diff --git a/src/utils/retryable_errors.py b/src/utils/retryable_errors.py new file mode 100644 index 00000000..357f594d --- /dev/null +++ b/src/utils/retryable_errors.py @@ -0,0 +1,86 @@ +"""Classify exceptions as transient (safe to retry) or terminal. + +Imports only exception taxonomies, so it is importable from anywhere and +unit-testable without a DB. +""" + +import asyncio +from collections.abc import Iterator + +import httpx +from sqlalchemy.exc import DBAPIError + +__all__ = ["is_retryable_db_error", "is_retryable_error"] + +_RETRYABLE_SQLSTATES = frozenset( + { + "40001", # serialization_failure + "40P01", # deadlock_detected + "55P03", # lock_not_available (lock_timeout / NOWAIT) + "57014", # query_canceled (statement_timeout) + "08000", # connection_exception family + "08001", + "08003", + "08004", + "08006", + } +) + +# Provider/network transport failures. SDK wrappers (anthropic/openai +# APIConnectionError etc.) chain to these via __cause__. +_TRANSPORT_ERRORS = ( + httpx.TransportError, + ConnectionError, + asyncio.TimeoutError, + TimeoutError, +) + + +def _iter_cause_chain(exc: BaseException) -> Iterator[BaseException]: + seen: set[int] = set() + current: BaseException | None = exc + while current is not None and id(current) not in seen: + seen.add(id(current)) + yield current + current = current.__cause__ + + +def _sqlstate(exc: DBAPIError) -> str | None: + """Extract the SQLSTATE off ``DBAPIError.orig``, driver-agnostically.""" + orig = getattr(exc, "orig", None) + for candidate in (orig, getattr(orig, "__cause__", None)): + code = getattr(candidate, "sqlstate", None) + if isinstance(code, str): + return code + return None + + +def is_retryable_db_error(exc: BaseException) -> bool: + """True for transient DB failures: deadlock, serialization failure, + lock/statement timeout, or a lost connection. + + Integrity (23xxx), data (22xxx), and programming (42xxx) errors are + deliberately terminal. + """ + for current in _iter_cause_chain(exc): + if not isinstance(current, DBAPIError): + continue + if current.connection_invalidated: + return True + if _sqlstate(current) in _RETRYABLE_SQLSTATES: + return True + return False + + +def is_retryable_error(exc: BaseException) -> bool: + """Superset of ``is_retryable_db_error``: also transient network/provider + transport failures (timeouts, connection refused/reset). + + Auth failures (401 from a rotated key) are deliberately terminal: they + never self-heal, so retrying only delays the burn. + """ + if is_retryable_db_error(exc): + return True + return any( + isinstance(current, _TRANSPORT_ERRORS) for current in _iter_cause_chain(exc) + ) diff --git a/tests/crud/test_document.py b/tests/crud/test_document.py index 5f03c80b..f216f2a1 100644 --- a/tests/crud/test_document.py +++ b/tests/crud/test_document.py @@ -1,10 +1,13 @@ +import asyncio import datetime +from typing import Any from unittest.mock import AsyncMock, patch import pytest from nanoid import generate as generate_nanoid from sqlalchemy import select -from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.exc import OperationalError +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker from src import crud, models, schemas from src.crud.document import SemanticRejectionResult, is_rejected_duplicate @@ -1334,3 +1337,477 @@ class TestSessionPurityInvariant: ) assert rejected is SemanticRejectionResult.NOT_DUPLICATE mock_query.assert_not_awaited() + + +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. + """ + + N_DOCS: int = 20 + N_ROUNDS: int = 5 + + async def _setup( + self, + db_session: AsyncSession, + test_workspace: models.Workspace, + test_peer: models.Peer, + ) -> tuple[models.Peer, models.Session]: + """Create an observed peer, session, and collection, committed so + they are visible to independent concurrent sessions.""" + test_peer2 = 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([test_peer2, test_session]) + await db_session.flush() + collection = models.Collection( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + db_session.add(collection) + await db_session.commit() + return test_peer2, test_session + + def _batch(self, session_name: str) -> list[schemas.DocumentCreate]: + return [ + schemas.DocumentCreate( + content=f"user fact number {i}", + embedding=[0.1] * 1536, + session_name=session_name, + metadata=schemas.DocumentMetadata( + message_ids=[i], + message_created_at="2026-01-01T00:00:00Z", + ), + ) + for i in range(self.N_DOCS) + ] + + @staticmethod + def _chain(exc: BaseException) -> str: + parts: list[str] = [] + seen: set[int] = set() + e: BaseException | None = exc + while e is not None and id(e) not in seen: + seen.add(id(e)) + parts.append(f"{type(e).__name__}: {e}") + e = e.__cause__ or e.__context__ + return " <- ".join(parts) + + @pytest.mark.asyncio + async def test_concurrent_reinforcement_does_not_deadlock( + self, + db_engine: "AsyncEngine", + 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).""" + test_workspace, test_peer = sample_data + test_peer2, test_session = await self._setup( + db_session, test_workspace, test_peer + ) + + # Seed the rows both writers will reinforce. + await crud.create_documents( + db_session, + self._batch(test_session.name), + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + + session_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False) + + for round_num in range(self.N_ROUNDS): + forward = self._batch(test_session.name) + backward = list(reversed(self._batch(test_session.name))) + + async def _run(batch: list[schemas.DocumentCreate]) -> None: + async with session_factory() as db: + await crud.create_documents( + db, + batch, + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + + results = await asyncio.gather( + _run(forward), _run(backward), return_exceptions=True + ) + errors = [r for r in results if isinstance(r, BaseException)] + assert not errors, ( + f"round {round_num}: concurrent create_documents failed: " + + "; ".join(self._chain(e) for e in errors) + ) + + # Every round reinforced the same rows: 1 seed + 2 per round. + docs = ( + ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == test_workspace.name, + models.Document.observer == test_peer.name, + models.Document.observed == test_peer2.name, + models.Document.deleted_at.is_(None), + ) + ) + ) + .scalars() + .all() + ) + assert len(docs) == self.N_DOCS + assert all(d.times_derived == 1 + 2 * self.N_ROUNDS for d in docs) + + +class TestCreateDocumentsErrorHandling: + """A dead transaction aborts the batch loudly; per-document failures + still skip just that document (DEV-1975).""" + + async def _setup( + self, + db_session: AsyncSession, + test_workspace: models.Workspace, + test_peer: models.Peer, + ) -> tuple[models.Peer, models.Session]: + test_peer2 = 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([test_peer2, test_session]) + await db_session.flush() + collection = models.Collection( + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + db_session.add(collection) + await db_session.commit() + return test_peer2, test_session + + def _doc(self, content: str, session_name: str) -> schemas.DocumentCreate: + return schemas.DocumentCreate( + content=content, + embedding=[0.1] * 1536, + session_name=session_name, + metadata=schemas.DocumentMetadata( + message_ids=[1], + message_created_at="2026-01-01T00:00:00Z", + ), + ) + + @pytest.mark.asyncio + async def test_db_error_in_loop_aborts_batch( + self, + 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.""" + test_workspace, test_peer = sample_data + test_peer2, test_session = await self._setup( + db_session, test_workspace, test_peer + ) + # Plain strings: the rollback below expires ORM objects in the session. + workspace_name = test_workspace.name + observer = test_peer.name + observed = test_peer2.name + session_name = test_session.name + + await crud.create_documents( + db_session, + [self._doc("existing fact", session_name)], + workspace_name=workspace_name, + observer=observer, + observed=observed, + ) + + class FakePGError(Exception): + 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), + ): + await crud.create_documents( + db_session, + [ + self._doc("existing fact", session_name), + self._doc("a brand new fact", session_name), + ], + workspace_name=workspace_name, + observer=observer, + observed=observed, + ) + + docs = ( + ( + await db_session.execute( + select(models.Document).where( + models.Document.workspace_name == workspace_name, + models.Document.observer == observer, + models.Document.observed == observed, + ) + ) + ) + .scalars() + .all() + ) + assert [d.content for d in docs] == ["existing fact"] + assert docs[0].times_derived == 1 + + @pytest.mark.asyncio + async def test_per_document_error_still_skips_only_that_document( + self, + db_session: AsyncSession, + sample_data: tuple[models.Workspace, models.Peer], + ): + """Non-DB per-document failures keep their skip semantics.""" + test_workspace, test_peer = sample_data + test_peer2, test_session = await self._setup( + db_session, test_workspace, test_peer + ) + + from src.crud import document as document_module + + real_dedup_key = document_module._dedup_key # pyright: ignore[reportPrivateUsage] + + def flaky_dedup_key( + content: str, level: str, session_name: str | None + ) -> tuple[str, str, str | None]: + if content == "poison": + raise ValueError("bad content") + return real_dedup_key(content, level, session_name) + + with patch.object(document_module, "_dedup_key", flaky_dedup_key): + result = await crud.create_documents( + db_session, + [ + self._doc("good fact one", test_session.name), + self._doc("poison", test_session.name), + self._doc("good fact two", test_session.name), + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=test_peer2.name, + ) + + assert sorted(d.content for d in result.created_documents) == [ + "good fact one", + "good fact two", + ] + + +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.""" + + 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) + ] + test_session = models.Session( + name=str(generate_nanoid()), workspace_name=test_workspace.name + ) + db_session.add_all([*observed_peers, 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, + ) + ) + await db_session.commit() + return observed_peers, test_session + + def _doc(self, content: str, session_name: str) -> schemas.DocumentCreate: + return schemas.DocumentCreate( + content=content, + embedding=[0.1] * 1536, + session_name=session_name, + metadata=schemas.DocumentMetadata( + message_ids=[1], + message_created_at="2026-01-01T00:00:00Z", + ), + ) + + @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( + 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( + db_session, test_workspace, test_peer + ) + monkeypatch.setattr(settings.VECTOR_STORE, "TYPE", "turbopuffer") + monkeypatch.setattr(settings.VECTOR_STORE, "MIGRATED", True) + + events: list[str] = [] + 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") + return await real_execute(statement, *args, **kwargs) + + async def fake_resolve(*_args: Any, **_kwargs: Any) -> list[str]: + events.append("resolve") + return [] + + with ( + patch.object(db_session, "execute", side_effect=spying_execute), + patch( + "src.crud.document.query_external_vector_document_ids", + side_effect=fake_resolve, + ), + ): + result = await crud.create_documents( + db_session, + [ + self._doc("fact one", test_session.name), + self._doc("fact two", test_session.name), + ], + workspace_name=test_workspace.name, + observer=test_peer.name, + observed=observed_peer.name, + deduplicate=True, + ) + + assert len(result.created_documents) == 2 + assert events == ["resolve", "resolve", "lock"] diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 81dd0632..f37067ff 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -7,6 +7,7 @@ from unittest.mock import patch import pytest from nanoid import generate as generate_nanoid from sqlalchemy import select +from sqlalchemy.exc import OperationalError from sqlalchemy.ext.asyncio import AsyncSession from src import models @@ -1874,3 +1875,235 @@ class TestPollingJitter: qm.shutdown_event.set() # A shutdown already signalled must short-circuit the (long) jitter sleep. await asyncio.wait_for(qm._sleep_startup_jitter(), timeout=1.0) # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +class TestQueueRetry: + """Bounded retry of transient errors in process_work_unit (DEV-1975). + + A transient failure (deadlock, lost connection, provider transport) must + leave the batch's queue items unprocessed and release the work unit for + re-claim, up to MAX_RETRYABLE_ATTEMPTS per work unit; terminal failures + keep today's burn-one-item behavior. + """ + + async def _seed_work_unit( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + n_messages: int = 1, + ) -> tuple[QueueManager, str, str, list[models.QueueItem]]: + """Seed a claimed representation work unit owned by a test worker.""" + session, peers = sample_session_with_peers + peer = peers[0] + + messages: list[models.Message] = [] + for index in range(n_messages): + message = models.Message( + session_name=session.name, + workspace_name=session.workspace_name, + peer_name=peer.name, + content=f"Message {index}", + token_count=10, + seq_in_session=index + 1, + ) + db_session.add(message) + messages.append(message) + await db_session.commit() + for message in messages: + await db_session.refresh(message) + + queue_items: list[models.QueueItem] = [] + work_unit_key = "" + for message in messages: + payload = create_queue_payload( + message=message, + task_type="representation", + observed=peer.name, + observer=peer.name, + ) + work_unit_key = work_unit_key or construct_work_unit_key( + session.workspace_name, payload + ) + queue_item = models.QueueItem( + session_id=session.id, + task_type="representation", + work_unit_key=work_unit_key, + payload=payload, + processed=False, + workspace_name=session.workspace_name, + message_id=message.id, + ) + db_session.add(queue_item) + queue_items.append(queue_item) + await db_session.commit() + for queue_item in queue_items: + await db_session.refresh(queue_item) + + qm = QueueManager() + worker_id = "test_worker" + claimed_units = await qm.claim_work_units(db_session, [work_unit_key]) + qm.worker_ownership[worker_id] = WorkerOwnership( + work_unit_key=work_unit_key, aqs_id=claimed_units[work_unit_key] + ) + await db_session.commit() + return qm, work_unit_key, worker_id, queue_items + + @staticmethod + def _retryable_error() -> OperationalError: + class FakePGError(Exception): + sqlstate: str = "40P01" + + return OperationalError("UPDATE documents", {}, FakePGError()) + + async def _fetch_items( + self, db_session: AsyncSession, work_unit_key: str + ) -> list[models.QueueItem]: + db_session.expire_all() + return list( + ( + await db_session.execute( + select(models.QueueItem) + .where(models.QueueItem.work_unit_key == work_unit_key) + .order_by(models.QueueItem.id) + ) + ) + .scalars() + .all() + ) + + async def _aqs_rows(self, db_session: AsyncSession, work_unit_key: str) -> int: + return len( + ( + await db_session.execute( + select(models.ActiveQueueSession).where( + models.ActiveQueueSession.work_unit_key == work_unit_key + ) + ) + ) + .scalars() + .all() + ) + + async def test_retryable_error_leaves_items_unprocessed( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A transient error stops the work unit after ONE batch fetch (no + tight loop), leaves items unprocessed with no error, and releases + the ActiveQueueSession row.""" + monkeypatch.setattr("src.deriver.queue_manager.RETRY_BACKOFF_SECONDS", 0.0) + qm, work_unit_key, worker_id, _ = await self._seed_work_unit( + db_session, sample_session_with_peers, create_queue_payload, n_messages=2 + ) + initial_semaphore_value = qm.semaphore._value + + batch_fetches = 0 + original_get_batch = qm.get_queue_item_batch + + async def counting_get_batch(*args: Any, **kwargs: Any) -> Any: + nonlocal batch_fetches + batch_fetches += 1 + return await original_get_batch(*args, **kwargs) + + with ( + patch.object(qm, "get_queue_item_batch", side_effect=counting_get_batch), + patch( + "src.deriver.queue_manager.process_representation_batch", + side_effect=self._retryable_error(), + ), + ): + await qm.process_work_unit(work_unit_key, worker_id) + + assert batch_fetches == 1 + items = await self._fetch_items(db_session, work_unit_key) + assert all(not item.processed for item in items) + assert all(item.error is None for item in items) + assert await self._aqs_rows(db_session, work_unit_key) == 0 + assert qm._retry_attempts[work_unit_key] == 1 # pyright: ignore[reportPrivateUsage] + assert qm.semaphore._value == initial_semaphore_value + + async def test_retry_exhaustion_is_terminal( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """At the attempt cap a transient error burns the first item exactly + like today's terminal path and clears the counter.""" + from src.deriver.queue_manager import MAX_RETRYABLE_ATTEMPTS + + monkeypatch.setattr("src.deriver.queue_manager.RETRY_BACKOFF_SECONDS", 0.0) + qm, work_unit_key, worker_id, _ = await self._seed_work_unit( + db_session, sample_session_with_peers, create_queue_payload + ) + qm._retry_attempts[work_unit_key] = MAX_RETRYABLE_ATTEMPTS - 1 # pyright: ignore[reportPrivateUsage] + + with patch( + "src.deriver.queue_manager.process_representation_batch", + side_effect=self._retryable_error(), + ): + await qm.process_work_unit(work_unit_key, worker_id) + + items = await self._fetch_items(db_session, work_unit_key) + assert len(items) == 1 + assert items[0].processed + assert items[0].error is not None + assert "OperationalError" in items[0].error + assert work_unit_key not in qm._retry_attempts # pyright: ignore[reportPrivateUsage] + + async def test_non_retryable_error_burns_immediately( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + ) -> None: + """A non-retryable error keeps today's behavior verbatim: the first + item is marked errored on the first attempt.""" + qm, work_unit_key, worker_id, _ = await self._seed_work_unit( + db_session, sample_session_with_peers, create_queue_payload + ) + + with patch( + "src.deriver.queue_manager.process_representation_batch", + side_effect=ValueError("bad batch"), + ): + await qm.process_work_unit(work_unit_key, worker_id) + + items = await self._fetch_items(db_session, work_unit_key) + assert len(items) == 1 + assert items[0].processed + assert items[0].error is not None + assert "ValueError" in items[0].error + assert work_unit_key not in qm._retry_attempts # pyright: ignore[reportPrivateUsage] + + async def test_counter_cleared_after_success( + self, + db_session: AsyncSession, + sample_session_with_peers: tuple[models.Session, list[models.Peer]], + create_queue_payload: Callable[..., Any], + ) -> None: + """A success wipes the accumulated attempt count for the work unit.""" + qm, work_unit_key, worker_id, _ = await self._seed_work_unit( + db_session, sample_session_with_peers, create_queue_payload + ) + qm._retry_attempts[work_unit_key] = 1 # pyright: ignore[reportPrivateUsage] + + async def noop_batch(*_args: Any, **_kwargs: Any) -> None: + return None + + with patch( + "src.deriver.queue_manager.process_representation_batch", + side_effect=noop_batch, + ): + await qm.process_work_unit(work_unit_key, worker_id) + + items = await self._fetch_items(db_session, work_unit_key) + assert all(item.processed for item in items) + assert all(item.error is None for item in items) + assert work_unit_key not in qm._retry_attempts # pyright: ignore[reportPrivateUsage] diff --git a/tests/utils/test_retryable_errors.py b/tests/utils/test_retryable_errors.py new file mode 100644 index 00000000..bead2d2e --- /dev/null +++ b/tests/utils/test_retryable_errors.py @@ -0,0 +1,117 @@ +"""DB-free unit tests for src/utils/retryable_errors.py.""" + +import asyncio +from typing import cast + +import httpx +import pytest +from sqlalchemy.exc import DBAPIError, OperationalError + +from src.utils.retryable_errors import is_retryable_db_error, is_retryable_error + + +class FakePGError(Exception): + """Stands in for a driver exception carrying a SQLSTATE.""" + + sqlstate: str | None + + def __init__(self, sqlstate: str | None) -> None: + super().__init__(f"fake pg error ({sqlstate})") + self.sqlstate = sqlstate + + +def _dbapi_error( + sqlstate: str | None, + *, + orig: BaseException | None = None, + connection_invalidated: bool = False, +) -> DBAPIError: + if orig is None and sqlstate is not None: + orig = FakePGError(sqlstate) + return OperationalError( + "SELECT 1", + {}, + cast(BaseException, orig), + connection_invalidated=connection_invalidated, + ) + + +@pytest.mark.parametrize( + ("sqlstate", "expected"), + [ + ("40P01", True), # deadlock_detected + ("40001", True), # serialization_failure + ("55P03", True), # lock_not_available + ("57014", True), # query_canceled + ("08006", True), # connection_failure + ("23505", False), # unique_violation + ("42P01", False), # undefined_table + ("22P02", False), # invalid_text_representation + ], +) +def test_sqlstate_classification(sqlstate: str, expected: bool): + exc = _dbapi_error(sqlstate) + assert is_retryable_db_error(exc) is expected + assert is_retryable_error(exc) is expected + + +def test_orig_none_is_terminal(): + assert not is_retryable_db_error(_dbapi_error(None)) + + +def test_sqlstate_on_orig_cause(): + """SQLSTATE found by walking orig.__cause__ when orig itself has none.""" + wrapper = Exception("driver wrapper") + wrapper.__cause__ = FakePGError("40P01") + assert is_retryable_db_error(_dbapi_error(None, orig=wrapper)) + + +def test_connection_invalidated_is_retryable(): + exc = _dbapi_error(None, connection_invalidated=True) + assert is_retryable_db_error(exc) + + +def test_dbapi_error_nested_in_cause_chain(): + outer = RuntimeError("save failed") + outer.__cause__ = _dbapi_error("40P01") + assert is_retryable_db_error(outer) + assert is_retryable_error(outer) + + +def test_non_db_exceptions_are_not_db_retryable(): + assert not is_retryable_db_error(ValueError("bad input")) + assert not is_retryable_db_error(httpx.ConnectTimeout("timed out")) + + +@pytest.mark.parametrize( + ("exc", "expected"), + [ + (httpx.ConnectTimeout("timed out"), True), + (httpx.ReadTimeout("timed out"), True), + (httpx.ConnectError("connection refused"), True), + (ConnectionResetError("reset"), True), + (asyncio.TimeoutError(), True), + (TimeoutError(), True), + (ValueError("bad input"), False), + (httpx.HTTPStatusError("401", request=None, response=None), False), # pyright: ignore[reportArgumentType] + ], +) +def test_transport_classification(exc: BaseException, expected: bool): + assert is_retryable_error(exc) is expected + assert not is_retryable_db_error(exc) + + +def test_transport_error_nested_in_cause_chain(): + """SDK wrappers (e.g. APIConnectionError) chain to httpx via __cause__.""" + wrapper = RuntimeError("provider call failed") + wrapper.__cause__ = httpx.ConnectError("connection refused") + assert is_retryable_error(wrapper) + assert not is_retryable_db_error(wrapper) + + +def test_cause_cycle_terminates(): + a = RuntimeError("a") + b = RuntimeError("b") + a.__cause__ = b + b.__cause__ = a + assert not is_retryable_error(a)