feat(observability): export database connection pool stats on /metrics

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:

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",

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"

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