Concept: Rust database access via Python database connection pool (via runInteraction) (LLM splat v1)

changelog.d/19846.misc

@@ -0,0 +1 @@
+Allow Rust code to have database access via Python database connection pool.

rust/Cargo.toml

@@ -22,6 +22,7 @@ crate-type = ["lib", "cdylib"]
 name = "synapse.synapse_rust"
 
 [dependencies]
+async-trait = "0.1.89"
 anyhow = "1.0.63"
 base64 = "0.22.1"
 bytes = "1.6.0"
@@ -64,6 +65,12 @@ tokio = { version = "1.44.2", features = ["rt", "rt-multi-thread"] }
 once_cell = "1.18.0"
 itertools = "0.14.0"
 
+# TODO: Remove: These are just used to make sure a tokio-postgres backed database pool makes sense
+# with our interfaces
+bb8 = "0.8.3"
+bb8-postgres = "0.8.1"
+postgres-native-tls = "0.5.0"
+
 [build-dependencies]
 blake2 = "0.10.4"
 hex = "0.4.3"

rust/src/config/mod.rs

@@ -0,0 +1,103 @@
+/*
+ * This file is licensed under the Affero General Public License (AGPL) version 3.
+ *
+ * Copyright (C) 2026 Element Creations Ltd
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * See the GNU Affero General Public License for more details:
+ * <https://www.gnu.org/licenses/agpl-3.0.html>.
+ *
+ */
+
+use std::collections::BTreeSet;
+use std::str::FromStr;
+
+use pyo3::{exceptions::PyRuntimeError, prelude::*};
+
+/// A Rust-side view of Synapse's Python `HomeServerConfig`.
+///
+/// This only mirrors the subset of config that the Rust handlers need, rather
+/// than the whole thing. Thanks to `#[derive(FromPyObject)]`, each field is
+/// pulled directly off the corresponding attribute of the Python `config`
+/// object, so you can populate it in one shot with
+/// `homeserver.getattr("config")?.extract()?`.
+#[derive(FromPyObject, Clone)]
+pub struct SynapseHomeServerConfig {
+    pub room: RoomConfig,
+    pub auth: AuthConfig,
+    pub server: ServerConfig,
+    pub experimental: ExperimentalConfig,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub enum RoomCreationPreset {
+    PrivateChat,
+    PublicChat,
+    TrustedPrivateChat,
+}
+
+impl FromStr for RoomCreationPreset {
+    type Err = PyErr;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        Ok(match s {
+            "private_chat" => RoomCreationPreset::PrivateChat,
+            "public_chat" => RoomCreationPreset::PublicChat,
+            "trusted_private_chat" => RoomCreationPreset::TrustedPrivateChat,
+            other => {
+                return Err(PyRuntimeError::new_err(format!(
+                    "Unknown variant {other:?} does not translate to `RoomCreationPreset`. \
+                     This is a Synapse programming error."
+                )))
+            }
+        })
+    }
+}
+
+impl<'a, 'py> FromPyObject<'a, 'py> for RoomCreationPreset {
+    type Error = PyErr;
+
+    fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
+        value.extract::<&str>()?.parse()
+    }
+}
+
+#[derive(FromPyObject, Clone)]
+pub struct RoomConfig {
+    pub encryption_enabled_by_default_for_room_presets: BTreeSet<RoomCreationPreset>,
+}
+
+#[derive(FromPyObject, Clone)]
+pub struct AuthConfig {
+    pub login_via_existing_enabled: bool,
+}
+#[derive(FromPyObject, Clone)]
+pub struct ServerConfig {
+    pub max_event_delay_ms: Option<u64>,
+}
+
+#[derive(FromPyObject, Clone)]
+pub struct ExperimentalConfig {
+    pub msc3026_enabled: bool,
+    pub msc3773_enabled: bool,
+    pub msc2815_enabled: bool,
+    pub msc3881_enabled: bool,
+    pub msc3874_enabled: bool,
+    pub msc3912_enabled: bool,
+    pub msc3391_enabled: bool,
+    pub msc4069_profile_inhibit_propagation: bool,
+    pub msc4028_push_encrypted_events: bool,
+    pub msc4108_enabled: bool,
+    pub msc4108_delegation_endpoint: Option<String>,
+    pub msc3575_enabled: bool,
+    pub msc4133_enabled: bool,
+    pub msc4155_enabled: bool,
+    pub msc4306_enabled: bool,
+    pub msc4169_enabled: bool,
+    pub msc4354_enabled: bool,
+    pub msc4222_enabled: bool,
+}

