diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs index 4d115d27f3..911291727c 100644 --- a/rust/src/database/postgres/connection.rs +++ b/rust/src/database/postgres/connection.rs @@ -22,7 +22,7 @@ //! A [`Cursor`] is therefore cheap: it holds a clone of the owning //! [`Connection`] (an `Arc`) plus its own result-set state //! ([`CursorQueryState`]). It borrows the client only for the brief moment it -//! takes to *start* a query (`prepare` + `query_raw`); the resulting row stream +//! takes to *start* a query (a cached `prepare` + `query_raw`); the resulting row stream //! is self-contained (`'static`), so once a query has been issued the cursor //! reads rows from its own state without touching the connection again. Many //! cursors can share one connection this way, though in practice Synapse's @@ -74,10 +74,11 @@ use pyo3::{ prelude::*, types::{PyDict, PyInt, PyList, PyTuple}, }; -use tokio_postgres::Client; use crate::database::postgres::{ - cursor_state::CursorQueryState, helpers::BlockingPostgresResult, pool::PooledConnection, + cursor_state::CursorQueryState, + helpers::BlockingPostgresResult, + pool::{PooledClient, PooledConnection}, value::PgValue, }; @@ -253,7 +254,7 @@ impl Connection { fn with_client( &self, py: Python<'_>, - f: impl FnOnce(&Client) -> PyResult, + f: impl FnOnce(&PooledClient) -> PyResult, ) -> PyResult { let mut guard = self.lock()?; @@ -261,11 +262,12 @@ impl Connection { // matching psycopg2. We set `in_txn` *after* a successful `BEGIN` but // before running `f`, so that if `f` (the user's statement) fails the // open-but-aborted transaction is still tracked and `rollback()` knows - // to clean it up. + // to clean it up. `batch_execute` uses the simple-query protocol: one + // round-trip, no prepare. if !guard.autocommit && !guard.in_txn { { let client = client_ref(&guard)?; - client.execute("BEGIN", &[]).block_on_result(py)?; + client.batch_execute("BEGIN").block_on_result(py)?; } guard.in_txn = true; } @@ -301,7 +303,8 @@ impl Connection { let result = { let client = client_ref(&guard)?; - client.execute(stmt, &[]).block_on_result(py) + // Simple-query protocol: one round-trip, no prepare. + client.batch_execute(stmt).block_on_result(py) }; match result { @@ -321,7 +324,7 @@ impl Connection { /// Borrow the live client out of a locked inner state, or error if the /// connection has been closed. -fn client_ref(guard: &ConnInner) -> PyResult<&Client> { +fn client_ref(guard: &ConnInner) -> PyResult<&PooledClient> { guard .client .as_deref() @@ -515,17 +518,16 @@ impl Cursor { *state = CursorQueryState::new() }) } -} -#[pymethods] -impl Cursor { - /// Execute `query`, optionally with positional `params` bound to `$1`, - /// `$2`, ... placeholders. - /// - /// Any previous result set is discarded. After this returns, rows (if any) - /// can be read with `fetch_one`/`fetch_all`/`fetch_next_batch`. - #[pyo3(signature = (query, params = None))] - fn execute(&self, py: Python<'_>, query: &str, params: Option>) -> PyResult<()> { + /// One attempt at [`Cursor::execute`]'s body; `refresh` forces a fresh + /// prepare (bypassing and replacing the cached statement). + fn execute_once( + &self, + py: Python<'_>, + query: &str, + params: Vec, + refresh: bool, + ) -> PyResult<()> { // Drop any previous result set before starting the new query. self.lock_state()?.new_query(); @@ -533,11 +535,9 @@ impl Cursor { // the query. `query_raw` returns a `'static` `RowStream`, so the borrow // ends here and the cursor owns the stream from now on. let (stream, description) = self.connection.with_client(py, |client| { - let statement = client.prepare(query).block_on_result(py)?; + let statement = client.prepare_cached(query, refresh).block_on_result(py)?; - let stream = client - .query_raw(&statement, params.unwrap_or_default()) - .block_on_result(py)?; + let stream = client.query_raw(&statement, params).block_on_result(py)?; // The column names back the (future) PEP-249 `Cursor.description`; // pull them out of the prepared statement here so `cursor_state` @@ -568,6 +568,59 @@ impl Cursor { Ok(()) } +} + +/// Whether `err` is Postgres's "cached plan must not change result type" +/// (SQLSTATE `0A000`), raised at Bind time when concurrent DDL has invalidated +/// a cached prepared statement's result shape. +fn is_stale_plan_err(py: Python<'_>, err: &PyErr) -> bool { + err.value(py) + .getattr("pgcode") + .ok() + .and_then(|code| code.extract::>().ok().flatten()) + .is_some_and(|code| code == "0A000") +} + +#[pymethods] +impl Cursor { + /// Execute `query`, optionally with positional `params` bound to `$1`, + /// `$2`, ... placeholders. + /// + /// Any previous result set is discarded. After this returns, rows (if any) + /// can be read with `fetch_one`/`fetch_all`/`fetch_next_batch`. + /// + /// Prepared statements are cached per pooled connection (see + /// [`PooledClient::prepare_cached`]). If concurrent DDL invalidates a + /// cached plan (SQLSTATE `0A000`, "cached plan must not change result + /// type"), the statement is re-prepared and — outside a transaction, where + /// the error hasn't aborted anything — retried once. Inside a transaction + /// the error propagates (the transaction is aborted anyway); the fresh + /// prepare means the caller's retry gets a valid plan. + #[pyo3(signature = (query, params = None))] + fn execute(&self, py: Python<'_>, query: &str, params: Option>) -> PyResult<()> { + let params = params.unwrap_or_default(); + + match self.execute_once(py, query, params.clone(), false) { + Err(err) if is_stale_plan_err(py, &err) => { + if self.connection.lock()?.in_txn { + // The transaction is aborted, so nothing (not even a fresh + // prepare) can run on it; drop the stale cache entry — no + // I/O — and propagate. The caller's next transaction + // re-prepares a valid plan. + self.connection.with_client(py, |client| { + client.invalidate(query); + Ok(()) + })?; + Err(err) + } else { + // Autocommit: nothing is aborted, so re-prepare fresh and + // retry once. + self.execute_once(py, query, params, true) + } + } + other => other, + } + } /// Execute `query` once for each parameter set in `params_seq`. /// @@ -610,7 +663,7 @@ impl Cursor { // Prepare once, then build a future per parameter set. Driving them // concurrently is what makes `tokio_postgres` pipeline them onto the // connection; blocking on the joined future runs the whole batch. - let statement = client.prepare(query).block_on_result(py)?; + let statement = client.prepare_cached(query, false).block_on_result(py)?; let counts = try_join_all( params_seq diff --git a/rust/src/database/postgres/pool.rs b/rust/src/database/postgres/pool.rs index 8fb63ed865..8b075babdd 100644 --- a/rust/src/database/postgres/pool.rs +++ b/rust/src/database/postgres/pool.rs @@ -7,17 +7,22 @@ //! [`super::fixup_config_defaults`]) and drives the connection task on the shared //! runtime. //! -//! 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. +//! The pooled item is a [`PooledClient`]: a [`tokio_postgres::Client`] plus a +//! per-connection prepared-statement cache (see [`PooledClient::prepare_cached`]). +//! Rust-native code takes one from the pool and uses it with the standard +//! `tokio_postgres` async query functions (it derefs to the client); 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 std::collections::HashMap; +use std::sync::Mutex; 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 tokio_postgres::{Client, Config, NoTls, Statement}; use crate::database::postgres::connection::Connection; use crate::database::postgres::errors::{pg_err_to_py, OperationalError}; @@ -91,11 +96,74 @@ impl ConnectionManager { } } +/// A pooled [`Client`] plus its cache of prepared statements. +/// +/// The cache lives *with* the pooled connection (not with a checkout): named +/// prepared statements are per-session server state, so they survive being +/// returned to the pool, and a repeated query on any later checkout of the same +/// connection skips the prepare round-trip. Dereferences to the [`Client`] so +/// existing call sites are unaffected. +pub struct PooledClient { + client: Client, + /// Prepared statements keyed by their SQL. Guarded by a plain [`Mutex`]: + /// the critical sections are lookups/inserts (never awaits). + statements: Mutex>, +} + +/// The cache is cleared once it holds this many statements. Synapse's hot set +/// of distinct SQL strings is far smaller; the bound exists so one-off SQL +/// (e.g. `execute_values`'s literal-spliced statements) can't grow the map +/// without limit on a long-lived connection. A clear-all is crude but simple, +/// and the hot set re-warms in one round of queries. +const STATEMENT_CACHE_CAP: usize = 512; + +impl PooledClient { + /// Fetch the prepared [`Statement`] for `sql`, preparing and caching it on + /// a miss. With `refresh`, any cached entry is discarded and re-prepared — + /// used to recover when concurrent DDL invalidates a cached plan (SQLSTATE + /// `0A000`, "cached plan must not change result type"). + pub async fn prepare_cached( + &self, + sql: &str, + refresh: bool, + ) -> Result { + if refresh { + self.statements.lock().expect("not poisoned").remove(sql); + } else if let Some(statement) = self.statements.lock().expect("not poisoned").get(sql) { + return Ok(statement.clone()); + } + + let statement = self.client.prepare(sql).await?; + + let mut statements = self.statements.lock().expect("not poisoned"); + if statements.len() >= STATEMENT_CACHE_CAP { + statements.clear(); + } + statements.insert(sql.to_owned(), statement.clone()); + Ok(statement) + } + + /// Drop the cached statement for `sql`, if any. No I/O: used to shed a + /// stale plan from inside an aborted transaction, where re-preparing is + /// impossible until the transaction ends. + pub fn invalidate(&self, sql: &str) { + self.statements.lock().expect("not poisoned").remove(sql); + } +} + +impl std::ops::Deref for PooledClient { + type Target = Client; + + fn deref(&self) -> &Client { + &self.client + } +} + impl Manager for ConnectionManager { - type Type = Client; + type Type = PooledClient; type Error = tokio_postgres::Error; - async fn create(&self) -> Result { + async fn create(&self) -> Result { // 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 @@ -112,10 +180,13 @@ impl Manager for ConnectionManager { // equivalent of the engine's `on_new_connection`). client.batch_execute(&self.session.setup_sql()).await?; - Ok(client) + Ok(PooledClient { + client, + statements: Mutex::new(HashMap::new()), + }) } - async fn recycle(&self, client: &mut Client, _: &Metrics) -> RecycleResult { + async fn recycle(&self, client: &mut PooledClient, _: &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