dandri/synapse
1
0
Fork 0
mirror of https://github.com/element-hq/synapse.git synced 2026-07-21 11:11:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

LLM attempt at switching back to Store with dyn DatabasePool

Browse Source
This commit is contained in:
Eric Eastwood
2026-06-24 14:34:47 -05:00
parent ce68c0f73b
commit 2860e4ea01
6 changed files with 240 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/handlers/mod.rs
+3 -1
View File
@@ -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),
rust/src/handlers/versions.rs
+2 -3
View File
@@ -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<UnstableFeatureMap>,
pub store: Arc<Store<PythonDatabasePoolWrapper>>,
pub store: Arc<Store>,
/// The Twisted reactor, used to bridge our `async` response back into a
/// Twisted deferred that Python can `await`.
pub reactor: Py<PyAny>,
@@ -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<PythonDatabasePoolWrapper>,
store: &Store,
global_unstable_feature_map: &UnstableFeatureMap,
user_id: Option<&str>,
) -> Result<VersionsResponse, anyhow::Error> {
rust/src/storage/db/mod.rs
+90 -3
View File
@@ -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<dyn Any + Send>` 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<dyn for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, ErasedResult> + Send + Sync>;
/// The type-erased result of an [`ErasedInteraction`]: the concrete `R` boxed up
/// as `Box<dyn Any + Send>` so it can pass through the object-safe
/// [`DatabasePool::run_interaction_dyn`].
pub type ErasedResult = anyhow::Result<Box<dyn Any + Send>>;
/// 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<P: DatabasePool>`) rather than using dynamic dispatch.
/// is held behind a trait object (e.g. `Box<dyn DatabasePool>`), 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<R>>
+ Send
+ Sync
+ 'static;
+ 'static,
{
async move {
// Box the concrete result `R` up as `Box<dyn Any + Send>` 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::<R>().expect(
"run_interaction_dyn returned a value of an unexpected type; \
this is a programming error in the DatabasePool implementation",
))
}
}
}
impl<T: DatabasePool + ?Sized> DatabasePoolExt for T {}
/// Wrap a strongly-typed transaction callback so its result is boxed as
/// `Box<dyn Any + Send>`, producing the type-erased [`ErasedInteraction`] that
/// [`DatabasePool::run_interaction_dyn`] accepts.
fn erase_interaction<R, F>(func: F) -> ErasedInteraction
where
R: Send + 'static,
F: for<'txn> Fn(&'txn mut dyn Transaction) -> BoxFuture<'txn, anyhow::Result<R>>
+ 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<dyn Fn..>` doesn't reliably
// infer that lifetime.
fn constrain<C>(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<dyn Any + Send>) }.boxed()
}))
}
/// A [`tokio_postgres::Transaction`] looking thing that we can use on the Rust side to
rust/src/storage/db/python_db_pool.rs
+97 -93
View File
@@ -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<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.
//
// 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));
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<Mutex<Option<ErasedResult>>> = 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<Py<PyAny>> {
let py = args.py();
let txn_py = args.get_item(0)?;
let mut txn = txn_py.extract::<LoggingTransactionWrapper>()?;
// 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<Py<PyAny>> {
let py = args.py();
let txn_py = args.get_item(0)?;
let mut txn = txn_py.extract::<LoggingTransactionWrapper>()?;
// 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()
}
}
rust/src/storage/db/rust_db_pool.rs
+44 -41
View File
@@ -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<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 {
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()
}
}
rust/src/storage/store.rs
+4 -4
View File
@@ -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<P: DatabasePool> {
pub db_pool: P,
pub struct Store {
pub db_pool: Box<dyn DatabasePool>,
}
impl<P: DatabasePool> Store<P> {
impl Store {
/// Checks whether a given feature is enabled/disabled for this user
///
/// If there is no entry, returns None