From a5c43b8dc18b3c05e0ce73c15c9844d790805005 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Sun, 31 May 2026 23:59:50 -0400 Subject: [PATCH] feat(db): add acquisition counter and in-flight query gauge MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- src/db.py | 89 +++++++++++++++- src/deriver/__main__.py | 3 +- src/main.py | 3 +- src/telemetry/prometheus/metrics.py | 29 +++++ tests/test_db_resilience.py | 159 ++++++++++++++++++++++++++++ 5 files changed, 280 insertions(+), 3 deletions(-) create mode 100644 tests/test_db_resilience.py diff --git a/src/db.py b/src/db.py index f1c65887..6c7d5866 100644 --- a/src/db.py +++ b/src/db.py @@ -1,7 +1,9 @@ import contextvars +import logging +from typing import Any import sentry_sdk -from sqlalchemy import MetaData, text +from sqlalchemy import MetaData, event, text from sqlalchemy.exc import OperationalError from sqlalchemy.exc import TimeoutError as SQLAlchemyTimeoutError from sqlalchemy.ext.asyncio import ( @@ -19,6 +21,12 @@ from tenacity import ( ) 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} @@ -63,6 +71,18 @@ SessionLocal = async_sessionmaker( # 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. @@ -106,6 +126,7 @@ async def acquire_connection_with_retry(db: AsyncSession, context: str) -> None: if not settings.DB.CONNECTION_RETRY_ENABLED: await db.connection() return + attempts = 0 try: async for attempt in AsyncRetrying( wait=wait_exponential_jitter( @@ -117,12 +138,78 @@ async def acquire_connection_with_retry(db: AsyncSession, context: str) -> None: reraise=True, ): with attempt: + attempts += 1 await db.connection() 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 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 + + +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. Safe to call once per process. + """ + global _db_instance_type, _inflight_tracker + _db_instance_type = instance_type + if not settings.METRICS.ENABLED: + 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) # Define your naming convention diff --git a/src/deriver/__main__.py b/src/deriver/__main__.py index 0807fc11..c56ed6a0 100644 --- a/src/deriver/__main__.py +++ b/src/deriver/__main__.py @@ -6,7 +6,7 @@ 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, @@ -24,6 +24,7 @@ def start_metrics_server() -> None: 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/main.py b/src/main.py index bb604b32..6edd9d4c 100644 --- a/src/main.py +++ b/src/main.py @@ -19,7 +19,7 @@ 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, @@ -134,6 +134,7 @@ async def lifespan(_: FastAPI): # 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 diff --git a/src/telemetry/prometheus/metrics.py b/src/telemetry/prometheus/metrics.py index 20ee7ac1..d883816b 100644 --- a/src/telemetry/prometheus/metrics.py +++ b/src/telemetry/prometheus/metrics.py @@ -127,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: @@ -277,6 +294,18 @@ 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() diff --git a/tests/test_db_resilience.py b/tests/test_db_resilience.py new file mode 100644 index 00000000..0177373e --- /dev/null +++ b/tests/test_db_resilience.py @@ -0,0 +1,159 @@ +"""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 + +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 + + async def connection(self) -> None: + self.calls += 1 + if self.always_fail or self.calls <= self.fail_times: + raise _make_operational_error() + + +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 + 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 + + +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_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