Move per-homeserver Rust state into a RustRuntime object on the HomeServer (#20011)

Previously the tokio runtime was stashed in a hidden attribute on the
reactor object, installed lazily by whichever Rust code first needed it,
and started via `callWhenRunning`.

Instead, we create a `RustRuntime` (accessible via
`HomeServer.get_rust_runtime()`) that holds any per-reactor Rust state,
such as the tokio runtime. It is constructed lazily on use. Rust
consumers (`HttpClient`, `VersionsHandler`, the Python DB pool wrapper)
now receive the runtime or reactor handle explicitly, and the
`reactor.run()` / manual-startup workarounds in tests are no longer
needed.

We also add helper wrappers in Rust for `Reactor` and `HomeServer` that
exposes the needed functionality.

The aim is to allow us to have a Rust-side clock (mainly to get the
current time), that respects the unit test per-reactor time management.

---------

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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,
 }
 

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<Bound<'py, PyAny>>
 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<F>(
-    reactor: Py<PyAny>,
+    reactor: Reactor,
     make_awaitable: F,
 ) -> PyResult<Py<PyAny>>
 where
@@ -265,9 +265,7 @@ where
             },
         )?;
 
-        reactor
-            .bind(py)
-            .call_method1(intern!(py, "callFromThread"), (starter,))?;
+        reactor.call_from_thread(py, (starter,))?;
 
         Ok(())
     })?;

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<RustHandlers> {
-        let config: SynapseHomeServerConfig = homeserver.getattr("config")?.extract()?;
+    pub fn py_new(py: Python<'_>, homeserver: HomeServer) -> PyResult<RustHandlers> {
+        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(),
             },
         )?;
 

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<UnstableFeatureMap>,
     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>,
+    /// 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| {

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:
+ * <https://www.gnu.org/licenses/agpl-3.0.html>.
+ *
+ */
+
+//! 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<PyAny>);
+
+impl<'a, 'py> FromPyObject<'a, 'py> for HomeServer {
+    type Error = PyErr;
+
+    fn extract(obj: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
+        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<Reactor> {
+        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<RustRuntime> {
+        Ok(self
+            .0
+            .bind(py)
+            .call_method0(intern!(py, "get_rust_runtime"))?
+            .cast::<RustRuntime>()?
+            .get()
+            .clone())
+    }
+
+    /// The Rust-side view of `hs.config`.
+    pub fn config(&self, py: Python<'_>) -> PyResult<SynapseHomeServerConfig> {
+        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<Py<PyAny>> {
+        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<Bound<'py, PyAny>> {
+        self.0
+            .bind(py)
+            .call_method0(intern!(py, "get_datastores"))?
+            .getattr(intern!(py, "main"))?
+            .getattr(intern!(py, "db_pool"))
+    }
+}

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<PyAny>,
+    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<PyAny>,
+        runtime: &RustRuntime,
         user_agent: &str,
         http2_only: bool,
     ) -> PyResult<HttpClient> {
-        // 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<Bound<'a, PyAny>> {
-        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();

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)?;

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<Py<Self>> {
-        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_)
     }

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:
+ * <https://www.gnu.org/licenses/agpl-3.0.html>.
+ *
+ */
+
+//! 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<PyAny>);
+
+impl<'a, 'py> FromPyObject<'a, 'py> for Reactor {
+    type Error = PyErr;
+
+    fn extract(obj: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
+        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))
+    }
+}

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<Py<Self>> {
-        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_)
     }

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:
+ * <https://www.gnu.org/licenses/agpl-3.0.html>.
+ *
+ */
+
+//! 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<RustRuntimeInner>`] 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<TokioState>,
+    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<Handle> {
+        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<RustRuntimeInner>,
+}
+
+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<Self> {
+        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<RustRuntimeInner>,
+}
+
+#[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::<RustRuntime>()?;
+
+    m.add_submodule(&child_module)?;
+
+    py.import("sys")?
+        .getattr("modules")?
+        .set_item("synapse.synapse_rust.runtime", child_module)?;
+
+    Ok(())
+}

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<PyAny>,
+    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<PyAny>) -> PyResult<Self> {
+    pub fn new(database_pool: &Bound<'_, PyAny>, reactor: Reactor) -> PyResult<Self> {
         Ok(Self {
             database_pool_py_ref: PyWeakrefReference::new(database_pool)?.unbind(),
             reactor,

rust/src/tokio_runtime.rs

@@ -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:
- * <https://www.gnu.org/licenses/agpl-3.0.html>.
- *
- */
-
-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<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.
-    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<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())
-}