diff --git a/rust/src/handlers/mod.rs b/rust/src/handlers/mod.rs index e42952ec1a..e5243f87e1 100644 --- a/rust/src/handlers/mod.rs +++ b/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( diff --git a/rust/src/handlers/versions.rs b/rust/src/handlers/versions.rs index d929759939..78a80116e1 100644 --- a/rust/src/handlers/versions.rs +++ b/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, + pub store: Arc>, /// The Twisted reactor, used to bridge our `async` response back into a /// Twisted deferred that Python can `await`. pub reactor: Py, @@ -77,7 +78,7 @@ impl VersionsHandler { /// Assemble a `/versions` response body. async fn build_versions_response( - store: &Store, + store: &Store, config: &SynapseConfig, user_id: Option<&str>, ) -> Result { diff --git a/rust/src/http_client.rs b/rust/src/http_client.rs index ecd5c7b55e..be3bf1c1bf 100644 --- a/rust/src/http_client.rs +++ b/rust/src/http_client.rs @@ -12,14 +12,22 @@ * . */ -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(reactor: Py, make_deferred: F) -> PyResult> +where + F: for<'py> Fn(Python<'py>) -> PyResult> + Send + 'static, +{ + // Resolves when the deferred fires; carries the resolved value or the error. + let (tx, rx) = oneshot::channel::>>(); + // 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> { + 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> { + 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> { + 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(|_| "".to_owned()), + ), + } +} + static MAKE_DEFERRED_YIELDABLE: OnceLock> = OnceLock::new(); /// Given a deferred, make it follow the Synapse logcontext rules diff --git a/rust/src/storage/db/mod.rs b/rust/src/storage/db/mod.rs index 804c8ac937..bdbf43a18e 100644 --- a/rust/src/storage/db/mod.rs +++ b/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; -/// 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`) 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>; - -/// 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 Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, AnyResult> + Send>; - /// A database connection pool. /// -/// We use a `Box` 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`) 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( &self, name: &'static str, @@ -94,29 +53,9 @@ pub trait DatabasePoolExt: DatabasePool { F: for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, anyhow::Result> + Send + Sync - + 'static, - { - // Erase the concrete return type `R` into `Box` 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) - }) - }); - - async move { - let boxed = self.run_interaction_erased(name, erased).await?; - Ok(*boxed - .downcast::() - .expect("run_interaction return type mismatch (this is a Synapse programming error)")) - } - } + + 'static; } -impl 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] diff --git a/rust/src/storage/db/python_db_pool.rs b/rust/src/storage/db/python_db_pool.rs index 1939fe3749..dcd537c3d2 100644 --- a/rust/src/storage/db/python_db_pool.rs +++ b/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::(); - + 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. - let result_slot: Arc>> = 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>>> = 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::()?; - 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> { - 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> { - 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(|_| "".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> { - 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)), + }, + } } } diff --git a/rust/src/storage/db/rust_db_pool.rs b/rust/src/storage/db/rust_db_pool.rs index addddd489d..284bd18924 100644 --- a/rust/src/storage/db/rust_db_pool.rs +++ b/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>, } -#[async_trait::async_trait] impl DatabasePool for RustDatabasePool { - async fn run_interaction_erased(&self, _name: &'static str, func: InteractionFn) -> AnyResult { + 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 { diff --git a/rust/src/storage/store.rs b/rust/src/storage/store.rs index 76d092efc0..9d1d3f2534 100644 --- a/rust/src/storage/store.rs +++ b/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 { pub config: SynapseConfig, - pub db_pool: Box, + pub db_pool: P, } -impl Store { +impl Store

{ pub async fn is_feature_enabled( &self, user_id: &str,