diff --git a/docs/deployment.md b/docs/deployment.md index 60bdc7e257..2992ef36c7 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -19,6 +19,29 @@ A durable standalone deployment should: The default SQLite database is intended for evaluation and single-node local use. +### Watch the connection pool + +On PostgreSQL the gateway serves requests from a fixed pool of database +connections, and running out of them makes every request fail at once. `/metrics` +reports the pool directly, labeled by `pool` (`request` for request traffic, +`log` for the usage-log writer): + +| Metric | Meaning | +| --- | --- | +| `gateway_db_pool_connections_checked_out` | Connections in use right now | +| `gateway_db_pool_connections_idle` | Connections available to hand out | +| `gateway_db_pool_overflow_connections` | Connections open beyond the pool's base size | +| `gateway_db_pool_capacity` | Ceiling on connections the pool hands out at once | + +The `request` pool is sized by `db_pool_size` and may open `db_max_overflow` +connections beyond it, so its capacity is the two added together. The `log` +pool is sized by `db_log_pool_size` and has no overflow, so its capacity is +that value. + +Alert on checked-out connections approaching capacity for a sustained period. +These metrics do not appear on SQLite, which opens a connection per use and +keeps no pool. + ## Docker Compose The repository Compose stack runs Otari and PostgreSQL: diff --git a/src/gateway/core/database.py b/src/gateway/core/database.py index 1c9f282943..2767620b65 100644 --- a/src/gateway/core/database.py +++ b/src/gateway/core/database.py @@ -4,10 +4,11 @@ import asyncio import contextlib -from collections.abc import AsyncGenerator, AsyncIterator +from collections.abc import AsyncGenerator, AsyncIterator, Iterator from contextlib import asynccontextmanager +from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Any, Protocol, runtime_checkable from alembic import command from alembic.config import Config @@ -19,6 +20,7 @@ from gateway.core.config import GatewayConfig from gateway.log_config import logger +from gateway.metrics import REGISTRY, Collector, GaugeMetricFamily _engine: AsyncEngine | None = None _SessionLocal: async_sessionmaker[AsyncSession] | None = None @@ -196,6 +198,122 @@ async def release_session(session: AsyncSession | None) -> bool: return True +REQUEST_POOL = "request" +LOG_POOL = "log" + +# The configured ``max_overflow`` of each live pool, by pool name. ``QueuePool`` +# has no public accessor for it, and the gateway is the side that chose the +# value, so it is recorded here at :func:`init_db` time rather than read back +# off the pool. +_pool_max_overflow: dict[str, int] = {} + + +@runtime_checkable +class _CountingPool(Protocol): + """The counters a ``QueuePool`` keeps and a ``NullPool`` does not.""" + + def checkedout(self) -> int: ... + def checkedin(self) -> int: ... + def overflow(self) -> int: ... + def size(self) -> int: ... + + +@dataclass(frozen=True, slots=True) +class PoolStats: + """A point-in-time reading of one engine's connection pool.""" + + checked_out: int + checked_in: int + overflow: int + size: int + max_overflow: int + + @property + def capacity(self) -> int: + """The most connections this pool will ever hand out at once.""" + return self.size + self.max_overflow + + +def _engine_pool_stats(engine: AsyncEngine | None, max_overflow: int) -> PoolStats | None: + """Read *engine*'s pool. + + Returns ``None`` when there is no pool to read: no engine yet, or a pool + with no counters (``NullPool``, which SQLite uses). + """ + if engine is None: + return None + pool = engine.pool + if not isinstance(pool, _CountingPool): + return None + # ``QueuePool`` reports a negative overflow until it has opened its base + # connections. + overflow = max(pool.overflow(), 0) + return PoolStats( + checked_out=pool.checkedout(), + checked_in=pool.checkedin(), + overflow=overflow, + size=pool.size(), + max_overflow=max_overflow, + ) + + +def pool_stats() -> dict[str, PoolStats]: + """Pool stats for every active engine, keyed by pool name. + + A pool with nothing to report is omitted, so the result is empty before + :func:`init_db` runs and on SQLite. + """ + readings = { + REQUEST_POOL: _engine_pool_stats(_engine, _pool_max_overflow.get(REQUEST_POOL, 0)), + LOG_POOL: _engine_pool_stats(_log_engine, _pool_max_overflow.get(LOG_POOL, 0)), + } + return {name: stats for name, stats in readings.items() if stats is not None} + + +class _PoolCollector(Collector): + """Publish the connection-pool counters at scrape time. + + A collector rather than gauges refreshed on a timer: the counters are + already maintained by SQLAlchemy, so every scrape reads the live pool with + no way to lag it, and a pool that reports nothing emits no series at all + rather than leaving a stale value behind. + """ + + def collect(self) -> Iterator[GaugeMetricFamily]: + checked_out = GaugeMetricFamily( + "gateway_db_pool_connections_checked_out", + "Pooled database connections currently checked out", + labels=["pool"], + ) + idle = GaugeMetricFamily( + "gateway_db_pool_connections_idle", + "Pooled database connections checked in and available", + labels=["pool"], + ) + overflow = GaugeMetricFamily( + "gateway_db_pool_overflow_connections", + "Database connections open beyond the base pool size", + labels=["pool"], + ) + capacity = GaugeMetricFamily( + "gateway_db_pool_capacity", + "Most database connections the pool will hand out at once (size plus max overflow)", + labels=["pool"], + ) + for name, stats in pool_stats().items(): + checked_out.add_metric([name], stats.checked_out) + idle.add_metric([name], stats.checked_in) + overflow.add_metric([name], stats.overflow) + capacity.add_metric([name], stats.capacity) + yield checked_out + yield idle + yield overflow + yield capacity + + +REGISTRY.register(_PoolCollector()) + + def engine_kwargs( config: GatewayConfig, *, @@ -266,6 +384,7 @@ def init_db(config: GatewayConfig) -> None: ) _install_timeout_translation(_engine) _SessionLocal = async_sessionmaker(_engine, expire_on_commit=False) + _pool_max_overflow[REQUEST_POOL] = config.db_max_overflow if is_sqlite: _configure_sqlite_pragmas(_engine) @@ -298,6 +417,7 @@ def init_db(config: GatewayConfig) -> None: ) _install_timeout_translation(_log_engine) _LogSessionLocal = async_sessionmaker(_log_engine, expire_on_commit=False) + _pool_max_overflow[LOG_POOL] = 0 if config.auto_migrate: _run_migrations(database_url) @@ -352,6 +472,7 @@ def _take_engines() -> list[AsyncEngine]: _SessionLocal = None _log_engine = None _LogSessionLocal = None + _pool_max_overflow.clear() return engines @@ -387,12 +508,16 @@ async def _dispose_all() -> None: __all__ = [ "DATABASE_ERRORS", + "LOG_POOL", + "REQUEST_POOL", + "PoolStats", "create_log_session", "create_session", "dispose_db", "engine_kwargs", "get_db", "init_db", + "pool_stats", "release_session", "reset_db", "translate_timeout_error", diff --git a/src/gateway/metrics.py b/src/gateway/metrics.py index 882cda1882..839e29633d 100644 --- a/src/gateway/metrics.py +++ b/src/gateway/metrics.py @@ -16,6 +16,8 @@ ProcessCollector, generate_latest, ) +from prometheus_client.core import GaugeMetricFamily +from prometheus_client.registry import Collector from starlette.responses import Response from gateway.core.config import API_ROOT, API_VERSION @@ -24,7 +26,16 @@ from starlette.requests import Request from starlette.types import ASGIApp, Message, Receive, Scope, Send -__all__ = ["Counter", "Gauge", "Histogram", "MetricsMiddleware", "REGISTRY", "metrics_endpoint"] +__all__ = [ + "REGISTRY", + "Collector", + "Counter", + "Gauge", + "GaugeMetricFamily", + "Histogram", + "MetricsMiddleware", + "metrics_endpoint", +] REGISTRY = CollectorRegistry() @@ -54,7 +65,6 @@ registry=REGISTRY, ) - _PROMETHEUS_CONTENT_TYPE = "text/plain; version=0.0.4; charset=utf-8" _UNMATCHED_ENDPOINT = "unmatched" diff --git a/tests/unit/test_database_pool_stats.py b/tests/unit/test_database_pool_stats.py new file mode 100644 index 0000000000..d0733f60c4 --- /dev/null +++ b/tests/unit/test_database_pool_stats.py @@ -0,0 +1,142 @@ +"""Unit tests for the connection-pool stats helper and its Prometheus collector.""" + +from collections.abc import Iterator +from types import SimpleNamespace +from typing import Any + +import pytest + +from gateway.core import database +from gateway.core.database import LOG_POOL, REQUEST_POOL, PoolStats, pool_stats +from gateway.metrics import REGISTRY + + +class _QueuePool: + """Stand-in for ``AsyncAdaptedQueuePool`` exposing the counters read here.""" + + def __init__(self, checked_out: int, checked_in: int, overflow: int, size: int) -> None: + self._checked_out = checked_out + self._checked_in = checked_in + self._overflow = overflow + self._size = size + + def checkedout(self) -> int: + return self._checked_out + + def checkedin(self) -> int: + return self._checked_in + + def overflow(self) -> int: + return self._overflow + + def size(self) -> int: + return self._size + + +class _NullPool: + """Stand-in for ``NullPool``, which implements none of those counters.""" + + def status(self) -> str: + return "NullPool" + + +def _engine(pool: Any) -> Any: + return SimpleNamespace(pool=pool) + + +@pytest.fixture(autouse=True) +def _restore_engines() -> Iterator[None]: + engine = database._engine + log_engine = database._log_engine + overflow = dict(database._pool_max_overflow) + yield + database._engine = engine + database._log_engine = log_engine + database._pool_max_overflow.clear() + database._pool_max_overflow.update(overflow) + + +def test_pool_stats_reads_the_queue_pool_counters() -> None: + database._engine = _engine(_QueuePool(checked_out=7, checked_in=3, overflow=2, size=10)) + database._log_engine = None + database._pool_max_overflow[REQUEST_POOL] = 20 + + readings = pool_stats() + + assert readings[REQUEST_POOL] == PoolStats(checked_out=7, checked_in=3, overflow=2, size=10, max_overflow=20) + assert readings[REQUEST_POOL].capacity == 30 + + +def test_capacity_uses_the_configured_overflow_not_the_pool_object() -> None: + """The gateway chose ``max_overflow``, so capacity comes from the config, not a private field.""" + database._engine = _engine(_QueuePool(checked_out=1, checked_in=1, overflow=0, size=10)) + database._log_engine = None + database._pool_max_overflow[REQUEST_POOL] = 5 + + assert pool_stats()[REQUEST_POOL].capacity == 15 + + +def test_pool_stats_clamps_a_negative_overflow() -> None: + """A pool that has not created its base connections yet reports a negative overflow.""" + database._engine = _engine(_QueuePool(checked_out=1, checked_in=0, overflow=-9, size=10)) + database._log_engine = None + + assert pool_stats()[REQUEST_POOL].overflow == 0 + + +def test_pool_stats_omits_a_null_pool() -> None: + database._engine = _engine(_NullPool()) + database._log_engine = None + + assert pool_stats() == {} + + +def test_pool_stats_is_empty_when_the_database_is_not_initialized() -> None: + database._engine = None + database._log_engine = None + + assert pool_stats() == {} + + +def test_pool_stats_covers_both_engines() -> None: + database._engine = _engine(_QueuePool(checked_out=1, checked_in=2, overflow=0, size=10)) + database._log_engine = _engine(_QueuePool(checked_out=0, checked_in=1, overflow=0, size=2)) + database._pool_max_overflow[REQUEST_POOL] = 20 + database._pool_max_overflow[LOG_POOL] = 0 + + readings = pool_stats() + + assert set(readings) == {REQUEST_POOL, LOG_POOL} + assert readings[LOG_POOL].capacity == 2 + + +def test_pool_stats_omits_an_engine_without_a_pool() -> None: + database._engine = _engine(_QueuePool(checked_out=1, checked_in=2, overflow=0, size=10)) + database._log_engine = _engine(_NullPool()) + + assert set(pool_stats()) == {REQUEST_POOL} + + +def test_the_collector_publishes_the_live_pool_on_every_scrape() -> None: + database._engine = _engine(_QueuePool(checked_out=4, checked_in=6, overflow=1, size=10)) + database._log_engine = None + database._pool_max_overflow[REQUEST_POOL] = 20 + + labels = {"pool": REQUEST_POOL} + assert REGISTRY.get_sample_value("gateway_db_pool_connections_checked_out", labels) == 4.0 + assert REGISTRY.get_sample_value("gateway_db_pool_connections_idle", labels) == 6.0 + assert REGISTRY.get_sample_value("gateway_db_pool_overflow_connections", labels) == 1.0 + assert REGISTRY.get_sample_value("gateway_db_pool_capacity", labels) == 30.0 + + database._engine = _engine(_QueuePool(checked_out=9, checked_in=1, overflow=0, size=10)) + + assert REGISTRY.get_sample_value("gateway_db_pool_connections_checked_out", labels) == 9.0 + + +def test_the_collector_publishes_no_series_without_a_pool() -> None: + """SQLite runs on NullPool, and an uninitialized engine has no pool at all.""" + database._engine = None + database._log_engine = None + + assert REGISTRY.get_sample_value("gateway_db_pool_capacity", {"pool": REQUEST_POOL}) is None + assert REGISTRY.get_sample_value("gateway_db_pool_capacity", {"pool": LOG_POOL}) is None diff --git a/tests/unit/test_gateway_metrics.py b/tests/unit/test_gateway_metrics.py index d6529c8f2b..9d8951498d 100644 --- a/tests/unit/test_gateway_metrics.py +++ b/tests/unit/test_gateway_metrics.py @@ -30,6 +30,10 @@ def _sample(name: str, labels: dict[str, str] | None = None) -> float: ("gateway_active_requests", "gauge", ()), ("gateway_auth_failures", "counter", ("reason",)), ("gateway_budget_exceeded", "counter", ()), + ("gateway_db_pool_capacity", "gauge", ("pool",)), + ("gateway_db_pool_connections_checked_out", "gauge", ("pool",)), + ("gateway_db_pool_connections_idle", "gauge", ("pool",)), + ("gateway_db_pool_overflow_connections", "gauge", ("pool",)), ("gateway_inline_cost_settlements", "counter", ("outcome",)), ("gateway_rate_limit_hits", "counter", ()), ("gateway_request_cost_dollars", "histogram", ("provider", "model")), @@ -47,7 +51,8 @@ def test_scrape_exposes_the_pinned_families() -> None: """The set of gateway metric families, with their types and label names, is fixed. A labeled family with no series yet shows only its HELP and TYPE lines in a - scrape, so the label names are read off the collectors rather than the text. + scrape, so the label names are read off the collector, or off the family it + yields where the collector is a custom one that keeps no label names. """ import gateway.main # noqa: F401 # imports every module that registers a metric @@ -56,7 +61,8 @@ def test_scrape_exposes_the_pinned_families() -> None: describe = getattr(collector, "describe", collector.collect) for metric in describe(): if metric.name.startswith("gateway_"): - families.add((metric.name, metric.type, tuple(getattr(collector, "_labelnames", ())))) + labelnames = getattr(collector, "_labelnames", ()) or getattr(metric, "_labelnames", ()) + families.add((metric.name, metric.type, tuple(labelnames))) assert families == _EXPOSED_FAMILIES