diff --git a/src/config.py b/src/config.py index 5ae93fa4..c6bfe65e 100644 --- a/src/config.py +++ b/src/config.py @@ -622,6 +622,21 @@ class DBSettings(HonchoSettings): SQL_DEBUG: bool = False TRACING: bool = False + # Bounded exponential-backoff retry around connection acquisition. Guards + # against transient transaction-pooler saturation (e.g. Supavisor rejecting + # with "too many clients") by retrying the pool checkout instead of failing + # the request immediately. Applied to both the API and background paths. + 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 + class AuthSettings(HonchoSettings): model_config = SettingsConfigDict(env_prefix="AUTH_", extra="ignore") # pyright: ignore @@ -737,6 +752,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..f1c65887 100644 --- a/src/db.py +++ b/src/db.py @@ -1,9 +1,22 @@ import contextvars +import sentry_sdk from sqlalchemy import MetaData, text -from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine +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 @@ -45,6 +58,73 @@ SessionLocal = async_sessionmaker( bind=engine, ) +# 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) + + +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: + return { + "checked_out": pool.checkedout(), + "checked_in": pool.checkedin(), + "size": pool.size(), + "overflow": 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. + + Retrying the same session is safe: no connection is bound until checkout + succeeds, so each attempt re-attempts the checkout cleanly. + """ + with sentry_sdk.start_span(op="db.pool.acquire", name=context): + if not settings.DB.CONNECTION_RETRY_ENABLED: + await db.connection() + return + 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: + await db.connection() + except RETRYABLE_DB_CONNECTION_ERRORS as e: + if settings.SENTRY.ENABLED: + sentry_sdk.set_context("db_pool", get_pool_stats()) + sentry_sdk.capture_exception(e) + raise + + # 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..710715a7 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -6,7 +6,7 @@ from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession from src.config import settings -from src.db import SessionLocal, request_context +from src.db import SessionLocal, acquire_connection_with_retry, request_context async def get_db(): @@ -16,6 +16,7 @@ async def get_db(): db: AsyncSession = SessionLocal() try: + await acquire_connection_with_retry(db, context) if settings.DB.TRACING: await db.execute( text("SELECT set_config('application_name', :name, false)"), @@ -50,6 +51,7 @@ async def tracked_db(operation_name: str | None = None): db = SessionLocal() try: + await acquire_connection_with_retry(db, context or f"task:{operation_name}") if settings.DB.TRACING: await db.execute( text("SELECT set_config('application_name', :name, false)"), diff --git a/src/deriver/__main__.py b/src/deriver/__main__.py index 20ffbb66..0807fc11 100644 --- a/src/deriver/__main__.py +++ b/src/deriver/__main__.py @@ -8,7 +8,11 @@ from prometheus_client import start_http_server from src.config import settings from src.db import engine 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,8 @@ 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") 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..2fd141c9 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,6 +387,21 @@ 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") @@ -385,11 +409,15 @@ class QueueManager: 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) + # Sleep the already-grown interval; the backoff is advanced + # once per empty-detection below, not here. + await asyncio.sleep(self._current_poll_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 +427,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 +443,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..bb604b32 100644 --- a/src/main.py +++ b/src/main.py @@ -13,6 +13,7 @@ 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 @@ -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,9 @@ 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") + # 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..20ee7ac1 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 @@ -279,6 +281,46 @@ class PrometheusMetrics: 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]: + # Lazy import to avoid an import cycle at module load (db imports config, + # telemetry is imported widely). Reads the async engine.pool directly. + from src.db import get_pool_stats + + namespace = settings.METRICS.NAMESPACE or "" + gauge = GaugeMetricFamily( + "db_pool_connections", + "DB connections held by this instance, by pool state", + labels=["namespace", "instance_type", "state"], + ) + for state, value in get_pool_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_dependencies.py b/tests/test_dependencies.py index c90643ec..511e2a2e 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: + # acquire_connection_with_retry forces the (otherwise lazy) pool + # checkout via this call before any query runs. + self.connection_calls += 1 async def execute(self, statement: Any, params: Any = None) -> None: self.execute_calls.append((statement, params)) @@ -44,6 +50,7 @@ async def test_get_db_sets_application_name_when_tracing_enabled( try: db = await anext(dep_gen) assert db is fake_db + assert fake_db.connection_calls == 1 # eager pool checkout with retry assert len(fake_db.execute_calls) == 1 stmt, params = fake_db.execute_calls[0] assert "set_config" in str(stmt)