diff --git a/rust/src/handlers/mod.rs b/rust/src/handlers/mod.rs index 1f04b65e10..9911443721 100644 --- a/rust/src/handlers/mod.rs +++ b/rust/src/handlers/mod.rs @@ -52,7 +52,9 @@ impl RustHandlers { let db_pool = PythonDatabasePoolWrapper::new(db_pool_py, reactor.clone_ref(py)); // Store is shared across all of the handlers so let's use an `Arc` - let store = Arc::new(Store { db_pool }); + let store = Arc::new(Store { + db_pool: Box::new(db_pool), + }); let global_unstable_feature_map = Arc::new( versions::synapse_config_to_global_unstable_feature_map(&config), diff --git a/rust/src/handlers/versions.rs b/rust/src/handlers/versions.rs index 99d28a4773..029f5478b8 100644 --- a/rust/src/handlers/versions.rs +++ b/rust/src/handlers/versions.rs @@ -21,7 +21,6 @@ use serde::{Deserialize, Serialize}; use crate::config::{RoomCreationPreset, SynapseConfig}; use crate::deferred::create_deferred; -use crate::storage::db::python_db_pool::PythonDatabasePoolWrapper; use crate::storage::store::{PerUserExperimentalFeature, Store}; /// `GET /_matrix/client/versions` response @@ -45,7 +44,7 @@ impl<'py> IntoPyObject<'py> for VersionsResponse { #[pyclass] pub struct VersionsHandler { pub global_unstable_feature_map: Arc, - pub store: Arc>, + pub store: Arc, /// The Twisted reactor, used to bridge our `async` response back into a /// Twisted deferred that Python can `await`. pub reactor: Py, @@ -83,7 +82,7 @@ impl VersionsHandler { /// * global_unstable_feature_map: The global values before any per-user overrides /// * user_id: The user making the request async fn build_versions_response( - store: &Store, + store: &Store, global_unstable_feature_map: &UnstableFeatureMap, user_id: Option<&str>, ) -> Result { diff --git a/rust/src/storage/db/mod.rs b/rust/src/storage/db/mod.rs index bfbf081fa7..51029161e2 100644 --- a/rust/src/storage/db/mod.rs +++ b/rust/src/storage/db/mod.rs @@ -13,23 +13,66 @@ * */ +use std::any::Any; use std::future::Future; use futures::future::BoxFuture; +use futures::FutureExt; pub mod python_db_pool; pub mod rust_db_pool; +/// A type-erased transaction callback, as accepted by +/// [`DatabasePool::run_interaction_dyn`]. +/// +/// This is the object-safe form of the `func` passed to +/// [`DatabasePoolExt::run_interaction`]: the concrete result type `R` is boxed up +/// as `Box` so the trait can stay dyn-compatible. The ergonomic +/// [`DatabasePoolExt::run_interaction`] wrapper takes care of boxing the result +/// and downcasting it back to `R`. +pub type ErasedInteraction = + Box Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, ErasedResult> + Send + Sync>; + +/// The type-erased result of an [`ErasedInteraction`]: the concrete `R` boxed up +/// as `Box` so it can pass through the object-safe +/// [`DatabasePool::run_interaction_dyn`]. +pub type ErasedResult = anyhow::Result>; + /// A database connection pool. /// /// 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`) rather than using dynamic dispatch. +/// is held behind a trait object (e.g. `Box`), so a single +/// `Store` type can wrap whichever pool the binary happens to use. +/// +/// Callers don't use this trait directly; they go through the ergonomic, +/// strongly-typed [`run_interaction`](DatabasePoolExt::run_interaction) on +/// [`DatabasePoolExt`] (blanket-implemented for every `DatabasePool`, including +/// `dyn DatabasePool`). This trait only carries the object-safe, type-erased +/// [`run_interaction_dyn`](Self::run_interaction_dyn) that the wrapper dispatches +/// to. /// /// `Send + Sync` so it can be stored in a `#[pyclass]` and shared across threads. pub trait DatabasePool: Send + Sync { + /// Object-safe, type-erased core of [`DatabasePoolExt::run_interaction`]. + /// + /// Implementors implement this; callers should prefer + /// [`DatabasePoolExt::run_interaction`], which boxes up the result and + /// downcasts it back to the concrete type for you. See that method for the + /// semantics of `name` and `func`. + fn run_interaction_dyn<'a>( + &'a self, + name: &'static str, + func: ErasedInteraction, + ) -> BoxFuture<'a, ErasedResult>; +} + +/// Ergonomic, strongly-typed access to a [`DatabasePool`]. +/// +/// Blanket-implemented for every `T: DatabasePool + ?Sized`, so it is available +/// both on concrete pools and on `dyn DatabasePool` trait objects. +pub trait DatabasePoolExt: DatabasePool { /// Starts a transaction on the database and runs the given function, /// returning its result. /// @@ -79,7 +122,51 @@ pub trait DatabasePool: Send + Sync { F: for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, anyhow::Result> + Send + Sync - + 'static; + + 'static, + { + async move { + // Box the concrete result `R` up as `Box` so we can + // route through the object-safe `run_interaction_dyn`, then downcast + // back to `R` once it returns. + let erased = erase_interaction(func); + + let result = self.run_interaction_dyn(name, erased).await?; + Ok(*result.downcast::().expect( + "run_interaction_dyn returned a value of an unexpected type; \ + this is a programming error in the DatabasePool implementation", + )) + } + } +} + +impl DatabasePoolExt for T {} + +/// Wrap a strongly-typed transaction callback so its result is boxed as +/// `Box`, producing the type-erased [`ErasedInteraction`] that +/// [`DatabasePool::run_interaction_dyn`] accepts. +fn erase_interaction(func: F) -> ErasedInteraction +where + R: Send + 'static, + F: for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, anyhow::Result> + + Send + + Sync + + 'static, +{ + // Routing the closure through a generic fn parameter that carries the + // `for<'txn>` bound pins its higher-ranked signature; assigning a closure + // that borrows its argument straight into a `Box` doesn't reliably + // infer that lifetime. + fn constrain(c: C) -> C + where + C: for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, ErasedResult>, + { + c + } + + Box::new(constrain(move |txn| { + let fut = func(txn); + async move { Ok(Box::new(fut.await?) as Box) }.boxed() + })) } /// A [`tokio_postgres::Transaction`] looking thing that we can use on the Rust side to diff --git a/rust/src/storage/db/python_db_pool.rs b/rust/src/storage/db/python_db_pool.rs index b9fc142e3f..7388bc0765 100644 --- a/rust/src/storage/db/python_db_pool.rs +++ b/rust/src/storage/db/python_db_pool.rs @@ -23,6 +23,7 @@ use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use futures::future::BoxFuture; +use futures::FutureExt; use pyo3::{ exceptions::{PyRuntimeError, PyTypeError}, intern, @@ -31,7 +32,9 @@ use pyo3::{ }; use crate::deferred::run_python_awaitable; -use crate::storage::db::{DatabasePool, DbValue, Row, Transaction}; +use crate::storage::db::{ + DatabasePool, DbValue, ErasedInteraction, ErasedResult, Row, Transaction, +}; /// The database engines we support in the Python side of Synapse #[derive(Copy, Clone, Debug)] @@ -73,103 +76,104 @@ impl PythonDatabasePoolWrapper { } impl DatabasePool for PythonDatabasePoolWrapper { - async fn run_interaction(&self, name: &'static str, func: F) -> anyhow::Result - where - R: Send + 'static, - F: for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, anyhow::Result> - + 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. - // - // 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>>> = Arc::new(Mutex::new(None)); + fn run_interaction_dyn<'a>( + &'a self, + name: &'static str, + func: ErasedInteraction, + ) -> BoxFuture<'a, ErasedResult> { + async move { + // `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 result back out through Python, the + // callback stashes the result here and we pick it up once the deferred + // fires. + // + // 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>> = Arc::new(Mutex::new(None)); - // 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, - None, - move |args, _kwargs| -> PyResult> { - let py = args.py(); - let txn_py = args.get_item(0)?; - let mut txn = txn_py.extract::()?; + // 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, + None, + move |args, _kwargs| -> PyResult> { + let py = args.py(); + let txn_py = args.get_item(0)?; + let mut txn = txn_py.extract::()?; - // Since we expect people to only call `.await` on [`Transaction`] - // related methods (mentioned in the [`Transaction`] docstring) AND - // because there is no async work to suspend on in the Python - // [`Transaction`] synchronously, we can get away with polling once - // as it should immediately resolve to [`Poll::Ready`]. Getting - // [`Poll::Pending`] would be considered a programming error. - // - // Alternatively, we could just use `futures::executor::block_on` - // which is probably cleaner but a single-shot poll is more - // enforcing of the concept we want to represent. - match poll_once(func(&mut txn)) { - Poll::Ready(Ok(value)) => { - *callback_slot.lock().unwrap() = Some(Ok(value)); - Ok(py.None()) + // Since we expect people to only call `.await` on [`Transaction`] + // related methods (mentioned in the [`Transaction`] docstring) AND + // because there is no async work to suspend on in the Python + // [`Transaction`] synchronously, we can get away with polling once + // as it should immediately resolve to [`Poll::Ready`]. Getting + // [`Poll::Pending`] would be considered a programming error. + // + // Alternatively, we could just use `futures::executor::block_on` + // which is probably cleaner but a single-shot poll is more + // enforcing of the concept we want to represent. + match poll_once(func(&mut txn)) { + Poll::Ready(Ok(value)) => { + *callback_slot.lock().unwrap() = Some(Ok(value)); + Ok(py.None()) + } + Poll::Ready(Err(err)) => { + // Re-raise into Python so `runInteraction` rolls the + // transaction back (and can apply its retry logic for + // serialization/deadlock errors). + let py_err = anyhow_to_pyerr(&err); + *callback_slot.lock().unwrap() = Some(Err(err)); + Err(py_err) + } + Poll::Pending => unreachable!( + "The `run_interaction` transaction callback future returned `Poll::Pending`, \ + but we expect Synapse Python database work to resolve synchronously. \ + This is a Synapse programming error: genuine async work is \ + not supported here.", + ), } - Poll::Ready(Err(err)) => { - // Re-raise into Python so `runInteraction` rolls the - // transaction back (and can apply its retry logic for - // serialization/deadlock errors). - let py_err = anyhow_to_pyerr(&err); - *callback_slot.lock().unwrap() = Some(Err(err)); - Err(py_err) - } - Poll::Pending => unreachable!( - "The `run_interaction` transaction callback future returned `Poll::Pending`, \ - but we expect Synapse Python database work to resolve synchronously. \ - This is a Synapse programming error: genuine async work is \ - not supported here.", - ), - } + }, + )? + .unbind(); + + Ok(( + callback, + self.database_pool_py.clone_ref(py), + self.reactor.clone_ref(py), + )) + }) + .map_err(anyhow::Error::from)?; + + // Use `runInteraction` directly + let outcome = run_python_awaitable(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)), }, - )? - .unbind(); - - Ok(( - callback, - self.database_pool_py.clone_ref(py), - self.reactor.clone_ref(py), - )) - }) - .map_err(anyhow::Error::from)?; - - // Use `runInteraction` directly - let outcome = run_python_awaitable(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)), - }, + } } + .boxed() } } diff --git a/rust/src/storage/db/rust_db_pool.rs b/rust/src/storage/db/rust_db_pool.rs index 94d701643f..8fb13b658e 100644 --- a/rust/src/storage/db/rust_db_pool.rs +++ b/rust/src/storage/db/rust_db_pool.rs @@ -20,9 +20,12 @@ use anyhow::Context; use bb8_postgres::tokio_postgres::{self, types::ToSql, IsolationLevel}; use bb8_postgres::PostgresConnectionManager; use futures::future::BoxFuture; +use futures::FutureExt; use postgres_native_tls::MakeTlsConnector; -use crate::storage::db::{DatabasePool, DbValue, Row, Transaction}; +use crate::storage::db::{ + DatabasePool, DbValue, ErasedInteraction, ErasedResult, Row, Transaction, +}; /// Native Rust database access backed by `tokio-postgres` (for use in synapse-rust-apps) pub struct RustDatabasePool { @@ -30,52 +33,52 @@ pub struct RustDatabasePool { } impl DatabasePool for RustDatabasePool { - async fn run_interaction(&self, _name: &'static str, func: F) -> anyhow::Result - where - R: Send + 'static, - F: for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, anyhow::Result> - + Send - + Sync - + 'static, - { - // Like Synapse's `runInteraction`, retry the whole transaction on - // serialization/deadlock errors (which can happen under repeatable-read). - loop { - let mut conn = self - .db_pool - .get() - .await - .context("Failed to acquire database connection")?; + fn run_interaction_dyn<'a>( + &'a self, + _name: &'static str, + func: ErasedInteraction, + ) -> BoxFuture<'a, ErasedResult> { + async move { + // Like Synapse's `runInteraction`, retry the whole transaction on + // serialization/deadlock errors (which can happen under repeatable-read). + loop { + let mut conn = self + .db_pool + .get() + .await + .context("Failed to acquire database connection")?; - // Repeatable-read isolation level (like Synapse). - let txn = conn - .build_transaction() - .isolation_level(IsolationLevel::RepeatableRead) - .start() - .await - .context("Failed to start transaction")?; + // Repeatable-read isolation level (like Synapse). + let txn = conn + .build_transaction() + .isolation_level(IsolationLevel::RepeatableRead) + .start() + .await + .context("Failed to start transaction")?; - let mut wrapper = TokioPostgresTransaction { txn }; - match func(&mut wrapper).await { - Ok(value) => { - wrapper - .txn - .commit() - .await - .context("Failed to commit transaction")?; - return Ok(value); - } - Err(err) => { - // The transaction is rolled back implicitly when dropped, but - // be explicit about it before deciding whether to retry. - let _ = wrapper.txn.rollback().await; - if is_retryable(&err) { - continue; + let mut wrapper = TokioPostgresTransaction { txn }; + match func(&mut wrapper).await { + Ok(value) => { + wrapper + .txn + .commit() + .await + .context("Failed to commit transaction")?; + return Ok(value); + } + Err(err) => { + // The transaction is rolled back implicitly when dropped, but + // be explicit about it before deciding whether to retry. + let _ = wrapper.txn.rollback().await; + if is_retryable(&err) { + continue; + } + return Err(err); } - return Err(err); } } } + .boxed() } } diff --git a/rust/src/storage/store.rs b/rust/src/storage/store.rs index be6a577467..f2bf8db3fa 100644 --- a/rust/src/storage/store.rs +++ b/rust/src/storage/store.rs @@ -16,7 +16,7 @@ use futures::FutureExt; use serde::Serialize; -use crate::storage::db::{DatabasePool, RowExt}; +use crate::storage::db::{DatabasePool, DatabasePoolExt, RowExt}; /// Currently supported per-user features #[derive(Serialize, Debug)] @@ -46,11 +46,11 @@ impl std::fmt::Display for PerUserExperimentalFeature { } } -pub struct Store { - pub db_pool: P, +pub struct Store { + pub db_pool: Box, } -impl Store

{ +impl Store { /// Checks whether a given feature is enabled/disabled for this user /// /// If there is no entry, returns None