From cb7b622656cadebb9118edd06a8d63756502555a Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 2 Jul 2026 15:03:24 +0000 Subject: [PATCH] Make the Postgres Connection shim pool-backed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Python-facing `Connection`/`Cursor` shim now always wraps a connection checked out of the `deadpool` pool. `ConnInner` holds a `PooledConnection` and disposes of it correctly when the last reference (the `Connection` and all its cursors) goes away, via a `Drop` impl plus `release`/`discard` helpers: - a clean connection is returned to the pool for reuse; - a connection is *discarded* (detached with `Object::take`, which also shrinks the pool) whenever reuse would be unsafe — a failed COMMIT/ROLLBACK, a poisoned mutex, or being dropped with a transaction still open (which can't be rolled back synchronously from `Drop`, so the socket close makes the server do it). A plain query error still does *not* throw the connection away — it stays open+aborted for the driver to `rollback()`, exactly as psycopg2 behaves. The pool is the only way to obtain a connection: the standalone `connect(dsn)` free function is replaced by a Python-facing `ConnectionPool` class (Rust `PyConnectionPool`, exposed as `postgres.ConnectionPool`). Build it once from a DSN, then check connections out with `pool.connect()`. Checkout failures map onto the DBAPI2 hierarchy — a backend connect error reuses the query-error mapping (and its `pgcode`), while a timeout / closed pool becomes `OperationalError` — so `connect()` behaves like psycopg2's for callers. Adds live-Postgres tests (gated on SYNAPSE_TEST_POSTGRES_DSN) asserting which connections end up back in the pool, and drives the Python test suite through a pool in `setUp`. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d --- rust/src/database/postgres/connection.rs | 223 ++++++++++++++++--- rust/src/database/postgres/mod.rs | 38 +--- rust/src/database/postgres/pool.rs | 58 ++++- tests/synapse_rust/test_database_postgres.py | 47 ++-- 4 files changed, 284 insertions(+), 82 deletions(-) diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs index ace581ac72..47e38a8cf3 100644 --- a/rust/src/database/postgres/connection.rs +++ b/rust/src/database/postgres/connection.rs @@ -39,21 +39,34 @@ //! `commit`/`rollback` end the transaction and clear the flag; with no //! transaction open they are no-ops, just like psycopg2. //! -//! ## Dropping the `Client` on error +//! ## Returning vs discarding the connection +//! +//! Every [`Connection`] wraps a [`PooledConnection`] checked out of the +//! [`super::pool`]. Dropping it normally *returns it to the pool* for reuse. +//! Where that reuse would be unsafe we instead **discard** it: the connection is +//! detached from the pool with [`Object::take`] and dropped, which both closes +//! the socket and shrinks the pool so the bad connection is never handed out +//! again. //! //! A *query* error (bad SQL, a constraint violation, an integer out of range, //! …) leaves the connection open with its transaction in the aborted state, //! exactly as psycopg2 does: the error propagates to Python and the driver is //! expected to `rollback()`. We do **not** throw the connection away for these. //! -//! The transaction-control statements are different. If `COMMIT` or `ROLLBACK` -//! itself fails we no longer know what state the server-side session is in, so -//! we drop the `Client` (closing the socket) rather than hand a possibly-broken -//! connection back for reuse. Likewise, [`Connection::close`] drops the client; -//! the server rolls back any transaction left open when the socket closes. +//! Three situations do force a discard, because the session state is unknown or +//! unclean and must not reach the next caller: +//! - a failed `COMMIT`/`ROLLBACK` — we no longer know the session state; +//! - a poisoned mutex — a panic happened mid-operation; +//! - the connection being dropped with a transaction still open — it can't be +//! rolled back synchronously from `Drop`, so the server does it for us when +//! the socket closes. +//! +//! [`Connection::close`], by contrast, returns a clean connection to the pool +//! (discarding it only if a transaction was left open). use std::sync::{Arc, Mutex, MutexGuard, TryLockError}; +use deadpool::managed::Object; use futures::future::try_join_all; use pyo3::{ exceptions::PyRuntimeError, @@ -63,7 +76,8 @@ use pyo3::{ use tokio_postgres::Client; use crate::database::postgres::{ - cursor_state::CursorQueryState, helpers::BlockingPostgresResult, value::PgValue, + cursor_state::CursorQueryState, helpers::BlockingPostgresResult, pool::PooledConnection, + value::PgValue, }; /// `try_lock` a mutex that is single-threaded by contract, mapping its two @@ -94,10 +108,10 @@ fn try_lock_or_reset<'a, T>( /// A single Postgres connection exposed to Python. /// -/// Owns the [`tokio_postgres::Client`] for its whole life and is the authority -/// on transaction state. The `Arc>` lets cursors hold a cheap clone -/// (so they can reach the client to start a query) while keeping all access to -/// the client serialised. +/// Wraps a connection checked out of the pool for its whole life and is the +/// authority on transaction state. The `Arc>` lets cursors hold a +/// cheap clone (so they can reach the client to start a query) while keeping all +/// access to the client serialised. #[pyclass(frozen, skip_from_py_object)] #[derive(Clone)] pub struct Connection { @@ -106,9 +120,10 @@ pub struct Connection { /// The mutable guts of a [`Connection`], behind its mutex. struct ConnInner { - /// The driver client. `None` once the connection has been closed (or thrown - /// away after a transaction-control error); any further use is an error. - client: Option, + /// The pooled connection. `None` once the connection has been closed, + /// returned to the pool, or discarded after an error; any further use is an + /// error. + client: Option, /// Whether a transaction is currently open (a `BEGIN` has been issued and /// not yet matched by a `COMMIT`/`ROLLBACK`). Drives the lazy `BEGIN`. in_txn: bool, @@ -118,12 +133,51 @@ struct ConnInner { autocommit: bool, } +impl ConnInner { + /// Give up the connection cleanly. + /// + /// The connection is returned to the pool for reuse — **unless** a + /// transaction is still open, in which case it can't be rolled back from + /// here, so we discard it (detach and drop) rather than hand a + /// mid-transaction connection to the next caller. + fn release(&mut self) { + if let Some(conn) = self.client.take() { + if self.in_txn { + let _ = Object::take(conn); // detach + drop: not returned to the pool + } + // else: `conn` dropped here → returned to the pool for reuse. + } + self.in_txn = false; + } + + /// Discard the connection: the session state is unknown or unclean, so it + /// must never be reused. It is detached from the pool (shrinking it). + fn discard(&mut self) { + if let Some(conn) = self.client.take() { + let _ = Object::take(conn); + } + self.in_txn = false; + } +} + +impl Drop for ConnInner { + fn drop(&mut self) { + // Return the connection to the pool (or discard it) when the last + // reference — the `Connection` and all its cursors — goes away. + self.release(); + } +} + impl Connection { - /// Wrap a freshly-established `Client` in a `Connection`. - pub fn new(client: Client) -> Self { + /// Wrap a connection checked out of the pool in a `Connection`. + /// + /// The connection is returned to the pool when this `Connection` (and every + /// cursor cloned from it) is dropped, unless it is discarded first (see the + /// module docs). + pub fn new(conn: PooledConnection) -> Self { Self { inner: Arc::new(Mutex::new(ConnInner { - client: Some(client), + client: Some(conn), in_txn: false, autocommit: false, })), @@ -140,10 +194,9 @@ impl Connection { /// know the session state — and errors. fn lock(&self) -> PyResult> { try_lock_or_reset(&self.inner, "connection", |inner| { - // On poison we no longer know the session state, so close the - // connection: drop the client and clear the transaction flag. - inner.client = None; - inner.in_txn = false; + // On poison we no longer know the session state, so discard the + // connection (never returning it to the pool). + inner.discard(); }) } @@ -208,9 +261,9 @@ impl Connection { Ok(()) } Err(err) => { - // Unknown session state: drop the connection rather than reuse it. - guard.client = None; - guard.in_txn = false; + // Unknown session state: discard the connection rather than + // reuse it (or return it to the pool). + guard.discard(); Err(err) } } @@ -222,7 +275,7 @@ impl Connection { fn client_ref(guard: &ConnInner) -> PyResult<&Client> { guard .client - .as_ref() + .as_deref() .ok_or_else(|| PyRuntimeError::new_err("connection already closed")) } @@ -247,15 +300,14 @@ impl Connection { self.end_txn(py, "ROLLBACK") } - /// Close the connection, dropping the underlying client. + /// Close the connection, releasing the underlying client. /// - /// Dropping the client closes the socket; the server rolls back any - /// transaction that was still open. Idempotent: closing an + /// A standalone client's socket is closed; a pooled connection is returned + /// to the pool for reuse (or discarded if a transaction was left open, in + /// which case the server rolls it back). Idempotent: closing an /// already-closed connection is fine. fn close(&self) -> PyResult<()> { - let mut guard = self.lock()?; - guard.client = None; - guard.in_txn = false; + self.lock()?.release(); Ok(()) } @@ -556,3 +608,112 @@ impl Cursor { Ok(()) } } + +#[cfg(test)] +mod tests { + //! These exercise the pool-backed `Connection` against a live Postgres, so + //! they only run when `SYNAPSE_TEST_POSTGRES_DSN` is set (e.g. to + //! `host=postgres user=postgres password=postgres dbname=postgres`); + //! otherwise they no-op. They assert *which* connections end up back in the + //! pool — the transaction/value logic itself is covered by the Python test + //! suite that drives these classes end to end. + + use super::*; + use crate::database::postgres::pool::create_pool; + use crate::database::runtime::runtime; + + fn test_dsn() -> Option { + std::env::var("SYNAPSE_TEST_POSTGRES_DSN").ok() + } + + /// A clean connection (its transaction committed) is returned to the pool + /// when the `Connection` is dropped. + #[test] + fn pooled_connection_returns_to_pool_after_commit() { + let Some(dsn) = test_dsn() else { + eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run"); + return; + }; + + let pool = create_pool(&dsn, 1).unwrap(); + let obj = runtime().block_on(async { pool.get().await.unwrap() }); + assert_eq!(pool.status().size, 1); + + Python::initialize(); + Python::attach(move |py| { + let conn = Connection::new(obj); + let cursor = conn.cursor(); + cursor.execute(py, "SELECT 1", None).unwrap(); + conn.commit(py).unwrap(); + // Dropping both references releases the pooled connection. + drop(cursor); + drop(conn); + }); + + // The (clean) connection went back to the pool rather than being torn + // down, so it's available for the next caller. + assert_eq!(pool.status().size, 1); + assert_eq!(pool.status().available, 1); + } + + /// A connection dropped with a transaction still open can't be rolled back + /// from `Drop`, so it is discarded (detached from the pool) rather than + /// handed to the next caller mid-transaction. + #[test] + fn pooled_connection_discarded_when_dropped_mid_transaction() { + let Some(dsn) = test_dsn() else { + eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run"); + return; + }; + + let pool = create_pool(&dsn, 1).unwrap(); + let obj = runtime().block_on(async { pool.get().await.unwrap() }); + assert_eq!(pool.status().size, 1); + + Python::initialize(); + Python::attach(move |py| { + let conn = Connection::new(obj); + let cursor = conn.cursor(); + // Opens a transaction lazily (BEGIN) but never commits/rolls back. + cursor.execute(py, "SELECT 1", None).unwrap(); + drop(cursor); + drop(conn); + }); + + // Detached: the pool shrank rather than accepting a mid-transaction + // connection back. + assert_eq!(pool.status().size, 0); + assert_eq!(pool.status().available, 0); + } + + /// A plain query error does *not* poison the connection: after the caller + /// rolls back, the (now-clean) connection returns to the pool. + #[test] + fn pooled_connection_returns_to_pool_after_query_error_and_rollback() { + let Some(dsn) = test_dsn() else { + eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run"); + return; + }; + + let pool = create_pool(&dsn, 1).unwrap(); + let obj = runtime().block_on(async { pool.get().await.unwrap() }); + + Python::initialize(); + Python::attach(move |py| { + let conn = Connection::new(obj); + let cursor = conn.cursor(); + // A bad statement aborts the transaction but leaves the connection + // usable, exactly as psycopg2 does. + cursor + .execute(py, "SELECT * FROM does_not_exist", None) + .unwrap_err(); + // The driver's job on failure: roll back, which clears the txn. + conn.rollback(py).unwrap(); + drop(cursor); + drop(conn); + }); + + assert_eq!(pool.status().size, 1); + assert_eq!(pool.status().available, 1); + } +} diff --git a/rust/src/database/postgres/mod.rs b/rust/src/database/postgres/mod.rs index fa115d99c5..d231dd1dd6 100644 --- a/rust/src/database/postgres/mod.rs +++ b/rust/src/database/postgres/mod.rs @@ -4,14 +4,9 @@ //! shared multi-thread tokio runtime (see `super::runtime`). use anyhow::Error; -use log::warn; -use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use pyo3::types::PyModule; -use crate::database::postgres::helpers::BlockingPostgresResult; -use crate::database::runtime::runtime; - mod connection; mod cursor_state; mod errors; @@ -21,15 +16,15 @@ pub mod pool; pub(crate) mod query; mod value; -/// Register the `postgres` submodule (the `Connection` / `Cursor` classes, the -/// DBAPI2 exception hierarchy and the `connect` factory) under the parent +/// Register the `postgres` submodule (the `ConnectionPool`, `Connection` and +/// `Cursor` classes and the DBAPI2 exception hierarchy) under the parent /// `database` module. pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { let child = PyModule::new(py, "postgres")?; + child.add_class::()?; child.add_class::()?; child.add_class::()?; - child.add_function(wrap_pyfunction!(connect, &child)?)?; errors::register_exceptions(py, &child)?; m.add_submodule(&child)?; @@ -43,33 +38,6 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> Ok(()) } -/// Open a new Postgres connection from a libpq-style DSN. -/// -/// Blocks until the connection is established, then spawns the long-lived -/// connection task (which drives the socket) onto the shared runtime and -/// hands back a `Connection` wrapping the client. -#[pyfunction] -fn connect<'py>(py: Python<'py>, dsn: &str) -> PyResult> { - let config = fixup_default_host(dsn) - .map_err(|e| PyRuntimeError::new_err(format!("Failed to prepare DSN: {e}")))?; - - // TLS is not yet supported: unlike libpq (whose default is - // `sslmode=prefer`), we never negotiate TLS regardless of the DSN's - // sslmode. Supporting it is left to a follow-up. - let (client, connection) = config.connect(tokio_postgres::NoTls).block_on_result(py)?; - - // Spawn the connection task on the runtime. - runtime().spawn(async move { - if let Err(e) = connection.await { - warn!("postgres connection error: {e}"); - } - }); - - let conn = connection::Connection::new(client); - - Bound::new(py, conn) -} - /// Fix up a DSN to ensure it has a host, using libpq's default host if /// necessary. /// diff --git a/rust/src/database/postgres/pool.rs b/rust/src/database/postgres/pool.rs index f29d4802d1..19fafb93eb 100644 --- a/rust/src/database/postgres/pool.rs +++ b/rust/src/database/postgres/pool.rs @@ -13,11 +13,16 @@ //! *same* pool, so both share a single set of connections rather than running //! two pools that could exhaust the server's connection limit between them. -use deadpool::managed::{Manager, Metrics, Object, Pool, RecycleError, RecycleResult}; +use deadpool::managed::{Manager, Metrics, Object, Pool, PoolError, RecycleError, RecycleResult}; use log::warn; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; use tokio_postgres::{Client, Config, NoTls}; +use crate::database::postgres::connection::Connection; +use crate::database::postgres::errors::{pg_err_to_py, OperationalError}; use crate::database::postgres::fixup_default_host; +use crate::database::postgres::helpers::BlockingPostgres; use crate::database::runtime::runtime; /// Creates and recycles [`tokio_postgres`] connections for a [`ConnectionPool`]. @@ -82,6 +87,57 @@ pub fn create_pool(dsn: &str, max_size: usize) -> Result PyResult { + let pool = create_pool(dsn, max_size).map_err(|e| { + PyRuntimeError::new_err(format!("failed to build connection pool: {e}")) + })?; + Ok(Self { pool }) + } + + /// Check a connection out of the pool. + /// + /// Blocks (releasing the GIL) until a connection is available, opening a new + /// one if the pool is below `max_size`. A failure to establish the + /// connection surfaces through the same DBAPI2 exception hierarchy as a + /// query error, so callers can treat it like psycopg2's `connect`. + fn connect(&self, py: Python<'_>) -> PyResult { + let conn = self.pool.get().block_on(py).map_err(pool_err_to_py)?; + Ok(Connection::new(conn)) + } +} + +/// Map a `deadpool` checkout failure onto the DBAPI2 exception hierarchy. +fn pool_err_to_py(err: PoolError) -> PyErr { + match err { + // The backend failed to establish the connection: reuse the exact + // mapping (and `pgcode` tagging) a query error gets. + PoolError::Backend(e) => pg_err_to_py(&e), + // Timed out waiting for a slot, pool closed, no runtime, or a + // post-create hook failure: all connection-level problems, which + // Synapse treats as retryable operational errors. + other => OperationalError::new_err(format!("failed to acquire connection: {other}")), + } +} + #[cfg(test)] mod tests { //! These tests need a live Postgres, so they only run when diff --git a/tests/synapse_rust/test_database_postgres.py b/tests/synapse_rust/test_database_postgres.py index 41f3a5edc1..81c7ca3ade 100644 --- a/tests/synapse_rust/test_database_postgres.py +++ b/tests/synapse_rust/test_database_postgres.py @@ -84,23 +84,31 @@ class PostgresConnectionTestCase(unittest.TestCase): """Tests for the Rust Postgres ``Connection`` / ``Cursor``.""" def setUp(self) -> None: - self.conn = postgres.connect(_build_dsn()) + # Connections are only ever obtained from a pool; check one out for the + # duration of the test. + self.pool = postgres.ConnectionPool(_build_dsn()) + self.conn = self.pool.connect() def tearDown(self) -> None: - # Explicitly drop the connection to ensure that the underlying Rust - # object is dropped before the Python interpreter shuts down. Otherwise, - # the open connection will block us tearing down the test database. + # Explicitly drop the connection (returning it to the pool) and then the + # pool, so the underlying Rust objects are dropped before the Python + # interpreter shuts down. Otherwise the open connection would block us + # tearing down the test database. del self.conn + del self.pool # ------------------------------------------------------------------ - # connect() + # ConnectionPool.connect() # ------------------------------------------------------------------ def test_connect_bad_dsn_raises(self) -> None: # A syntactically valid but unconnectable DSN should raise one of our - # DBAPI2 errors rather than return a half-open connection. + # DBAPI2 errors rather than return a half-open connection. The pool + # parses the DSN eagerly but only dials on checkout, so the failure + # surfaces from connect(). + pool = postgres.ConnectionPool("host=127.0.0.1 port=1 dbname=does_not_exist") with self.assertRaises(postgres.Error): - postgres.connect("host=127.0.0.1 port=1 dbname=does_not_exist") + pool.connect() # ------------------------------------------------------------------ # execute() / fetch_one() / fetch_all() @@ -783,12 +791,15 @@ class PostgresConnectionDrivenTestCase(unittest.TestCase): """ def setUp(self) -> None: - self.conn = postgres.connect(_build_dsn()) + self.pool = postgres.ConnectionPool(_build_dsn()) + self.conn = self.pool.connect() def tearDown(self) -> None: - # Drop the connection before the interpreter shuts down (see the note - # in PostgresConnectionTestCase.tearDown). + # Drop the connection (returning it to the pool) and the pool before the + # interpreter shuts down (see the note in + # PostgresConnectionTestCase.tearDown). del self.conn + del self.pool # -- small helpers ------------------------------------------------------ @@ -1046,10 +1057,12 @@ class PostgresErrorMappingTestCase(unittest.TestCase): """ def setUp(self) -> None: - self.conn = postgres.connect(_build_dsn()) + self.pool = postgres.ConnectionPool(_build_dsn()) + self.conn = self.pool.connect() def tearDown(self) -> None: del self.conn + del self.pool def _exec_commit(self, sql: str) -> None: """Run a single statement and commit it (its own transaction).""" @@ -1102,12 +1115,14 @@ class PostgresErrorMappingTestCase(unittest.TestCase): "requires Postgres reachable on libpq's default host", ) class PostgresDefaultHostTestCase(unittest.TestCase): - """Covers the libpq default-host fixup in ``connect``. + """Covers the libpq default-host fixup. When the DSN omits ``host=``, ``tokio-postgres`` would default to localhost, but Synapse wants libpq's default (honouring ``PGHOST`` / the compiled-in - socket dir). This only runs when the test Postgres is actually reachable on - that default host, so it's guarded separately from the main suite. + socket dir). The fixup lives in the pool's connection manager, so checking a + connection out of a pool built from a host-less DSN exercises it. This only + runs when the test Postgres is actually reachable on that default host, so + it's guarded separately from the main suite. """ def test_connect_without_host_uses_libpq_default(self) -> None: @@ -1120,13 +1135,15 @@ class PostgresDefaultHostTestCase(unittest.TestCase): parts.append(f"port={POSTGRES_PORT}") if POSTGRES_PASSWORD is not None: parts.append(f"password={POSTGRES_PASSWORD}") - conn = postgres.connect(" ".join(parts)) + pool = postgres.ConnectionPool(" ".join(parts)) + conn = pool.connect() try: self.assertEqual( run_interaction(conn, lambda cursor: _select_one(cursor)), (1,) ) finally: del conn + del pool def _select_one(cursor: Any) -> Optional[list[Any]]: