LLM attempt at simplifying

rust/src/handlers/mod.rs

@@ -54,7 +54,7 @@ impl RustHandlers {
         // Store is shared across all of the handlers so let's use an `Arc`
         let store = Arc::new(Store {
             config: config.clone(),
-            db_pool: Box::new(db_pool),
+            db_pool,
         });
 
         let versions = Py::new(

rust/src/handlers/versions.rs

@@ -21,6 +21,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::config::SynapseConfig;
 use crate::http_client::create_deferred;
+use crate::storage::db::python_db_pool::PythonDatabasePoolWrapper;
 use crate::storage::store::{PerUserExperimentalFeature, Store};
 
 /// `GET /_matrix/client/versions` response
@@ -44,7 +45,7 @@ impl<'py> IntoPyObject<'py> for VersionsResponse {
 #[pyclass]
 pub struct VersionsHandler {
     pub config: SynapseConfig,
-    pub store: Arc<Store>,
+    pub store: Arc<Store<PythonDatabasePoolWrapper>>,
     /// The Twisted reactor, used to bridge our `async` response back into a
     /// Twisted deferred that Python can `await`.
     pub reactor: Py<PyAny>,
@@ -77,7 +78,7 @@ impl VersionsHandler {
 
 /// Assemble a `/versions` response body.
 async fn build_versions_response(
-    store: &Store,
+    store: &Store<PythonDatabasePoolWrapper>,
     config: &SynapseConfig,
     user_id: Option<&str>,
 ) -> Result<VersionsResponse, anyhow::Error> {

rust/src/http_client.rs

@@ -12,14 +12,22 @@
  * <https://www.gnu.org/licenses/agpl-3.0.html>.
  */
 
-use std::{collections::HashMap, future::Future, sync::OnceLock};
+use std::{
+    collections::HashMap,
+    future::Future,
+    sync::{Arc, Mutex, OnceLock},
+};
 
 use anyhow::Context;
 use http_body_util::BodyExt;
 use once_cell::sync::OnceCell;
-use pyo3::{create_exception, exceptions::PyException, prelude::*};
+use pyo3::{
+    create_exception, exceptions::PyException, exceptions::PyRuntimeError, intern, prelude::*,
+    types::PyCFunction,
+};
 use reqwest::RequestBuilder;
 use tokio::runtime::Runtime;
+use tokio::sync::oneshot;
 
 use crate::errors::HttpResponseException;
 
@@ -322,6 +330,111 @@ where
     make_deferred_yieldable(py, &deferred)
 }
 
+/// Runs `make_deferred` on the Twisted reactor thread to obtain a Deferred (or
+/// coroutine), then resolves once that Deferred fires.
+///
+/// This is the inverse of [`create_deferred`]: where that turns a Rust future
+/// into a Twisted Deferred, this turns a Twisted Deferred into an awaitable Rust
+/// future.
+///
+/// We're called on a Tokio worker thread, but Twisted `Deferred`s (and the
+/// coroutine that `ensureDeferred` drives) are not thread-safe and Synapse's
+/// logcontext is thread-local, so the coroutine must both start and resume on
+/// the reactor thread. The `callFromThread` hop is what gets us there for the
+/// kickoff; the deferred's own callbacks then fire on the reactor thread too.
+/// (Note this is unrelated to offloading the DB work onto a thread — that's
+/// handled internally by whatever `make_deferred` calls, e.g. `runInteraction`.)
+///
+/// `make_deferred` is invoked on the reactor thread and may return either a
+/// coroutine or a `Deferred`; `ensureDeferred` normalises both to a `Deferred`.
+pub(crate) async fn await_deferred<F>(reactor: Py<PyAny>, make_deferred: F) -> PyResult<Py<PyAny>>
+where
+    F: for<'py> Fn(Python<'py>) -> PyResult<Bound<'py, PyAny>> + Send + 'static,
+{
+    // Resolves when the deferred fires; carries the resolved value or the error.
+    let (tx, rx) = oneshot::channel::<PyResult<Py<PyAny>>>();
+    // Shared between the success and error callbacks (only one ever fires).
+    let sender = Arc::new(Mutex::new(Some(tx)));
+
+    Python::attach(|py| -> PyResult<()> {
+        let success_sender = Arc::clone(&sender);
+        let on_success = PyCFunction::new_closure(
+            py,
+            None,
+            None,
+            move |args, _kwargs| -> PyResult<Py<PyAny>> {
+                let value = args.get_item(0)?.unbind();
+                if let Some(tx) = success_sender.lock().unwrap().take() {
+                    let _ = tx.send(Ok(value));
+                }
+                Ok(args.py().None())
+            },
+        )?
+        .unbind();
+
+        let error_sender = Arc::clone(&sender);
+        let on_error = PyCFunction::new_closure(
+            py,
+            None,
+            None,
+            move |args, _kwargs| -> PyResult<Py<PyAny>> {
+                let err = failure_to_pyerr(&args.get_item(0)?);
+                if let Some(tx) = error_sender.lock().unwrap().take() {
+                    let _ = tx.send(Err(err));
+                }
+                Ok(args.py().None())
+            },
+        )?
+        .unbind();
+
+        let starter = PyCFunction::new_closure(
+            py,
+            None,
+            None,
+            move |args, _kwargs| -> PyResult<Py<PyAny>> {
+                let py = args.py();
+                let deferred = defer(py)?
+                    .call_method1(intern!(py, "ensureDeferred"), (make_deferred(py)?,))?;
+                deferred.call_method1(
+                    intern!(py, "addCallbacks"),
+                    (on_success.bind(py), on_error.bind(py)),
+                )?;
+                Ok(py.None())
+            },
+        )?;
+
+        reactor
+            .bind(py)
+            .call_method1(intern!(py, "callFromThread"), (starter,))?;
+
+        Ok(())
+    })?;
+
+    match rx.await {
+        Ok(result) => result,
+        Err(_) => Err(PyRuntimeError::new_err(
+            "await_deferred channel closed before the deferred fired",
+        )),
+    }
+}
+
+/// Convert a Twisted `Failure` (as passed to an errback) into a [`PyErr`].
+///
+/// A `Failure` carries the original exception instance in its `.value`
+/// attribute, which we re-raise so callers see the real error. If that can't be
+/// reached, fall back to the `Failure`'s textual representation.
+fn failure_to_pyerr(failure: &Bound<'_, PyAny>) -> PyErr {
+    match failure.getattr(intern!(failure.py(), "value")) {
+        Ok(value) => PyErr::from_value(value),
+        Err(_) => PyRuntimeError::new_err(
+            failure
+                .str()
+                .map(|s| s.to_string_lossy().into_owned())
+                .unwrap_or_else(|_| "<unknown failure>".to_owned()),
+        ),
+    }
+}
+
 static MAKE_DEFERRED_YIELDABLE: OnceLock<pyo3::Py<pyo3::PyAny>> = OnceLock::new();
 
 /// Given a deferred, make it follow the Synapse logcontext rules

rust/src/storage/db/mod.rs

@@ -13,7 +13,6 @@
  *
  */
 
-use std::any::Any;
 use std::future::Future;
 
 use futures::future::BoxFuture;
@@ -29,61 +28,21 @@ pub mod rust_db_pool;
 /// parsing the strings into richer types as needed.
 pub type Row = Vec<String>;
 
-/// The type-erased result of a `run_interaction` callback.
-///
-/// We box the result as `dyn Any` so that the [`DatabasePool`] trait can stay
-/// object-safe (and therefore usable as `Box<dyn DatabasePool>`) while still
-/// allowing callbacks to return an arbitrary `R`. The ergonomic, generic
-/// [`DatabasePoolExt::run_interaction`] downcasts this back to the concrete
-/// type for the caller.
-pub type AnyResult = anyhow::Result<Box<dyn Any + Send>>;
-
-/// A type-erased `run_interaction` callback.
-///
-/// The callback is given a [`Transaction`] and returns a boxed future
-/// resolving to a type-erased result. It may be invoked multiple times under
-/// certain failure modes (serialization and deadlock errors), so it is `Fn`
-/// rather than `FnOnce`.
-pub type InteractionFn =
-    Box<dyn for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, AnyResult> + Send>;
-
 /// A database connection pool.
 ///
-/// We use a `Box<dyn DatabasePool>` so the same code can run against either the
-/// Python-backed pool (in Synapse, see [`python_db_pool`]) or a native
-/// `tokio-postgres` pool (in `synapse-rust-apps`, see [`rust_db_pool`]). To keep
-/// the trait object-safe, the only required method is the type-erased
-/// [`Self::run_interaction_erased`]; prefer the generic
-/// [`DatabasePoolExt::run_interaction`] at call sites.
+/// Code is written against this trait so the same store can run against either
+/// the Python-backed pool (in Synapse, see [`python_db_pool`]) or a native
+/// `tokio-postgres` pool (in `synapse-rust-apps`, see [`rust_db_pool`]). The pool
+/// type is fixed within any given binary, so callers are generic over the pool
+/// (e.g. `Store<P: DatabasePool>`) rather than using dynamic dispatch.
 ///
 /// `Send + Sync` so it can be stored in a `#[pyclass]` and shared across threads.
-#[async_trait::async_trait]
 pub trait DatabasePool: Send + Sync {
-    /// Starts a transaction on the database and runs the given (type-erased)
-    /// function.
-    ///
-    /// The given `func` may be called multiple times under certain failure
-    /// modes (like serialization and deadlock errors).
-    async fn run_interaction_erased(
-        &self,
-        name: &'static str,
-        func: InteractionFn,
-    ) -> AnyResult;
-}
-
-/// Ergonomic, generic extension to [`DatabasePool`].
-///
-/// This is automatically implemented for every `DatabasePool` (including
-/// `dyn DatabasePool`) via the blanket impl below, and provides the typed
-/// `run_interaction` that callers actually use. It lives in a separate trait
-/// (rather than on `DatabasePool` directly) because a generic method would make
-/// `DatabasePool` no longer object-safe.
-pub trait DatabasePoolExt: DatabasePool {
     /// Starts a transaction on the database and runs the given function,
     /// returning its result.
     ///
-    /// The given `func` may be called multiple times under certain failure
-    /// modes (like serialization and deadlock errors).
+    /// `func` may be called multiple times under certain failure modes (like
+    /// serialization and deadlock errors), so it is `Fn` rather than `FnOnce`.
     fn run_interaction<R, F>(
         &self,
         name: &'static str,
@@ -94,29 +53,9 @@ pub trait DatabasePoolExt: DatabasePool {
         F: for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, anyhow::Result<R>>
             + Send
             + Sync
-            + 'static,
-    {
-        // Erase the concrete return type `R` into `Box<dyn Any>` so we can call
-        // through the object-safe `run_interaction_erased`.
-        let erased: InteractionFn = Box::new(move |txn| {
-            let fut = func(txn);
-            Box::pin(async move {
-                let value = fut.await?;
-                Ok(Box::new(value) as Box<dyn Any + Send>)
-            })
-        });
-
-        async move {
-            let boxed = self.run_interaction_erased(name, erased).await?;
-            Ok(*boxed
-                .downcast::<R>()
-                .expect("run_interaction return type mismatch (this is a Synapse programming error)"))
-        }
-    }
+            + 'static;
 }
 
-impl<T: DatabasePool + ?Sized> DatabasePoolExt for T {}
-
 /// A [`tokio_postgres::Transaction`] looking thing that we can use on the Rust side to
 /// interact with the database
 #[async_trait::async_trait]

rust/src/storage/db/python_db_pool.rs

@@ -19,10 +19,11 @@
 
 use std::sync::{Arc, Mutex};
 
+use futures::future::BoxFuture;
 use pyo3::{exceptions::PyRuntimeError, intern, prelude::*, types::PyCFunction, types::PyList};
-use tokio::sync::oneshot;
 
-use crate::storage::db::{AnyResult, DatabasePool, InteractionFn, Row, Transaction};
+use crate::http_client::await_deferred;
+use crate::storage::db::{DatabasePool, Row, Transaction};
 
 /// The database engines we support in the Python side of Synapse
 #[derive(Copy, Clone, Debug)]
@@ -63,26 +64,33 @@ impl PythonDatabasePoolWrapper {
     }
 }
 
-#[async_trait::async_trait]
 impl DatabasePool for PythonDatabasePoolWrapper {
-    async fn run_interaction_erased(&self, name: &'static str, func: InteractionFn) -> AnyResult {
-        // We bridge the Python-side `runInteraction` (a coroutine, run on the
-        // Twisted reactor + thread pool) back into our `async` Rust world using a
-        // oneshot channel that resolves when the resulting deferred fires.
-        let (tx, rx) = oneshot::channel::<AnyResult>();
-
+    async fn run_interaction<R, F>(&self, name: &'static str, func: F) -> anyhow::Result<R>
+    where
+        R: Send + 'static,
+        F: for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, anyhow::Result<R>>
+            + Send
+            + Sync
+            + 'static,
+    {
         // `runInteraction` calls `func` with a `LoggingTransaction` on a DB
         // thread and expects a synchronous return value. Since we can't
         // round-trip an arbitrary Rust `R` back out through Python, the callback
         // stashes the result here and we pick it up once the deferred fires.
-        let result_slot: Arc<Mutex<Option<AnyResult>>> = Arc::new(Mutex::new(None));
+        //
+        // Note the callback may run more than once (`runInteraction` retries on
+        // serialization/deadlock errors), so we only trust this slot once the
+        // deferred has fired, i.e. once the transaction has finally committed or
+        // failed.
+        let result_slot: Arc<Mutex<Option<anyhow::Result<R>>>> = Arc::new(Mutex::new(None));
 
-        Python::attach(|py| -> PyResult<()> {
-            // (1) The callback that Python's `runInteraction` invokes on a DB
-            // thread with a `LoggingTransaction`. We drive `func` to completion
-            // here. The Python query path is synchronous under the hood, so it's
-            // safe to block this dedicated DB thread until the future resolves.
-            let callback_slot = Arc::clone(&result_slot);
+        // Build the callback that Python's `runInteraction` invokes on a DB
+        // thread with a `LoggingTransaction`, plus owned handles we can move onto
+        // the reactor thread. We drive `func` to completion in the callback; the
+        // Python query path is synchronous under the hood, so it's safe to block
+        // this dedicated DB thread until the future resolves.
+        let callback_slot = Arc::clone(&result_slot);
+        let (callback, database_pool_py, reactor) = Python::attach(|py| -> PyResult<_> {
             let callback = PyCFunction::new_closure(
                 py,
                 None,
@@ -92,9 +100,7 @@ impl DatabasePool for PythonDatabasePoolWrapper {
                     let txn_py = args.get_item(0)?;
                     let mut txn = txn_py.extract::<LoggingTransactionWrapper>()?;
 
-                    let result = futures::executor::block_on(func(&mut txn));
-
-                    match result {
+                    match futures::executor::block_on(func(&mut txn)) {
                         Ok(value) => {
                             *callback_slot.lock().unwrap() = Some(Ok(value));
                             Ok(py.None())
@@ -112,91 +118,37 @@ impl DatabasePool for PythonDatabasePoolWrapper {
             )?
             .unbind();
 
-            // The oneshot sender, shared between the success and error callbacks
-            // (only one of which ever fires).
-            let sender = Arc::new(Mutex::new(Some(tx)));
-
-            // (2a) Fired when the transaction succeeds: hand the stashed result
-            // back to the awaiting task.
-            let success_slot = Arc::clone(&result_slot);
-            let success_sender = Arc::clone(&sender);
-            let on_success = PyCFunction::new_closure(
-                py,
-                None,
-                None,
-                move |args, _kwargs| -> PyResult<Py<PyAny>> {
-                    let result = success_slot.lock().unwrap().take().unwrap_or_else(|| {
-                        Err(anyhow::anyhow!("run_interaction produced no result"))
-                    });
-                    if let Some(tx) = success_sender.lock().unwrap().take() {
-                        let _ = tx.send(result);
-                    }
-                    Ok(args.py().None())
-                },
-            )?
-            .unbind();
-
-            // (2b) Fired when the transaction fails. Prefer the original error
-            // captured in the callback (it carries the Rust context); otherwise
-            // fall back to the Twisted `Failure` text.
-            let error_slot = Arc::clone(&result_slot);
-            let error_sender = Arc::clone(&sender);
-            let on_error = PyCFunction::new_closure(
-                py,
-                None,
-                None,
-                move |args, _kwargs| -> PyResult<Py<PyAny>> {
-                    let result = error_slot.lock().unwrap().take().unwrap_or_else(|| {
-                        let description = args
-                            .get_item(0)
-                            .and_then(|failure| failure.str())
-                            .map(|s| s.to_string_lossy().into_owned())
-                            .unwrap_or_else(|_| "<unknown failure>".to_owned());
-                        Err(anyhow::anyhow!("run_interaction failed: {description}"))
-                    });
-                    if let Some(tx) = error_sender.lock().unwrap().take() {
-                        let _ = tx.send(result);
-                    }
-                    Ok(args.py().None())
-                },
-            )?
-            .unbind();
-
-            // (3) Kick off `runInteraction` on the reactor thread. It's a
-            // coroutine, so we wrap it with `ensureDeferred` and attach our
-            // callbacks.
-            let database_pool_py = self.database_pool_py.clone_ref(py);
-            let starter = PyCFunction::new_closure(
-                py,
-                None,
-                None,
-                move |args, _kwargs| -> PyResult<Py<PyAny>> {
-                    let py = args.py();
-                    let coro = database_pool_py.bind(py).call_method1(
-                        intern!(py, "runInteraction"),
-                        (name, callback.bind(py)),
-                    )?;
-                    let deferred = py
-                        .import("twisted.internet.defer")?
-                        .call_method1(intern!(py, "ensureDeferred"), (coro,))?;
-                    deferred.call_method1(
-                        intern!(py, "addCallbacks"),
-                        (on_success.bind(py), on_error.bind(py)),
-                    )?;
-                    Ok(py.None())
-                },
-            )?;
-
-            self.reactor
-                .bind(py)
-                .call_method1(intern!(py, "callFromThread"), (starter,))?;
-
-            Ok(())
+            Ok((
+                callback,
+                self.database_pool_py.clone_ref(py),
+                self.reactor.clone_ref(py),
+            ))
         })
         .map_err(anyhow::Error::from)?;
 
-        rx.await
-            .map_err(|_| anyhow::anyhow!("run_interaction channel closed before completing"))?
+        // Await `runInteraction` directly. `runInteraction` offloads the actual
+        // DB work onto a thread itself, so we don't; `await_deferred` only has to
+        // start the coroutine on the reactor thread and bridge its deferred back
+        // into our `async` world.
+        let outcome = await_deferred(reactor, move |py| {
+            database_pool_py
+                .bind(py)
+                .call_method1(intern!(py, "runInteraction"), (name, callback.bind(py)))
+        })
+        .await;
+
+        // Prefer the result captured by the callback (it carries the Rust
+        // context). If the slot is empty, `runInteraction` failed before ever
+        // invoking the callback (e.g. it couldn't acquire a connection), so
+        // surface the deferred's outcome instead.
+        let captured = result_slot.lock().unwrap().take();
+        match captured {
+            Some(result) => result,
+            None => match outcome {
+                Ok(_) => Err(anyhow::anyhow!("run_interaction produced no result")),
+                Err(py_err) => Err(anyhow::Error::from(py_err)),
+            },
+        }
     }
 }
 

rust/src/storage/db/rust_db_pool.rs

@@ -19,18 +19,25 @@
 use anyhow::Context;
 use bb8_postgres::tokio_postgres::{self, types::ToSql, IsolationLevel};
 use bb8_postgres::PostgresConnectionManager;
+use futures::future::BoxFuture;
 use postgres_native_tls::MakeTlsConnector;
 
-use crate::storage::db::{AnyResult, DatabasePool, InteractionFn, Row, Transaction};
+use crate::storage::db::{DatabasePool, Row, Transaction};
 
 /// Native Rust database access backed by `tokio-postgres` (for use in synapse-rust-apps)
 pub struct RustDatabasePool {
     db_pool: bb8::Pool<PostgresConnectionManager<MakeTlsConnector>>,
 }
 
-#[async_trait::async_trait]
 impl DatabasePool for RustDatabasePool {
-    async fn run_interaction_erased(&self, _name: &'static str, func: InteractionFn) -> AnyResult {
+    async fn run_interaction<R, F>(&self, _name: &'static str, func: F) -> anyhow::Result<R>
+    where
+        R: Send + 'static,
+        F: for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, anyhow::Result<R>>
+            + Send
+            + Sync
+            + 'static,
+    {
         // Like Synapse's `runInteraction`, retry the whole transaction on
         // serialization/deadlock errors (which can happen under repeatable-read).
         loop {

rust/src/storage/store.rs

@@ -16,10 +16,7 @@
 use futures::FutureExt;
 use serde::Serialize;
 
-use crate::{
-    config::SynapseConfig,
-    storage::db::{DatabasePool, DatabasePoolExt},
-};
+use crate::{config::SynapseConfig, storage::db::DatabasePool};
 
 /// Currently supported per-user features
 #[derive(Serialize, Debug)]
@@ -59,13 +56,12 @@ impl std::fmt::Display for PerUserExperimentalFeature {
     }
 }
 
-
-pub struct Store {
+pub struct Store<P: DatabasePool> {
     pub config: SynapseConfig,
-    pub db_pool: Box<dyn DatabasePool>,
+    pub db_pool: P,
 }
 
-impl Store {
+impl<P: DatabasePool> Store<P> {
     pub async fn is_feature_enabled(
         &self,
         user_id: &str,