diff --git a/changelog.d/20011.misc b/changelog.d/20011.misc new file mode 100644 index 0000000000..9582106675 --- /dev/null +++ b/changelog.d/20011.misc @@ -0,0 +1 @@ +Refactor the Rust code to have a single place to store per-homeserver state. diff --git a/rust/clippy.toml b/rust/clippy.toml new file mode 100644 index 0000000000..bf2f715c8e --- /dev/null +++ b/rust/clippy.toml @@ -0,0 +1,10 @@ +# Blocking work spawned onto the tokio threadpool is leaked if it is still +# running when the per-homeserver `RustRuntime` is shut down (see +# `rust/src/runtime.rs`). We should avoid spawning blocking work until we figure +# out how to cancel/stop such work on shutdown. +disallowed-methods = [ + { path = "tokio::task::spawn_blocking", reason = "leaks past RustRuntime shutdown; see rust/src/runtime.rs" }, + { path = "tokio::runtime::Runtime::spawn_blocking", reason = "leaks past RustRuntime shutdown; see rust/src/runtime.rs" }, + { path = "tokio::runtime::Handle::spawn_blocking", reason = "leaks past RustRuntime shutdown; see rust/src/runtime.rs" }, + { path = "tokio::task::block_in_place", reason = "blocks a tokio worker thread; see rust/src/runtime.rs" }, +] diff --git a/rust/src/config/mod.rs b/rust/src/config/mod.rs index db00197107..aba29046db 100644 --- a/rust/src/config/mod.rs +++ b/rust/src/config/mod.rs @@ -48,6 +48,7 @@ pub struct AuthConfig { #[derive(FromPyObject, Clone)] pub struct ServerConfig { pub msc4140_enabled: bool, + pub public_baseurl: String, pub include_profile_updates_in_sync: bool, } diff --git a/rust/src/deferred.rs b/rust/src/deferred.rs index 48ce283959..9f438993fb 100644 --- a/rust/src/deferred.rs +++ b/rust/src/deferred.rs @@ -26,7 +26,8 @@ use pyo3::{ use tokio::sync::oneshot; use crate::logging::context::with_logcontext; -use crate::tokio_runtime::runtime; +use crate::reactor::Reactor; +use crate::runtime::RustRuntime; create_exception!( synapse.synapse_rust.http_client, @@ -81,7 +82,7 @@ fn logging_context_module(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> { /// but the work itself is wasted. pub fn create_deferred<'py, F, O>( py: Python<'py>, - reactor: &Bound<'py, PyAny>, + runtime: &RustRuntime, fut: F, ) -> PyResult> where @@ -98,12 +99,12 @@ where // current when the caller invoked us. See `crate::logging::context`. let logcontext = crate::logging::context::LogContextHandle::capture(py); - let rt = runtime(reactor)?; - let handle = rt.handle()?; + let handle = runtime.tokio_handle()?; let task = handle.spawn(logcontext.scope(fut)); - // Unbind the reactor so that we can pass it to the task - let reactor = reactor.clone().unbind(); + // Keep the runtime state (and, through it, the reactor) alive while the + // task is in flight. + let runtime = runtime.clone(); handle.spawn(async move { let res = task.await; @@ -117,19 +118,18 @@ where }, }; - // 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) + runtime + .reactor() + .call_from_thread(py, (deferred_callback, obj)) .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) + runtime + .reactor() + .call_from_thread(py, (deferred_errback, err)) .expect("callFromThread should not fail"); // There's nothing we can really do with errors here } } @@ -150,7 +150,7 @@ where /// the Twisted reactor and runs to completion regardless of whether the returned Rust /// future is ever polled; awaiting it only observes the result. pub(crate) async fn run_python_awaitable( - reactor: Py, + reactor: Reactor, make_awaitable: F, ) -> PyResult> where @@ -265,9 +265,7 @@ where }, )?; - reactor - .bind(py) - .call_method1(intern!(py, "callFromThread"), (starter,))?; + reactor.call_from_thread(py, (starter,))?; Ok(()) })?; diff --git a/rust/src/handlers/mod.rs b/rust/src/handlers/mod.rs index 5df019adff..3119113a0e 100644 --- a/rust/src/handlers/mod.rs +++ b/rust/src/handlers/mod.rs @@ -21,7 +21,7 @@ use pyo3::{ Bound, PyResult, Python, }; -use crate::config::SynapseHomeServerConfig; +use crate::homeserver::HomeServer; use crate::storage::db::python_db_pool::PythonDatabasePoolWrapper; use crate::storage::store::Store; @@ -36,19 +36,17 @@ struct RustHandlers { impl RustHandlers { #[new] #[pyo3(signature = (homeserver))] - pub fn py_new(py: Python<'_>, homeserver: &Bound<'_, PyAny>) -> PyResult { - let config: SynapseHomeServerConfig = homeserver.getattr("config")?.extract()?; + pub fn py_new(py: Python<'_>, homeserver: HomeServer) -> PyResult { + let config = homeserver.config(py)?; - // The Twisted reactor, used both to drive our Tokio runtime and to - // marshal database work back onto the reactor thread. - let reactor = homeserver.call_method0("get_reactor")?.unbind(); + // The per-homeserver Rust state, which gives us the tokio runtime + // and the Twisted reactor. + let runtime = homeserver.get_rust_runtime(py)?; - // hs.get_datastores().main.db_pool - let db_pool_py = homeserver - .call_method0("get_datastores")? - .getattr("main")? - .getattr("db_pool")?; - let db_pool = PythonDatabasePoolWrapper::new(&db_pool_py, reactor.clone_ref(py))?; + let db_pool = PythonDatabasePoolWrapper::new( + &homeserver.main_database_pool(py)?, + runtime.reactor().clone_ref(py), + )?; // Store is shared across all of the handlers so let's use an `Arc` let store = Arc::new(Store { @@ -64,7 +62,7 @@ impl RustHandlers { versions::VersionsHandler { global_unstable_feature_map: Arc::clone(&global_unstable_feature_map), store: Arc::clone(&store), - reactor: reactor.clone_ref(py), + runtime: runtime.clone(), }, )?; diff --git a/rust/src/handlers/versions.rs b/rust/src/handlers/versions.rs index ab4bd9a160..f3e6ad8f20 100644 --- a/rust/src/handlers/versions.rs +++ b/rust/src/handlers/versions.rs @@ -21,6 +21,7 @@ use serde::Serialize; use crate::config::{types::RoomCreationPreset, SynapseHomeServerConfig}; use crate::deferred::create_deferred; +use crate::runtime::RustRuntime; use crate::storage::store::{PerUserExperimentalFeature, Store}; /// `GET /_matrix/client/versions` response @@ -45,9 +46,9 @@ impl<'py> IntoPyObject<'py> for VersionsResponse { pub struct VersionsHandler { pub global_unstable_feature_map: 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, + /// The per-homeserver Rust state, used to bridge our `async` response + /// back into a Twisted deferred that Python can `await`. + pub runtime: RustRuntime, } #[pymethods] @@ -63,7 +64,7 @@ impl VersionsHandler { let store = Arc::clone(&self.store); let global_unstable_feature_map = Arc::clone(&self.global_unstable_feature_map); - create_deferred(py, self.reactor.bind(py), async move { + create_deferred(py, &self.runtime, async move { build_versions_response(&store, &global_unstable_feature_map, user_id.as_deref()) .await .map_err(|err| { diff --git a/rust/src/homeserver.rs b/rust/src/homeserver.rs new file mode 100644 index 0000000000..44f8ba905e --- /dev/null +++ b/rust/src/homeserver.rs @@ -0,0 +1,102 @@ +/* + * 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: + * . + * + */ + +//! A typed wrapper around the Python `HomeServer`. + +use pyo3::{intern, prelude::*, types::PyDict}; + +use crate::config::SynapseHomeServerConfig; +use crate::reactor::Reactor; +use crate::runtime::RustRuntime; + +/// The Python `HomeServer`, as seen from Rust. +/// +/// This is a wrapper around a foreign Python object. No validation is +/// performed, on the assumption the Python side is correctly typed. +pub struct HomeServer(Py); + +impl<'a, 'py> FromPyObject<'a, 'py> for HomeServer { + type Error = PyErr; + + fn extract(obj: Borrowed<'a, 'py, PyAny>) -> PyResult { + Ok(HomeServer(obj.to_owned().unbind())) + } +} + +impl HomeServer { + /// Fetch the Twisted reactor in use by this HomeServer. + pub fn get_reactor(&self, py: Python<'_>) -> PyResult { + self.0 + .bind(py) + .call_method0(intern!(py, "get_reactor"))? + .extract() + } + + /// Register a system event trigger with the HomeServer so it can be cleanly + /// removed when the HomeServer is shutdown. + pub fn register_sync_shutdown_handler( + &self, + py: Python<'_>, + callable: &Bound<'_, PyAny>, + ) -> PyResult<()> { + let kwargs = PyDict::new(py); + kwargs.set_item(intern!(py, "phase"), intern!(py, "after"))?; + kwargs.set_item(intern!(py, "eventType"), intern!(py, "shutdown"))?; + kwargs.set_item(intern!(py, "shutdown_func"), callable)?; + self.0.bind(py).call_method( + intern!(py, "register_sync_shutdown_handler"), + (), + Some(&kwargs), + )?; + + Ok(()) + } + + /// The per-homeserver Rust state (`hs.get_rust_runtime()`), which gives + /// access to the tokio runtime and the reactor. + pub fn get_rust_runtime(&self, py: Python<'_>) -> PyResult { + Ok(self + .0 + .bind(py) + .call_method0(intern!(py, "get_rust_runtime"))? + .cast::()? + .get() + .clone()) + } + + /// The Rust-side view of `hs.config`. + pub fn config(&self, py: Python<'_>) -> PyResult { + self.0.bind(py).getattr(intern!(py, "config"))?.extract() + } + + /// The Synapse `Clock` (`hs.get_clock()`). + // TODO: give the clock a typed wrapper of its own. + pub fn get_clock(&self, py: Python<'_>) -> PyResult> { + Ok(self + .0 + .bind(py) + .call_method0(intern!(py, "get_clock"))? + .unbind()) + } + + /// The main database pool (`hs.get_datastores().main.db_pool`). + pub fn main_database_pool<'py>(&self, py: Python<'py>) -> PyResult> { + self.0 + .bind(py) + .call_method0(intern!(py, "get_datastores"))? + .getattr(intern!(py, "main"))? + .getattr(intern!(py, "db_pool")) + } +} diff --git a/rust/src/http_client.rs b/rust/src/http_client.rs index aaa1066a67..802f93b2d0 100644 --- a/rust/src/http_client.rs +++ b/rust/src/http_client.rs @@ -21,7 +21,7 @@ use reqwest::RequestBuilder; use crate::deferred::create_deferred; use crate::errors::HttpResponseException; -use crate::tokio_runtime::runtime; +use crate::runtime::RustRuntime; /// Called when registering modules with python. pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -42,21 +42,18 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> #[pyclass] struct HttpClient { client: reqwest::Client, - reactor: Py, + runtime: RustRuntime, } #[pymethods] impl HttpClient { #[new] - #[pyo3(signature = (reactor, user_agent, http2_only = false))] + #[pyo3(signature = (runtime, user_agent, http2_only = false))] pub fn py_new( - reactor: Bound, + runtime: &RustRuntime, user_agent: &str, http2_only: bool, ) -> PyResult { - // Make sure the runtime gets installed - let _ = runtime(&reactor)?; - let mut builder = reqwest::Client::builder().user_agent(user_agent); if http2_only { @@ -69,7 +66,7 @@ impl HttpClient { Ok(HttpClient { client, - reactor: reactor.unbind(), + runtime: runtime.clone(), }) } @@ -107,7 +104,7 @@ impl HttpClient { builder: RequestBuilder, response_limit: usize, ) -> PyResult> { - create_deferred(py, self.reactor.bind(py), async move { + create_deferred(py, &self.runtime, async move { let response = builder.send().await.context("sending request")?; let status = response.status(); diff --git a/rust/src/lib.rs b/rust/src/lib.rs index f68b7c17d7..3d965daf51 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -12,6 +12,7 @@ pub mod duration; pub mod errors; pub mod events; pub mod handlers; +pub mod homeserver; pub mod http; pub mod http_client; pub mod identifier; @@ -20,11 +21,12 @@ pub mod logging; pub mod matrix_const; pub mod msc4388_rendezvous; pub mod push; +pub mod reactor; pub mod rendezvous; pub mod room_versions; +pub mod runtime; pub mod segmenter; pub mod storage; -pub mod tokio_runtime; pub mod types; lazy_static! { @@ -78,6 +80,7 @@ fn synapse_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { handlers::register_module(py, m)?; http_client::register_module(py, m)?; rendezvous::register_module(py, m)?; + runtime::register_module(py, m)?; msc4388_rendezvous::register_module(py, m)?; segmenter::register_module(py, m)?; room_versions::register_module(py, m)?; diff --git a/rust/src/msc4388_rendezvous/mod.rs b/rust/src/msc4388_rendezvous/mod.rs index 2f1b004ec1..ffd4f240d9 100644 --- a/rust/src/msc4388_rendezvous/mod.rs +++ b/rust/src/msc4388_rendezvous/mod.rs @@ -21,7 +21,7 @@ use http::StatusCode; use pyo3::{ pyclass, pymethods, types::{PyAnyMethods, PyModule, PyModuleMethods}, - Bound, IntoPyObject, Py, PyAny, PyResult, Python, + Bound, Py, PyAny, PyResult, Python, }; use serde::Deserialize; use ulid::Ulid; @@ -30,9 +30,9 @@ use self::session::Session; use crate::{ duration::SynapseDuration, errors::{NotFoundError, SynapseError}, + homeserver::HomeServer, http::http_request_from_twisted, msc4388_rendezvous::session::{GetResponse, PostResponse, PutResponse}, - UnwrapInfallible, }; mod session; @@ -92,25 +92,21 @@ impl MSC4388RendezvousHandler { #[pyo3(signature = (homeserver, /, soft_limit=100, hard_limit=200,max_content_length=4*1024, eviction_interval=60*1000, ttl=2*60*1000))] fn new( py: Python<'_>, - homeserver: &Bound<'_, PyAny>, + homeserver: HomeServer, soft_limit: usize, hard_limit: usize, max_content_length: u64, eviction_interval: u64, ttl: u64, ) -> PyResult> { - let clock = homeserver - .call_method0("get_clock")? - .into_pyobject(py) - .unwrap_infallible() - .unbind(); + let clock = homeserver.get_clock(py)?; // Construct a Python object so that we can get a reference to the // evict method and schedule it to run. let self_ = Py::new( py, Self { - clock, + clock: clock.clone_ref(py), sessions: BTreeMap::new(), soft_limit, hard_limit, @@ -122,11 +118,9 @@ impl MSC4388RendezvousHandler { let eviction_duration = SynapseDuration::from_milliseconds(eviction_interval); let evict = self_.getattr(py, "_evict")?; - homeserver.call_method0("get_clock")?.call_method( - "looping_call", - (evict, &eviction_duration), - None, - )?; + clock + .bind(py) + .call_method("looping_call", (evict, &eviction_duration), None)?; Ok(self_) } diff --git a/rust/src/reactor.rs b/rust/src/reactor.rs new file mode 100644 index 0000000000..7522cd3390 --- /dev/null +++ b/rust/src/reactor.rs @@ -0,0 +1,55 @@ +/* + * 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: + * . + * + */ + +//! A typed wrapper around the Twisted reactor. + +use pyo3::{call::PyCallArgs, intern, prelude::*}; + +/// The Twisted reactor, as seen from Rust. +/// +/// This is a wrapper around a foreign Python object. No validation is +/// performed, on the assumption the Python side is correctly typed. +pub struct Reactor(Py); + +impl<'a, 'py> FromPyObject<'a, 'py> for Reactor { + type Error = PyErr; + + fn extract(obj: Borrowed<'a, 'py, PyAny>) -> PyResult { + Ok(Reactor(obj.to_owned().unbind())) + } +} + +impl Reactor { + /// `reactor.callFromThread(f, *args)`: schedule a call on the reactor + /// thread. This is the only reactor method that is safe to call from + /// other threads (e.g. tokio workers). + /// + /// `args` is the full argument tuple, starting with the callable itself. + pub fn call_from_thread<'py>( + &self, + py: Python<'py>, + args: impl PyCallArgs<'py>, + ) -> PyResult<()> { + self.0 + .bind(py) + .call_method1(intern!(py, "callFromThread"), args)?; + + Ok(()) + } + + pub fn clone_ref(&self, py: Python<'_>) -> Reactor { + Reactor(self.0.clone_ref(py)) + } +} diff --git a/rust/src/rendezvous/mod.rs b/rust/src/rendezvous/mod.rs index 9a6da9fcc3..64a5123f31 100644 --- a/rust/src/rendezvous/mod.rs +++ b/rust/src/rendezvous/mod.rs @@ -29,7 +29,7 @@ use pyo3::{ exceptions::PyValueError, pyclass, pymethods, types::{PyAnyMethods, PyModule, PyModuleMethods}, - Bound, IntoPyObject, Py, PyAny, PyResult, Python, + Bound, Py, PyAny, PyResult, Python, }; use ulid::Ulid; @@ -37,8 +37,8 @@ use self::session::Session; use crate::{ duration::SynapseDuration, errors::{NotFoundError, SynapseError}, + homeserver::HomeServer, http::{http_request_from_twisted, http_response_to_twisted, HeaderMapPyExt}, - UnwrapInfallible, }; mod session; @@ -113,25 +113,17 @@ impl RendezvousHandler { #[pyo3(signature = (homeserver, /, capacity=100, max_content_length=4*1024, eviction_interval=60*1000, ttl=60*1000))] fn new( py: Python<'_>, - homeserver: &Bound<'_, PyAny>, + homeserver: HomeServer, capacity: usize, max_content_length: u64, eviction_interval: u64, ttl: u64, ) -> PyResult> { - let base: String = homeserver - .getattr("config")? - .getattr("server")? - .getattr("public_baseurl")? - .extract()?; + let base = homeserver.config(py)?.server.public_baseurl; let base = Uri::try_from(format!("{base}_synapse/client/rendezvous")) .map_err(|_| PyValueError::new_err("Invalid base URI"))?; - let clock = homeserver - .call_method0("get_clock")? - .into_pyobject(py) - .unwrap_infallible() - .unbind(); + let clock = homeserver.get_clock(py)?; let eviction_duration = SynapseDuration::from_milliseconds(eviction_interval); @@ -141,7 +133,7 @@ impl RendezvousHandler { py, Self { base, - clock, + clock: clock.clone_ref(py), sessions: BTreeMap::new(), capacity, max_content_length, @@ -150,11 +142,9 @@ impl RendezvousHandler { )?; let evict = self_.getattr(py, "_evict")?; - homeserver.call_method0("get_clock")?.call_method( - "looping_call", - (evict, &eviction_duration), - None, - )?; + clock + .bind(py) + .call_method("looping_call", (evict, &eviction_duration), None)?; Ok(self_) } diff --git a/rust/src/runtime.rs b/rust/src/runtime.rs new file mode 100644 index 0000000000..46ceeec1a8 --- /dev/null +++ b/rust/src/runtime.rs @@ -0,0 +1,217 @@ +/* + * 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: + * . + * + */ + +//! The per-homeserver state for the Rust side of Synapse. +//! +//! A [`RustRuntime`] is created once per homeserver (`hs.get_rust_runtime()`) +//! and holds everything the Rust side keeps for the lifetime of that +//! homeserver: currently the tokio thread pool and a handle to the Twisted +//! reactor. Rust consumers (e.g. the HTTP client) clone the inner +//! [`Arc`] at construction time and don't need the GIL (or +//! the Python-facing object) to reach it afterwards. +//! +//! The tokio runtime is shut down with the homeserver, via a handler +//! registered with `HomeServer.register_sync_shutdown_handler`. + +use std::ops::Deref; +use std::sync::{Arc, Mutex, Weak}; +use std::time::Duration; + +use anyhow::Context; +use pyo3::{exceptions::PyRuntimeError, prelude::*}; +use tokio::runtime::{Handle, Runtime}; + +use crate::homeserver::HomeServer; +use crate::reactor::Reactor; + +/// How long to wait for in-flight tokio tasks to be cancelled when shutting +/// down with the reactor. +/// +/// Note that any [`Runtime::spawn_blocking`] work that is still running when +/// the timeout expires is leaked, along with the worker thread running it. +/// +/// See [`tokio::runtime`] for details. +const SHUTDOWN_TIMEOUT: Duration = Duration::from_millis(100); + +/// State of the lazily-started tokio runtime. +enum TokioState { + /// Not started yet; the runtime is built on first use. + NotStarted, + Running(Runtime), + /// Shut down with the homeserver. Cannot be restarted. + Shutdown, +} + +/// The state shared between the Python-facing [`RustRuntime`] handle and any +/// Rust-side consumers holding an `Arc` of this. +pub struct RustRuntimeInner { + reactor: Reactor, + tokio: Mutex, + worker_threads: usize, +} + +impl RustRuntimeInner { + /// The Twisted reactor this homeserver runs on. + pub fn reactor(&self) -> &Reactor { + &self.reactor + } + + /// Get a handle to the tokio runtime, starting the runtime if it hasn't + /// been started yet. + pub fn tokio_handle(&self) -> PyResult { + let mut state = self + .tokio + .lock() + .map_err(|_| PyRuntimeError::new_err("tokio runtime lock poisoned"))?; + + match &*state { + TokioState::Running(runtime) => Ok(runtime.handle().clone()), + TokioState::NotStarted => { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(self.worker_threads) + .enable_all() + .build() + .context("building tokio runtime")?; + let handle = runtime.handle().clone(); + *state = TokioState::Running(runtime); + Ok(handle) + } + TokioState::Shutdown => Err(PyRuntimeError::new_err( + "the tokio runtime has been shut down", + )), + } + } + + /// Shut the tokio runtime down, cancelling all in-flight tasks and waiting + /// for up to [`SHUTDOWN_TIMEOUT`] for them to finish. Called via + /// [`ShutdownHook`] when the reactor shuts down. + /// + /// Note that any [`Runtime::spawn_blocking`] work is leaked until it is + /// finished, along with the worker thread running it. + fn shutdown(&self, py: Python<'_>) -> PyResult<()> { + let mut state = self + .tokio + .lock() + .map_err(|_| PyRuntimeError::new_err("tokio runtime lock poisoned"))?; + let previous_state = std::mem::replace(&mut *state, TokioState::Shutdown); + // Don't hold the lock while blocking on the shutdown below. + drop(state); + + if let TokioState::Running(runtime) = previous_state { + // Shutdown the runtime, waiting for a small grace period for + // in-flight tasks to be cancelled. + // + // See [`tokio::runtime`] for details. + py.detach(|| runtime.shutdown_timeout(SHUTDOWN_TIMEOUT)); + } + + Ok(()) + } +} + +impl Drop for RustRuntimeInner { + fn drop(&mut self) { + // Backstop for homeservers whose shutdown trigger never fires (e.g. in + // tests). We use `shutdown_background` rather than a blocking shutdown, + // because the last `Arc` may be dropped from a task running on this + // very runtime, where blocking would panic. + if let Ok(state) = self.tokio.get_mut() { + if let TokioState::Running(runtime) = std::mem::replace(state, TokioState::Shutdown) { + runtime.shutdown_background(); + } + } + } +} + +/// A cheaply-clonable handle to the per-homeserver Rust state, and the +/// Python-facing class for it. +/// +/// One instance is constructed per homeserver by +/// `HomeServer.get_rust_runtime()`. Rust classes that need it take it as a +/// constructor argument and store their own clone, which is just an `Arc` +/// refcount bump. Derefs to [`RustRuntimeInner`]. +#[pyclass(frozen, skip_from_py_object)] +#[derive(Clone)] +pub struct RustRuntime { + inner: Arc, +} + +impl Deref for RustRuntime { + type Target = RustRuntimeInner; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +#[pymethods] +impl RustRuntime { + #[new] + #[pyo3(signature = (hs, worker_threads = 4))] + fn py_new(py: Python<'_>, hs: HomeServer, worker_threads: usize) -> PyResult { + let inner = Arc::new(RustRuntimeInner { + reactor: hs.get_reactor(py)?, + tokio: Mutex::new(TokioState::NotStarted), + worker_threads, + }); + + // Shut the tokio runtime down when the homeserver is shut down. The + // trigger holds only a `Weak` reference, as otherwise we risk a + // reference cycle passing through a Rust field that Python's GC cannot + // see into. + let hook = Py::new( + py, + ShutdownHook { + inner: Arc::downgrade(&inner), + }, + )?; + hs.register_sync_shutdown_handler(py, hook.bind(py).as_any())?; + + Ok(RustRuntime { inner }) + } +} + +/// The callable registered with `HomeServer.register_sync_shutdown_handler`, +/// which runs it on `HomeServer.shutdown()` or when the reactor stops. +#[pyclass(frozen)] +struct ShutdownHook { + inner: Weak, +} + +#[pymethods] +impl ShutdownHook { + fn __call__(&self, py: Python<'_>) -> PyResult<()> { + if let Some(inner) = self.inner.upgrade() { + inner.shutdown(py)?; + } + + Ok(()) + } +} + +/// Called when registering modules with python. +pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + let child_module = PyModule::new(py, "runtime")?; + + child_module.add_class::()?; + + m.add_submodule(&child_module)?; + + py.import("sys")? + .getattr("modules")? + .set_item("synapse.synapse_rust.runtime", child_module)?; + + Ok(()) +} diff --git a/rust/src/storage/db/python_db_pool.rs b/rust/src/storage/db/python_db_pool.rs index 89003d574b..b683a9c00e 100644 --- a/rust/src/storage/db/python_db_pool.rs +++ b/rust/src/storage/db/python_db_pool.rs @@ -40,6 +40,7 @@ use pyo3::{ }; use crate::deferred::run_python_awaitable; +use crate::reactor::Reactor; use crate::storage::db::{ DatabasePool, DbRow, DbValue, ErasedInteraction, ErasedResult, Transaction, }; @@ -112,13 +113,13 @@ pub struct PythonDatabasePoolWrapper { /// never gets garbage collected and never points back at the homeserver, so it is /// not part of any reference cycle. Ideally, we could worry about it but /// practically probably doesn't matter. - reactor: Py, + reactor: Reactor, } impl PythonDatabasePoolWrapper { /// Build a wrapper around the Python `DatabasePool` (e.g. /// `hs.get_datastores().main.db_pool`) and the Twisted `reactor`. - pub fn new(database_pool: &Bound<'_, PyAny>, reactor: Py) -> PyResult { + pub fn new(database_pool: &Bound<'_, PyAny>, reactor: Reactor) -> PyResult { Ok(Self { database_pool_py_ref: PyWeakrefReference::new(database_pool)?.unbind(), reactor, diff --git a/rust/src/tokio_runtime.rs b/rust/src/tokio_runtime.rs deleted file mode 100644 index f239040f30..0000000000 --- a/rust/src/tokio_runtime.rs +++ /dev/null @@ -1,107 +0,0 @@ -/* - * 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()) -} diff --git a/synapse/api/auth/mas.py b/synapse/api/auth/mas.py index ed0427d6f3..d094410ff9 100644 --- a/synapse/api/auth/mas.py +++ b/synapse/api/auth/mas.py @@ -109,7 +109,7 @@ class MasDelegatedAuth(BaseAuth): self._http_client = hs.get_proxied_http_client() self._rust_http_client = HttpClient( - reactor=hs.get_reactor(), + runtime=hs.get_rust_runtime(), user_agent=self._http_client.user_agent.decode("utf8"), http2_only=self._config.force_http2, ) diff --git a/synapse/server.py b/synapse/server.py index 9bea208974..266440534c 100644 --- a/synapse/server.py +++ b/synapse/server.py @@ -177,6 +177,7 @@ from synapse.streams.events import EventSources from synapse.synapse_rust.handlers import RustHandlers from synapse.synapse_rust.msc4388_rendezvous import MSC4388RendezvousHandler from synapse.synapse_rust.rendezvous import RendezvousHandler +from synapse.synapse_rust.runtime import RustRuntime from synapse.types import DomainSpecificString, ISynapseReactor from synapse.util import SYNAPSE_VERSION from synapse.util.caches import CACHE_METRIC_REGISTRY @@ -972,6 +973,18 @@ class HomeServer(metaclass=abc.ABCMeta): def get_rust_handlers(self) -> RustHandlers: return RustHandlers(self) + @cache_in_self + def get_rust_runtime(self) -> RustRuntime: + """The per-homeserver state for the Rust side of Synapse: the tokio + thread pool, plus anything else Rust code keeps for the lifetime of the + homeserver. + + The tokio runtime is started lazily on first use, and shut down when + this homeserver is shut down. + """ + # TODO: make the number of worker threads configurable + return RustRuntime(hs=self, worker_threads=4) + @cache_in_self def get_event_sources(self) -> EventSources: return EventSources(self) diff --git a/synapse/synapse_rust/http_client.pyi b/synapse/synapse_rust/http_client.pyi index 1814daec73..e6468bc46b 100644 --- a/synapse/synapse_rust/http_client.pyi +++ b/synapse/synapse_rust/http_client.pyi @@ -14,7 +14,7 @@ from typing import Mapping from twisted.internet.defer import Deferred -from synapse.types import ISynapseReactor +from synapse.synapse_rust.runtime import RustRuntime class HttpClient: """ @@ -23,7 +23,7 @@ class HttpClient: def __init__( self, - reactor: ISynapseReactor, + runtime: RustRuntime, user_agent: str, http2_only: bool = False, ) -> None: @@ -31,7 +31,7 @@ class HttpClient: Create a new HTTP client backed by reqwest. Args: - reactor: The Twisted reactor to coordinate with + runtime: The per-homeserver Rust state (`hs.get_rust_runtime()`) user_agent: The user agent to use for requests http2_only: Whether to use HTTP/2 only, even on unencrypted connections. By default, it will always use HTTP/1.1 over unencrypted connections, and diff --git a/synapse/synapse_rust/runtime.pyi b/synapse/synapse_rust/runtime.pyi new file mode 100644 index 0000000000..32ae77ec28 --- /dev/null +++ b/synapse/synapse_rust/runtime.pyi @@ -0,0 +1,15 @@ +from synapse.server import HomeServer + +class RustRuntime: + """The per-homeserver state for the Rust side of Synapse. + + Holds the tokio thread pool (started lazily on first use, shut down with the + Homeserver). Rust classes that need them take this as a constructor + argument. Get it from `hs.get_rust_runtime()`. + """ + + def __init__( + self, + hs: HomeServer, + worker_threads: int = 4, + ) -> None: ... diff --git a/tests/handlers/test_oauth_delegation.py b/tests/handlers/test_oauth_delegation.py index 71f83607ff..0fc1e17d6b 100644 --- a/tests/handlers/test_oauth_delegation.py +++ b/tests/handlers/test_oauth_delegation.py @@ -216,8 +216,6 @@ class MasAuthDelegation(HomeserverTestCase): def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: self.server = FakeMasServer() hs = self.setup_test_homeserver() - # This triggers the server startup hooks, which starts the Tokio thread pool - reactor.run() self._auth = checked_cast(MasDelegatedAuth, hs.get_auth()) return hs diff --git a/tests/push/test_http.py b/tests/push/test_http.py index 47521a773f..bd2ada9155 100644 --- a/tests/push/test_http.py +++ b/tests/push/test_http.py @@ -31,7 +31,6 @@ from synapse.rest import admin from synapse.rest.admin.experimental_features import ExperimentalFeature from synapse.rest.client import login, push_rule, pusher, receipts, room, versions from synapse.server import HomeServer -from synapse.synapse_rust.http_client import HttpClient from synapse.types import JsonDict from synapse.util.clock import Clock @@ -1233,24 +1232,6 @@ class MSC3881VersionsTestCase(HomeserverTestCase): versions.register_servlets, ] - def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: - hs = self.setup_test_homeserver() - - # XXX: We must create the Rust HTTP client before we call `reactor.run()` below. - # Twisted's `MemoryReactor` doesn't invoke `callWhenRunning` callbacks if it's - # already running and we rely on that to start the Tokio thread pool in Rust. In - # the future, this may not matter, see https://github.com/twisted/twisted/pull/12514 - self._http_client = hs.get_proxied_http_client() - _ = HttpClient( - reactor=hs.get_reactor(), - user_agent=self._http_client.user_agent.decode("utf8"), - ) - - # This triggers the server startup hooks, which starts the Tokio thread pool - reactor.run() - - return hs - def tearDown(self) -> None: # MemoryReactor doesn't trigger the shutdown phases, and we want the # Tokio thread pool to be stopped diff --git a/tests/rest/client/test_delayed_events.py b/tests/rest/client/test_delayed_events.py index 3af7d13858..644b373f36 100644 --- a/tests/rest/client/test_delayed_events.py +++ b/tests/rest/client/test_delayed_events.py @@ -24,7 +24,6 @@ from synapse.api.errors import Codes from synapse.rest import admin from synapse.rest.client import delayed_events, login, room, sync, versions from synapse.server import HomeServer -from synapse.synapse_rust.http_client import HttpClient from synapse.types import JsonDict from synapse.util.clock import Clock from synapse.util.duration import Duration @@ -41,24 +40,6 @@ _EVENT_TYPE = "com.example.test" class DelayedEventsUnstableSupportTestCase(HomeserverTestCase): servlets = [versions.register_servlets] - def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: - hs = self.setup_test_homeserver() - - # XXX: We must create the Rust HTTP client before we call `reactor.run()` below. - # Twisted's `MemoryReactor` doesn't invoke `callWhenRunning` callbacks if it's - # already running and we rely on that to start the Tokio thread pool in Rust. In - # the future, this may not matter, see https://github.com/twisted/twisted/pull/12514 - self._http_client = hs.get_proxied_http_client() - _ = HttpClient( - reactor=hs.get_reactor(), - user_agent=self._http_client.user_agent.decode("utf8"), - ) - - # This triggers the server startup hooks, which starts the Tokio thread pool - reactor.run() - - return hs - def tearDown(self) -> None: # MemoryReactor doesn't trigger the shutdown phases, and we want the # Tokio thread pool to be stopped diff --git a/tests/rest/client/test_login_token_request.py b/tests/rest/client/test_login_token_request.py index d6b2cf054e..b387cb98fb 100644 --- a/tests/rest/client/test_login_token_request.py +++ b/tests/rest/client/test_login_token_request.py @@ -24,7 +24,6 @@ from twisted.internet.testing import MemoryReactor from synapse.rest import admin from synapse.rest.client import login, login_token_request, versions from synapse.server import HomeServer -from synapse.synapse_rust.http_client import HttpClient from synapse.util.clock import Clock from tests import unittest @@ -48,19 +47,6 @@ class LoginTokenRequestServletTestCase(unittest.HomeserverTestCase): self.hs.config.registration.auto_join_rooms = [] self.hs.config.captcha.enable_registration_captcha = False - # XXX: We must create the Rust HTTP client before we call `reactor.run()` below. - # Twisted's `MemoryReactor` doesn't invoke `callWhenRunning` callbacks if it's - # already running and we rely on that to start the Tokio thread pool in Rust. In - # the future, this may not matter, see https://github.com/twisted/twisted/pull/12514 - self._http_client = self.hs.get_proxied_http_client() - _ = HttpClient( - reactor=self.hs.get_reactor(), - user_agent=self._http_client.user_agent.decode("utf8"), - ) - - # This triggers the server startup hooks, which starts the Tokio thread pool - reactor.run() - return self.hs def tearDown(self) -> None: diff --git a/tests/rest/client/test_matrixrtc.py b/tests/rest/client/test_matrixrtc.py index 63c7632a88..545a5a8e8c 100644 --- a/tests/rest/client/test_matrixrtc.py +++ b/tests/rest/client/test_matrixrtc.py @@ -27,7 +27,6 @@ from synapse.config.matrixrtc import TransportConfigModel from synapse.rest import admin from synapse.rest.client import login, matrixrtc, register, room, versions from synapse.server import HomeServer -from synapse.synapse_rust.http_client import HttpClient from synapse.util.clock import Clock from tests import unittest @@ -157,24 +156,6 @@ class MatrixRtcVersionsTestCase(HomeserverTestCase): servlets = [versions.register_servlets] - def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: - hs = self.setup_test_homeserver() - - # XXX: We must create the Rust HTTP client before we call `reactor.run()` below. - # Twisted's `MemoryReactor` doesn't invoke `callWhenRunning` callbacks if it's - # already running and we rely on that to start the Tokio thread pool in Rust. In - # the future, this may not matter, see https://github.com/twisted/twisted/pull/12514 - self._http_client = hs.get_proxied_http_client() - _ = HttpClient( - reactor=hs.get_reactor(), - user_agent=self._http_client.user_agent.decode("utf8"), - ) - - # This triggers the server startup hooks, which starts the Tokio thread pool - reactor.run() - - return hs - def tearDown(self) -> None: # MemoryReactor doesn't trigger the shutdown phases, and we want the # Tokio thread pool to be stopped diff --git a/tests/rest/client/test_versions.py b/tests/rest/client/test_versions.py index 1ed6bb145b..6ecff9de6f 100644 --- a/tests/rest/client/test_versions.py +++ b/tests/rest/client/test_versions.py @@ -17,7 +17,6 @@ from twisted.internet.testing import MemoryReactor from synapse.rest import admin from synapse.rest.client import login, versions from synapse.server import HomeServer -from synapse.synapse_rust.http_client import HttpClient from synapse.types import JsonDict from synapse.util.clock import Clock @@ -37,24 +36,6 @@ class VersionsTestCase(unittest.HomeserverTestCase): versions.register_servlets, ] - def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: - hs = self.setup_test_homeserver() - - # XXX: We must create the Rust HTTP client before we call `reactor.run()` below. - # Twisted's `MemoryReactor` doesn't invoke `callWhenRunning` callbacks if it's - # already running and we rely on that to start the Tokio thread pool in Rust. In - # the future, this may not matter, see https://github.com/twisted/twisted/pull/12514 - self._http_client = hs.get_proxied_http_client() - _ = HttpClient( - reactor=hs.get_reactor(), - user_agent=self._http_client.user_agent.decode("utf8"), - ) - - # This triggers the server startup hooks, which starts the Tokio thread pool - reactor.run() - - return hs - def tearDown(self) -> None: # MemoryReactor doesn't trigger the shutdown phases, and we want the # Tokio thread pool to be stopped diff --git a/tests/synapse_rust/test_http_client.py b/tests/synapse_rust/test_http_client.py index 845fe2b503..330fbeeba6 100644 --- a/tests/synapse_rust/test_http_client.py +++ b/tests/synapse_rust/test_http_client.py @@ -86,25 +86,13 @@ class StubServer(HTTPServer): class HttpClientTestCase(HomeserverTestCase): - def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: - hs = self.setup_test_homeserver() - - # XXX: We must create the Rust HTTP client before we call `reactor.run()` below. - # Twisted's `MemoryReactor` doesn't invoke `callWhenRunning` callbacks if it's - # already running and we rely on that to start the Tokio thread pool in Rust. In - # the future, this may not matter, see https://github.com/twisted/twisted/pull/12514 + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: self._http_client = hs.get_proxied_http_client() self._rust_http_client = HttpClient( - reactor=hs.get_reactor(), + runtime=hs.get_rust_runtime(), user_agent=self._http_client.user_agent.decode("utf8"), ) - # This triggers the server startup hooks, which starts the Tokio thread pool - reactor.run() - - return hs - - def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: self.server = StubServer() def tearDown(self) -> None: diff --git a/tests/synapse_rust/test_logcontext.py b/tests/synapse_rust/test_logcontext.py index 0be82f0caa..5e1cefaff9 100644 --- a/tests/synapse_rust/test_logcontext.py +++ b/tests/synapse_rust/test_logcontext.py @@ -82,7 +82,7 @@ class RustLogContextTestCase(HomeserverTestCase): # thread pool in Rust. self._http_client = hs.get_proxied_http_client() self._rust_http_client = HttpClient( - reactor=hs.get_reactor(), + runtime=hs.get_rust_runtime(), user_agent=self._http_client.user_agent.decode("utf8"), )