diff --git a/changelog.d/19868.misc b/changelog.d/19868.misc new file mode 100644 index 00000000000..cd182ea15cb --- /dev/null +++ b/changelog.d/19868.misc @@ -0,0 +1 @@ +Split out `deferred` and `tokio_runtime` to their own Rust modules. diff --git a/rust/src/deferred.rs b/rust/src/deferred.rs new file mode 100644 index 00000000000..91d5a08f7db --- /dev/null +++ b/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: + * . + * + */ + +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 + } else { + "unknown error" + }; + Self::new_err(panic_message.to_owned()) + } +} + +/// A reference to the `twisted.internet.defer` module. +static DEFER: OnceCell> = 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> +where + F: Future> + 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> = 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> { + 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(()) +} diff --git a/rust/src/http_client.rs b/rust/src/http_client.rs index 398ba9041f2..aaa1066a676 100644 --- a/rust/src/http_client.rs +++ b/rust/src/http_client.rs @@ -12,145 +12,22 @@ * . */ -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 - } 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, -} - -#[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> { - 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) -> 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> { - // 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 = reactor.getattr(TOKIO_RUNTIME_ATTR)?.extract()?; - Ok(runtime.borrow()) -} - -/// A reference to the `twisted.internet.defer` module. -static DEFER: OnceCell> = 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::()?; - // 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> -where - F: Future> + 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> = 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> { - 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) -} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index bf96422cbb0..eb2d8cecdde 100644 --- a/rust/src/lib.rs +++ b/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)?; diff --git a/rust/src/tokio_runtime.rs b/rust/src/tokio_runtime.rs new file mode 100644 index 00000000000..f239040f305 --- /dev/null +++ b/rust/src/tokio_runtime.rs @@ -0,0 +1,107 @@ +/* + * 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: + * . + * + */ + +use anyhow::Context; +use pyo3::prelude::*; +use tokio::runtime::Runtime; + +/// 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] +pub struct PyTokioRuntime { + runtime: Option, +} + +#[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. + pub 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. +pub fn runtime<'a>(reactor: &Bound<'a, PyAny>) -> PyResult> { + 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) -> 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> { + // 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 = reactor.getattr(TOKIO_RUNTIME_ATTR)?.extract()?; + Ok(runtime.borrow()) +}