Port LogContext to Rust

changelog.d/19979.misc

@@ -0,0 +1 @@
+Port the logcontext machinery (`LoggingContext`, `ContextResourceUsage` and the current-context storage) to Rust.

docs/log_contexts.md

@@ -2,9 +2,15 @@
 
 To help track the processing of individual requests, synapse uses a
 '`log context`' to track which request it is handling at any given
-moment. This is done via a thread-local variable; a `logging.Filter` is
-then used to fish the information back out of the thread-local variable
-and add it to each log record.
+moment. The "current" log context is stored in the Rust extension
+(`synapse.synapse_rust.logcontext`), which resolves it from the current
+tokio task (if we are running inside one) and otherwise from the current OS
+thread; a `logging.Filter` is then used to fish the information back out and
+add it to each log record. Storing it in Rust means a single source of truth is
+visible from both Python (the reactor and its thread pool) and Rust (tokio
+tasks), so log records emitted from either — including from Rust code polled on
+a worker thread — are attributed to the right request. See
+[the Rust side](#the-rust-side) below.
 
 Logcontexts are also used for CPU and database accounting, so that we
 can track which requests were responsible for high CPU use or database
@@ -550,6 +556,29 @@ actually happen too much. Unfortunately, when it does happen, it will
 lead to leaked logcontexts which are incredibly hard to track down.
 
 
+## The Rust side
+
+The "current" logcontext is stored in the Rust extension rather than in a Python
+thread-local, so that it is visible from both worlds. `current_context()` and
+`set_current_context()` are imported from `synapse.logging.context` as usual —
+the Rust storage is an implementation detail that Python code does not need to
+care about.
+
+The switch itself (`set_current_context`) only ever runs on the reactor (or its
+thread pool) — the Python side — where it does the `getrusage` CPU accounting.
+It is never driven from a tokio worker thread.
+
+What Rust code *does* need to be aware of: when you spawn a future onto the tokio
+runtime, the current logcontext must be captured and carried along so that log
+records emitted while the future is polled (including any `log::` records from
+dependencies, and any Python invoked back from Rust) are attributed correctly.
+Use the provided helper — `LogContextHandle::capture(py)` plus `LogContextHandle::scope` in
+`rust/src/logging/context.rs` — which captures the caller's logcontext at the FFI
+boundary and scopes it onto the spawned task (this is what `create_deferred`
+does), rather than a bare `tokio::spawn`. `current_context()` resolves the task's captured
+context first, so `LoggingContextFilter` — and therefore `pyo3-log` — resolves the
+right context on worker threads, with no per-log-record stamping.
+
 ## Debugging logcontext issues
 
 Debugging logcontext issues can be tricky as leaking or losing a logcontext will surface
@@ -564,3 +593,18 @@ loggers:
     synapse.logging.context.debug:
         level: DEBUG
 ```
+
+Note that some of these traces (`LoggingContext(...).__enter__`/`__exit__`) are
+emitted from Rust via `pyo3-log`, which caches logger levels for performance.
+Configuring the logger in the logging config (as above) works — the cache is
+flushed whenever the config is (re)loaded, including on `SIGHUP` — but enabling
+the logger *at runtime* (e.g. `setLevel(logging.DEBUG)` from the manhole) will
+not surface the Rust-emitted traces until you also flush the cache:
+
+```python
+import logging
+from synapse.synapse_rust import reset_logging_config
+
+logging.getLogger("synapse.logging.context.debug").setLevel(logging.DEBUG)
+reset_logging_config()
+```

rust/Cargo.toml

@@ -29,6 +29,7 @@ bytes = "1.6.0"
 headers = "0.4.0"
 http = "1.1.0"
 lazy_static = "1.4.0"
+libc = "0.2.174"
 log = "0.4.17"
 mime = "0.3.17"
 pyo3 = { version = "0.28.3", features = [

rust/src/deferred.rs

@@ -18,13 +18,15 @@ use std::{
     sync::{Arc, Mutex},
 };
 
+use log::{debug, log_enabled, Level};
 use once_cell::sync::OnceCell;
 use pyo3::{
     create_exception, exceptions::PyException, exceptions::PyRuntimeError, intern, prelude::*,
     types::PyCFunction,
 };
 use tokio::sync::oneshot;
 
+use crate::logging::context::{with_logcontext, DEBUG_LOGGER_NAME};
 use crate::tokio_runtime::runtime;
 
 create_exception!(
@@ -71,7 +73,13 @@ fn logging_context_module(py: Python<'_>) -> PyResult<&Bound<'_, PyAny>> {
 /// Creates a twisted deferred from the given future, spawning the task on the
 /// tokio runtime.
 ///
-/// Does not handle deferred cancellation or contextvars.
+/// Does not handle contextvars.
+///
+/// TODO: propagate deferred cancellation to the tokio task (via
+/// `JoinHandle::abort`). Until then a cancelled request leaves its task
+/// running, so the task can outlive the request's logcontext —
+/// `run_python_awaitable` defends against the resulting finished-context case,
+/// but the work itself is wasted.
 pub fn create_deferred<'py, F, O>(
     py: Python<'py>,
     reactor: &Bound<'py, PyAny>,
@@ -85,9 +93,16 @@ where
     let deferred_callback = deferred.getattr("callback")?.unbind();
     let deferred_errback = deferred.getattr("errback")?.unbind();
 
+    // Capture the caller's logcontext at the boundary (GIL held, on the reactor
+    // thread) and scope it onto the spawned task, so that logging emitted while
+    // the future is polled — and any `run_python_awaitable` callbacks back into
+    // Python — are attributed to the context that was 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 task = handle.spawn(fut);
+    let task = handle.spawn(logcontext.scope(fut));
 
     // Unbind the reactor so that we can pass it to the task
     let reactor = reactor.clone().unbind();
@@ -148,7 +163,15 @@ where
     // Shared between the success and error callbacks (only one ever fires).
     let sender = Arc::new(Mutex::new(Some(tx)));
 
-    Python::attach(|py| -> PyResult<()> {
+    // Capture the logcontext of the calling tokio task (if any). We restore it on
+    // the reactor thread before driving the awaitable, so Python code invoked from
+    // Rust (e.g. `DatabasePool.runInteraction`) runs in the same logcontext that was
+    // current when Python originally called into Rust — its logging and DB-metrics
+    // accounting are then attributed to the right request. `None` (called outside a
+    // scoped task) falls back to the sentinel.
+    let logcontext = crate::logging::context::LogContextHandle::current();
+
+    Python::attach(move |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);
@@ -216,21 +239,61 @@ where
             move |args, _kwargs| -> PyResult<Py<PyAny>> {
                 let py = args.py();
 
-                // We fire-and-forget using `run_in_background`. Re-using
-                // `run_in_background` also makes sure the awaitable gets run with the
-                // current logcontext while following the logcontext rules.
+                // Choose the logcontext to drive the awaitable in: the captured
+                // one, restored on the reactor thread — the one thread where the
+                // context's `main_thread` affinity check passes.
                 //
-                // FIXME: Currently runs in the sentinel logcontext because we don't manage it here
-                let deferred = logging_context_module(py)?.call_method1(
-                    intern!(py, "run_in_background"),
-                    (awaitable_factory.bind(py),),
-                );
-
-                let deferred = deferred?;
-                deferred.call_method1(
-                    intern!(py, "addCallbacks"),
-                    (on_success.bind(py), on_error.bind(py)),
-                )?;
+                // Never re-start a context that has already finished: the request may
+                // have completed (or been cancelled — `create_deferred` does not
+                // propagate cancellation) while this task was still running. Restoring
+                // it would trip the "Re-starting finished log context" abuse check and
+                // account our work against a context whose metrics are already
+                // finalised, so such work runs in the sentinel instead. Both
+                // `__exit__` (which sets `finished`) and this check run on the
+                // reactor thread, so the check cannot race.
+                let context = match &logcontext {
+                    Some(handle) => {
+                        let finished = handle
+                            .logging_context()
+                            .is_some_and(|ctx| ctx.borrow(py).is_finished());
+                        if finished {
+                            if log_enabled!(target: DEBUG_LOGGER_NAME, Level::Debug) {
+                                // Only a real context can be finished, so
+                                // `logging_context()` is `Some` here.
+                                if let Some(ctx) = handle.logging_context() {
+                                    debug!(
+                                        target: DEBUG_LOGGER_NAME,
+                                        "run_python_awaitable: captured logcontext {} has \
+                                         finished; running in the sentinel",
+                                        ctx.bind(py).str()?
+                                    );
+                                }
+                            }
+                            None
+                        } else {
+                            handle.logging_context().map(|ctx| ctx.clone_ref(py))
+                        }
+                    }
+                    // Called from outside any scoped task: the sentinel. (The
+                    // reactor thread is normally at the sentinel already, in which
+                    // case the switch below is a no-op.)
+                    None => None,
+                };
+
+                // Kick off the awaitable, fire-and-forget, via `run_in_background`:
+                // it calls the factory in the current logcontext and follows the
+                // logcontext rules from there — in particular, it arranges for the
+                // reactor to be back at the sentinel when the awaitable later
+                // completes.
+                with_logcontext(py, context, || {
+                    let deferred = run_in_background(py, awaitable_factory.bind(py))?;
+                    deferred.call_method1(
+                        intern!(py, "addCallbacks"),
+                        (on_success.bind(py), on_error.bind(py)),
+                    )?;
+                    Ok(())
+                })?;
+
                 Ok(py.None())
             },
         )?;
@@ -270,6 +333,26 @@ fn failure_to_pyerr(failure: &Bound<'_, PyAny>) -> PyErr {
     }
 }
 
+/// A reference to `synapse.logging.context.run_in_background`.
+static RUN_IN_BACKGROUND: OnceCell<Py<PyAny>> = OnceCell::new();
+
+/// Call `synapse.logging.context.run_in_background(f)`: call `f` in the current
+/// logcontext and drive the awaitable it returns to completion, following the
+/// logcontext rules. Returns the resulting `Deferred`.
+fn run_in_background<'py>(py: Python<'py>, f: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
+    let run_in_background = RUN_IN_BACKGROUND.get_or_try_init(|| {
+        logging_context_module(py)?
+            .getattr("run_in_background")
+            .map(Into::into)
+    })?;
+
+    run_in_background
+        .call1(py, (f,))?
+        .extract(py)
+        .map_err(Into::into)
+}
+
+/// A reference to `synapse.logging.context.make_deferred_yieldable`.
 static MAKE_DEFERRED_YIELDABLE: OnceCell<Py<PyAny>> = OnceCell::new();
 
 /// Given a deferred, make it follow the Synapse logcontext rules

rust/src/lib.rs

@@ -16,6 +16,7 @@ pub mod http;
 pub mod http_client;
 pub mod identifier;
 pub mod json;
+pub mod logging;
 pub mod matrix_const;
 pub mod msc4388_rendezvous;
 pub mod push;
@@ -70,6 +71,7 @@ fn synapse_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(reset_logging_config, m)?)?;
 
     acl::register_module(py, m)?;
+    logging::context::register_module(py, m)?;
     deferred::register_module(py, m)?;
     push::register_module(py, m)?;
     events::register_module(py, m)?;