From 065c5cb4c88cc64d6e50595b44fd004ce74eb059 Mon Sep 17 00:00:00 2001 From: Eric Eastwood Date: Tue, 23 Jun 2026 19:32:29 -0500 Subject: [PATCH] Non-working: Try to save/restore logcontext --- rust/src/deferred.rs | 121 +++++++++++++++++++++----- rust/src/handlers/versions.rs | 6 ++ rust/src/storage/db/python_db_pool.rs | 8 +- 3 files changed, 111 insertions(+), 24 deletions(-) diff --git a/rust/src/deferred.rs b/rust/src/deferred.rs index 1399a9f481..46f4fb6d7f 100644 --- a/rust/src/deferred.rs +++ b/rust/src/deferred.rs @@ -58,10 +58,46 @@ fn defer(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> { .bind(py)) } +/// A reference to the `synapse.logging.context` module. +static LOGGING_CONTEXT_MODULE: OnceCell> = OnceCell::new(); + +/// Access to the `synapse.logging.context` module. +fn logging_context_module(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> { + Ok(LOGGING_CONTEXT_MODULE + .get_or_try_init(|| py.import("synapse.logging.context").map(Into::into))? + .bind(py)) +} + +/// Set the Synapse logcontext active on the current (reactor) thread, returning +/// the context that was previously active so it can be restored afterwards. +fn set_current_logcontext<'py>( + py: Python<'py>, + context: &Bound<'py, PyAny>, +) -> PyResult> { + logging_context_module(py)?.call_method1(intern!(py, "set_current_context"), (context,)) +} + +tokio::task_local! { + /// The Synapse `LoggingContext` that was active on the reactor thread when + /// the Rust future was created (see [`create_deferred`]). + /// + /// Synapse attributes per-request CPU/DB usage via a thread-local + /// `LoggingContext`, but our work runs on a Tokio worker thread that is + /// detached from it. We stash the originating context here so that whenever + /// we hop back onto the reactor thread to drive Python work (see + /// [`run_awaitable`]) we can re-activate it; otherwise that work runs in + /// the sentinel context and its resource usage (e.g. database time) is lost + /// instead of being charged to the request. + static LOGGING_CONTEXT: Py; +} + /// Creates a twisted deferred from the given future, spawning the task on the /// tokio runtime. /// -/// Does not handle deferred cancellation or contextvars. +/// Captures the Synapse logcontext active on the reactor thread (so work driven +/// by the future is attributed to the originating request, see +/// [`LOGGING_CONTEXT`]), but does not handle deferred cancellation or +/// contextvars. pub fn create_deferred<'py, F, O>( py: Python<'py>, reactor: &Bound<'py, PyAny>, @@ -75,9 +111,17 @@ where let deferred_callback = deferred.getattr("callback")?.unbind(); let deferred_errback = deferred.getattr("errback")?.unbind(); + // Capture the logcontext active on the reactor thread now, while we're still + // on it, so the spawned task can re-apply it when it hops back to drive + // Python work. If there's no real context this is the sentinel, and + // re-applying it later is a no-op. + let logging_context = logging_context_module(py)? + .call_method0(intern!(py, "current_context"))? + .unbind(); + let rt = runtime(reactor)?; let handle = rt.handle()?; - let task = handle.spawn(fut); + let task = handle.spawn(LOGGING_CONTEXT.scope(logging_context, fut)); // Unbind the reactor so that we can pass it to the task let reactor = reactor.clone().unbind(); @@ -117,33 +161,53 @@ where make_deferred_yieldable(py, &deferred) } -/// Runs `make_deferred` on the Twisted reactor thread to obtain a Deferred (or -/// coroutine), then resolves once that Deferred fires. +/// Runs a Python awaitable to completion on the Twisted reactor and resolves +/// with its result. /// /// This is the inverse of [`create_deferred`]: where that turns a Rust future -/// into a Twisted Deferred, this turns a Twisted Deferred into an awaitable Rust -/// future. +/// into a Twisted `Deferred`, this turns a Python awaitable into a Rust future. /// -/// We're called on a Tokio worker thread, but Twisted `Deferred`s (and the -/// coroutine that `ensureDeferred` drives) are not thread-safe and Synapse's -/// logcontext is thread-local, so the coroutine must both start and resume on -/// the reactor thread. The `callFromThread` hop is what gets us there for the -/// kickoff; the deferred's own callbacks then fire on the reactor thread too. -/// (Note this is unrelated to offloading the DB work onto a thread — that's -/// handled internally by whatever `make_deferred` calls, e.g. `runInteraction`.) +/// `make_awaitable` is invoked on the reactor thread to produce the awaitable. +/// In Synapse that's typically a coroutine (e.g. from calling an `async def` +/// like `runInteraction`), though a `Deferred` works too — both are awaitable, +/// but a coroutine is not itself a `Deferred`. /// -/// `make_deferred` is invoked on the reactor thread and may return either a -/// coroutine or a `Deferred`; `ensureDeferred` normalises both to a `Deferred`. -pub(crate) async fn await_deferred(reactor: Py, make_deferred: F) -> PyResult> +/// Despite returning a future, the awaitable is kicked off in the background and +/// runs to completion regardless of whether the returned Rust future is ever +/// polled; awaiting it only observes the result. +/// +/// We're called on a Tokio worker thread, but Python awaitables and Synapse's +/// logcontext (thread-local) are not thread-safe, so the awaitable must be both +/// started and driven on the reactor thread. The `callFromThread` hop is what +/// gets us there for the kickoff; the resulting deferred's callbacks then fire +/// on the reactor thread too. (Note this is unrelated to offloading the DB work +/// onto a thread — that's handled internally by whatever `make_awaitable` +/// produces, e.g. `runInteraction`.) +/// +/// If a logcontext was captured when the task was created (see +/// [`create_deferred`]), it is re-activated while the awaitable is kicked off so +/// that its resource usage (e.g. database time) is attributed to the originating +/// request. We drive it with `run_coroutine_in_background` so this follows the +/// Synapse logcontext rules (see `docs/log_contexts.md`): the awaitable runs in +/// the request context but the logcontext is reset back to the sentinel once the +/// work completes, so the request context doesn't leak into the reactor. +pub(crate) async fn run_awaitable(reactor: Py, make_awaitable: F) -> PyResult> where F: for<'py> Fn(Python<'py>) -> PyResult> + Send + 'static, { - // Resolves when the deferred fires; carries the resolved value or the error. + // Resolves when the awaitable completes; carries the resolved value or error. let (tx, rx) = oneshot::channel::>>(); // Shared between the success and error callbacks (only one ever fires). let sender = Arc::new(Mutex::new(Some(tx))); + // The logcontext that is active for this thread + let calling_logging_context = LOGGING_CONTEXT + .try_with(|ctx| Python::attach(|py| ctx.clone_ref(py))) + .ok(); + Python::attach(|py| -> PyResult<()> { + // Create some deferred success/error callback functions that we will use to get + // the result from Python to Rust. let success_sender = Arc::clone(&sender); let on_success = PyCFunction::new_closure( py, @@ -174,14 +238,31 @@ where )? .unbind(); + // Create a function that we will run with the Twisted reactor that will drive + // the Python awaitable. let starter = PyCFunction::new_closure( py, None, None, move |args, _kwargs| -> PyResult> { let py = args.py(); - let deferred = defer(py)? - .call_method1(intern!(py, "ensureDeferred"), (make_deferred(py)?,))?; + + // Re-activate the logcontext while we kick the awaitable off. + let previous_context = calling_logging_context + .as_ref() + .map(|ctx| set_current_logcontext(py, ctx.bind(py))) + .transpose()?; + + // We use `run_in_background` to ensure the logcontext rules are being followed + let deferred = logging_context_module(py)? + .call_method1(intern!(py, "run_in_background"), (make_awaitable,)); + + // Put the reactor thread back to the logcontext we found it in + if let Some(previous_context) = previous_context { + set_current_logcontext(py, &previous_context)?; + } + + let deferred = deferred?; deferred.call_method1( intern!(py, "addCallbacks"), (on_success.bind(py), on_error.bind(py)), @@ -200,7 +281,7 @@ where match rx.await { Ok(result) => result, Err(_) => Err(PyRuntimeError::new_err( - "await_deferred channel closed before the deferred fired", + "run_awaitable channel closed before the awaitable completed", )), } } diff --git a/rust/src/handlers/versions.rs b/rust/src/handlers/versions.rs index a3fc4bbd85..aeb2efaaef 100644 --- a/rust/src/handlers/versions.rs +++ b/rust/src/handlers/versions.rs @@ -121,6 +121,12 @@ async fn build_versions_response( None => global_unstable_feature_map.msc3575, }; + log::info!( + "asdf msc3881_enabled={:?} msc3575_enabled={:?}", + msc3881_enabled, + msc3575_enabled + ); + let unstable_feature_map = UnstableFeatureMap { msc3575: msc3575_enabled, msc3881: msc3881_enabled, diff --git a/rust/src/storage/db/python_db_pool.rs b/rust/src/storage/db/python_db_pool.rs index 781b4b8473..79140042e5 100644 --- a/rust/src/storage/db/python_db_pool.rs +++ b/rust/src/storage/db/python_db_pool.rs @@ -30,7 +30,7 @@ use pyo3::{ types::{PyBool, PyCFunction, PyFloat, PyInt, PyList, PyString}, }; -use crate::deferred::await_deferred; +use crate::deferred::run_awaitable; use crate::storage::db::{DatabasePool, Row, Transaction, Value}; /// The database engines we support in the Python side of Synapse @@ -151,10 +151,10 @@ impl DatabasePool for PythonDatabasePoolWrapper { .map_err(anyhow::Error::from)?; // Await `runInteraction` directly. `runInteraction` offloads the actual - // DB work onto a thread itself, so we don't; `await_deferred` only has to - // start the coroutine on the reactor thread and bridge its deferred back + // DB work onto a thread itself, so we don't; `run_awaitable` only has to + // start the coroutine on the reactor thread and bridge its result back // into our `async` world. - let outcome = await_deferred(reactor, move |py| { + let outcome = run_awaitable(reactor, move |py| { database_pool_py .bind(py) .call_method1(intern!(py, "runInteraction"), (name, callback.bind(py)))