From 868c3a7c74d376d2102a0e16d304aa7d30f927c2 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Wed, 2 Sep 2026 10:41:00 -0400 Subject: [PATCH] fix: add test and simplify queue query --- CLAUDE.md | 10 +++++ src/deriver/consumer.py | 7 +++- src/deriver/queue_manager.py | 48 +++++++++++----------- src/utils/queue_payload.py | 8 ++++ tests/deriver/test_queue_processing.py | 56 ++++++++++++++++++++------ 5 files changed, 90 insertions(+), 39 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 6a83066d..ab842f99 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -118,6 +118,16 @@ cd sdks/typescript && bun run tsc --noEmit - **Never hold a DB session during external calls** (LLM, embedding, HTTP). If a function needs both a DB session and an external call result, compute the external result first and pass it as a parameter. This avoids tying up DB connections during slow network I/O. Use `tracked_db` for short-lived, DB-only operations; pass a shared session when multiple DB-only calls can reuse one connection. - **Never write through a read-only session** (`tracked_db(..., read_only=True)`, `get_read_db`, `ReadSessionLocal`). These run in AUTOCOMMIT mode with no transaction: writes are NOT blocked by the database — they silently commit immediately, and `begin_nested()` savepoints break. There is no runtime guard; this is enforced by convention only. Use `read_only=True` strictly for SELECT-only windows; anything that mutates (including get-or-create paths) must use a regular write session. +#### Multi-row locking and deadlocks + +Tables written concurrently by more than one worker — `documents` (deriver, dreamer, scope backfill/removal, reconciler) and `queue` (every deriver replica) — deadlock when two writers touch an overlapping row set in different orders. Rules: + +- **A multi-row `SELECT ... FOR UPDATE` MUST carry an explicit `ORDER BY `.** Without it Postgres locks in scan order, which differs per plan, so two writers with overlapping sets can cycle. `_apply_document_row_updates` in `src/crud/document.py` is the reference implementation. +- **`WHERE id IN (...)` does NOT impose an order**, so sorting the Python list is a no-op — the list order is discarded and the planner picks `Bitmap Heap Scan` (ctid order), `Index Scan` (id order), or `Seq Scan` per invocation. Deterministic ordering requires either a preceding `SELECT ... ORDER BY id FOR UPDATE` or `WHERE id IN (SELECT id ... ORDER BY id FOR UPDATE)`. +- **`Document.id` is a random nanoid** (`models.py`), so id order is uncorrelated with physical order — an unordered predicate `UPDATE`/`DELETE` is roughly a coin flip against an id-ordered locker per row pair, not a rare edge case. (`QueueItem.id` is an integer identity, so there id order is also chronological.) +- **Prefer no lock at all.** A single `UPDATE ... WHERE ` acquires row locks as it writes and has no separate lock phase to get wrong. Reach for `FOR UPDATE` only when a value must be read, computed in Python, and written back — that read-modify-write is the only reason `_apply_document_row_updates` locks (it replaced a server-side `func.greatest()`), and `populate_existing=True` is required with it so the identity map doesn't serve a stale pre-lock value. Server-side expressions (`func.greatest`, the JSONB `-` operator) avoid the lock entirely; see `_clear_work_unit_retry_attempts` in `src/deriver/queue_manager.py`. +- `FOR UPDATE SKIP LOCKED` (the reconciler's claim pattern) never waits, so it cannot be a deadlock partner — but holding those locks across an external call still stalls other writers. See the "never hold a DB session during external calls" rule above. + #### Auth scoping - **`allow_member_read=True` (in `require_auth(...)`) is read-only — NEVER set it on a route that mutates state.** It lets a peer-scoped key reach a session route when its peer is an active member of the session, so on a mutating route it would hand any session member write access (message injection, config mutation, deletion). HTTP method is not a reliable read/write signal here (some read routes use POST for a richer body), so this is enforced by an explicit allowlist in `tests/routes/test_auth_route_policy.py` — adding the flag to a new route fails that test until you consciously add the route to `EXPECTED_MEMBER_READ_ROUTES`, and you must never add a mutating method there. diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index 8335798b..6251dc73 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -27,6 +27,7 @@ from src.telemetry.events import ( from src.telemetry.logging import log_performance_metrics from src.utils import summarizer from src.utils.queue_payload import ( + RETRY_ATTEMPTS_PAYLOAD_KEY, DeletionPayload, DreamPayload, ReconcilerPayload, @@ -44,9 +45,11 @@ logging.getLogger("sqlalchemy.engine.Engine").disabled = True async def process_item(queue_item: models.QueueItem) -> None: """Process a single item from the queue.""" task_type = queue_item.task_type - # Drop the work-unit retry counter before payload validation. + # Drop the work-unit retry counter before payload validation: every payload + # model sets extra="forbid", so leaving it in burns the item as + # extra_forbidden on the reclaim that was supposed to retry it. queue_payload = dict(queue_item.payload or {}) - queue_payload.pop("_retry_attempts", None) + queue_payload.pop(RETRY_ATTEMPTS_PAYLOAD_KEY, None) workspace_name = queue_item.workspace_name # Handle reconciler first - it's the only task type that doesn't require workspace_name diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 58532472..b98c0ef6 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -15,7 +15,7 @@ from dotenv import load_dotenv from nanoid import generate as generate_nanoid from sentry_sdk.integrations.asyncio import AsyncioIntegration from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration -from sqlalchemy import and_, delete, or_, select, update +from sqlalchemy import Text, and_, delete, literal, or_, select, update from sqlalchemy.dialects.postgresql import insert from sqlalchemy.engine import CursorResult from sqlalchemy.ext.asyncio import AsyncSession @@ -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.queue_payload import RETRY_ATTEMPTS_PAYLOAD_KEY from src.utils.retryable_errors import is_retryable_error from src.utils.work_unit import parse_work_unit_key from src.webhooks.events import ( @@ -59,7 +60,6 @@ load_dotenv(override=True) # shares one budget. MAX_RETRYABLE_ATTEMPTS = 3 RETRY_BACKOFF_SECONDS = 1.0 -_RETRY_ATTEMPTS_PAYLOAD_KEY = "_retry_attempts" class WorkerOwnership(NamedTuple): @@ -1152,7 +1152,7 @@ class QueueManager: item = await self._oldest_unprocessed_item(db, work_unit_key) if item is None: return 0 - raw = (item.payload or {}).get(_RETRY_ATTEMPTS_PAYLOAD_KEY, 0) + raw = (item.payload or {}).get(RETRY_ATTEMPTS_PAYLOAD_KEY, 0) try: return max(0, int(raw)) except (TypeError, ValueError): @@ -1170,7 +1170,7 @@ class QueueManager: await db.commit() return new_payload = dict(item.payload or {}) - new_payload[_RETRY_ATTEMPTS_PAYLOAD_KEY] = attempts + new_payload[RETRY_ATTEMPTS_PAYLOAD_KEY] = attempts await db.execute( update(models.QueueItem) .where(models.QueueItem.id == item.id) @@ -1178,34 +1178,32 @@ class QueueManager: ) await db.commit() - @staticmethod - def _payload_without_retry_attempts( - payload: dict[str, Any] | None, - ) -> dict[str, Any]: - cleaned = dict(payload or {}) - cleaned.pop(_RETRY_ATTEMPTS_PAYLOAD_KEY, None) - return cleaned - async def _clear_work_unit_retry_attempts(self, work_unit_key: str) -> None: - """Drop the shared attempt count from remaining unprocessed items.""" + """Drop the shared attempt count from remaining unprocessed items. + + One statement on purpose: a multi-row ``SELECT ... FOR UPDATE`` here + would take locks on ``queue`` in scan order, which is a deadlock partner + for any other multi-row writer on the same table. The JSONB ``-`` + operator does the strip server-side, so no rows are locked ahead of the + write and there is no lock order to get wrong. + """ async with tracked_db("clear_work_unit_retry_attempts") as db: - result = await db.execute( - select(models.QueueItem) + await db.execute( + update(models.QueueItem) .where( models.QueueItem.work_unit_key == work_unit_key, models.QueueItem.processed.is_(False), + models.QueueItem.payload.has_key(RETRY_ATTEMPTS_PAYLOAD_KEY), + ) + .values( + # literal(..., Text) is required: an untyped bind leaves + # Postgres unable to pick between jsonb - text and its + # integer/array siblings. + payload=models.QueueItem.payload.op("-")( + literal(RETRY_ATTEMPTS_PAYLOAD_KEY, Text) + ) ) - .with_for_update() ) - for item in result.scalars(): - payload = item.payload or {} - if _RETRY_ATTEMPTS_PAYLOAD_KEY not in payload: - continue - await db.execute( - update(models.QueueItem) - .where(models.QueueItem.id == item.id) - .values(payload=self._payload_without_retry_attempts(payload)) - ) await db.commit() async def mark_queue_items_as_processed( diff --git a/src/utils/queue_payload.py b/src/utils/queue_payload.py index 59815ca5..6f3f1105 100644 --- a/src/utils/queue_payload.py +++ b/src/utils/queue_payload.py @@ -5,6 +5,14 @@ from pydantic import BaseModel, ConfigDict from src.schemas import DreamType, ReconcilerType, ResolvedConfiguration +# Queue mechanics, not task data: the deriver stores a per-work-unit transient +# failure count under this key so a retry budget survives work-unit reclaim. +# Every payload model below forbids extras, so anything that reads a raw +# QueueItem.payload must strip this key before validating. Lives here rather +# than in the deriver because both the writer (queue_manager) and the stripper +# (consumer) need it, and queue_manager imports consumer. +RETRY_ATTEMPTS_PAYLOAD_KEY = "_retry_attempts" + class BasePayload(BaseModel): """Base payload with common fields.""" diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py index 99f44931..ce8e3e13 100644 --- a/tests/deriver/test_queue_processing.py +++ b/tests/deriver/test_queue_processing.py @@ -2,17 +2,20 @@ import asyncio from collections.abc import Callable from datetime import UTC, datetime, timedelta from typing import Any -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import pytest from nanoid import generate as generate_nanoid +from pydantic import ValidationError from sqlalchemy import select from sqlalchemy.exc import OperationalError from sqlalchemy.ext.asyncio import AsyncSession from src import models from src.config import settings +from src.deriver.consumer import process_item from src.deriver.queue_manager import QueueManager, WorkerOwnership +from src.utils.queue_payload import RETRY_ATTEMPTS_PAYLOAD_KEY, SummaryPayload from src.utils.work_unit import construct_work_unit_key @@ -2170,16 +2173,23 @@ class TestQueueRetry: assert items[0].error is not None assert "OperationalError" in items[0].error - async def test_summary_payload_forbids_retry_attempts_key(self) -> None: - from pydantic import ValidationError + async def test_process_item_strips_retry_counter_before_validation(self) -> None: + """A reclaimed non-representation item must survive its own retry counter. - from src.deriver.queue_manager import QueueManager - from src.utils.queue_payload import SummaryPayload - - raw = { + The counter is written onto an *unprocessed* item so the budget outlives + a work-unit reclaim -- which means the next claim re-reads it. Every + payload model sets ``extra="forbid"``, so without the strip in + ``process_item`` the reclaim raises extra_forbidden -> ValueError -> + not retryable -> the item is burned terminally on the very attempt that + was supposed to retry it. Representation tasks never hit this: their + batch path reads the payload with ``.get()`` instead of validating, + which is why the rest of this class cannot catch it. + """ + raw: dict[str, Any] = { "task_type": "summary", "session_name": "s", "message_seq_in_session": 1, + "message_public_id": "msg-public-id", "configuration": { "reasoning": {"enabled": True}, "peer_card": {"use": True, "create": True}, @@ -2190,10 +2200,32 @@ class TestQueueRetry: }, "dream": {"enabled": True}, }, - "_retry_attempts": 1, + RETRY_ATTEMPTS_PAYLOAD_KEY: 1, } - with pytest.raises(ValidationError) as ei: + + # Pin the premise: the payload model must keep rejecting the key, so + # this fails loudly if someone "fixes" the burn with extra="allow" + # instead of stripping. + with pytest.raises(ValidationError) as exc_info: SummaryPayload.model_validate(raw) - assert any(err["type"] == "extra_forbidden" for err in ei.value.errors()) - cleaned = QueueManager._payload_without_retry_attempts(raw) # pyright: ignore[reportPrivateUsage] - SummaryPayload.model_validate(cleaned) + assert any(err["type"] == "extra_forbidden" for err in exc_info.value.errors()) + + queue_item = models.QueueItem( + task_type="summary", + work_unit_key="summary:test-workspace:test-session", + payload=raw, + processed=False, + workspace_name="test-workspace", + message_id=1, + ) + + with patch( + "src.deriver.consumer.summarizer.summarize_if_needed", + new_callable=AsyncMock, + ) as mock_summarize: + await process_item(queue_item) + + mock_summarize.assert_awaited_once() + # The strip must happen on a copy: the counter has to stay on the row so + # the budget still advances if this attempt fails again. + assert raw[RETRY_ATTEMPTS_PAYLOAD_KEY] == 1