From 396976db34fb6a49f5491b54c5e31bbe04649f34 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Mon, 1 Jun 2026 12:57:07 -0400 Subject: [PATCH] Connection Exponential Backoff (#758) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(db): add connection retry, adaptive deriver polling, and pool metrics Add resilience and visibility for DB connection handling under transaction- pooler (Supavisor) saturation, where client-connection limits get exhausted across many tenants. - get_db/tracked_db now force an eager pool checkout with bounded exponential backoff (tenacity), retrying SQLAlchemy TimeoutError + OperationalError so transient pooler rejections degrade gracefully instead of 500ing. Toggle via DB_CONNECTION_RETRY_ENABLED (+ delay/backoff knobs); ~10s default budget. - Deriver polling backs off when idle or erroring (base -> max, x2 each cycle) and snaps back to base on claimed work, cutting steady-state query load. Toggle via DERIVER_POLLING_BACKOFF_ENABLED (+ max/multiplier). - Add scrape-time db_pool_connections Prometheus gauge (checked_out/checked_in/ size/overflow, labeled api|deriver), registered in both the API lifespan and the deriver metrics server. - Make SqlalchemyIntegration explicit in both Sentry inits; wrap connection acquisition in a db.pool.acquire span and capture live pool stats on retry-exhaustion. * feat(db): add acquisition counter and in-flight query gauge Build on the pool-connection metrics with two signals that turn detection into diagnosis under transaction-pooler saturation: - db_connection_acquisitions{outcome=ok|retried|exhausted}: counts how often connection checkout retries through pooler rejection — the alertable early warning before requests start failing. - db_queries_in_flight: statements actually executing on the wire (via SQLAlchemy cursor-execute events, drift-proof across query errors). Pairs with checked_out: the gap reveals connections held but parked (the "idle in transaction during an external call" antipattern). Labeled namespace + instance_type only; gated on METRICS.ENABLED for zero overhead when off. Add DB-free unit tests for retry outcomes, polling backoff, and in-flight gauge drift handling. * fix: address CodeRabbit review on PR #758 - db: roll back the session on a retryable checkout failure before retrying — a failed autobegin can leave it pending-rollback, making the next db.connection() raise instead of re-checking-out cleanly. Cheap Python-side cleanup when no connection was bound. - metrics: guard DBPoolCollector.collect() so a pool-read/import hiccup can't raise and abort the whole /metrics scrape (Prometheus drops ALL metrics if any collector raises) — log and fall back to empty. * fix(db): lazy retrying session + review fixes for connection backoff Address Codex/CodeRabbit review on PR #758. - Replace eager checkout with HonchoAsyncSession: a lazy AsyncSession that checks out its connection (with retry) on the first DB-touching call, not at construction. Request handlers doing non-DB work (embedding/file/LLM) before their first query no longer pin a connection across it, while the API path still gets checkout retry. Only the checkout is retried — the statement runs once via super(), so writes are never duplicated. Tracing's set_config moves into the same lazy acquire hook. - Roll the session back on a retryable checkout failure before retrying, so a failed autobegin can't leave it pending-rollback. - Lower default POOL_TIMEOUT to 5s and validate it stays under the retry budget for pooled (non-null) POOL_CLASS; update config.toml.example and v2/v3 docs. - Clamp pool overflow gauge to >= 0 (was negative before the pool fills). - Remove double-sleep in the deriver idle poll (true backoff cap, not 2x); make in-flight instrumentation registration idempotent. - Tests: HonchoAsyncSession lazy/idempotent acquire, statement-runs-once, tracing, commit/rollback flag reset, get_db no-acquire-at-entry, polling-loop single-sleep, and the POOL_TIMEOUT/retry-budget validator. * fix(db): cover all DB-touching session methods; clear flag on close/reset Address Codex follow-up review on PR #758 (polish, no behavior-critical bug). - HonchoAsyncSession: wrap get/get_one/stream/stream_scalars/delete in addition to execute/scalar/scalars/flush/merge/refresh/commit, so the "lazy checkout with retry on first DB use" guarantee has no holes. connection() stays unwrapped (acquire_connection_with_retry calls it — wrapping would recurse). - Reset the acquired flag on close()/reset() too, so a session reused after close/reset re-acquires (and re-wraps retry) on its next DB use. - Fix stale comments: connection retry now applies lazily to the request path via HonchoAsyncSession (config.py), and the FakeSession helper note. - Tests: close/reset flag reset, and get/delete route through acquisition. --- config.toml.example | 9 +- docs/v2/contributing/configuration.mdx | 2 +- docs/v3/contributing/configuration.mdx | 8 +- src/config.py | 59 +++- src/db.py | 314 ++++++++++++++++++++- src/dependencies.py | 34 +-- src/deriver/__main__.py | 11 +- src/deriver/queue_manager.py | 45 ++- src/main.py | 11 +- src/telemetry/__init__.py | 7 +- src/telemetry/prometheus/__init__.py | 2 + src/telemetry/prometheus/metrics.py | 79 ++++++ tests/test_db_resilience.py | 370 +++++++++++++++++++++++++ tests/test_dependencies.py | 33 +-- 14 files changed, 924 insertions(+), 60 deletions(-) create mode 100644 tests/test_db_resilience.py diff --git a/config.toml.example b/config.toml.example index 0f4b92c8..258172a7 100644 --- a/config.toml.example +++ b/config.toml.example @@ -26,11 +26,18 @@ POOL_CLASS = "default" POOL_PRE_PING = true POOL_SIZE = 10 MAX_OVERFLOW = 20 -POOL_TIMEOUT = 30 # seconds +POOL_TIMEOUT = 5 # seconds; must stay under CONNECTION_RETRY_MAX_DELAY_SECONDS POOL_RECYCLE = 300 # seconds POOL_USE_LIFO = true SQL_DEBUG = false TRACING = false +# Bounded retry around connection checkout (used by background/tracked_db and +# the lazy request session). For a pooled (non-null) POOL_CLASS, POOL_TIMEOUT +# must be < CONNECTION_RETRY_MAX_DELAY_SECONDS or startup validation fails. +CONNECTION_RETRY_ENABLED = true +CONNECTION_RETRY_MAX_DELAY_SECONDS = 10.0 +CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS = 0.1 +CONNECTION_RETRY_BACKOFF_MAX_SECONDS = 2.0 # Authentication settings [auth] diff --git a/docs/v2/contributing/configuration.mdx b/docs/v2/contributing/configuration.mdx index c172369c..f2548f94 100644 --- a/docs/v2/contributing/configuration.mdx +++ b/docs/v2/contributing/configuration.mdx @@ -161,7 +161,7 @@ DB_CONNECTION_URI=postgresql+psycopg://honcho_user:secure_password@db.example.co DB_SCHEMA=public DB_POOL_SIZE=10 DB_MAX_OVERFLOW=20 -DB_POOL_TIMEOUT=30 +DB_POOL_TIMEOUT=5 DB_POOL_RECYCLE=300 DB_POOL_PRE_PING=true DB_SQL_DEBUG=false diff --git a/docs/v3/contributing/configuration.mdx b/docs/v3/contributing/configuration.mdx index 582bde3a..d7a71b6c 100644 --- a/docs/v3/contributing/configuration.mdx +++ b/docs/v3/contributing/configuration.mdx @@ -459,10 +459,16 @@ DB_SCHEMA=public DB_POOL_PRE_PING=true DB_POOL_SIZE=10 DB_MAX_OVERFLOW=20 -DB_POOL_TIMEOUT=30 +DB_POOL_TIMEOUT=5 DB_POOL_RECYCLE=300 DB_POOL_USE_LIFO=true DB_SQL_DEBUG=false +# Bounded retry around connection checkout. For a pooled (non-null) DB_POOL_CLASS, +# DB_POOL_TIMEOUT must be < DB_CONNECTION_RETRY_MAX_DELAY_SECONDS or startup fails. +DB_CONNECTION_RETRY_ENABLED=true +DB_CONNECTION_RETRY_MAX_DELAY_SECONDS=10.0 +DB_CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS=0.1 +DB_CONNECTION_RETRY_BACKOFF_MAX_SECONDS=2.0 ``` ### Authentication diff --git a/src/config.py b/src/config.py index 5ae93fa4..35d72d3e 100644 --- a/src/config.py +++ b/src/config.py @@ -612,8 +612,9 @@ class DBSettings(HonchoSettings): POOL_PRE_PING: bool = True POOL_SIZE: Annotated[int, Field(default=10, gt=0, le=1000)] = 10 MAX_OVERFLOW: Annotated[int, Field(default=20, ge=0, le=1000)] = 20 - POOL_TIMEOUT: Annotated[int, Field(default=30, gt=0, le=300)] = ( - 30 # seconds (max 5 minutes) + POOL_TIMEOUT: Annotated[int, Field(default=5, gt=0, le=300)] = ( + 5 # seconds; kept under CONNECTION_RETRY_MAX_DELAY_SECONDS so a pool + # checkout fails fast enough to allow a retry within the budget ) POOL_RECYCLE: Annotated[int, Field(default=300, gt=0, le=7200)] = ( 300 # seconds (max 2 hours) @@ -622,6 +623,48 @@ class DBSettings(HonchoSettings): SQL_DEBUG: bool = False TRACING: bool = False + # Bounded exponential-backoff retry around connection acquisition. Applied + # lazily on the first DB use of any session (HonchoAsyncSession) — both the + # request path and background/tracked_db scopes — without forcing an eager + # checkout. Guards against transient transaction-pooler saturation (e.g. + # Supavisor rejecting with "too many clients") by retrying the checkout + # instead of failing immediately. CONNECTION_RETRY_MAX_DELAY_SECONDS is the + # TOTAL retry + # budget; with a real QueuePool, a single checkout can block up to + # POOL_TIMEOUT, so POOL_TIMEOUT must stay below the budget for a retry to be + # possible (enforced below). With NullPool (the transaction-pooler setup) + # there is no local queue wait, so saturation is a fast OperationalError and + # POOL_TIMEOUT does not apply. + CONNECTION_RETRY_ENABLED: bool = True + CONNECTION_RETRY_MAX_DELAY_SECONDS: Annotated[ + float, Field(default=10.0, gt=0.0, le=120.0) + ] = 10.0 + CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS: Annotated[ + float, Field(default=0.1, gt=0.0, le=10.0) + ] = 0.1 + CONNECTION_RETRY_BACKOFF_MAX_SECONDS: Annotated[ + float, Field(default=2.0, gt=0.0, le=30.0) + ] = 2.0 + + @model_validator(mode="after") + def _validate_retry_budget_vs_pool_timeout(self) -> "DBSettings": + # Only meaningful for a real local pool: NullPool has no queue wait, so + # POOL_TIMEOUT is unused there. For a QueuePool, a checkout can block up + # to POOL_TIMEOUT, so it must be < the total retry budget or the first + # attempt consumes the whole budget and no retry ever happens. + if ( + self.POOL_CLASS != "null" + and self.CONNECTION_RETRY_ENABLED + and self.POOL_TIMEOUT >= self.CONNECTION_RETRY_MAX_DELAY_SECONDS + ): + raise ValueError( + f"DB_POOL_TIMEOUT ({self.POOL_TIMEOUT}s) must be less than " + + "DB_CONNECTION_RETRY_MAX_DELAY_SECONDS " + + f"({self.CONNECTION_RETRY_MAX_DELAY_SECONDS}s) so a pooled checkout " + + "can fail fast enough to be retried within the budget." + ) + return self + class AuthSettings(HonchoSettings): model_config = SettingsConfigDict(env_prefix="AUTH_", extra="ignore") # pyright: ignore @@ -737,6 +780,18 @@ class DeriverSettings(HonchoSettings): POLLING_SLEEP_INTERVAL_SECONDS: Annotated[ float, Field(default=1.0, gt=0.0, le=60.0) ] = 1.0 + # Adaptive polling: when the queue is idle (or the loop is erroring) the + # sleep interval grows from POLLING_SLEEP_INTERVAL_SECONDS toward + # POLLING_SLEEP_MAX_INTERVAL_SECONDS by POLLING_BACKOFF_MULTIPLIER each + # cycle, then snaps back to the base interval as soon as work is found. + # Reduces steady-state query load against the (shared) DB/pooler. + POLLING_BACKOFF_ENABLED: bool = True + POLLING_SLEEP_MAX_INTERVAL_SECONDS: Annotated[ + float, Field(default=30.0, gt=0.0, le=300.0) + ] = 30.0 + POLLING_BACKOFF_MULTIPLIER: Annotated[ + float, Field(default=2.0, ge=1.0, le=10.0) + ] = 2.0 STALE_SESSION_TIMEOUT_MINUTES: Annotated[int, Field(default=5, gt=0, le=1440)] = 5 # Retention window (seconds) for keeping errored items in the queue diff --git a/src/db.py b/src/db.py index 4b656175..0a133c05 100644 --- a/src/db.py +++ b/src/db.py @@ -1,11 +1,32 @@ import contextvars +import logging +from typing import Any -from sqlalchemy import MetaData, text -from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +import sentry_sdk +from sqlalchemy import MetaData, event, text +from sqlalchemy.exc import OperationalError +from sqlalchemy.exc import TimeoutError as SQLAlchemyTimeoutError +from sqlalchemy.ext.asyncio import ( + AsyncSession, + async_sessionmaker, + create_async_engine, +) from sqlalchemy.orm import declarative_base -from sqlalchemy.pool import NullPool +from sqlalchemy.pool import NullPool, QueuePool +from tenacity import ( + AsyncRetrying, + retry_if_exception_type, + stop_after_delay, + wait_exponential_jitter, +) from src.config import settings +from src.telemetry.prometheus.metrics import ( + db_queries_in_flight_gauge, + prometheus_metrics, +) + +logger = logging.getLogger(__name__) connect_args = {"prepare_threshold": None} @@ -38,13 +59,300 @@ engine = create_async_engine( **engine_kwargs, ) +# NOTE: SessionLocal is defined further down, after HonchoAsyncSession (its +# session class) and acquire_connection_with_retry (which that class calls). + +# Errors worth retrying when acquiring a pooled connection: SQLAlchemy's local +# pool-checkout timeout, and OperationalError (how a saturated transaction +# pooler surfaces "too many clients" / connection refusals). +RETRYABLE_DB_CONNECTION_ERRORS = (SQLAlchemyTimeoutError, OperationalError) + +# Identifies this process ("api" | "deriver") on DB metrics. Set once at startup +# by register_db_query_instrumentation; stays "unknown" if metrics are disabled. +_db_instance_type: str = "unknown" + + +def _record_acquisition_outcome(outcome: str) -> None: + """Record a connection-acquisition outcome (no-op when metrics disabled).""" + if settings.METRICS.ENABLED: + prometheus_metrics.record_db_connection_acquisition( + instance_type=_db_instance_type, outcome=outcome + ) + + +def get_pool_stats() -> dict[str, int]: + """Return live connection-pool stats for this process. + + ``engine.pool`` is the AsyncEngine's pool (the same object as + ``engine.sync_engine.pool``); its stat methods are synchronous counter + reads with no I/O, so they are safe to call without ``await``. Returns + zeros for pools that do not track connections (e.g. ``NullPool``). + """ + zeros = {"checked_out": 0, "checked_in": 0, "size": 0, "overflow": 0} + pool = engine.pool + # Only QueuePool (and its AsyncAdaptedQueuePool subclass) tracks connection + # counts; NullPool and others have no meaningful stats. + if not isinstance(pool, QueuePool): + return zeros + try: + # overflow() is negative until the base pool fills (it starts at + # -pool_size); clamp to the count of overflow connections actually open. + return { + "checked_out": pool.checkedout(), + "checked_in": pool.checkedin(), + "size": pool.size(), + "overflow": max(0, pool.overflow()), + } + except Exception: + return zeros + + +async def acquire_connection_with_retry(db: AsyncSession, context: str) -> None: + """Force pool checkout (which ``SessionLocal()`` defers) with bounded backoff. + + ``SessionLocal()`` is lazy: the pool checkout — and any pooler rejection — + happens on the first query. We force it here inside a retry block so that + transient saturation of the transaction pooler is retried with exponential + backoff rather than surfacing as an immediate error. The checkout is wrapped + in a Sentry span so wait time is visible in traces; on budget exhaustion the + original error is reraised after capturing live pool stats to Sentry. + + Each attempt rolls the session back on a retryable failure before retrying: + a failed checkout can leave the autobegun transaction in a pending-rollback + state, which would make the next ``db.connection()`` raise instead of + re-checking-out cleanly. The rollback is pure Python-side state cleanup when + no connection was bound, so it is cheap and safe. + """ + with sentry_sdk.start_span(op="db.pool.acquire", name=context): + if not settings.DB.CONNECTION_RETRY_ENABLED: + await db.connection() + return + attempts = 0 + try: + async for attempt in AsyncRetrying( + wait=wait_exponential_jitter( + initial=settings.DB.CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS, + max=settings.DB.CONNECTION_RETRY_BACKOFF_MAX_SECONDS, + ), + stop=stop_after_delay(settings.DB.CONNECTION_RETRY_MAX_DELAY_SECONDS), + retry=retry_if_exception_type(RETRYABLE_DB_CONNECTION_ERRORS), + reraise=True, + ): + with attempt: + attempts += 1 + try: + await db.connection() + except RETRYABLE_DB_CONNECTION_ERRORS: + # Reset session state so the next attempt starts clean. + try: + await db.rollback() + except Exception: + logger.debug( + "rollback after failed checkout failed", + exc_info=True, + ) + raise + except RETRYABLE_DB_CONNECTION_ERRORS as e: + _record_acquisition_outcome("exhausted") + if settings.SENTRY.ENABLED: + sentry_sdk.set_context("db_pool", get_pool_stats()) + sentry_sdk.capture_exception(e) + raise + # "ok" on first try, "retried" if backoff was needed before success. + _record_acquisition_outcome("ok" if attempts <= 1 else "retried") + + +class HonchoAsyncSession(AsyncSession): + """AsyncSession that lazily checks out its connection, with retry. + + The pool checkout — and any pooler-rejection retry via + ``acquire_connection_with_retry`` — happens on the FIRST DB-touching call, + not at construction. So a request handler that does non-DB work (embedding, + file processing, an LLM call) before its first query does NOT pin a + connection across that work, while still getting checkout retry on the + request path. + + Only the checkout is retried; the SQL statement itself runs exactly once + (we never retry ``super().execute`` after a broad OperationalError), so + writes are never duplicated. The context for tracing is read from the + ``request_context`` ContextVar, which the request/task scope has already set. + """ + + # Class-level default; per-instance assignment shadows it (the subclass has + # a __dict__ even though AsyncSession declares __slots__). + _honcho_acquired: bool = False + + async def _ensure_acquired(self) -> None: + if self._honcho_acquired: + return + context = request_context.get() or "unknown" + await acquire_connection_with_retry(self, context) + self._honcho_acquired = True + if settings.DB.TRACING: + # Forced checkout already happened above; this rides the same + # connection. super() to avoid re-entering _ensure_acquired. + await super().execute( + text("SELECT set_config('application_name', :name, false)"), + {"name": context}, + ) + + # The overrides below are thin: ensure the connection is checked out (once, + # with retry) before delegating to AsyncSession. They cover every public + # DB-touching async method so the "lazy retry on first DB use" guarantee has + # no holes. Signatures are widened to *args/**kwargs because we only forward; + # call sites are typed against the AsyncSession base, so this does not weaken + # type-checking elsewhere. (connection() is intentionally NOT wrapped — + # acquire_connection_with_retry calls it, so wrapping would recurse.) + async def execute(self, *args: Any, **kwargs: Any) -> Any: + await self._ensure_acquired() + return await super().execute(*args, **kwargs) + + async def scalar(self, *args: Any, **kwargs: Any) -> Any: + await self._ensure_acquired() + return await super().scalar(*args, **kwargs) + + async def scalars(self, *args: Any, **kwargs: Any) -> Any: + await self._ensure_acquired() + return await super().scalars(*args, **kwargs) + + async def get(self, *args: Any, **kwargs: Any) -> Any: + await self._ensure_acquired() + return await super().get(*args, **kwargs) + + async def get_one(self, *args: Any, **kwargs: Any) -> Any: + await self._ensure_acquired() + return await super().get_one(*args, **kwargs) + + async def stream(self, *args: Any, **kwargs: Any) -> Any: + await self._ensure_acquired() + return await super().stream(*args, **kwargs) + + async def stream_scalars(self, *args: Any, **kwargs: Any) -> Any: + await self._ensure_acquired() + return await super().stream_scalars(*args, **kwargs) + + async def flush(self, *args: Any, **kwargs: Any) -> None: + await self._ensure_acquired() + await super().flush(*args, **kwargs) + + async def merge(self, *args: Any, **kwargs: Any) -> Any: + await self._ensure_acquired() + return await super().merge(*args, **kwargs) + + async def refresh(self, *args: Any, **kwargs: Any) -> None: + await self._ensure_acquired() + await super().refresh(*args, **kwargs) + + async def delete(self, *args: Any, **kwargs: Any) -> None: + await self._ensure_acquired() + await super().delete(*args, **kwargs) + + async def commit(self) -> None: + # Ensures the add()->commit() path (autoflush on commit) also retries. + await self._ensure_acquired() + try: + await super().commit() + finally: + # Transaction ended; a later op must re-acquire (and re-wrap retry). + self._honcho_acquired = False + + async def rollback(self) -> None: + try: + await super().rollback() + finally: + self._honcho_acquired = False + + async def close(self) -> None: + try: + await super().close() + finally: + # The connection is released; a reused session must re-acquire. + self._honcho_acquired = False + + async def reset(self) -> None: + try: + await super().reset() + finally: + self._honcho_acquired = False + + SessionLocal = async_sessionmaker( autocommit=False, autoflush=False, expire_on_commit=False, bind=engine, + class_=HonchoAsyncSession, ) + +class DBQueryInflightTracker: + """Tracks statements executing on the wire via SQLAlchemy cursor events. + + Drift-proof: marks ``Connection.info`` when a statement starts and clears it + on completion OR error, so the gauge can't leak upward (an errored statement + skips ``after_cursor_execute``) or go negative (a connect-time error has no + matching start). Bound to a pre-resolved labeled gauge child so the + per-statement hot path does no label resolution. + """ + + # Marker on Connection.info recording that we incremented for the current + # statement, so we decrement exactly once on completion or error. + INFLIGHT_KEY: str = "_honcho_inflight" + + def __init__(self, gauge_child: Any) -> None: + self._child: Any = gauge_child + + def on_before(self, conn: Any, *_: Any) -> None: + try: + conn.info[self.INFLIGHT_KEY] = True + self._child.inc() + except Exception: + logger.debug("in-flight gauge inc failed", exc_info=True) + + def on_after(self, conn: Any, *_: Any) -> None: + try: + if conn.info.pop(self.INFLIGHT_KEY, False): + self._child.dec() + except Exception: + logger.debug("in-flight gauge dec failed", exc_info=True) + + def on_error(self, exception_context: Any) -> None: + try: + conn = exception_context.connection + if conn is not None and conn.info.pop(self.INFLIGHT_KEY, False): + self._child.dec() + except Exception: + logger.debug("in-flight gauge error-path dec failed", exc_info=True) + + +# Process-wide tracker, created at registration (None until then / if metrics off). +_inflight_tracker: DBQueryInflightTracker | None = None + + +_db_query_instrumentation_registered = False + + +def register_db_query_instrumentation(instance_type: str) -> None: + """Attach per-statement in-flight tracking to the engine (no-op if off). + + Gated on METRICS.ENABLED so there is zero overhead — not even attached event + listeners — when metrics are disabled. Idempotent: repeated calls (e.g. a + re-run lifespan or test startup) won't attach duplicate listeners, which + would double-count in-flight statements. + """ + global _db_instance_type, _inflight_tracker, _db_query_instrumentation_registered + _db_instance_type = instance_type + if not settings.METRICS.ENABLED or _db_query_instrumentation_registered: + return + child = db_queries_in_flight_gauge.labels(instance_type=instance_type) + _inflight_tracker = DBQueryInflightTracker(child) + sync_engine = engine.sync_engine + event.listen(sync_engine, "before_cursor_execute", _inflight_tracker.on_before) + event.listen(sync_engine, "after_cursor_execute", _inflight_tracker.on_after) + event.listen(sync_engine, "handle_error", _inflight_tracker.on_error) + _db_query_instrumentation_registered = True + + # Define your naming convention convention = { "ix": "ix_%(table_name)s_%(column_0_N_name)s", # Index - supports multi-column diff --git a/src/dependencies.py b/src/dependencies.py index 66b1931e..cd7007e0 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -2,25 +2,22 @@ import uuid from contextlib import asynccontextmanager from fastapi import Depends -from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession -from src.config import settings from src.db import SessionLocal, request_context async def get_db(): - """FastAPI Dependency Generator for Database""" - - context = request_context.get() or "unknown" + """FastAPI Dependency Generator for Database. + The session is lazy: it does NOT check out a pooled connection here. + HonchoAsyncSession acquires it (with retry) on the first DB-touching call, + so a handler doing non-DB work (embedding/file/LLM) before its first query + does not pin a connection across it. Tracing's application_name and the + checkout retry both live in the session, keyed off request_context. + """ db: AsyncSession = SessionLocal() try: - if settings.DB.TRACING: - await db.execute( - text("SELECT set_config('application_name', :name, false)"), - {"name": context}, - ) yield db except Exception: await db.rollback() @@ -30,14 +27,19 @@ async def get_db(): # is closed before the TCP connection drops. Supavisor v2 does NOT # clean up orphaned transactions on client disconnect in transaction- # pooling mode, so relying on `in_transaction()` (Python-side state) - # can leave the backend pinned with an open BEGIN. + # can leave the backend pinned with an open BEGIN. (Cheap no-op if the + # lazy session never checked out a connection.) await db.rollback() await db.close() @asynccontextmanager async def tracked_db(operation_name: str | None = None): - """Context manager for tracked database sessions""" + """Context manager for tracked database sessions. + + Sets a task-scoped request_context so the lazy session picks it up for + tracing/attribution, then yields a lazy session (see get_db). + """ # Get request ID if available, or create operation-specific one context = request_context.get() token = None @@ -46,16 +48,8 @@ async def tracked_db(operation_name: str | None = None): context = f"task:{operation_name}:{str(uuid.uuid4())[:8]}" token = request_context.set(context) - # Create session with tracking info db = SessionLocal() - try: - if settings.DB.TRACING: - await db.execute( - text("SELECT set_config('application_name', :name, false)"), - {"name": context or f"task:{operation_name}"}, - ) - yield db except Exception: await db.rollback() diff --git a/src/deriver/__main__.py b/src/deriver/__main__.py index 20ffbb66..c56ed6a0 100644 --- a/src/deriver/__main__.py +++ b/src/deriver/__main__.py @@ -6,9 +6,13 @@ import uvloop from prometheus_client import start_http_server from src.config import settings -from src.db import engine +from src.db import engine, register_db_query_instrumentation from src.startup import validate_embedding_schema -from src.telemetry import initialize_telemetry_async, shutdown_telemetry +from src.telemetry import ( + initialize_telemetry_async, + register_db_pool_collector, + shutdown_telemetry, +) from .queue_manager import main @@ -18,6 +22,9 @@ logger = logging.getLogger(__name__) def start_metrics_server() -> None: """Start the Prometheus metrics HTTP server on port 9090.""" start_http_server(9090) + # Expose DB connection-pool stats for this deriver instance. + register_db_pool_collector("deriver") + register_db_query_instrumentation("deriver") logger.info("Prometheus metrics server started on port 9090") diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py index 8c2b5850..22716cba 100644 --- a/src/deriver/queue_manager.py +++ b/src/deriver/queue_manager.py @@ -11,6 +11,7 @@ import sentry_sdk 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.dialects.postgresql import insert from sqlalchemy.engine import CursorResult @@ -125,6 +126,12 @@ class QueueManager: self.worker_ownership: dict[str, WorkerOwnership] = {} self.queue_empty_flag: asyncio.Event = asyncio.Event() + # 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 = ( + settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS + ) + # Initialize from settings self.workers: int = settings.DERIVER.WORKERS self.semaphore: asyncio.Semaphore = asyncio.Semaphore(self.workers) @@ -147,7 +154,9 @@ class QueueManager: # Initialize Sentry if enabled, using settings if settings.SENTRY.ENABLED: - initialize_sentry(integrations=[AsyncioIntegration()]) + initialize_sentry( + integrations=[AsyncioIntegration(), SqlalchemyIntegration()] + ) def add_task(self, task: asyncio.Task[None]) -> None: """Track a new task""" @@ -378,18 +387,36 @@ class QueueManager: ) return claimed_mapping + def _reset_poll_interval(self) -> None: + """Snap the polling interval back to the base after finding work.""" + self._current_poll_interval = settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS + + def _advance_poll_interval(self) -> float: + """Return the current idle/backoff sleep, then grow it toward the cap.""" + interval = self._current_poll_interval + if settings.DERIVER.POLLING_BACKOFF_ENABLED: + self._current_poll_interval = min( + self._current_poll_interval + * settings.DERIVER.POLLING_BACKOFF_MULTIPLIER, + settings.DERIVER.POLLING_SLEEP_MAX_INTERVAL_SECONDS, + ) + return interval + async def polling_loop(self) -> None: """Main polling loop to find and process new work units""" logger.debug("Starting polling loop") try: while not self.shutdown_event.is_set(): if self.queue_empty_flag.is_set(): - # logger.debug("Queue empty flag set, waiting") - await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS) + # The empty-poll branch below already slept this cycle's + # interval; just clear the flag and re-query (no second + # sleep — that would double the effective idle interval). self.queue_empty_flag.clear() continue - # Check if we have capacity before querying + # Check if we have capacity before querying. There is work to do + # (workers are busy), so keep the base interval for fast pickup + # when capacity frees rather than backing off. if self.semaphore.locked(): # logger.debug("All workers busy, waiting") await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS) @@ -399,6 +426,7 @@ class QueueManager: await self.cleanup_stale_work_units() claimed_work_units = await self.get_and_claim_work_units() if claimed_work_units: + self._reset_poll_interval() for work_unit_key, aqs_id in claimed_work_units.items(): # Create a new task for processing this work unit if not self.shutdown_event.is_set(): @@ -414,15 +442,14 @@ class QueueManager: self.add_task(task) else: self.queue_empty_flag.set() - await asyncio.sleep( - settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS - ) + await asyncio.sleep(self._advance_poll_interval()) except Exception as e: logger.exception("Error in polling loop") if settings.SENTRY.ENABLED: sentry_sdk.capture_exception(e) - # Note: rollback is handled by tracked_db dependency - await asyncio.sleep(settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS) + # Note: rollback is handled by tracked_db dependency. + # Back off so a down/saturated DB isn't hammered every cycle. + await asyncio.sleep(self._advance_poll_interval()) finally: logger.info("Polling loop stopped") diff --git a/src/main.py b/src/main.py index d08d1164..6edd9d4c 100644 --- a/src/main.py +++ b/src/main.py @@ -13,12 +13,13 @@ from fastapi.responses import JSONResponse from fastapi_pagination import add_pagination from pydantic import ValidationError from sentry_sdk.integrations.fastapi import FastApiIntegration +from sentry_sdk.integrations.sqlalchemy import SqlalchemyIntegration from sentry_sdk.integrations.starlette import StarletteIntegration from src._version import HONCHO_VERSION from src.cache.client import close_cache, init_cache from src.config import settings -from src.db import engine, request_context +from src.db import engine, register_db_query_instrumentation, request_context from src.exceptions import HonchoException from src.routers import ( conclusions, @@ -34,6 +35,7 @@ from src.telemetry import ( initialize_telemetry_async, metrics_endpoint, prometheus_metrics, + register_db_pool_collector, shutdown_telemetry, ) from src.telemetry.logging import get_route_template @@ -117,6 +119,9 @@ if SENTRY_ENABLED: FastApiIntegration( transaction_style="endpoint", ), + # Explicit so DB-query spans (and our db.pool.acquire span) are not + # reliant on auto-enabling. + SqlalchemyIntegration(), ], before_send=before_send, ) @@ -127,6 +132,10 @@ async def lifespan(_: FastAPI): # Initialize CloudEvents telemetry await initialize_telemetry_async() + # Expose DB connection-pool stats for this API instance (no-op if metrics off) + register_db_pool_collector("api") + register_db_query_instrumentation("api") + # Validate embedding schema before serving any traffic. Fails closed: if # the configured EMBEDDING_VECTOR_DIMENSIONS does not match the physical # pgvector columns, the process refuses to start rather than silently diff --git a/src/telemetry/__init__.py b/src/telemetry/__init__.py index 401be523..abbc5df0 100644 --- a/src/telemetry/__init__.py +++ b/src/telemetry/__init__.py @@ -12,13 +12,18 @@ This module consolidates all telemetry, metrics, and observability functionality """ from src.telemetry.events import emit -from src.telemetry.prometheus import metrics_endpoint, prometheus_metrics +from src.telemetry.prometheus import ( + metrics_endpoint, + prometheus_metrics, + register_db_pool_collector, +) __all__ = [ "emit", "initialize_telemetry_async", "metrics_endpoint", "prometheus_metrics", + "register_db_pool_collector", "shutdown_telemetry", ] diff --git a/src/telemetry/prometheus/__init__.py b/src/telemetry/prometheus/__init__.py index ab0a2616..876ea110 100644 --- a/src/telemetry/prometheus/__init__.py +++ b/src/telemetry/prometheus/__init__.py @@ -14,6 +14,7 @@ from src.telemetry.prometheus.metrics import ( TokenTypes, metrics_endpoint, prometheus_metrics, + register_db_pool_collector, ) __all__ = [ @@ -23,4 +24,5 @@ __all__ = [ "TokenTypes", "metrics_endpoint", "prometheus_metrics", + "register_db_pool_collector", ] diff --git a/src/telemetry/prometheus/metrics.py b/src/telemetry/prometheus/metrics.py index 90082b85..1a79c634 100644 --- a/src/telemetry/prometheus/metrics.py +++ b/src/telemetry/prometheus/metrics.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from collections.abc import Iterator from enum import Enum from typing import cast, final @@ -14,6 +15,7 @@ from prometheus_client import ( disable_created_metrics, generate_latest, ) +from prometheus_client.core import GaugeMetricFamily from starlette.requests import Request from starlette.responses import Response @@ -125,6 +127,23 @@ telemetry_buffer_size_gauge = NamespacedGauge( ["namespace"], ) +# DB connection-pool health. The acquisitions counter measures how often we hit +# (and retry through) transaction-pooler saturation; the in-flight gauge counts +# statements actually executing on the wire, so checked_out minus in_flight +# reveals connections held but parked (the "idle in transaction during an +# external call" antipattern). +db_connection_acquisitions_counter = NamespacedCounter( + "db_connection_acquisitions", + "DB connection acquisitions by outcome (ok=first try, retried, exhausted)", + ["namespace", "instance_type", "outcome"], +) + +db_queries_in_flight_gauge = NamespacedGauge( + "db_queries_in_flight", + "DB statements currently executing on a connection for this instance", + ["namespace", "instance_type"], +) + @final class PrometheusMetrics: @@ -275,10 +294,70 @@ class PrometheusMetrics: except Exception as e: self._handle_metric_error("set_telemetry_buffer_size", e) + def record_db_connection_acquisition( + self, *, instance_type: str, outcome: str + ) -> None: + # outcome is one of "ok" | "retried" | "exhausted". + try: + db_connection_acquisitions_counter.labels( + instance_type=instance_type, + outcome=outcome, + ).inc() + except Exception as e: + self._handle_metric_error("record_db_connection_acquisition", e) + prometheus_metrics = PrometheusMetrics() +class DBPoolCollector: + """Scrape-time collector for SQLAlchemy connection-pool stats. + + Computed live on each /metrics scrape from the async engine's pool, so it + is always current with no background task or sampling lag. One instance is + registered per process (the API server or a deriver worker). + """ + + def __init__(self, instance_type: str) -> None: + # instance_type: "api" | "deriver" + self.instance_type: str = instance_type + + def collect(self) -> Iterator[GaugeMetricFamily]: + namespace = settings.METRICS.NAMESPACE or "" + gauge = GaugeMetricFamily( + "db_pool_connections", + "DB connections held by this instance, by pool state", + labels=["namespace", "instance_type", "state"], + ) + # Fail soft: Prometheus aborts the entire scrape (dropping ALL metrics) + # if any collector raises, so never let a pool/import hiccup here sink + # the whole /metrics response. + try: + # Lazy import to avoid an import cycle at module load (db imports + # config, telemetry is imported widely). Reads engine.pool directly. + from src.db import get_pool_stats + + stats = get_pool_stats() + except Exception: + logger.warning("Failed to collect DB pool stats", exc_info=True) + stats = {} + for state, value in stats.items(): + gauge.add_metric([namespace, self.instance_type, state], value) + yield gauge + + +_db_pool_collector_registered = False + + +def register_db_pool_collector(instance_type: str) -> None: + """Register the DB pool collector once per process (no-op if metrics off).""" + global _db_pool_collector_registered + if _db_pool_collector_registered or not settings.METRICS.ENABLED: + return + REGISTRY.register(DBPoolCollector(instance_type)) + _db_pool_collector_registered = True + + async def metrics_endpoint(_request: Request) -> Response: if not settings.METRICS.ENABLED: return Response("Metrics are disabled", status_code=404) diff --git a/tests/test_db_resilience.py b/tests/test_db_resilience.py new file mode 100644 index 00000000..a82f307a --- /dev/null +++ b/tests/test_db_resilience.py @@ -0,0 +1,370 @@ +"""Unit tests for DB connection resilience + observability. + +These are DB-free: they exercise the retry helper against a fake session, the +deriver polling backoff math, and the in-flight gauge listeners directly. +""" + +from types import SimpleNamespace +from typing import Any + +import pytest +from sqlalchemy.exc import OperationalError +from sqlalchemy.ext.asyncio import AsyncSession + +import src.db as db_module +from src.config import settings +from src.db import DBQueryInflightTracker, acquire_connection_with_retry +from src.telemetry.prometheus.metrics import ( + db_connection_acquisitions_counter, + db_queries_in_flight_gauge, +) + + +def _make_operational_error() -> OperationalError: + """A stand-in for how a saturated pooler surfaces ('too many clients').""" + return OperationalError("SELECT 1", {}, Exception("too many clients")) + + +class _FlakyConnSession: + """Fake AsyncSession whose connection() fails N times then succeeds.""" + + def __init__(self, fail_times: int, *, always_fail: bool = False) -> None: + self.fail_times: int = fail_times + self.always_fail: bool = always_fail + self.calls: int = 0 + self.rollback_calls: int = 0 + + async def connection(self) -> None: + self.calls += 1 + if self.always_fail or self.calls <= self.fail_times: + raise _make_operational_error() + + async def rollback(self) -> None: + self.rollback_calls += 1 + + +def _acq_count(outcome: str) -> float: + child = db_connection_acquisitions_counter.labels( + instance_type="api", outcome=outcome + ) + return float(child._value.get()) # pyright: ignore[reportPrivateUsage, reportUnknownArgumentType] + + +@pytest.fixture +def metrics_on(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(settings.METRICS, "ENABLED", True) + monkeypatch.setattr(settings.METRICS, "NAMESPACE", "test") + monkeypatch.setattr(db_module, "_db_instance_type", "api") + # Fast, deterministic backoff so retry tests don't sleep. + monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_ENABLED", True) + monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS", 0.001) + monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_BACKOFF_MAX_SECONDS", 0.01) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("metrics_on") +async def test_acquire_succeeds_first_try_records_ok() -> None: + before = _acq_count("ok") + session = _FlakyConnSession(fail_times=0) + await acquire_connection_with_retry(session, "request:test") # pyright: ignore[reportArgumentType] + assert session.calls == 1 + assert _acq_count("ok") == before + 1 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("metrics_on") +async def test_acquire_retries_then_succeeds_records_retried() -> None: + before = _acq_count("retried") + session = _FlakyConnSession(fail_times=2) + await acquire_connection_with_retry(session, "request:test") # pyright: ignore[reportArgumentType] + assert session.calls == 3 # two failures then success + # Session is reset after each failed checkout before the next attempt. + assert session.rollback_calls == 2 + assert _acq_count("retried") == before + 1 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("metrics_on") +async def test_acquire_exhausts_budget_reraises_and_records( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Tiny budget so the loop gives up quickly under sustained failure. + monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_MAX_DELAY_SECONDS", 0.05) + before = _acq_count("exhausted") + session = _FlakyConnSession(fail_times=0, always_fail=True) + with pytest.raises(OperationalError): + await acquire_connection_with_retry(session, "request:test") # pyright: ignore[reportArgumentType] + assert session.calls >= 1 + assert _acq_count("exhausted") == before + 1 + + +@pytest.mark.asyncio +async def test_acquire_disabled_calls_once(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings.DB, "CONNECTION_RETRY_ENABLED", False) + session = _FlakyConnSession(fail_times=0) + await acquire_connection_with_retry(session, "request:test") # pyright: ignore[reportArgumentType] + assert session.calls == 1 + + +# --- HonchoAsyncSession: lazy checkout with retry on first DB use ------------ + + +@pytest.mark.asyncio +async def test_session_lazy_acquires_on_first_db_use( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """No checkout at construction; acquisition (with retry) happens once on the + first execute, and the statement itself runs exactly once.""" + monkeypatch.setattr(settings.DB, "TRACING", False) + acquired: list[str] = [] + executed: list[Any] = [] + + async def fake_acquire(_session: Any, context: str) -> None: + acquired.append(context) + + async def fake_super_execute(_self: Any, *args: Any, **_kw: Any) -> str: + executed.append(args[0] if args else None) + return "result" + + monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire) + monkeypatch.setattr(AsyncSession, "execute", fake_super_execute) + + session = db_module.SessionLocal() + assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage] + assert acquired == [] # constructing the session does NOT check out + + result = await session.execute("SELECT 1") + assert result == "result" + assert acquired == ["unknown"] # acquired exactly once, on first use + assert session._honcho_acquired is True # pyright: ignore[reportPrivateUsage] + assert executed == ["SELECT 1"] # statement ran once (not retried) + + await session.execute("SELECT 2") + assert acquired == ["unknown"] # still once — no re-acquire on later use + + +@pytest.mark.asyncio +async def test_session_tracing_sets_application_name_on_acquire( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(settings.DB, "TRACING", True) + calls: list[tuple[Any, ...]] = [] + + async def fake_acquire(_session: Any, _context: str) -> None: + return None + + async def fake_super_execute(_self: Any, *args: Any, **_kw: Any) -> None: + calls.append(args) + + monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire) + monkeypatch.setattr(AsyncSession, "execute", fake_super_execute) + + token = db_module.request_context.set("request:trace-ctx") + try: + session = db_module.SessionLocal() + await session.execute("SELECT 1") + finally: + db_module.request_context.reset(token) + + # First super().execute is the set_config, then the real statement. + assert any("set_config" in str(c[0]) for c in calls) + set_config_call = next(c for c in calls if "set_config" in str(c[0])) + assert set_config_call[1] == {"name": "request:trace-ctx"} + + +@pytest.mark.asyncio +async def test_session_lifecycle_methods_reset_acquired_flag( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def fake_acquire(_session: Any, _context: str) -> None: + return None + + async def noop(_self: Any) -> None: + return None + + monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire) + for method in ("commit", "rollback", "close", "reset"): + monkeypatch.setattr(AsyncSession, method, noop) + + session = db_module.SessionLocal() + await session.commit() # ensures acquired, commits, then resets + assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage] + + # rollback/close/reset must each clear the flag so a reused session + # re-acquires (and re-wraps retry) on its next DB use. + for method in ("rollback", "close", "reset"): + session._honcho_acquired = True # pyright: ignore[reportPrivateUsage] + await getattr(session, method)() + assert session._honcho_acquired is False # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.asyncio +async def test_session_get_and_delete_also_acquire( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The lazy-acquire guarantee covers get/delete, not just execute.""" + monkeypatch.setattr(settings.DB, "TRACING", False) + acquired: list[str] = [] + + async def fake_acquire(_session: Any, context: str) -> None: + acquired.append(context) + + async def fake_get(_self: Any, *_a: Any, **_k: Any) -> str: + return "row" + + async def fake_delete(_self: Any, *_a: Any, **_k: Any) -> None: + return None + + monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire) + monkeypatch.setattr(AsyncSession, "get", fake_get) + monkeypatch.setattr(AsyncSession, "delete", fake_delete) + + session = db_module.SessionLocal() + await session.get(object, 1) + await session.delete(object()) + assert acquired == ["unknown"] # acquired once on the first DB-touching call + + +@pytest.mark.asyncio +async def test_get_db_does_not_acquire_at_dependency_entry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from src.dependencies import get_db as real_get_db + + acquired: list[str] = [] + + async def fake_acquire(_session: Any, context: str) -> None: + acquired.append(context) + + monkeypatch.setattr(db_module, "acquire_connection_with_retry", fake_acquire) + + dep_gen = real_get_db() + await anext(dep_gen) + assert acquired == [] # yielding the session must not check out + await dep_gen.aclose() + assert acquired == [] # finally rollback/close must not check out either + + +def test_polling_backoff_sequence_and_reset(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", True) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 30.0) + + from src.deriver.queue_manager import QueueManager + + qm = QueueManager() + seq = [qm._advance_poll_interval() for _ in range(8)] # pyright: ignore[reportPrivateUsage] + assert seq == [1.0, 2.0, 4.0, 8.0, 16.0, 30.0, 30.0, 30.0] # caps at max + + qm._reset_poll_interval() # pyright: ignore[reportPrivateUsage] + assert qm._advance_poll_interval() == 1.0 # pyright: ignore[reportPrivateUsage] + + +def test_polling_backoff_disabled_stays_constant( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", False) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0) + + from src.deriver.queue_manager import QueueManager + + qm = QueueManager() + assert [qm._advance_poll_interval() for _ in range(3)] == [1.0, 1.0, 1.0] # pyright: ignore[reportPrivateUsage] + + +def test_pool_timeout_must_be_under_retry_budget() -> None: + """A pooled QueuePool checkout that can block past the retry budget is a + contradiction (no retry ever happens) and must fail config validation.""" + from src.config import DBSettings + + # Contradiction: pooled + retry on, but POOL_TIMEOUT >= budget. + with pytest.raises(ValueError, match="DB_POOL_TIMEOUT"): + DBSettings( + POOL_CLASS="default", + POOL_TIMEOUT=30, + CONNECTION_RETRY_ENABLED=True, + CONNECTION_RETRY_MAX_DELAY_SECONDS=10.0, + ) + + # NullPool has no queue wait, so POOL_TIMEOUT is irrelevant -> allowed. + DBSettings( + POOL_CLASS="null", + POOL_TIMEOUT=30, + CONNECTION_RETRY_ENABLED=True, + CONNECTION_RETRY_MAX_DELAY_SECONDS=10.0, + ) + + +@pytest.mark.asyncio +async def test_polling_loop_idle_sleeps_once_per_cycle( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Drive the real loop on an empty queue: exactly one (growing, capped) + sleep per empty poll — no double-sleep from the queue_empty_flag branch.""" + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_ENABLED", True) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_INTERVAL_SECONDS", 1.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_BACKOFF_MULTIPLIER", 2.0) + monkeypatch.setattr(settings.DERIVER, "POLLING_SLEEP_MAX_INTERVAL_SECONDS", 8.0) + + import asyncio + + from src.deriver import queue_manager as qm_mod + + qm = qm_mod.QueueManager() + sleeps: list[float] = [] + polls = {"n": 0} + + async def fake_cleanup() -> None: + return None + + async def fake_claim() -> dict[str, str]: + polls["n"] += 1 + if polls["n"] >= 5: + qm.shutdown_event.set() # stop after 5 empty polls + return {} + + async def fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + + monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup) + monkeypatch.setattr(qm, "get_and_claim_work_units", fake_claim) + # queue_manager calls asyncio.sleep on the stdlib module; patch it there. + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + await qm.polling_loop() + + # One sleep per empty poll, growing 1->2->4->8 then capped at 8 (not doubled). + assert sleeps == [1.0, 2.0, 4.0, 8.0, 8.0] + + +def test_inflight_gauge_no_drift(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(settings.METRICS, "NAMESPACE", "test") + child: Any = db_queries_in_flight_gauge.labels(instance_type="api") + tracker = DBQueryInflightTracker(child) + key = DBQueryInflightTracker.INFLIGHT_KEY + + def value() -> float: + return float(child._value.get()) + + start = value() + conn = SimpleNamespace(info={}) + + # Normal execute: before -> after returns to baseline. + tracker.on_before(conn) + assert value() == start + 1 + assert conn.info[key] is True + tracker.on_after(conn) + assert value() == start + assert key not in conn.info + + # Errored execute: before -> on_error decrements (after never fires). + tracker.on_before(conn) + assert value() == start + 1 + tracker.on_error(SimpleNamespace(connection=conn)) + assert value() == start + + # on_error without a matching before (e.g. connect error) must not push the + # gauge negative. + tracker.on_error(SimpleNamespace(connection=SimpleNamespace(info={}))) + assert value() == start diff --git a/tests/test_dependencies.py b/tests/test_dependencies.py index c90643ec..2996cddc 100644 --- a/tests/test_dependencies.py +++ b/tests/test_dependencies.py @@ -16,6 +16,12 @@ class FakeSession: self.execute_calls: list[tuple[Any, ...]] = [] self.rollback_calls: int = 0 self.close_calls: int = 0 + self.connection_calls: int = 0 + + async def connection(self) -> None: + # Tracks checkout attempts so tests can assert get_db/tracked_db stay + # lazy (they should never force a checkout themselves). + self.connection_calls += 1 async def execute(self, statement: Any, params: Any = None) -> None: self.execute_calls.append((statement, params)) @@ -31,26 +37,24 @@ class FakeSession: @pytest.mark.asyncio -async def test_get_db_sets_application_name_when_tracing_enabled( +async def test_get_db_yields_lazily_without_checkout_or_tracing( monkeypatch: pytest.MonkeyPatch, ) -> None: + # get_db must NOT touch the connection or run set_config itself — those now + # happen lazily inside HonchoAsyncSession on first DB use, so a handler doing + # non-DB work before its first query never pins a connection. fake_db = FakeSession() monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) - monkeypatch.setattr(settings.DB, "TRACING", True) + monkeypatch.setattr(settings.DB, "TRACING", True) # still no set_config here - context_token = request_context.set("request:test-ctx") dep_gen = real_get_db() - try: db = await anext(dep_gen) assert db is fake_db - assert len(fake_db.execute_calls) == 1 - stmt, params = fake_db.execute_calls[0] - assert "set_config" in str(stmt) - assert params == {"name": "request:test-ctx"} + assert fake_db.connection_calls == 0 # no eager checkout + assert fake_db.execute_calls == [] # no set_config in get_db finally: await dep_gen.aclose() - request_context.reset(context_token) assert fake_db.rollback_calls == 1 # unconditional rollback in finally assert fake_db.close_calls == 1 @@ -80,7 +84,6 @@ async def test_tracked_db_creates_and_resets_task_context( ) -> None: fake_db = FakeSession() monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) - monkeypatch.setattr(settings.DB, "TRACING", True) monkeypatch.setattr( uuid, "uuid4", @@ -90,15 +93,12 @@ async def test_tracked_db_creates_and_resets_task_context( clear_token = request_context.set(None) try: async with real_tracked_db("cleanup_job"): + # tracked_db sets the task context so the lazy session can read it. assert request_context.get() == "task:cleanup_job:12345678" finally: request_context.reset(clear_token) assert request_context.get() is None - assert len(fake_db.execute_calls) == 1 - stmt, params = fake_db.execute_calls[0] - assert "set_config" in str(stmt) - assert params == {"name": "task:cleanup_job:12345678"} assert fake_db.rollback_calls == 1 # unconditional rollback in finally assert fake_db.close_calls == 1 @@ -109,7 +109,6 @@ async def test_tracked_db_preserves_existing_request_context( ) -> None: fake_db = FakeSession() monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db) - monkeypatch.setattr(settings.DB, "TRACING", True) context_token = request_context.set("request:existing") try: @@ -118,10 +117,6 @@ async def test_tracked_db_preserves_existing_request_context( finally: request_context.reset(context_token) - assert len(fake_db.execute_calls) == 1 - stmt, params = fake_db.execute_calls[0] - assert "set_config" in str(stmt) - assert params == {"name": "request:existing"} assert fake_db.rollback_calls == 1 # unconditional rollback in finally assert fake_db.close_calls == 1