From 99472ba3cf2d72442bf07f90f22aaf44979a45ec Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Sun, 18 Oct 2026 19:28:42 +0000 Subject: [PATCH] Count native pool statement-cache misses and retired connections Each connection of the native Postgres pool caches its prepared statements. A miss costs a round trip to prepare, and the cache is emptied whole when it reaches 512 statements. Nothing showed how often either happens, or how often `txn_limit` replaces a connection and loses its cache. The pool now counts cache hits, prepares, cache clears, and connections opened and retired. Python reports them per pool as `synapse_rust_db_pool_statement_cache_hits_total`, `..._statement_prepares_total`, `..._statement_cache_clears_total`, `..._connections_opened_total` and `..._connections_retired_total`, carried over when the native pool is closed. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01PVLsNh6akt4ifEizY4D1ta --- changelog.d/db-roundtrips.misc | 1 + rust/src/database/postgres/pool.rs | 57 +++++++++++- scripts-dev/mypy_synapse_plugin.py | 1 + synapse/storage/rust_pool.py | 102 ++++++++++++++++++++- synapse/synapse_rust/database/postgres.pyi | 10 ++ tests/storage/test_rust_pool.py | 41 +++++++++ 6 files changed, 209 insertions(+), 3 deletions(-) create mode 100644 changelog.d/db-roundtrips.misc diff --git a/changelog.d/db-roundtrips.misc b/changelog.d/db-roundtrips.misc new file mode 100644 index 0000000000..c3ad597ef6 --- /dev/null +++ b/changelog.d/db-roundtrips.misc @@ -0,0 +1 @@ +Make fewer database round trips on the native Postgres driver. diff --git a/rust/src/database/postgres/pool.rs b/rust/src/database/postgres/pool.rs index 1691ef667a..88500762df 100644 --- a/rust/src/database/postgres/pool.rs +++ b/rust/src/database/postgres/pool.rs @@ -16,7 +16,8 @@ //! could exhaust the server's connection limit between them. use std::collections::HashMap; -use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use std::time::Duration; use deadpool::managed::{ @@ -93,6 +94,42 @@ impl SessionConfig { } } +/// Counts of what a pool's connections did, for metrics. Shared by the +/// [`ConnectionManager`] and every [`PooledClient`] it creates. +#[derive(Debug, Default)] +pub struct PoolStats { + /// Statements found in a connection's prepared-statement cache. + statement_cache_hits: AtomicU64, + /// Statements prepared, each costing a round trip: cache misses and + /// refreshes. + statement_prepares: AtomicU64, + /// Times a connection's cache was emptied for reaching + /// [`STATEMENT_CACHE_CAP`]. + statement_cache_clears: AtomicU64, + /// Connections opened. + connections_opened: AtomicU64, + /// Connections closed for reaching their use limit (`txn_limit`). + connections_retired: AtomicU64, +} + +impl PoolStats { + fn bump(counter: &AtomicU64) { + counter.fetch_add(1, Ordering::Relaxed); + } + + /// `(statement_cache_hits, statement_prepares, statement_cache_clears, + /// connections_opened, connections_retired)`. + pub fn snapshot(&self) -> (u64, u64, u64, u64, u64) { + ( + self.statement_cache_hits.load(Ordering::Relaxed), + self.statement_prepares.load(Ordering::Relaxed), + self.statement_cache_clears.load(Ordering::Relaxed), + self.connections_opened.load(Ordering::Relaxed), + self.connections_retired.load(Ordering::Relaxed), + ) + } +} + /// Creates and recycles [`tokio_postgres`] connections for a [`ConnectionPool`]. pub struct ConnectionManager { /// The resolved connection config, parsed once from the DSN (with libpq's @@ -111,6 +148,8 @@ pub struct ConnectionManager { /// If set, a connection is closed rather than reused once it has been /// checked out this many times. See [`PoolConfig::max_uses`]. max_uses: Option, + /// What the pool's connections did, for metrics. + stats: Arc, } impl ConnectionManager { @@ -131,6 +170,7 @@ impl ConnectionManager { session, tls: connector, max_uses: None, + stats: Arc::default(), }) } } @@ -162,6 +202,8 @@ pub struct PooledClient { /// Prepared statements keyed by their SQL. Guarded by a plain [`Mutex`]: /// the critical sections are lookups/inserts (never awaits). statements: Mutex>, + /// The pool's counters. + stats: Arc, } /// The cache is cleared once it holds this many statements. Synapse's hot set @@ -184,13 +226,16 @@ impl PooledClient { if refresh { self.statements.lock().expect("not poisoned").remove(sql); } else if let Some(statement) = self.statements.lock().expect("not poisoned").get(sql) { + PoolStats::bump(&self.stats.statement_cache_hits); return Ok(statement.clone()); } + PoolStats::bump(&self.stats.statement_prepares); let statement = self.client.prepare(sql).await?; let mut statements = self.statements.lock().expect("not poisoned"); if statements.len() >= STATEMENT_CACHE_CAP { + PoolStats::bump(&self.stats.statement_cache_clears); statements.clear(); } statements.insert(sql.to_owned(), statement.clone()); @@ -242,9 +287,11 @@ impl Manager for ConnectionManager { client.batch_execute(&setup).await?; } + PoolStats::bump(&self.stats.connections_opened); Ok(PooledClient { client, statements: Mutex::new(HashMap::new()), + stats: Arc::clone(&self.stats), }) } @@ -266,6 +313,7 @@ impl Manager for ConnectionManager { // first, so this runs when the `max_uses`th checkout would reuse it. if let Some(max_uses) = self.max_uses { if metrics.recycle_count + 1 >= max_uses { + PoolStats::bump(&self.stats.connections_retired); return Err(RecycleError::message("connection reached its use limit")); } } @@ -434,6 +482,13 @@ impl PyConnectionPool { self.pool.close(); } + /// What the pool's connections have done since it was built, for + /// metrics: `(statement_cache_hits, statement_prepares, + /// statement_cache_clears, connections_opened, connections_retired)`. + fn stats(&self) -> (u64, u64, u64, u64, u64) { + self.pool.manager().stats.snapshot() + } + /// The pool's occupancy, for metrics: `(max_size, size, available, /// waiting)`. `size` is the number of open connections, `available` how /// many of those are idle in the pool, and `waiting` how many checkouts diff --git a/scripts-dev/mypy_synapse_plugin.py b/scripts-dev/mypy_synapse_plugin.py index a57f07af56..2c52bcb2b3 100644 --- a/scripts-dev/mypy_synapse_plugin.py +++ b/scripts-dev/mypy_synapse_plugin.py @@ -151,6 +151,7 @@ prometheus_metric_fullname_to_label_arg_map: Mapping[str, ArgLocation | None] = "synapse.metrics.jemalloc.JemallocCollector": None, "synapse.util.metrics.DynamicCollectorRegistry": None, "synapse.metrics.background_process_metrics._Collector": None, + "synapse.storage.rust_pool._RustPoolStatsCollector": None, # # `Metric` subclasses: "prometheus_client.metrics_core.Metric": None, diff --git a/synapse/storage/rust_pool.py b/synapse/storage/rust_pool.py index a6ba3287c1..82b13bdbd8 100644 --- a/synapse/storage/rust_pool.py +++ b/synapse/storage/rust_pool.py @@ -32,15 +32,17 @@ configured with ``use_rust_driver``. import logging import threading import weakref -from typing import TYPE_CHECKING, Any, Callable, TypeVar +from typing import TYPE_CHECKING, Any, Callable, Iterable, TypeVar from prometheus_client import Gauge, Histogram +from prometheus_client.core import REGISTRY, CounterMetricFamily, Metric from typing_extensions import Concatenate, ParamSpec from twisted.internet import threads from twisted.python.threadpool import ThreadPool from synapse.metrics import SERVER_NAME_LABEL +from synapse.metrics._types import Collector from synapse.storage.rust_dbapi import Connection as DBAPI2Connection from synapse.synapse_rust.database import postgres @@ -90,6 +92,72 @@ rust_pool_checkout_wait_seconds = Histogram( ) +# The counters of `ConnectionPool.stats()`, in order. +_POOL_COUNTERS = ( + ( + "synapse_rust_db_pool_statement_cache_hits_total", + "Statements found in a native Rust pool connection's prepared-statement cache", + ), + ( + "synapse_rust_db_pool_statement_prepares_total", + "Statements a native Rust pool connection prepared, each a round trip to the " + "database", + ), + ( + "synapse_rust_db_pool_statement_cache_clears_total", + "Times a native Rust pool connection's prepared-statement cache was emptied " + "for being full", + ), + ( + "synapse_rust_db_pool_connections_opened_total", + "Connections the native Rust database pool opened", + ), + ( + "synapse_rust_db_pool_connections_retired_total", + "Connections the native Rust database pool closed for reaching `txn_limit`", + ), +) + + +class _RustPoolStatsCollector(Collector): + """Reports the counters of every live `RustConnectionPool`, summed over + pools with the same labels. Holds the pools weakly.""" + + def __init__(self) -> None: + self._pools: "weakref.WeakSet[RustConnectionPool]" = weakref.WeakSet() + # Pools are added on the reactor thread and read on the metrics + # thread. + self._lock = threading.Lock() + + def add(self, pool: "RustConnectionPool") -> None: + with self._lock: + self._pools.add(pool) + + def collect(self) -> "Iterable[Metric]": + with self._lock: + pools = list(self._pools) + totals: dict[tuple[str, str], list[int]] = {} + for pool in pools: + labels = tuple(pool._metric_labels.values()) + stats = pool.stats() + total = totals.setdefault(labels, [0] * len(stats)) # type: ignore[arg-type] + for index, value in enumerate(stats): + total[index] += value + + for index, (name, documentation) in enumerate(_POOL_COUNTERS): + family = CounterMetricFamily( + name, documentation, labels=["name", SERVER_NAME_LABEL] + ) + for labels, total in totals.items(): + family.add_metric(list(labels), total[index]) + yield family + + +# Its metrics carry `SERVER_NAME_LABEL`, which the collector adds per pool. +_pool_stats_collector = _RustPoolStatsCollector() # type: ignore[missing-server-name-label] +REGISTRY.register(_pool_stats_collector) + + class RustConnectionPool: """Runs blocking database functions against pooled Rust connections. @@ -185,6 +253,13 @@ class RustConnectionPool: self._metric_labels = {"name": name, SERVER_NAME_LABEL: server_name} self._register_gauges() + # What the native pools this replaced did, as `stats()` reports it, so + # the counters keep counting up across a `close()`/`start()` cycle. + self._closed_pool_stats = (0, 0, 0, 0, 0) + # Held while `close()` swaps the native pool for `_closed_pool_stats`, + # so a metrics scrape never counts the old pool twice. + self._stats_lock = threading.Lock() + _pool_stats_collector.add(self) def _register_gauges(self) -> None: """Point the occupancy gauges at this pool. @@ -205,6 +280,27 @@ class RustConnectionPool: for gauge, index in _POOL_GAUGES: gauge.labels(**self._metric_labels).set_function(status_field(index)) + def stats(self) -> tuple[int, int, int, int, int]: + """What this pool's connections have done, counted over every native + pool it has opened: ``(statement_cache_hits, statement_prepares, + statement_cache_clears, connections_opened, connections_retired)``.""" + with self._stats_lock: + return self._stats_unlocked() + + def _stats_unlocked(self) -> tuple[int, int, int, int, int]: + """`stats()`, for a caller holding `_stats_lock`.""" + closed = self._closed_pool_stats + if self._pool is None: + return closed + current = self._pool.stats() + return ( + closed[0] + current[0], + closed[1] + current[1], + closed[2] + current[2], + closed[3] + current[3], + closed[4] + current[4], + ) + def status(self) -> tuple[int, int, int, int]: """The native pool's occupancy: ``(max_size, size, available, waiting)``, all zero once the pool is closed.""" @@ -276,7 +372,9 @@ class RustConnectionPool: self._connections.clear() if self._pool is not None: self._pool.close() - self._pool = None + with self._stats_lock: + self._closed_pool_stats = self._stats_unlocked() + self._pool = None for gauge, _ in _POOL_GAUGES: # Stop reporting a closed pool; `start()` registers afresh. try: diff --git a/synapse/synapse_rust/database/postgres.pyi b/synapse/synapse_rust/database/postgres.pyi index 4a4af3188b..ac0b26dced 100644 --- a/synapse/synapse_rust/database/postgres.pyi +++ b/synapse/synapse_rust/database/postgres.pyi @@ -66,6 +66,16 @@ class ConnectionPool: """ def close(self) -> None: ... + def stats(self) -> tuple[int, int, int, int, int]: + """What the pool's connections have done since it was built: + ``(statement_cache_hits, statement_prepares, statement_cache_clears, + connections_opened, connections_retired)``. + + A prepare is a round trip to the database: a statement missing from the + connection's cache. A clear empties a full cache. A retired connection + reached its use limit (``max_uses``). + """ + def status(self) -> tuple[int, int, int, int]: """The pool's occupancy: ``(max_size, size, available, waiting)``. diff --git a/tests/storage/test_rust_pool.py b/tests/storage/test_rust_pool.py index 96308bbd89..4503b65ab7 100644 --- a/tests/storage/test_rust_pool.py +++ b/tests/storage/test_rust_pool.py @@ -23,6 +23,8 @@ import time from typing import TYPE_CHECKING, Any, cast from unittest.mock import Mock +from prometheus_client import REGISTRY + # The reactor the trial runner spins up; real, so threads and callFromThread work. from twisted.internet import reactor as _reactor from twisted.internet.defer import gatherResults @@ -299,6 +301,45 @@ class RustConnectionPoolTestCase(trial_unittest.TestCase): self.assertEqual(first, second) self.assertNotEqual(second, third) + async def test_stats_count_prepares_and_retirements(self) -> None: + """A statement is prepared once per connection, then found in the cache. + Connections retired by `txn_limit` are counted, the counts reach + Prometheus, and closing the pool keeps them.""" + pool = RustConnectionPool( + reactor, + _test_runtime(), + _build_dsn(), + name="test-rust-db-stats", + server_name="test", + threads=1, + txn_limit=2, + ) + pool.start() + self.addCleanup(pool.close) + + def query(conn: Any) -> None: + cursor = conn.cursor() + cursor.execute("SELECT 1 + $1::int", (1,)) + cursor.fetchall() + + # Two calls on the first connection, then one on its replacement. + for _ in range(3): + await pool.runWithConnection(query) + + hits, prepares, clears, opened, retired = pool.stats() + self.assertEqual((hits, prepares, clears), (1, 2, 0)) + self.assertEqual((opened, retired), (2, 1)) + + prepares_metric = REGISTRY.get_sample_value( + "synapse_rust_db_pool_statement_prepares_total", + {"name": "test-rust-db-stats", "server_name": "test"}, + ) + self.assertEqual(prepares_metric, 2) + + # Closing the native pool keeps what it counted. + pool.close() + self.assertEqual(pool.stats(), (1, 2, 0, 2, 1)) + def test_status_reports_pool_occupancy(self) -> None: # Nothing checked out yet: at most the configured number of # connections, none of them in use.