rust/src/deferred.rs

@@ -13,10 +13,17 @@
  *
  */
 
-use std::{future::Future, sync::OnceLock};
+use std::{
+    future::Future,
+    sync::{Arc, Mutex},
+};
 
 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 tokio::sync::oneshot;
 
 use crate::tokio_runtime::runtime;
 
@@ -51,10 +58,46 @@ fn defer(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> {
         .bind(py))
 }
 
+/// A reference to the `synapse.logging.context` module.
+static LOGGING_CONTEXT_MODULE: OnceCell<Py<PyAny>> = OnceCell::new();
+
+/// Access to the `synapse.logging.context` module.
+fn logging_context_module(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> {
+    Ok(LOGGING_CONTEXT_MODULE
+        .get_or_try_init(|| py.import("synapse.logging.context").map(Into::into))?
+        .bind(py))
+}
+
+/// Set the Synapse logcontext active on the current (reactor) thread, returning
+/// the context that was previously active so it can be restored afterwards.
+fn set_current_logging_context<'py>(
+    py: Python<'py>,
+    context: &Bound<'py, PyAny>,
+) -> PyResult<Bound<'py, PyAny>> {
+    logging_context_module(py)?.call_method1(intern!(py, "set_current_context"), (context,))
+}
+
+tokio::task_local! {
+    /// The Synapse `LoggingContext` that was active on the reactor thread when
+    /// the Rust future was created (see [`create_deferred`]).
+    ///
+    /// Synapse attributes per-request CPU/DB usage via a thread-local
+    /// `LoggingContext`, but our work runs on a Tokio worker thread that is
+    /// detached from it. We stash the originating context here so that whenever
+    /// we hop back onto the reactor thread to drive Python work (see
+    /// [`run_python_awaitable`]) we can re-activate it; otherwise that work runs in
+    /// the sentinel context and its resource usage (e.g. database time) is lost
+    /// instead of being charged to the request.
+    static LOGGING_CONTEXT: Py<PyAny>;
+}
+
 /// Creates a twisted deferred from the given future, spawning the task on the
 /// tokio runtime.
 ///
-/// Does not handle deferred cancellation or contextvars.
+/// Captures the Synapse logcontext active on the reactor thread (so work driven
+/// by the future is attributed to the originating request, see
+/// [`LOGGING_CONTEXT`]), but does not handle deferred cancellation or
+/// contextvars.
 pub fn create_deferred<'py, F, O>(
     py: Python<'py>,
     reactor: &Bound<'py, PyAny>,
@@ -68,9 +111,16 @@ where
     let deferred_callback = deferred.getattr("callback")?.unbind();
     let deferred_errback = deferred.getattr("errback")?.unbind();
 
+    // Capture the logcontext active on the reactor thread now, while we're still on it,
+    // so the spawned task can keep track of what logcontext its associated with and can
+    // re-apply it when it hops back to drive Python work.
+    let logging_context = logging_context_module(py)?
+        .call_method0(intern!(py, "current_context"))?
+        .unbind();
+
     let rt = runtime(reactor)?;
     let handle = rt.handle()?;
-    let task = handle.spawn(fut);
+    let task = handle.spawn(LOGGING_CONTEXT.scope(logging_context, fut));
 
     // Unbind the reactor so that we can pass it to the task
     let reactor = reactor.clone().unbind();
@@ -110,18 +160,171 @@ where
     make_deferred_yieldable(py, &deferred)
 }
 
