diff --git a/Cargo.lock b/Cargo.lock index caef3ebaad..099d51d38a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -212,6 +212,24 @@ dependencies = [ "cmov", ] +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" + [[package]] name = "digest" version = "0.10.7" @@ -471,6 +489,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "hex" version = "0.4.3" @@ -892,6 +916,16 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -1628,6 +1662,7 @@ dependencies = [ "base64", "blake2", "bytes", + "deadpool", "futures", "headers", "hex", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 180f32d866..a74e3c9ea5 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -62,6 +62,10 @@ http-body-util = "0.1.3" futures = "0.3.31" tokio = { version = "1.44.2", features = ["rt", "rt-multi-thread"] } tokio-postgres = "0.7" +# Connection pooling for the native Rust Postgres backend. We use the generic +# `deadpool::managed` pool with our own manager (rather than `deadpool-postgres`) +# so it reuses our `connect`/default-host logic and connection-task spawning. +deadpool = "0.12" once_cell = "1.18.0" itertools = "0.14.0" postgres-protocol = "0.6.10" diff --git a/rust/src/database/postgres/mod.rs b/rust/src/database/postgres/mod.rs index 28e30d5e28..5a210b0b60 100644 --- a/rust/src/database/postgres/mod.rs +++ b/rust/src/database/postgres/mod.rs @@ -17,6 +17,7 @@ mod cursor_state; mod errors; mod helpers; mod libpq; +pub mod pool; mod value; /// Register the `postgres` submodule (the `Connection` / `Cursor` classes, the @@ -75,7 +76,7 @@ fn connect<'py>(py: Python<'py>, dsn: &str) -> PyResult Result { +pub(crate) fn fixup_default_host(dsn: &str) -> Result { let mut config = dsn.parse::()?; // `tokio_postgres` parses only the DSN string (it does not consult `PGHOST` diff --git a/rust/src/database/postgres/pool.rs b/rust/src/database/postgres/pool.rs new file mode 100644 index 0000000000..08da9064fa --- /dev/null +++ b/rust/src/database/postgres/pool.rs @@ -0,0 +1,154 @@ +//! A `deadpool`-managed pool of `tokio_postgres` connections for the native +//! Rust Postgres backend. +//! +//! We use the generic `deadpool::managed` pool with our own [`ConnectionManager`] +//! rather than the `deadpool-postgres` crate, so that creating a connection +//! reuses the same DSN handling (libpq's default host, see +//! [`super::fixup_default_host`]) and connection-task spawning as +//! [`super::connect`]. +//! +//! The pooled item is a plain [`tokio_postgres::Client`]. Rust-native code takes +//! one from the pool and uses it with the standard `tokio_postgres` async query +//! functions; the Python-facing `Connection`/`Cursor` shim borrows from the +//! *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, Pool, RecycleError, RecycleResult}; +use log::warn; +use tokio_postgres::{Client, Config, NoTls}; + +use crate::database::postgres::fixup_default_host; +use crate::database::runtime::runtime; + +/// Creates and recycles [`tokio_postgres`] connections for a [`ConnectionPool`]. +pub struct ConnectionManager { + /// The resolved connection config, parsed once from the DSN (with libpq's + /// default host filled in if the DSN omitted one). + config: Config, +} + +impl ConnectionManager { + /// Build a manager from a libpq-style DSN. + pub fn from_dsn(dsn: &str) -> Result { + Ok(Self { + config: fixup_default_host(dsn)?, + }) + } +} + +impl Manager for ConnectionManager { + type Type = Client; + type Error = tokio_postgres::Error; + + async fn create(&self) -> Result { + // As in `super::connect`: establish the connection, then drive its + // long-lived connection task (which pumps the socket) on the shared + // runtime. The task ends on its own when the `Client` is dropped, i.e. + // when the pool discards this connection. + let (client, connection) = self.config.connect(NoTls).await?; + + runtime().spawn(async move { + if let Err(e) = connection.await { + warn!("postgres connection error: {e}"); + } + }); + + Ok(client) + } + + async fn recycle(&self, client: &mut Client, _: &Metrics) -> RecycleResult { + // Cheap liveness check before handing a pooled connection back out: if + // the connection task has ended (server closed the socket, fatal error, + // …) the client reports closed, so tell deadpool to drop it and create + // a fresh one rather than returning a dead connection. + if client.is_closed() { + return Err(RecycleError::message("connection closed")); + } + Ok(()) + } +} + +/// A pool of [`tokio_postgres`] connections. +pub type ConnectionPool = Pool; + +/// Build a [`ConnectionPool`] from a libpq-style DSN, capped at `max_size` +/// connections. +pub fn create_pool(dsn: &str, max_size: usize) -> Result { + let manager = ConnectionManager::from_dsn(dsn)?; + Ok(Pool::builder(manager).max_size(max_size).build()?) +} + +#[cfg(test)] +mod tests { + //! These tests need a live Postgres, so they only run when + //! `SYNAPSE_TEST_POSTGRES_DSN` is set (e.g. to + //! `host=postgres user=postgres password=postgres`); otherwise they no-op. + //! The state-machine / value / error logic is unit-tested elsewhere against + //! fakes — this is specifically the pooling behaviour against a real server. + + use super::*; + use crate::database::runtime::runtime; + + fn test_dsn() -> Option { + std::env::var("SYNAPSE_TEST_POSTGRES_DSN").ok() + } + + #[test] + fn pool_acquires_runs_a_query_and_reuses_the_connection() { + let Some(dsn) = test_dsn() else { + eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run"); + return; + }; + + let pool = create_pool(&dsn, 2).unwrap(); + + runtime().block_on(async { + // Acquire a connection and run a query with the standard async API. + let client = pool.get().await.unwrap(); + let row = client.query_one("SELECT 1::int", &[]).await.unwrap(); + assert_eq!(row.get::<_, i32>(0), 1); + + // Returning it to the pool (drop) and acquiring again reuses it. + drop(client); + assert_eq!(pool.status().size, 1); + + let client = pool.get().await.unwrap(); + let row = client.query_one("SELECT 2::int", &[]).await.unwrap(); + assert_eq!(row.get::<_, i32>(0), 2); + assert_eq!(pool.status().size, 1); + }); + } + + #[test] + fn pool_hands_out_up_to_max_size_connections_concurrently() { + let Some(dsn) = test_dsn() else { + eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run"); + return; + }; + + let pool = create_pool(&dsn, 2).unwrap(); + + runtime().block_on(async { + // Two live checkouts at once => two distinct connections created. + let a = pool.get().await.unwrap(); + let b = pool.get().await.unwrap(); + assert_eq!(pool.status().size, 2); + + // Both usable independently. + assert_eq!( + a.query_one("SELECT 10::int", &[]) + .await + .unwrap() + .get::<_, i32>(0), + 10 + ); + assert_eq!( + b.query_one("SELECT 20::int", &[]) + .await + .unwrap() + .get::<_, i32>(0), + 20 + ); + }); + } +}