Split out deferred and tokio_runtime to their own Rust modules

changelog.d/19868.misc

@@ -0,0 +1 @@
+Split out `deferred` and `tokio_runtime` to their own Rust modules.

rust/src/deferred.rs

@@ -0,0 +1,138 @@
+/*
+ * 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::{future::Future, sync::OnceLock};
+
+use once_cell::sync::OnceCell;
+use pyo3::{create_exception, exceptions::PyException, prelude::*};
+
+use crate::tokio_runtime::runtime;
+
+create_exception!(
+    synapse.synapse_rust.http_client,
+    RustPanicError,
+    PyException,
+    "A panic which happened in a Rust future"
+);
+
+impl RustPanicError {
+    fn from_panic(panic_err: &(dyn std::any::Any + Send + 'static)) -> PyErr {
+        // Apparently this is how you extract the panic message from a panic
+        let panic_message = if let Some(str_slice) = panic_err.downcast_ref::<&str>() {
+            str_slice
+        } else if let Some(string) = panic_err.downcast_ref::<String>() {
+            string
+        } else {
+            "unknown error"
+        };
+        Self::new_err(panic_message.to_owned())
+    }
+}
+
+/// A reference to the `twisted.internet.defer` module.
+static DEFER: OnceCell<Py<PyAny>> = OnceCell::new();
+
+/// Access to the `twisted.internet.defer` module.
+fn defer(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> {
+    Ok(DEFER
+        .get_or_try_init(|| py.import("twisted.internet.defer").map(Into::into))?
+        .bind(py))
+}
+
+/// Creates a twisted deferred from the given future, spawning the task on the
+/// tokio runtime.
+///
+/// Does not handle deferred cancellation or contextvars.
+pub fn create_deferred<'py, F, O>(
+    py: Python<'py>,
+    reactor: &Bound<'py, PyAny>,
+    fut: F,
+) -> PyResult<Bound<'py, PyAny>>
+where
+    F: Future<Output = PyResult<O>> + Send + 'static,
+    for<'a> O: IntoPyObject<'a> + Send + 'static,
+{
+    let deferred = defer(py)?.call_method0("Deferred")?;
+    let deferred_callback = deferred.getattr("callback")?.unbind();
+    let deferred_errback = deferred.getattr("errback")?.unbind();
+
+    let rt = runtime(reactor)?;
+    let handle = rt.handle()?;
+    let task = handle.spawn(fut);
+
+    // Unbind the reactor so that we can pass it to the task
+    let reactor = reactor.clone().unbind();
+    handle.spawn(async move {
+        let res = task.await;
+
+        Python::attach(move |py| {
+            // Flatten the panic into standard python error
+            let res = match res {
+                Ok(r) => r,
+                Err(join_err) => match join_err.try_into_panic() {
+                    Ok(panic_err) => Err(RustPanicError::from_panic(&panic_err)),
+                    Err(err) => Err(PyException::new_err(format!("Task cancelled: {err}"))),
+                },
+            };
+
+            // Re-bind the reactor
+            let reactor = reactor.bind(py);
+
+            // Send the result to the deferred, via `.callback(..)` or `.errback(..)`
+            match res {
+                Ok(obj) => {
+                    reactor
+                        .call_method("callFromThread", (deferred_callback, obj), None)
+                        .expect("callFromThread should not fail"); // There's nothing we can really do with errors here
+                }
+                Err(err) => {
+                    reactor
+                        .call_method("callFromThread", (deferred_errback, err), None)
+                        .expect("callFromThread should not fail"); // There's nothing we can really do with errors here
+                }
+            }
+        });
+    });
+
+    // Make the deferred follow the Synapse logcontext rules
+    make_deferred_yieldable(py, &deferred)
+}
+
+static MAKE_DEFERRED_YIELDABLE: OnceLock<pyo3::Py<pyo3::PyAny>> = OnceLock::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
+    });
+
+    make_deferred_yieldable
+        .call1(py, (deferred,))?
+        .extract(py)
+        .map_err(Into::into)
+}
+
+/// Called when registering modules with python.
+pub fn register_module(py: Python<'_>, _m: &Bound<'_, PyModule>) -> PyResult<()> {
+    // Make sure we fail early if we can't load some modules
+    defer(py)?;
+
+    Ok(())
+}

rust/src/http_client.rs

@@ -12,145 +12,22 @@
  * <https://www.gnu.org/licenses/agpl-3.0.html>.
  */
 
