Production stress-test harness (Locust, seed/cleanup CLI, observability, runner) + fix concurrent cold-evaluation race (#607)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: nofurtherinformation

backend/app/core/config.py

@@ -45,6 +45,9 @@ class Settings(BaseSettings):
     ACCESS_TOKEN_EXPIRE_MINUTES: int = 60 * 24 * 8
     DOMAIN: str = "localhost"
     ENVIRONMENT: Environment = Environment.local
+    # Env var so lowering it under Sentry quota pressure is a config flip,
+    # not a code change + deploy.
+    SENTRY_TRACES_SAMPLE_RATE: float = 1.0
 
     @computed_field  # type: ignore[misc]
     @property

backend/app/evaluation/main.py

@@ -7,9 +7,12 @@
 """
 
 import logging
+from time import monotonic, sleep
 from typing import Any
 
 from fastapi import BackgroundTasks
+from sqlalchemy import text
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.sql import func
 from sqlmodel import Session, select
 
@@ -26,24 +29,18 @@
 
 logger = logging.getLogger(__name__)
 
+EVAL_LOCK_POLL_SECONDS = 1.0
+EVAL_LOCK_WAIT_SECONDS = 120.0
 
-def update_or_select_document_evaluation(
-    background_tasks: BackgroundTasks,
-    session: Session,
-    document: Document,
-) -> MetricsEnvelope:
-    """Return the document's metrics, recomputing on cache miss or stale row.
 
-    A cached `Evaluation` row is considered fresh iff its `payload_version`
-    matches `CURRENT_PAYLOAD_VERSION` (registry hasn't changed) and its
-    `updated_at` is at least as new as the document's. Otherwise the metrics
-    are recomputed via `compute_metrics` and the row is upserted.
+def _cached_envelope(
+    evaluation: Evaluation | None, document: Document
+) -> MetricsEnvelope | None:
+    """Return the cached envelope iff the row is fresh, else None.
 
-    Failures are never cached — a cache hit always returns `failed=[]`.
+    Fresh iff `payload_version` matches `CURRENT_PAYLOAD_VERSION` (registry
+    hasn't changed) and `updated_at` is at least as new as the document's.
     """
-    evaluation = session.exec(
-        select(Evaluation).where(Evaluation.document_id == document.document_id)
-    ).one_or_none()
     if (
         evaluation
         and evaluation.payload_version == CURRENT_PAYLOAD_VERSION
@@ -56,6 +53,62 @@ def update_or_select_document_evaluation(
             metrics=evaluation.metrics,
             failed=[],
         )
+    return None
+
+
+def update_or_select_document_evaluation(
+    background_tasks: BackgroundTasks,
+    session: Session,
+    document: Document,
+) -> MetricsEnvelope:
+    """Return the document's metrics, recomputing on cache miss or stale row.
+
+    On a miss, a per-document Postgres advisory lock serializes computes
+    across all backend tasks so a thundering herd of cache-cold requests runs
+    one compute instead of N. Losers poll instead of blocking on the lock:
+    a parked waiter holds its pooled connection for the whole compute, so
+    ~15 cache-cold requests on one document would exhaust a task's pool
+    (5 + 10 overflow) and starve unrelated endpoints.
+
+    Failures are never cached — a cache hit always returns `failed=[]`.
+    """
+    evaluation = session.exec(
+        select(Evaluation).where(Evaluation.document_id == document.document_id)
+    ).one_or_none()
+    if cached := _cached_envelope(evaluation, document):
+        return cached
+
+    deadline = monotonic() + EVAL_LOCK_WAIT_SECONDS
+    while not session.execute(
+        text("SELECT pg_try_advisory_xact_lock(hashtextextended(:doc, 0))"),
+        {"doc": str(document.document_id)},
+    ).scalar_one():
+        # End the transaction so the pool reclaims our connection while we sleep.
+        session.rollback()
+        sleep(EVAL_LOCK_POLL_SECONDS)
+        evaluation = session.exec(
+            select(Evaluation).where(Evaluation.document_id == document.document_id)
+        ).one_or_none()
+        if cached := _cached_envelope(evaluation, document):
+            return cached
+        if monotonic() >= deadline:
+            # Winner is wedged; compute without the lock rather than fail —
+            # the IntegrityError path below tolerates concurrent commits.
+            logger.warning(
+                "evaluation lock wait timed out for %s; computing without it",
+                document.document_id,
+            )
+            break
+    else:
+        # A winner may have committed between our cache check and the
+        # acquire; expire so the re-read isn't served from the identity map.
+        session.expire_all()
+        evaluation = session.exec(
+            select(Evaluation).where(Evaluation.document_id == document.document_id)
+        ).one_or_none()
+        if cached := _cached_envelope(evaluation, document):
+            return cached
+
     envelope = compute_metrics(background_tasks, session, document.document_id)
 
     if evaluation:
@@ -70,7 +123,12 @@ def update_or_select_document_evaluation(
                 payload_version=envelope["payload_version"],
             )
         )
-    session.commit()
+    try:
+        session.commit()
+    except IntegrityError:
+        # A concurrent request computed and cached the same document first
+        # (both saw a cold cache); our freshly computed envelope is still valid.
+        session.rollback()
     return envelope
 
 

backend/app/main.py

@@ -107,7 +107,8 @@
 if settings.ENVIRONMENT in ("production", "qa"):
     sentry_sdk.init(
         dsn="https://b14aae02017e3a9c425de4b22af7dd0c@o4507623009091584.ingest.us.sentry.io/4507623009746944",
-        traces_sample_rate=1.0,
+        traces_sample_rate=settings.SENTRY_TRACES_SAMPLE_RATE,
+        # relative to traces_sample_rate, so one knob scales both
         profiles_sample_rate=1.0,
         environment=settings.ENVIRONMENT.value,
     )

backend/cli.py

@@ -35,6 +35,14 @@
 from os import environ
 from app.models import DistrictrMap, Overlay
 from datetime import datetime, timezone
+from stress_test.config import settings as stress_settings
+from stress_test.seed import (
+    NAME_PREFIX as STRESS_TEST_NAME_PREFIX,
+    delete_documents as _delete_stress_documents,
+    find_stress_documents as _find_stress_documents,
+    manifest_document_ids as _manifest_document_ids,
+    seed_documents as _seed_stress_documents,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -902,5 +910,87 @@ def check_missing_graphs(session: Session, skip_alert: bool):
     logger.info("Alert published to SNS topic %s", topic_arn)
 
 
+@cli.command("stress-test-seed")
+@click.option(
+    "--config-url",
+    help="Stress-test config JSON (URL or local path). Default: $STRESS_CONFIG_URL "
+    "or the CDN config (see stress_test/config.py).",
+)
+@click.option(
+    "--base-url",
+    help="API origin to create seed documents against (seeding goes over HTTP so "
+    "docs are created exactly as the app would). Default: $STRESS_BASE_URL or "
+    "http://localhost:8000.",
+)
+@click.option(
+    "--run-id",
+    help="Tags document names, User-Agent, and the manifest filename. "
+    "Default: $STRESS_RUN_ID.",
+)
+@click.option(
+    "--manifest",
+    help="Manifest output path (local or s3://). "
+    "Default: stress_test_manifest_<run-id>.json.",
+)
+@with_session
+def stress_test_seed(
+    session: Session,
+    config_url: str | None,
+    base_url: str | None,
+    run_id: str | None,
+    manifest: str | None,
+):
+    """Create the stress-test seed documents from the config JSON and write the
+    manifest of created document ids (consumed by the locustfile and by
+    `stress-test-cleanup`)."""
+    run_id = run_id or stress_settings.RUN_ID
+    documents = _seed_stress_documents(
+        session=session,
+        base_url=base_url or stress_settings.BASE_URL,
+        config_url=config_url or stress_settings.CONFIG_URL,
+        run_id=run_id,
+        manifest_path=manifest or f"stress_test_manifest_{run_id}.json",
+    )
+    logger.info(f"Seeded {len(documents)} documents")
+
+
+@cli.command("stress-test-cleanup")
+@click.option(
+    "--manifest",
+    "-m",
+    "manifests",
+    multiple=True,
+    help="Manifest JSON (seed or runtime shape; local path or s3://). Repeatable.",
+)
+@click.option(
+    "--yes",
+    "-y",
+    is_flag=True,
+    help="Delete leftover [STRESS-TEST]-named documents without prompting",
+)
+@with_session
+def stress_test_cleanup(session: Session, manifests: tuple[str, ...], yes: bool):
+    """Delete stress-test documents (rows + assignment partitions): everything
+    listed in the given manifests, then — belt and suspenders — any leftover
+    document whose metadata name starts with [STRESS-TEST], after confirmation."""
+    ids: list[str] = []
+    for path in manifests:
+        ids.extend(_manifest_document_ids(get_local_or_s3_path(file_path=path)))
+    if ids:
+        deleted = _delete_stress_documents(session, ids)
+        logger.info(f"Deleted {deleted} of {len(ids)} manifest-listed documents")
+
+    leftovers = _find_stress_documents(session)
+    if not leftovers:
+        logger.info(f"No {STRESS_TEST_NAME_PREFIX}-named documents remain")
+        return
+    for document_id, name in leftovers:
+        click.echo(f"  {document_id}  {name}")
+    if yes or click.confirm(
+        f"Delete these {len(leftovers)} {STRESS_TEST_NAME_PREFIX}-named document(s)?"
+    ):
+        _delete_stress_documents(session, [document_id for document_id, _ in leftovers])
+
+
 if __name__ == "__main__":
     cli()

backend/stress_test/.gitignore

@@ -0,0 +1,10 @@
+venv/
+artifacts/
+*.html
+*_stats.csv
+*_stats_history.csv
+*_failures.csv
+*_exceptions.csv
+stress_test_manifest_*.json
+stress_test_runtime_manifest_*.json
+.env