Cache prepared statements per pooled connection

Every `execute` did a fresh `prepare` + `query_raw` — two server round-trips
per query — and threw the `Statement` away; even BEGIN/COMMIT/ROLLBACK went
through a prepare. psycopg2 pays neither cost, which showed up as timing
margins under load (e.g. the worker-lock contention stress test timing out
under a parallel test run).

The pooled item is now a `PooledClient`: the `Client` plus a per-connection
statement cache keyed by SQL (`prepare_cached`). The cache lives with the
pooled connection — named prepared statements are per-session server state, so
they survive checkouts and a repeated query on any later checkout skips the
prepare round-trip. `PooledClient` derefs to `Client`, so existing call sites
are unchanged. The map is cleared at a size cap so one-off SQL (e.g.
execute_values' literal-spliced statements) can't grow it without bound; the
hot set re-warms in one round.

Concurrent DDL can invalidate a cached plan (SQLSTATE 0A000, "cached plan must
not change result type"). `execute` handles it: outside a transaction the
statement is re-prepared fresh and retried once; inside one (already aborted)
the stale cache entry is shed without I/O and the error propagates, so the
caller's next transaction gets a valid plan.

Transaction control (BEGIN/COMMIT/ROLLBACK) now uses `batch_execute` — the
simple-query protocol, one round-trip, nothing to prepare.

Verified: rust tests 158/158; storage/lock/worker-lock suites pass on the Rust
backend, including the worker-lock contention test under -j4 that previously
exceeded its wall-clock budget.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01W3G4M92AmwSSZCbmtMJU3d

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<R>(
         &self,
         py: Python<'_>,
-        f: impl FnOnce(&Client) -> PyResult<R>,
+        f: impl FnOnce(&PooledClient) -> PyResult<R>,
     ) -> PyResult<R> {
         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<Vec<PgValue>>) -> 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<PgValue>,
+        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::<Option<String>>().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<Vec<PgValue>>) -> 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

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<HashMap<String, Statement>>,
+}
+
+/// 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<Statement, tokio_postgres::Error> {
+        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<Client, Self::Error> {
+    async fn create(&self) -> Result<PooledClient, Self::Error> {
         // 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<Self::Error> {
+    async fn recycle(&self, client: &mut PooledClient, _: &Metrics) -> RecycleResult<Self::Error> {
         // 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