-use std::{collections::HashMap, future::Future, sync::OnceLock};
+use std::collections::HashMap;
 
 use anyhow::Context;
 use http_body_util::BodyExt;
-use once_cell::sync::OnceCell;
-use pyo3::{create_exception, exceptions::PyException, prelude::*};
+use pyo3::prelude::*;
 use reqwest::RequestBuilder;
-use tokio::runtime::Runtime;
 
+use crate::deferred::create_deferred;
 use crate::errors::HttpResponseException;
-
-create_exception!(
-    synapse.synapse_rust.http_client,
-    RustPanicError,
-    PyException,
-    "A panic which happened in a Rust future"
-);
-
-impl RustPanicError {
-    fn from_panic(panic_err: &(dyn std::any::Any + Send + 'static)) -> PyErr {
-        // Apparently this is how you extract the panic message from a panic
-        let panic_message = if let Some(str_slice) = panic_err.downcast_ref::<&str>() {
-            str_slice
-        } else if let Some(string) = panic_err.downcast_ref::<String>() {
-            string
-        } else {
-            "unknown error"
-        };
-        Self::new_err(panic_message.to_owned())
-    }
-}
-
-/// This is the name of the attribute where we store the runtime on the reactor
-static TOKIO_RUNTIME_ATTR: &str = "__synapse_rust_tokio_runtime";
-
-/// A Python wrapper around a Tokio runtime.
-///
-/// This allows us to 'store' the runtime on the reactor instance, starting it
-/// when the reactor starts, and stopping it when the reactor shuts down.
-#[pyclass]
-struct PyTokioRuntime {
-    runtime: Option<Runtime>,
-}
-
-#[pymethods]
-impl PyTokioRuntime {
-    fn start(&mut self) -> PyResult<()> {
-        // TODO: allow customization of the runtime like the number of threads
-        let runtime = tokio::runtime::Builder::new_multi_thread()
-            .worker_threads(4)
-            .enable_all()
-            .build()?;
-
-        self.runtime = Some(runtime);
-
-        Ok(())
-    }
-
-    fn shutdown(&mut self) -> PyResult<()> {
-        let runtime = self
-            .runtime
-            .take()
-            .context("Runtime was already shutdown")?;
-
-        // Dropping the runtime will shut it down
-        drop(runtime);
-
-        Ok(())
-    }
-}
-
-impl PyTokioRuntime {
-    /// Get the handle to the Tokio runtime, if it is running.
-    fn handle(&self) -> PyResult<&tokio::runtime::Handle> {
-        let handle = self
-            .runtime
-            .as_ref()
-            .context("Tokio runtime is not running")?
-            .handle();
-
-        Ok(handle)
-    }
-}
-
-/// Get a handle to the Tokio runtime stored on the reactor instance, or create
-/// a new one.
-fn runtime<'a>(reactor: &Bound<'a, PyAny>) -> PyResult<PyRef<'a, PyTokioRuntime>> {
-    if !reactor.hasattr(TOKIO_RUNTIME_ATTR)? {
-        install_runtime(reactor)?;
-    }
-
-    get_runtime(reactor)
-}
-
-/// Install a new Tokio runtime on the reactor instance.
-fn install_runtime(reactor: &Bound<PyAny>) -> PyResult<()> {
-    let py = reactor.py();
-    let runtime = PyTokioRuntime { runtime: None };
-    let runtime = runtime.into_pyobject(py)?;
-
-    // Attach the runtime to the reactor, starting it when the reactor is
-    // running, stopping it when the reactor is shutting down
-    reactor.call_method1("callWhenRunning", (runtime.getattr("start")?,))?;
-    reactor.call_method1(
-        "addSystemEventTrigger",
-        ("after", "shutdown", runtime.getattr("shutdown")?),
-    )?;
-    reactor.setattr(TOKIO_RUNTIME_ATTR, runtime)?;
-
-    Ok(())
-}
-
-/// Get a reference to a Tokio runtime handle stored on the reactor instance.
-fn get_runtime<'a>(reactor: &Bound<'a, PyAny>) -> PyResult<PyRef<'a, PyTokioRuntime>> {
-    // This will raise if `TOKIO_RUNTIME_ATTR` is not set or if it is
-    // not a `Runtime`. Careful that this could happen if the user sets it
-    // manually, or if multiple versions of `pyo3-twisted` are used!
-    let runtime: Bound<PyTokioRuntime> = reactor.getattr(TOKIO_RUNTIME_ATTR)?.extract()?;
-    Ok(runtime.borrow())
-}
-
-/// A reference to the `twisted.internet.defer` module.
-static DEFER: OnceCell<Py<PyAny>> = OnceCell::new();
-
-/// Access to the `twisted.internet.defer` module.
-fn defer(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> {
-    Ok(DEFER
-        .get_or_try_init(|| py.import("twisted.internet.defer").map(Into::into))?
-        .bind(py))
-}
+use crate::tokio_runtime::runtime;
 
 /// Called when registering modules with python.
 pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     let child_module: Bound<'_, PyModule> = PyModule::new(py, "http_client")?;
     child_module.add_class::<HttpClient>()?;
 
-    // Make sure we fail early if we can't load some modules
-    defer(py)?;
-
     m.add_submodule(&child_module)?;
 
     // We need to manually add the module to sys.modules to make `from
@@ -262,81 +139,3 @@ impl HttpClient {
         })
     }
 }