-static MAKE_DEFERRED_YIELDABLE: OnceLock<pyo3::Py<pyo3::PyAny>> = OnceLock::new();
+/// Runs a Python awaitable to completion on the Twisted reactor and resolves
+/// with its result.
+///
+/// This is the inverse of [`create_deferred`]: where that turns a Rust future
+/// into a Twisted `Deferred`, this turns a Python awaitable into a Rust future.
+///
+/// Despite returning a future, the awaitable is kicked off in the background running in
+/// the Twisted reactor and runs to completion regardless of whether the returned Rust
+/// future is ever polled; awaiting it only observes the result.
+///
+/// Calls `make_awaitable` with the current logcontext associated with this async Rust
+/// task. Follows the Synapse logcontext rules.
+pub(crate) async fn run_python_awaitable<F>(
+    reactor: Py<PyAny>,
+    make_awaitable: F,
+) -> PyResult<Py<PyAny>>
+where
+    F: for<'py> Fn(Python<'py>) -> PyResult<Bound<'py, PyAny>> + Send + 'static,
+{
+    // Resolves when the awaitable completes; carries the resolved value or 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)));
+
+    // The logcontext that is active for this async task
+    let current_rust_task_logging_context = LOGGING_CONTEXT
+        .try_with(|ctx| Python::attach(|py| ctx.clone_ref(py)))
+        .ok();
+
+    Python::attach(|py| -> PyResult<()> {
+        // Create some deferred success/error callback functions that we will use to get
+        // the result from Python to Rust.
+        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();
+
+        // Wrap `make_awaitable` as a Python callable so we can hand it to
+        // `run_in_background`, which calls it (in the active logcontext) to produce
+        // the awaitable it then drives.
+        let awaitable_factory = PyCFunction::new_closure(
+            py,
+            None,
+            None,
+            move |args, _kwargs| -> PyResult<Py<PyAny>> {
+                let py = args.py();
+                Ok(make_awaitable(py)?.unbind())
+            },
+        )?
+        .unbind();
+
+        // Create a function that we will run with the Twisted reactor that will drive
+        // the Python awaitable.
+        let starter = PyCFunction::new_closure(
+            py,
+            None,
+            None,
+            move |args, _kwargs| -> PyResult<Py<PyAny>> {
+                let py = args.py();
+
+                // Activate the captured Rust task-local logging context. This way we
+                // properly record metrics/logging for the thing being run.
+                let calling_logging_context = current_rust_task_logging_context
+                    .as_ref()
+                    .map(|ctx| set_current_logging_context(py, ctx.bind(py)))
+                    .transpose()?
+                    .expect("No `LoggingContext` returned from `set_current_logging_context(...)`. This is a Synapse programming error.");
+
+                // We fire-and-forget using `run_in_background`. Re-using
+                // `run_in_background` also makes sure the awaitable gets run with the
+                // current logcontext (the one we just activated) while following the
+                // logcontext rules.
+                let deferred = logging_context_module(py)?.call_method1(
+                    intern!(py, "run_in_background"),
+                    (awaitable_factory.bind(py),),
+                );
+
+                // Restore the `calling_logging_context` after we kick off the
+                // background task.
+                //
+                // Our goal is to have the caller logcontext unchanged after firing off
+                // the background task and returning.
+                set_current_logging_context(py, &calling_logging_context)?;
+
+                let deferred = deferred?;
+                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(
+            "run_python_awaitable channel closed before the awaitable completed",
+        )),
+    }
+}
+
+/// Convert a Twisted `Failure` (as passed to an Deferred errback) into a [`PyErr`].
+///
+/// A Twisted `Failure` carries the original exception instance in its `.value`
+/// attribute, which we re-raise so callers see the real error. If the `Failure` is
+/// mangled, we fallback to raising a generic [`PyRuntimeError`] explaining what we saw
+/// instead.
+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(format!(
+            "Expected Python object passed here to be a Twisted `Failure` with a `value` attribute \
+            but saw something else: {}",
+            failure
+                .str()
+                .map(|s| s.to_string_lossy().into_owned())
+                .unwrap_or_else(|_| "<failed to stringify Python object>".to_owned()),
+        )),
+    }
+}
+
+static MAKE_DEFERRED_YIELDABLE: OnceCell<Py<PyAny>> = OnceCell::new();
 
 /// Given a deferred, make it follow the Synapse logcontext rules
 fn make_deferred_yieldable<'py>(
     py: Python<'py>,
     deferred: &Bound<'py, PyAny>,
 ) -> PyResult<Bound<'py, PyAny>> {
-    let make_deferred_yieldable = MAKE_DEFERRED_YIELDABLE.get_or_init(|| {
-        let sys = PyModule::import(py, "synapse.logging.context").unwrap();
-        let func = sys.getattr("make_deferred_yieldable").unwrap().unbind();
-        func
-    });
+    let make_deferred_yieldable = MAKE_DEFERRED_YIELDABLE.get_or_try_init(|| {
+        logging_context_module(py)?
+            .getattr("make_deferred_yieldable")
+            .map(Into::into)
+    })?;
 
     make_deferred_yieldable
         .call1(py, (deferred,))?
@@ -133,6 +336,8 @@ fn make_deferred_yieldable<'py>(
 pub fn register_module(py: Python<'_>, _m: &Bound<'_, PyModule>) -> PyResult<()> {
     // Make sure we fail early if we can't load some modules
     defer(py)?;
+    // We can't check this here because of circular import issues
+    // logging_context_module(py)?;
 
     Ok(())
 }