diff --git a/changelog.d/20252.misc b/changelog.d/20252.misc new file mode 100644 index 0000000000..a8e1359f34 --- /dev/null +++ b/changelog.d/20252.misc @@ -0,0 +1 @@ +Avoid taking the GIL on Tokio worker threads when completing Rust futures. diff --git a/rust/src/deferred.rs b/rust/src/deferred.rs index 9f438993fb..1931862e79 100644 --- a/rust/src/deferred.rs +++ b/rust/src/deferred.rs @@ -108,7 +108,8 @@ where handle.spawn(async move { let res = task.await; - Python::attach(move |py| { + // Once done pass the result to the Twisted reactor thread for handling. + runtime.dispatch_to_twisted(move |py| { // Flatten the panic into standard python error let res = match res { Ok(r) => r, @@ -119,21 +120,18 @@ where }; // Send the result to the deferred, via `.callback(..)` or `.errback(..)` - match res { - Ok(obj) => { - 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) => { - runtime - .reactor() - .call_from_thread(py, (deferred_errback, err)) - .expect("callFromThread should not fail"); // There's nothing we can really do with errors here - } + let fired = match res { + Ok(obj) => deferred_callback.call1(py, (obj,)), + Err(err) => deferred_errback.call1(py, (err,)), + }; + + if let Err(err) = fired { + // There is nowhere to propagate this to. The closure runs from + // the dispatch reader's `doRead`, and an exception out of that + // makes Twisted drop the reader. Log it instead. + log::error!("Failed to fire a deferred from a Rust future: {err}"); } - }); + }) }); // Make the deferred follow the Synapse logcontext rules diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 3d965daf51..af1162709c 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -27,6 +27,7 @@ pub mod room_versions; pub mod runtime; pub mod segmenter; pub mod storage; +pub mod twisted_dispatch; pub mod types; lazy_static! { diff --git a/rust/src/logging/context.rs b/rust/src/logging/context.rs index d360f57cf6..d01223e7b7 100644 --- a/rust/src/logging/context.rs +++ b/rust/src/logging/context.rs @@ -975,23 +975,22 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> Ok(()) } +/// Helpers for Rust unit tests elsewhere in the crate that need a logcontext. #[cfg(test)] -mod tests { - use std::sync::Arc; - +pub(crate) mod testing { use pyo3::types::PyString; use super::*; /// A minimal `LoggingContext` for these tests, built directly (bypassing /// `__init__`, which would capture the current context and thread id). - fn test_context(py: Python<'_>, name: &str) -> Py { + pub(crate) fn test_context(py: Python<'_>, name: &str) -> Py { Py::new( py, LoggingContext { name: PyString::new(py, name).unbind(), server_name: PyString::new(py, "test_server").unbind(), - main_thread: 0, + main_thread: get_thread_id(), finished: false, usage_start: None, tag: Some(String::new()), @@ -1005,6 +1004,14 @@ mod tests { ) .expect("failed to allocate LoggingContext") } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::testing::test_context; + use super::*; #[test] fn thread_local_defaults_to_sentinel() { diff --git a/rust/src/reactor.rs b/rust/src/reactor.rs index 7522cd3390..329057af22 100644 --- a/rust/src/reactor.rs +++ b/rust/src/reactor.rs @@ -49,6 +49,30 @@ impl Reactor { Ok(()) } + /// `reactor.addReader(reader)`: have the reactor poll `reader.fileno()` + /// and call `reader.doRead()` when it is readable. + /// + /// Must be called on the reactor thread. + pub fn add_reader(&self, py: Python<'_>, reader: &Bound<'_, PyAny>) -> PyResult<()> { + self.0 + .bind(py) + .call_method1(intern!(py, "addReader"), (reader,))?; + + Ok(()) + } + + /// `reactor.removeReader(reader)`: stop polling a reader added with + /// [`Reactor::add_reader`]. A no-op if it was never added. + /// + /// Must be called on the reactor thread. + pub fn remove_reader(&self, py: Python<'_>, reader: &Bound<'_, PyAny>) -> PyResult<()> { + self.0 + .bind(py) + .call_method1(intern!(py, "removeReader"), (reader,))?; + + Ok(()) + } + pub fn clone_ref(&self, py: Python<'_>) -> Reactor { Reactor(self.0.clone_ref(py)) } diff --git a/rust/src/runtime.rs b/rust/src/runtime.rs index 46ceeec1a8..f89f24957d 100644 --- a/rust/src/runtime.rs +++ b/rust/src/runtime.rs @@ -17,8 +17,7 @@ //! //! 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 +//! homeserver. 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. //! @@ -35,6 +34,7 @@ use tokio::runtime::{Handle, Runtime}; use crate::homeserver::HomeServer; use crate::reactor::Reactor; +use crate::twisted_dispatch::{self, TwistedDispatchReader, TwistedDispatcher}; /// How long to wait for in-flight tokio tasks to be cancelled when shutting /// down with the reactor. @@ -60,6 +60,13 @@ pub struct RustRuntimeInner { reactor: Reactor, tokio: Mutex, worker_threads: usize, + + /// Runs closures on the Twisted reactor thread without taking the GIL on + /// the calling thread. See [`crate::twisted_dispatch`]. + dispatcher: Arc, + /// The reactor-facing half of `dispatcher`, registered with the Twisted + /// reactor via `addReader` until `shutdown`. + dispatch_reader: Py, } impl RustRuntimeInner { @@ -68,6 +75,18 @@ impl RustRuntimeInner { &self.reactor } + /// Queue `f` to run on the Twisted reactor thread with the GIL held, and + /// wake the reactor. Never takes the GIL itself, so a tokio task can call + /// it to hand a result back to Twisted. See [`crate::twisted_dispatch`]. + /// + /// Returns an error once the homeserver has shut down. + pub fn dispatch_to_twisted(&self, f: F) -> anyhow::Result<()> + where + F: FnOnce(Python<'_>) + Send + 'static, + { + self.dispatcher.dispatch(f) + } + /// Get a handle to the tokio runtime, starting the runtime if it hasn't /// been started yet. pub fn tokio_handle(&self) -> PyResult { @@ -117,6 +136,17 @@ impl RustRuntimeInner { py.detach(|| runtime.shutdown_timeout(SHUTDOWN_TIMEOUT)); } + // Unregister the wakeup socket from the Twisted reactor, then close + // the dispatcher and run the closures still queued so that their + // deferreds fire. + // + // All tokio tasks should be stopped by now, so only long running + // blocking threads may still be active. If they try to dispatch new + // work they get an error and should stop. + self.reactor + .remove_reader(py, self.dispatch_reader.bind(py).as_any())?; + self.dispatch_reader.get().close_and_drain(py); + Ok(()) } } @@ -161,10 +191,21 @@ impl RustRuntime { #[new] #[pyo3(signature = (hs, worker_threads = 4))] fn py_new(py: Python<'_>, hs: HomeServer, worker_threads: usize) -> PyResult { + let reactor = hs.get_reactor(py)?; + + // Register the read end of the dispatcher's wakeup socket with the + // Twisted reactor, so that closures dispatched by tokio tasks run on + // the reactor thread. `shutdown` removes it again. + let (dispatcher, dispatch_reader) = twisted_dispatch::new_pair()?; + let dispatch_reader = Py::new(py, dispatch_reader)?; + reactor.add_reader(py, dispatch_reader.bind(py).as_any())?; + let inner = Arc::new(RustRuntimeInner { - reactor: hs.get_reactor(py)?, + reactor, tokio: Mutex::new(TokioState::NotStarted), worker_threads, + dispatcher, + dispatch_reader, }); // Shut the tokio runtime down when the homeserver is shut down. The diff --git a/rust/src/twisted_dispatch.rs b/rust/src/twisted_dispatch.rs new file mode 100644 index 0000000000..158dd0b9c6 --- /dev/null +++ b/rust/src/twisted_dispatch.rs @@ -0,0 +1,335 @@ +/* + * 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: + * . + * + */ + +//! Run closures on the Twisted reactor thread from Tokio worker threads, +//! without taking the GIL on the worker. +//! +//! Taking the GIL on a worker would block the worker until the Python thread +//! releases it, which under load means waiting for the interpreter's switch +//! interval. This could severely impact the performance of the Tokio runtime. +//! +//! # Implementation +//! +//! The implementation combines a queue of closures with the self-pipe trick, +//! which is the same implementation as Twisted uses for `callFromThread`. +//! +//! ```text +//! Tokio worker Twisted reactor +//! │ │ +//! TwistedDispatcher | +//! ├─ Add a closure to a Rust queue │ +//! └─ Write one byte to socket A ───────► socket B becomes readable +//! │ +//! TwistedDispatchReader +//! └─ Reader.doRead() +//! ├─ Drain socket B +//! ├─ Take queued closures +//! └─ Execute them +//! ``` +//! +//! The [`TwistedDispatcher`] worker holds a queue of closures and the write end +//! of a unix socket pair. A worker pushes a closure and writes one byte. The +//! [`TwistedDispatchReader`] holds the read end and is registered with the +//! Twisted reactor via `addReader`. The Twisted reactor is woken up and calls +//! the `doRead` method on the Twisted reactor thread, which then reads from the +//! queue and runs any pending closures. +//! +//! One pair exists per homeserver, owned by [`crate::runtime::RustRuntime`], +//! which registers the reader on construction and removes it at shutdown. +//! +//! # Homeserver shutdown +//! +//! When the homeserver shuts down, the dispatcher is closed and `dispatch` will +//! return an error. The caller should stop what it is doing as its associated +//! homeserver has been shut down. + +use std::{ + any::Any, + io::{ErrorKind, Read, Write}, + os::{ + fd::{AsRawFd, RawFd}, + unix::net::UnixStream, + }, + panic::{catch_unwind, AssertUnwindSafe}, + sync::{Arc, Mutex}, +}; + +use anyhow::Context; +use log::error; +use pyo3::prelude::*; + +use crate::logging::context::with_logcontext; + +/// A queued closure. Runs on the Twisted reactor thread. +type DispatchedCall = Box) + Send + 'static>; + +/// The Tokio-facing half: a queue of closures plus the write end of the wakeup +/// socket. +pub struct TwistedDispatcher { + /// The queue of closures to run on the Twisted reactor thread. Or None if + /// the dispatcher has been closed, indicating the homeserver has been + /// shutdown. + queue: Mutex>>, + write_end: UnixStream, +} + +impl TwistedDispatcher { + /// Queue `f` to run on the Twisted reactor thread, and wake the reactor. + /// + /// Never takes the GIL, so this is safe to call from a Tokio worker. + /// + /// Returns an error if the dispatcher has been closed, which indicates the + /// homeserver has been shutdown. + pub fn dispatch(&self, f: F) -> anyhow::Result<()> + where + F: FnOnce(Python<'_>) + Send + 'static, + { + // First add the closure to the queue. + { + let mut queue = self.queue.lock().expect("dispatcher poisoned"); + let Some(queue) = &mut *queue else { + return Err(anyhow::anyhow!("dispatcher is closed")); + }; + + queue.push(Box::new(f)); + } + + // Second, write a byte to the socket to wakeup the Twisted reactor. + // + // Errors are ignored. `WouldBlock` means the socket buffer is full, so + // a wakeup is already pending. Nothing else can fail while the reader + // holds the other end. + let _ = (&self.write_end).write(b"x"); + + Ok(()) + } +} + +/// The Twisted-reactor-facing half of a [`TwistedDispatcher`]. +#[pyclass(frozen)] +pub struct TwistedDispatchReader { + dispatcher: Arc, + read_end: UnixStream, +} + +#[pymethods] +impl TwistedDispatchReader { + fn fileno(&self) -> RawFd { + self.read_end.as_raw_fd() + } + + #[pyo3(name = "logPrefix")] + fn log_prefix(&self) -> &'static str { + "synapse-rust-twisted-dispatch" + } + + /// Called by the Twisted reactor when the wakeup socket has something in it. + /// + /// Runs on the reactor thread. + #[pyo3(name = "doRead")] + fn do_read(&self, py: Python<'_>) { + // Drain the socket before taking the queue. This way around avoids + // races between the socket being drained and closures being added to + // the queue. + let mut buf = [0u8; 64]; + loop { + match (&self.read_end).read(&mut buf) { + Ok(0) => break, + Ok(_) => continue, + // Interrupted: try again. + Err(err) if err.kind() == ErrorKind::Interrupted => continue, + // WouldBlock: the socket buffer is empty, so we can stop + // draining. + Err(err) if err.kind() == ErrorKind::WouldBlock => break, + Err(err) => { + // This should not happen, as the socket is never closed + // except when this struct is dropped. + // + // We can't do much here except log the error and break out + // of the loop. + error!("Error reading from wakeup socket: {:?}", err); + break; + } + } + } + + // Take from the queue. Ensuring the lock is released before the + // closures run. This may be `None` if the dispatcher has been closed. + let jobs = self + .dispatcher + .queue + .lock() + .expect("dispatcher poisoned") + .as_mut() + .map(std::mem::take); + if let Some(jobs) = jobs { + self.run_dispatched(py, jobs); + } + + // Returning `None` keeps the reader registered. + } + + /// Called by the Twisted reactor at shutdown. The socket closes when the + /// owning `RustRuntime` drops the reader, so there is nothing to do here. + #[pyo3(name = "connectionLost")] + fn connection_lost(&self, _reason: &Bound<'_, PyAny>) {} +} + +impl TwistedDispatchReader { + /// Close the dispatcher so that no further closures can be queued, then + /// run the ones already queued. + pub fn close_and_drain(&self, py: Python<'_>) { + let Some(jobs) = self + .dispatcher + .queue + .lock() + .expect("dispatcher poisoned") + .take() + else { + return; + }; + + self.run_dispatched(py, jobs); + } + + /// Run closures taken from the queue. Must be called with the queue lock + /// released. + fn run_dispatched(&self, py: Python<'_>, jobs: Vec) { + for job in jobs { + // A panic here would leave `doRead` as an exception, and Twisted + // responds to that by dropping the reader. No later closure would + // ever run. Log it instead. + // + // `catch_unwind(AssertUnwindSafe(..))` is exactly the same as what + // pyo3 does when catching panics in Rust code called from Python. + // + // `UnwindSafe` guards against code after `catch_unwind` reading + // state that the panicking code was part way through mutating. + // Since we never read `job` after this, it is safe to catch the + // panic. `Python` is already `UnwindSafe`. + let ran = catch_unwind(AssertUnwindSafe(move || { + // Ensure the dispatched job gets run within the sentinel + // logcontext, and that the sentinel context is restored + // afterwards. + with_logcontext(py, None, move || { + job(py); + Ok(()) + }) + })); + + match ran { + Ok(Ok(())) => {} + Ok(Err(err)) => { + // Log an error if setting the logcontexts in the dispatched + // Rust completion fails. + log::error!( + "error setting logcontexts in dispatched Rust completion: {}", + err + ); + } + Err(err) => { + // Log an error if the dispatched Rust completion panicked. + log::error!( + "panic while running a dispatched Rust completion: {}", + panic_payload_as_str(&*err) + ); + } + } + } + } +} + +/// Convert a panic payload to a string for logging purposes. +/// +/// This is the object passed to [`panic!`] and returned from [`catch_unwind`]. +/// The vast majority of panics are either a `&'static str` or a `String`, but +/// technically they could be anything. +fn panic_payload_as_str(payload: &dyn Any) -> &str { + if let Some(&s) = payload.downcast_ref::<&'static str>() { + s + } else if let Some(s) = payload.downcast_ref::() { + s.as_str() + } else { + "" + } +} + +/// Build a connected [`TwistedDispatcher`]/[`TwistedDispatchReader`] pair. +pub fn new_pair() -> PyResult<(Arc, TwistedDispatchReader)> { + let (read_end, write_end) = + UnixStream::pair().context("creating the Tokio wakeup socket pair")?; + + // Neither end may block. The writer is a Tokio worker and the reader is + // the Twisted reactor thread. + read_end + .set_nonblocking(true) + .context("making the Tokio wakeup socket non-blocking")?; + write_end + .set_nonblocking(true) + .context("making the Tokio wakeup socket non-blocking")?; + + let dispatcher = Arc::new(TwistedDispatcher { + queue: Mutex::new(Some(Vec::new())), + write_end, + }); + let reader = TwistedDispatchReader { + dispatcher: Arc::clone(&dispatcher), + read_end, + }; + + Ok((dispatcher, reader)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::logging::context::{current_context, set_current_context, testing::test_context}; + + /// `close_and_drain` runs from the homeserver's shutdown handler, in the + /// caller's logcontext. Ensure that the logcontext is correctly preserved + /// across `close_and_drain` no matter what the dispatch job leaves it as. + #[test] + fn close_and_drain_runs_at_sentinel_and_restores_caller_logcontext() { + Python::initialize(); + Python::attach(|py| { + let (dispatcher, reader) = new_pair().expect("creating the dispatcher pair"); + + // Create and set the "shutdown" log context. + let caller = test_context(py, "shutdown"); + set_current_context(py, Some(caller.clone_ref(py))).expect("setting caller context"); + + // Dispatch a closure that changes the logcontext. + dispatcher + .dispatch(move |py| { + // Mimic the dispatch job not restoring the sentinel context. + let awaiter = test_context(py, "requester"); + set_current_context(py, Some(awaiter)).expect("switching to the awaiter"); + }) + .expect("dispatching while open"); + + // Call `close_and_drain`, which should restore the caller's + // context. + reader.close_and_drain(py); + + // Check the caller's context is current again. + let after = current_context(py).expect("caller's context was lost"); + assert!(after.is(&caller), "caller's context was not restored"); + + // Reset to the sentinel so we don't leak into another test that + // reuses this OS thread from the test harness's pool. + set_current_context(py, None).expect("resetting to the sentinel"); + }); + } +} diff --git a/tests/server.py b/tests/server.py index 5297e0b1ac..71f9da9c4c 100644 --- a/tests/server.py +++ b/tests/server.py @@ -24,6 +24,7 @@ import json import logging import os import os.path +import select import sqlite3 import time import uuid @@ -69,6 +70,7 @@ from twisted.internet.interfaces import ( IPushProducer, IReactorPluggableNameResolver, IReactorTime, + IReadDescriptor, IResolverSimple, ITCPTransport, ITransport, @@ -757,6 +759,29 @@ class ThreadedMemoryReactorClock(MemoryReactorClock): # main thread. super().advance(0) + # Now poll anything registered with `addReader`. A real reactor does + # this in its poll loop, but `MemoryReactor` only stores the readers, so + # results from Rust futures (see `TwistedDispatch`) would never reach + # their deferreds. Firing those deferreds can in turn queue more + # callbacks hence the recursive `advance(0)`. + readable = self._poll_readers() + if readable: + for reader in readable: + reader.doRead() + self.advance(0) + + def _poll_readers(self) -> list[IReadDescriptor]: + """The readers registered with `addReader` that have data waiting.""" + readers = {reader.fileno(): reader for reader in self.getReaders()} + if not readers: + return [] + + # Now poll the readers to see if any have data waiting. + poller = select.poll() + for fileno in readers: + poller.register(fileno, select.POLLIN) + return [readers[fileno] for fileno, _event in poller.poll(0)] + def cleanup_test_reactor_system_event_triggers( reactor: ThreadedMemoryReactorClock,