-
-/// Creates a twisted deferred from the given future, spawning the task on the
-/// tokio runtime.
-///
-/// Does not handle deferred cancellation or contextvars.
-fn create_deferred<'py, F, O>(
-    py: Python<'py>,
-    reactor: &Bound<'py, PyAny>,
-    fut: F,
-) -> PyResult<Bound<'py, PyAny>>
-where
-    F: Future<Output = PyResult<O>> + Send + 'static,
-    for<'a> O: IntoPyObject<'a> + Send + 'static,
-{
-    let deferred = defer(py)?.call_method0("Deferred")?;
-    let deferred_callback = deferred.getattr("callback")?.unbind();
-    let deferred_errback = deferred.getattr("errback")?.unbind();
-
-    let rt = runtime(reactor)?;
-    let handle = rt.handle()?;
-    let task = handle.spawn(fut);
-
-    // Unbind the reactor so that we can pass it to the task
-    let reactor = reactor.clone().unbind();
-    handle.spawn(async move {
-        let res = task.await;
-
-        Python::attach(move |py| {
-            // Flatten the panic into standard python error
-            let res = match res {
-                Ok(r) => r,
-                Err(join_err) => match join_err.try_into_panic() {
-                    Ok(panic_err) => Err(RustPanicError::from_panic(&panic_err)),
-                    Err(err) => Err(PyException::new_err(format!("Task cancelled: {err}"))),
-                },
-            };
-
-            // Re-bind the reactor
-            let reactor = reactor.bind(py);
-
-            // Send the result to the deferred, via `.callback(..)` or `.errback(..)`
-            match res {
-                Ok(obj) => {
-                    reactor
-                        .call_method("callFromThread", (deferred_callback, obj), None)
-                        .expect("callFromThread should not fail"); // There's nothing we can really do with errors here
-                }
-                Err(err) => {
-                    reactor
-                        .call_method("callFromThread", (deferred_errback, err), None)
-                        .expect("callFromThread should not fail"); // There's nothing we can really do with errors here
-                }
-            }
-        });
-    });
-
-    // Make the deferred follow the Synapse logcontext rules
-    make_deferred_yieldable(py, &deferred)
-}
-
-static MAKE_DEFERRED_YIELDABLE: OnceLock<pyo3::Py<pyo3::PyAny>> = OnceLock::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
-    });
-
-    make_deferred_yieldable
-        .call1(py, (deferred,))?
-        .extract(py)
-        .map_err(Into::into)
-}

rust/src/lib.rs

@@ -6,6 +6,7 @@ use pyo3_log::ResetHandle;
 
 pub mod acl;
 pub mod canonical_json;
+pub mod deferred;
 pub mod duration;
 pub mod errors;
 pub mod events;
@@ -19,6 +20,7 @@ pub mod push;
 pub mod rendezvous;
 pub mod room_versions;
 pub mod segmenter;
+pub mod tokio_runtime;
 pub mod types;
 
 lazy_static! {
@@ -65,6 +67,7 @@ fn synapse_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(reset_logging_config, m)?)?;
 
     acl::register_module(py, m)?;
+    deferred::register_module(py, m)?;
     push::register_module(py, m)?;
     events::register_module(py, m)?;
     http_client::register_module(py, m)?;