diff --git a/Cargo.lock b/Cargo.lock index 7ca2265aca9..7a44c58863f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1342,6 +1342,7 @@ dependencies = [ "icu_segmenter", "itertools", "lazy_static", + "libc", "log", "mime", "once_cell", diff --git a/changelog.d/19979.misc b/changelog.d/19979.misc new file mode 100644 index 00000000000..3210329a56f --- /dev/null +++ b/changelog.d/19979.misc @@ -0,0 +1 @@ +Port the logcontext machinery (`LoggingContext`, `ContextResourceUsage` and the current-context storage) to Rust. diff --git a/docs/log_contexts.md b/docs/log_contexts.md index 76710e10e09..ea60fde20ed 100644 --- a/docs/log_contexts.md +++ b/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() +``` diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 612ab09f6d0..a271e4cf90a 100644 --- a/rust/Cargo.toml +++ b/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 = [ diff --git a/rust/src/deferred.rs b/rust/src/deferred.rs index 62a1ce6b903..bd40c57518d 100644 --- a/rust/src/deferred.rs +++ b/rust/src/deferred.rs @@ -18,6 +18,7 @@ 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::*, @@ -25,6 +26,7 @@ use pyo3::{ }; 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> { 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> = 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> { + 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> = OnceCell::new(); /// Given a deferred, make it follow the Synapse logcontext rules diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 28783afbbac..f68b7c17d76 100644 --- a/rust/src/lib.rs +++ b/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)?; diff --git a/rust/src/logging/context.rs b/rust/src/logging/context.rs new file mode 100644 index 00000000000..6618729e7bd --- /dev/null +++ b/rust/src/logging/context.rs @@ -0,0 +1,1136 @@ +/* + * 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: + * . + * + */ + +//! Native storage for the Synapse "current logcontext". +//! +//! The storage lives in Rust rather than in a Python `threading.local` because a +//! Python thread-local is invisible to Rust: each tokio worker thread would see +//! its own slot, permanently at the sentinel, and logging emitted from Rust — +//! including from spawned tokio tasks — could not be attributed to the request +//! that caused it. +//! +//! This module holds that storage and unifies two sources of truth so that a +//! single [`current_context`] answer is correct from *both* worlds: +//! +//! 1. a per-OS-thread slot ([`THREAD_LOCAL_CONTEXT`]) — used by the reactor +//! thread and any reactor-managed threadpool threads; and +//! 2. a per-tokio-task slot ([`TASK_LOCAL_CONTEXT`]), which rides with a task as it +//! migrates between worker threads across `.await` points. +//! +//! Both slots hold an `Option>`: `None` means "no context" — +//! what Synapse calls the sentinel. The `_Sentinel` marker object itself is pure +//! Python (`synapse.logging.context.SENTINEL_CONTEXT`); the wrappers there +//! convert between it and `None` at the boundary, so no Rust code ever sees or +//! produces the sentinel object. +//! +//! [`current_context`] consults the task-local first (when called from inside a +//! runtime task) and falls back to the thread-local. Because +//! `LoggingContextFilter` (and therefore `pyo3-log`) resolves the context by +//! calling [`current_context`] at log-record time, log records emitted while a +//! task is being polled are attributed to the task's captured context with no +//! per-record stamping machinery. +//! +//! The accounting policy is native too: [`set_current_context`] reads the thread +//! rusage via libc, runs the `stop`/`start` bookkeeping, and uses +//! [`swap_current_context`] for the raw slot write. The switch primitive is only +//! ever driven on the reactor (or threadpool) threads — never on tokio worker +//! threads — so it always writes the thread-local, and the task-local (populated +//! only by [`LogContextHandle::scope`] at spawn time) takes read precedence during a +//! poll. [`swap_current_context`] checks that invariant rather than trusting it. + +use std::{cell::RefCell, future::Future}; + +use log::{debug, error, log_enabled, Level}; +use pyo3::call::PyCallArgs; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyString, PyTuple}; +use pyo3::{intern, PyTraverseError, PyVisit}; + +/// Name of the opt-in logger for logcontext switch tracing. +/// +/// This is the single source of truth for the logger name: it is used as the +/// `debug!` `target:` for the switch traces emitted below, and it is exported to +/// Python via [`register_module`] so that `synapse.logging.context` builds +/// *exactly* this logger (see `logcontext_debug_logger` there). Keeping one +/// constant stops the Rust `target:` and the Python `getLogger` name from +/// drifting apart. The messages only surface when this logger is explicitly +/// configured — see `ExplicitlyConfiguredLogger` on the Python side, whose +/// `isEnabledFor` pyo3-log honours, so the no-inherit opt-in works from Rust too. +pub const DEBUG_LOGGER_NAME: &str = "synapse.logging.context.debug"; + +thread_local! { + /// The current logcontext for this OS thread. The slot is typed: it holds a + /// [`LoggingContext`] (possibly a Python subclass instance), and `None` means + /// "no context set on this thread", i.e. the sentinel — a fresh thread (e.g. + /// a new threadpool worker) therefore starts in the sentinel. + static THREAD_LOCAL_CONTEXT: RefCell>> = const { RefCell::new(None) }; +} + +tokio::task_local! { + /// The logcontext captured for the current tokio task, set by + /// [`LogContextHandle::scope`] when the task is spawned. Only present inside a + /// scoped task; readable synchronously during any poll of that task, + /// regardless of which worker thread the poll runs on. + static TASK_LOCAL_CONTEXT: LogContextHandle; +} + +/// A cheap, clone-able, GIL-free handle on a captured logcontext, in the same +/// representation the storage slots use: a [`LoggingContext`] (possibly a +/// Python subclass instance), or `None` for the sentinel. +/// +/// `Py` is `Send + Sync`, so this can travel with a tokio task +/// across worker threads and be dropped on a detached thread (pyo3 defers the +/// decref). Cloning only needs the GIL for the underlying object, so we clone +/// the `Py` eagerly (with the GIL) at capture time and hand out clones of the +/// handle, which are GIL-free — see [`LogContextHandle::current`], called during +/// a poll where the GIL may not be held. +#[derive(Clone)] +pub struct LogContextHandle { + // Held behind an `Arc` so that cloning the handle (e.g. `LogContextHandle::current`, + // called during a poll where the GIL may not be held) and dropping it are + // both GIL-free; cloning a bare `Py` would require the GIL. + context: std::sync::Arc>>, +} + +impl LogContextHandle { + /// Capture the calling thread's current logcontext. + /// + /// Must be called with the GIL held, on the thread whose context we want + /// (i.e. at the FFI boundary, before spawning onto tokio). + pub fn capture(py: Python<'_>) -> Self { + LogContextHandle { + context: std::sync::Arc::new(current_context(py)), + } + } + + /// The logcontext of the current tokio task, if we are running inside one + /// that was spawned through [`LogContextHandle::scope`]. + pub fn current() -> Option { + TASK_LOCAL_CONTEXT.try_with(|c| c.clone()).ok() + } + + /// Run `fut` with this logcontext active (visible to [`current_context`] and + /// therefore to logging) for the duration of the task. + pub fn scope(self, fut: F) -> impl Future + where + F: Future, + { + TASK_LOCAL_CONTEXT.scope(self, fut) + } + + /// The captured [`LoggingContext`], or `None` if the sentinel was captured. + pub fn logging_context(&self) -> Option<&Py> { + self.context.as_ref().as_ref() + } +} + +/// Tracks the resources used by a log context. +/// +/// The public attribute surface, operators and `repr` are a compatibility +/// contract with the Python callers (Measure, request/background-process +/// metrics, the task scheduler, ...) — change both sides together. Native so the +/// switch machinery can do its rusage accounting without allocating a Python +/// object per operation. +// `skip_from_py_object`: pyo3 0.28 requires `Clone` pyclasses to explicitly +// opt in or out of a generated extract-by-clone `FromPyObject` (a bare +// `#[pyclass]` is a deprecation warning, and we build with `-D warnings`). +// Nothing extracts this type by value, so opt out. +#[pyclass(skip_from_py_object, get_all, set_all)] +#[derive(Clone, Default)] +pub struct ContextResourceUsage { + /// System CPU time, in seconds. + pub ru_stime: f64, + /// User CPU time, in seconds. + pub ru_utime: f64, + /// Number of database transactions done. + pub db_txn_count: i64, + /// Time spent doing database transactions (excluding scheduling), in seconds. + pub db_txn_duration_sec: f64, + /// Time spent waiting for a database connection, in seconds. + pub db_sched_duration_sec: f64, + /// Number of events requested from the database. + pub evt_db_fetch_count: i64, +} + +impl ContextResourceUsage { + fn add_assign(&mut self, other: &ContextResourceUsage) { + self.ru_utime += other.ru_utime; + self.ru_stime += other.ru_stime; + self.db_txn_count += other.db_txn_count; + self.db_txn_duration_sec += other.db_txn_duration_sec; + self.db_sched_duration_sec += other.db_sched_duration_sec; + self.evt_db_fetch_count += other.evt_db_fetch_count; + } + + fn sub_assign(&mut self, other: &ContextResourceUsage) { + self.ru_utime -= other.ru_utime; + self.ru_stime -= other.ru_stime; + self.db_txn_count -= other.db_txn_count; + self.db_txn_duration_sec -= other.db_txn_duration_sec; + self.db_sched_duration_sec -= other.db_sched_duration_sec; + self.evt_db_fetch_count -= other.evt_db_fetch_count; + } +} + +#[pymethods] +impl ContextResourceUsage { + /// `ContextResourceUsage(copy_from=None)` — if `copy_from` is given, copy its + /// stats; otherwise start at zero. + #[new] + #[pyo3(signature = (copy_from=None))] + fn new(copy_from: Option<&ContextResourceUsage>) -> Self { + copy_from.cloned().unwrap_or_default() + } + + /// Return a copy of this object. + fn copy(&self) -> ContextResourceUsage { + self.clone() + } + + /// Reset all stats to zero. + fn reset(&mut self) { + *self = ContextResourceUsage::default(); + } + + fn __repr__(&self) -> String { + // The single-quoted, `repr()`-style value formatting is the shape + // this class logs in, and scrapers may match on it — keep it stable. + // Rust's `{:?}` renders exponent-form floats as e.g. `1e-7` (Python + // `repr` writes `1e-07`); the only consumer of the string form is + // Measure's "Failed to save metrics!" warning, so we don't chase exact + // parity with Python there. + format!( + "", + self.ru_stime, + self.ru_utime, + self.db_txn_count, + self.db_txn_duration_sec, + self.db_sched_duration_sec, + self.evt_db_fetch_count, + ) + } + + /// `self += other`; mutate in place. pyo3 returns `self` for the in-place slot. + fn __iadd__(&mut self, other: &ContextResourceUsage) { + self.add_assign(other); + } + + /// `self -= other`; mutate in place. pyo3 returns `self` for the in-place slot. + fn __isub__(&mut self, other: &ContextResourceUsage) { + self.sub_assign(other); + } + + /// `self + other`, returning a new object. + fn __add__(&self, other: &ContextResourceUsage) -> ContextResourceUsage { + let mut res = self.clone(); + res.add_assign(other); + res + } + + /// `self - other`, returning a new object. + fn __sub__(&self, other: &ContextResourceUsage) -> ContextResourceUsage { + let mut res = self.clone(); + res.sub_assign(other); + res + } +} + +/// Call the (possibly test-patched) module-level `logcontext_error(msg)`. +fn logcontext_error(py: Python<'_>, msg: String) -> PyResult<()> { + let module = py.import("synapse.logging.context")?; + module.getattr("logcontext_error")?.call1((msg,))?; + Ok(()) +} + +extern "C" { + /// CPython's thread identifier (`pythread.h`) — the exact value + /// `threading.get_ident()` returns. Part of the stable ABI, but not bound + /// by pyo3-ffi, so declared here. + fn PyThread_get_thread_ident() -> std::os::raw::c_ulong; +} + +/// This thread's `threading.get_ident()` value, read natively (`get_ident` is +/// a thin wrapper around `PyThread_get_thread_ident`, which is `pthread_self()` +/// on POSIX) rather than by calling into Python. +/// +/// Note that `get_ident` is *not* an OS-level tid: on Linux it returns the same +/// value either side of a `fork()` call. Synapse forks in exactly one place, so +/// contexts created before the fork still pass the `main_thread` affinity check +/// after it. +fn get_thread_id() -> u64 { + // SAFETY: no preconditions; returns an identifier for the calling thread. + (unsafe { PyThread_get_thread_ident() }) as u64 +} + +/// Propagate a usage update to the parent context, if there is a (truthy) one. +/// +/// A plain base `LoggingContext` parent — the overwhelmingly common case, e.g. +/// every `Measure`-created nested context — is dispatched via `native` directly, +/// skipping the Python method call (attribute lookup, args tuple, call frame) +/// that `add_cputime` would otherwise pay on every switch away from a parented +/// context and `add_database_*` per DB operation. Anything else truthy is a +/// Python subclass, dispatched via `call_method1` so its overrides are +/// respected — the same split as `switch_context`. The truthiness guard matches +/// Python's `if self.parent_context:` and cannot be relaxed to `is_some()`: the +/// sentinel is falsy and implements no `add_cputime`, so that would call a +/// nonexistent method on it. +fn forward_to_parent<'py, N>( + parent: &Option>, + py: Python<'py>, + method: &str, + args: impl PyCallArgs<'py>, + native: N, +) -> PyResult<()> +where + N: FnOnce(&Bound<'py, LoggingContext>) -> PyResult<()>, +{ + if let Some(parent) = parent { + let parent = parent.bind(py); + if let Ok(base) = parent.cast_exact::() { + return native(base); + } + if parent.is_truthy()? { + parent.call_method1(method, args)?; + } + } + Ok(()) +} + +/// The current thread's CPU usage as `(ru_utime, ru_stime)` in seconds, read +/// directly via `getrusage(RUSAGE_THREAD)`. +/// +/// Returns `None` where per-thread rusage isn't available — which we take to be +/// any non-Linux target (`RUSAGE_THREAD` is Linux-only; macOS gets no per-context +/// CPU accounting). Reading it natively avoids allocating a Python +/// `resource.struct_rusage` object on every switch. +#[cfg(target_os = "linux")] +fn get_thread_rusage() -> Option<(f64, f64)> { + fn timeval_to_secs(tv: libc::timeval) -> f64 { + tv.tv_sec as f64 + tv.tv_usec as f64 / 1_000_000.0 + } + + // SAFETY: `getrusage` only writes into `usage`, and we only read it once it + // reports success. + let mut usage = std::mem::MaybeUninit::::uninit(); + let ret = unsafe { libc::getrusage(libc::RUSAGE_THREAD, usage.as_mut_ptr()) }; + if ret != 0 { + return None; + } + let usage = unsafe { usage.assume_init() }; + Some(( + timeval_to_secs(usage.ru_utime), + timeval_to_secs(usage.ru_stime), + )) +} + +#[cfg(not(target_os = "linux"))] +fn get_thread_rusage() -> Option<(f64, f64)> { + None +} + +/// The `(user, system)` CPU seconds elapsed between `start` and `current`. +/// +/// Guards against the clock going backwards +/// (clamping to zero and logging, as the accounting must never go negative). +fn cputime_delta(current: (f64, f64), start: (f64, f64)) -> (f64, f64) { + let mut utime_delta = current.0 - start.0; + let mut stime_delta = current.1 - start.1; + + // sanity check + if utime_delta < 0.0 { + error!("utime went backwards! {} < {}", current.0, start.0); + utime_delta = 0.0; + } + if stime_delta < 0.0 { + error!("stime went backwards! {} < {}", current.1, start.1); + stime_delta = 0.0; + } + + (utime_delta, stime_delta) +} + +/// Additional context for log formatting, tracking which request a unit of work +/// belongs to and accounting CPU/DB usage against it. Contexts are scoped within +/// a `with` block. +/// +/// The attribute surface, methods, error-message wording and abuse-detection +/// behaviour are a compatibility contract with Python callers and subclasses +/// (notably `BackgroundProcessLoggingContext`) — change both sides together. +/// +/// Construction is split between `__new__` (which allocates a blank +/// instance) and `__init__` (which does the real initialisation), matching how a +/// pure-Python class behaves. This lets Python subclasses — in particular +/// `synapse.metrics.background_process_metrics.BackgroundProcessLoggingContext`, +/// which composes a name and then calls `super().__init__(name=..., ...)` — work. +#[pyclass(subclass)] +pub struct LoggingContext { + /// Name for the context, used in logging. Stored as a Python string: + /// `LoggingContextFilter` calls `str(context)` on every log record, and a + /// `Py` getter is INCREF-only where a `String` getter would + /// allocate a fresh `str` per record. Rust-side reads only happen on cold + /// error/debug paths (see [`Self::name_string`]). + #[pyo3(get, set)] + name: Py, + /// The homeserver name this context is associated with. Stored as a Python + /// string for the same reason as `name` (read per log record). + #[pyo3(get, set)] + server_name: Py, + /// The `threading.get_ident()` value of the thread this context was created + /// on (see [`get_thread_id`] for why it is not a real OS tid); activity on + /// any other thread is an error. Settable only so tests can simulate + /// activity on the wrong thread. + #[pyo3(get, set)] + main_thread: u64, + /// Whether `__exit__` has run. Re-activating a finished context is an error. + #[pyo3(get, set)] + finished: bool, + /// The thread CPU usage `(ru_utime, ru_stime)` in seconds captured when this + /// context became active, or `None` if it is not currently active. Private + /// (native `(f64, f64)` rather than a Python `struct_rusage`) so the switch + /// path does no per-switch Python allocation; nothing outside this module + /// reads it. + usage_start: Option<(f64, f64)>, + /// A short human-readable tag (e.g. the sync type). Initialised to `""` and + /// treated as a `str` by everything in-tree, but `Option` so that assigning + /// `None` (which the sentinel's `tag` reports, and which out-of-tree callers + /// may assign) is accepted rather than raising `TypeError`. + #[pyo3(get, set)] + tag: Option, + /// The resources used by this context so far; mutated in place. Not + /// exposed as an attribute: Python reads it via `get_resource_usage()`, + /// which returns a copy. + resource_usage: Py, + /// The context that was current when this one was created; restored on exit. + /// `None` means the sentinel was current (or — only before `__init__` has + /// run — that nothing has been recorded yet; both restore to the sentinel), + /// so the getter reports `None` to Python where the old pure-Python + /// attribute held `SENTINEL_CONTEXT`. + #[pyo3(get, set)] + previous_context: Option>, + /// The parent context, if any; usage is propagated up to it. + #[pyo3(get, set)] + parent_context: Option>, + /// The `ContextRequest` this work belongs to, if any. + #[pyo3(get, set)] + request: Option>, + /// The opentracing scope associated with this context, if any. + #[pyo3(get, set)] + scope: Option>, +} + +#[pymethods] +impl LoggingContext { + /// Allocate a blank context. The real initialisation happens in `__init__`; + /// see the type docstring for why this is split. Extra positional/keyword + /// arguments are accepted and ignored so that subclasses passing their own + /// constructor arguments up through `type.__call__` (which feeds the same + /// arguments to both `__new__` and `__init__`) are not rejected here. + #[new] + #[pyo3(signature = (*_args, **_kwargs))] + fn __new__( + py: Python<'_>, + _args: &Bound<'_, PyTuple>, + _kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + Ok(LoggingContext { + name: intern!(py, "").clone().unbind(), + server_name: intern!(py, "").clone().unbind(), + main_thread: 0, + finished: false, + usage_start: None, + tag: Some(String::new()), + resource_usage: Py::new(py, ContextResourceUsage::default())?, + previous_context: None, + parent_context: None, + request: None, + scope: None, + }) + } + + #[pyo3(signature = (*, name, server_name, parent_context=None, request=None))] + fn __init__( + &mut self, + py: Python<'_>, + name: Bound<'_, PyString>, + server_name: Bound<'_, PyString>, + parent_context: Option>, + request: Option>, + ) -> PyResult<()> { + self.previous_context = current_context(py); + + // The resource-usage tracker was already allocated (zeroed) by `__new__`, + // which `type.__call__` runs immediately before this. + + // The thread resource usage when the logcontext became active. None if + // the context is not currently active. + self.usage_start = None; + + self.name = name.unbind(); + self.server_name = server_name.unbind(); + self.main_thread = get_thread_id(); + self.request = None; + self.tag = Some(String::new()); + self.scope = None; + + // keep track of whether we have hit the __exit__ block for this context + self.finished = false; + + // Inherit some fields from the parent context (read before we move it + // into `self`, so no borrow of `self.parent_context` is held). + if let Some(parent) = &parent_context { + let parent = parent.bind(py); + // which request this corresponds to + self.request = parent.getattr("request")?.extract()?; + // we also track the current scope + self.scope = parent.getattr("scope")?.extract()?; + } + + if let Some(request) = request { + // the request param overrides the request from the parent context + self.request = Some(request); + } + + self.parent_context = parent_context; + + Ok(()) + } + + /// Returns the stored name object itself (INCREF-only): this runs per log + /// record via `LoggingContextFilter`. + fn __str__(&self, py: Python<'_>) -> Py { + self.name.clone_ref(py) + } + + /// Enter this logging context, making it the current context. + fn __enter__<'py>(slf: Bound<'py, Self>) -> PyResult> { + let py = slf.py(); + // An owned handle rather than a held borrow: `set_current_context` + // re-enters `slf` (borrow_mut in `start_inner`). + let previous = slf + .borrow() + .previous_context + .as_ref() + .map(|p| p.clone_ref(py)); + + if log_enabled!(target: DEBUG_LOGGER_NAME, Level::Debug) { + // The name is only materialised when the opt-in debug logger is + // actually enabled: this runs on every context entry. + debug!( + target: DEBUG_LOGGER_NAME, + "LoggingContext({}).__enter__", + slf.borrow().name_string(py) + ); + } + + let old_context = set_current_context(py, Some(slf.clone().unbind()))?; + + if !slots_identical(&previous, &old_context) { + let previous_repr = slot_repr(py, &previous)?; + let old_repr = slot_repr(py, &old_context)?; + logcontext_error( + py, + format!("Expected previous context {previous_repr}, found {old_repr}"), + )?; + } + + Ok(slf) + } + + /// Restore the previous logging context. Returns `None` (does not suppress + /// exceptions). + fn __exit__( + slf: Bound<'_, Self>, + _exc_type: Bound<'_, PyAny>, + _exc_value: Bound<'_, PyAny>, + _traceback: Bound<'_, PyAny>, + ) -> PyResult<()> { + let py = slf.py(); + + let previous = slf + .borrow() + .previous_context + .as_ref() + .map(|p| p.clone_ref(py)); + + if log_enabled!(target: DEBUG_LOGGER_NAME, Level::Debug) { + // Match the Python `%s`: the str() of the previous context + // (`"sentinel"` for an empty slot). Computed (along with the name) + // only when the opt-in debug logger is actually enabled: this runs + // on every context exit. + let previous_str = match &previous { + Some(p) => p.bind(py).str()?.extract::()?, + None => "sentinel".to_owned(), + }; + debug!( + target: DEBUG_LOGGER_NAME, + "LoggingContext({}).__exit__ --> {previous_str}", + slf.borrow().name_string(py) + ); + } + + let current = set_current_context(py, previous)?; + + let restored_self = current.as_ref().is_some_and(|c| c.bind(py).is(&slf)); + if !restored_self { + // Cold path: the name is only materialised for the error message. + let name = slf.borrow().name_string(py); + match ¤t { + None => logcontext_error(py, format!("Expected logging context {name} was lost"))?, + Some(current) => { + let current_str: String = current.bind(py).str()?.extract()?; + logcontext_error( + py, + format!("Expected logging context {name} but found {current_str}"), + )?; + } + } + } + + // the fact that we are here suggests that the caller thinks everything is + // done and dusted for this logcontext, and further activity will not get + // recorded against the correct metrics. + slf.borrow_mut().finished = true; + + Ok(()) + } + + /// Record that this logcontext is currently running. + /// + /// This should not be called directly: use `set_current_context`. `rusage` is + /// the thread CPU usage `(ru_utime, ru_stime)` at the point of switching to + /// this context (`None` if the platform doesn't track it). + fn start(slf: Bound<'_, Self>, rusage: Option<(f64, f64)>) -> PyResult<()> { + Self::start_inner(&slf, rusage) + } + + /// Record that this logcontext is no longer running. + /// + /// This should not be called directly: use `set_current_context`. + fn stop(slf: Bound<'_, Self>, rusage: Option<(f64, f64)>) -> PyResult<()> { + Self::stop_inner(&slf, rusage) + } + + /// Get a *copy* of the resources used by this logcontext so far. + fn get_resource_usage(slf: Bound<'_, Self>) -> PyResult { + let py = slf.py(); + + // we always return a copy, for consistency + let mut res = slf.borrow().resource_usage.borrow(py).clone(); + + let (usage_start, main_thread) = { + let this = slf.borrow(); + (this.usage_start, this.main_thread) + }; + + // If we are on the correct thread and we're currently running then we can + // include resource usage so far. + if let Some(start) = usage_start { + if get_thread_id() == main_thread { + if let Some(current) = get_thread_rusage() { + let (utime_delta, stime_delta) = cputime_delta(current, start); + res.ru_utime += utime_delta; + res.ru_stime += stime_delta; + } + } + } + + Ok(res) + } + + /// Update the CPU time usage of this context (and any parents, recursively). + fn add_cputime(&self, py: Python<'_>, utime_delta: f64, stime_delta: f64) -> PyResult<()> { + { + let mut usage = self.resource_usage.borrow_mut(py); + usage.ru_utime += utime_delta; + usage.ru_stime += stime_delta; + } + forward_to_parent( + &self.parent_context, + py, + "add_cputime", + (utime_delta, stime_delta), + |p| p.borrow().add_cputime(py, utime_delta, stime_delta), + ) + } + + /// Record the use of a database transaction and how long it took. + fn add_database_transaction(&self, py: Python<'_>, duration_sec: f64) -> PyResult<()> { + if duration_sec < 0.0 { + return Err(PyValueError::new_err( + "DB txn time can only be non-negative", + )); + } + { + let mut usage = self.resource_usage.borrow_mut(py); + usage.db_txn_count += 1; + usage.db_txn_duration_sec += duration_sec; + } + forward_to_parent( + &self.parent_context, + py, + "add_database_transaction", + (duration_sec,), + |p| p.borrow().add_database_transaction(py, duration_sec), + ) + } + + /// Record a use of the database pool (the time taken to get a connection). + fn add_database_scheduled(&self, py: Python<'_>, sched_sec: f64) -> PyResult<()> { + if sched_sec < 0.0 { + return Err(PyValueError::new_err( + "DB scheduling time can only be non-negative", + )); + } + { + let mut usage = self.resource_usage.borrow_mut(py); + usage.db_sched_duration_sec += sched_sec; + } + forward_to_parent( + &self.parent_context, + py, + "add_database_scheduled", + (sched_sec,), + |p| p.borrow().add_database_scheduled(py, sched_sec), + ) + } + + /// Record a number of events being fetched from the db. + fn record_event_fetch(&self, py: Python<'_>, event_count: i64) -> PyResult<()> { + { + let mut usage = self.resource_usage.borrow_mut(py); + usage.evt_db_fetch_count += event_count; + } + forward_to_parent( + &self.parent_context, + py, + "record_event_fetch", + (event_count,), + |p| p.borrow().record_event_fetch(py, event_count), + ) + } + + /// Traverse referenced Python objects for the cyclic garbage collector. + /// `scope` and the context can reference each other, forming a real cycle. + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Some(previous_context) = &self.previous_context { + visit.call(previous_context)?; + } + if let Some(parent_context) = &self.parent_context { + visit.call(parent_context)?; + } + if let Some(request) = &self.request { + visit.call(request)?; + } + if let Some(scope) = &self.scope { + visit.call(scope)?; + } + Ok(()) + } + + fn __clear__(&mut self) { + self.previous_context = None; + self.parent_context = None; + self.request = None; + self.scope = None; + } +} + +impl LoggingContext { + /// Whether `__exit__` has run, for crate-internal callers (Python code reads + /// the `finished` attribute instead). + pub(crate) fn is_finished(&self) -> bool { + self.finished + } + + /// The context name as an owned Rust string. + /// + /// This copies the string data, so it is for cold error/debug paths only — + /// the switch fast path must not allocate. + fn name_string(&self, py: Python<'_>) -> String { + self.name.bind(py).to_string_lossy().into_owned() + } + + /// Native body of the `start` pymethod. Shared with the switch fast path in + /// [`set_current_context`], which calls this directly for a base + /// `LoggingContext` rather than dispatching through Python. Runs the same + /// thread-affinity and abuse checks on both paths. + /// + /// This (like [`Self::stop_inner`]) runs on every context switch: the error + /// branches materialise the name themselves so the fast path stays + /// allocation-free. + fn start_inner(slf: &Bound<'_, Self>, rusage: Option<(f64, f64)>) -> PyResult<()> { + let py = slf.py(); + let main_thread = slf.borrow().main_thread; + + if get_thread_id() != main_thread { + let name = slf.borrow().name_string(py); + logcontext_error(py, format!("Started logcontext {name} on different thread"))?; + return Ok(()); + } + + if slf.borrow().finished { + let name = slf.borrow().name_string(py); + logcontext_error(py, format!("Re-starting finished log context {name}"))?; + } + + // If we haven't already started, record the thread resource usage so far. + if slf.borrow().usage_start.is_some() { + let name = slf.borrow().name_string(py); + logcontext_error(py, format!("Re-starting already-active log context {name}"))?; + } else { + slf.borrow_mut().usage_start = rusage; + } + + Ok(()) + } + + /// Native body of the `stop` pymethod. Shared with the switch fast path. + fn stop_inner(slf: &Bound<'_, Self>, rusage: Option<(f64, f64)>) -> PyResult<()> { + let py = slf.py(); + let main_thread = slf.borrow().main_thread; + + // `finally`-style: `usage_start` must be cleared however we exit. + let result = (|| -> PyResult<()> { + if get_thread_id() != main_thread { + let name = slf.borrow().name_string(py); + logcontext_error(py, format!("Stopped logcontext {name} on different thread"))?; + return Ok(()); + } + + // `if not rusage: return` — no rusage means this platform doesn't + // track per-thread CPU, so there is nothing to account. + let Some(current) = rusage else { + return Ok(()); + }; + + // Record the cpu used since we started. + let Some(start) = slf.borrow().usage_start else { + let name = slf.borrow().name_string(py); + logcontext_error( + py, + format!("Called stop on logcontext {name} without recording a start rusage"), + )?; + return Ok(()); + }; + + let (utime_delta, stime_delta) = cputime_delta(current, start); + slf.borrow().add_cputime(py, utime_delta, stime_delta)?; + Ok(()) + })(); + + slf.borrow_mut().usage_start = None; + result + } +} + +/// Which way a [`switch_context`] dispatch goes. +#[derive(Clone, Copy)] +enum SwitchDirection { + Start, + Stop, +} + +/// Dispatch a `stop`/`start` to a slot value during a switch, respecting +/// subclass overrides. +/// +/// The sentinel (an empty slot) is a no-op. A base `LoggingContext` takes the +/// native accounting directly (the hot path — no Python dispatch, no +/// `struct_rusage` allocation). A Python subclass (e.g. +/// `BackgroundProcessLoggingContext`) goes through Python — materialising the +/// rusage as a `(utime, stime)` tuple only here, off the hot path — so its +/// overrides run. Both arms match on the same [`SwitchDirection`], so the +/// native and Python paths provably dispatch the same operation. +fn switch_context( + py: Python<'_>, + slot: Option<&Py>, + direction: SwitchDirection, + rusage: Option<(f64, f64)>, +) -> PyResult<()> { + let Some(ctx) = slot else { + return Ok(()); + }; + let ctx = ctx.bind(py); + if ctx.as_any().is_exact_instance_of::() { + match direction { + SwitchDirection::Start => LoggingContext::start_inner(ctx, rusage), + SwitchDirection::Stop => LoggingContext::stop_inner(ctx, rusage), + } + } else { + let method = match direction { + SwitchDirection::Start => intern!(py, "start"), + SwitchDirection::Stop => intern!(py, "stop"), + }; + ctx.call_method1(method, (rusage,))?; + Ok(()) + } +} + +/// Whether two slot values are the same context (or both the sentinel). +/// `LoggingContext` defines no `__eq__`, so identity is the comparison Python +/// callers got too. +fn slots_identical(a: &Option>, b: &Option>) -> bool { + match (a, b) { + (None, None) => true, + (Some(a), Some(b)) => a.is(b), + _ => false, + } +} + +/// `repr()` of a slot value, for error messages (cold paths only). An empty +/// slot renders as `None`, matching what the `previous_context` getter exposes +/// to Python. +fn slot_repr(py: Python<'_>, slot: &Option>) -> PyResult { + match slot { + Some(ctx) => Ok(ctx.bind(py).repr()?.extract()?), + None => Ok("None".to_owned()), + } +} + +/// Set the current logging context, returning the context that was previously +/// current. `None` means the sentinel, in both directions. +/// +/// `context` must be a [`LoggingContext`] (or subclass) or `None` — anything +/// else fails extraction with a `TypeError`, so the storage slots stay typed. +/// This is not the Python-facing API: `synapse.logging.context.set_current_context` +/// wraps this with the `SENTINEL_CONTEXT` <-> `None` mapping (and it, not this, +/// rejects `None` from callers — here `None` legitimately means the sentinel). +/// +/// Reads the thread rusage once (`getrusage` via libc), `stop`s the old context +/// and `start`s the new one; for the common base-`LoggingContext` case the +/// bookkeeping runs inline, with no per-switch Python dispatch or allocation. +#[pyfunction] +#[pyo3(signature = (context))] +pub fn set_current_context( + py: Python<'_>, + context: Option>, +) -> PyResult>> { + let current = current_context(py); + + if !slots_identical(¤t, &context) { + let rusage = get_thread_rusage(); + switch_context(py, current.as_ref(), SwitchDirection::Stop, rusage)?; + // Raw slot write; we already hold `current`, so ignore the previous value + // it returns. The clone_ref keeps a reference for the `start` below. + let new_ref = context.as_ref().map(|ctx| ctx.clone_ref(py)); + swap_current_context(context); + switch_context(py, new_ref.as_ref(), SwitchDirection::Start, rusage)?; + } + + Ok(current) +} + +/// Run `f` with `context` (a slot value: `None` is the sentinel) as the current +/// logcontext, restoring the previously-current context afterwards — the Rust +/// equivalent of Python's `with PreserveLoggingContext(context):`. +/// +/// The restore runs whether or not `f` fails: an error must not skip it, or +/// `context` would leak onto the calling thread, misattributing everything the +/// thread does next. If both `f` and the restore fail, `f`'s error is +/// reported. +pub(crate) fn with_logcontext( + py: Python<'_>, + context: Option>, + f: impl FnOnce() -> PyResult, +) -> PyResult { + let previous = set_current_context(py, context)?; + let result = f(); + let restored = set_current_context(py, previous); + + let value = result?; + restored?; + Ok(value) +} + +/// Get the current logging context, or `None` for the sentinel. +/// +/// Resolves the tokio task-local first (so logging emitted while a task is being +/// polled is attributed to the context that was current when the task was +/// spawned — even when that captured the sentinel), then this OS thread's slot. +/// This is not the Python-facing API: `synapse.logging.context.current_context` +/// wraps this and returns `SENTINEL_CONTEXT` instead of `None`. +#[pyfunction] +pub fn current_context(py: Python<'_>) -> Option> { + if let Some(handle) = LogContextHandle::current() { + return handle.logging_context().map(|ctx| ctx.clone_ref(py)); + } + + THREAD_LOCAL_CONTEXT.with(|slot| slot.borrow().as_ref().map(|ctx| ctx.clone_ref(py))) +} + +/// Set this OS thread's current logging context slot, returning the previous +/// slot value (`None` is the sentinel, in both directions). +/// +/// This is the raw slot write only — it does **not** do any resource-usage +/// accounting or thread-affinity checks; [`set_current_context`] wraps this with +/// the `getrusage` start/stop bookkeeping. +/// +/// Note this only touches the thread-local slot, never the tokio task-local: +/// the switch primitive is only ever driven on reactor/threadpool threads +/// (Python code), while the task-local is populated once at spawn time by +/// [`LogContextHandle::scope`]. +/// +/// Crate-internal: a raw slot write that bypasses the rusage accounting and +/// thread-affinity checks has no Python caller, so it is not exported (Python +/// uses [`set_current_context`]). +fn swap_current_context(context: Option>) -> Option> { + // Enforce the invariant above rather than trusting it: with a scoped + // task-local populated, `current_context` gives it read precedence, so this + // write would be invisible (and never restored) — everything that follows + // would be silently misattributed. `try_with` on an unset task-local is + // cheap on the normal (reactor-thread) path. + if TASK_LOCAL_CONTEXT.try_with(|_| ()).is_ok() { + error!( + "swap_current_context called during a tokio-scoped poll; the switch is \ + invisible to current_context() and will misattribute logs and metrics" + ); + } + + THREAD_LOCAL_CONTEXT.with(|slot| std::mem::replace(&mut *slot.borrow_mut(), context)) +} + +/// Called when registering modules with python. +pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + let child_module: Bound<'_, PyModule> = PyModule::new(py, "logcontext")?; + child_module.add_class::()?; + child_module.add_class::()?; + child_module.add_function(wrap_pyfunction!(current_context, &child_module)?)?; + child_module.add_function(wrap_pyfunction!(set_current_context, &child_module)?)?; + child_module.add("DEBUG_LOGGER_NAME", DEBUG_LOGGER_NAME)?; + + m.add_submodule(&child_module)?; + + // We need to manually add the module to sys.modules to make `from + // synapse.synapse_rust import logcontext` work. + py.import("sys")? + .getattr("modules")? + .set_item("synapse.synapse_rust.logcontext", child_module)?; + + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use pyo3::types::PyString; + + use super::*; + + /// A minimal `LoggingContext` for slot tests, built directly (bypassing + /// `__init__`, which would capture the current context and thread id). + 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, + finished: false, + usage_start: None, + tag: Some(String::new()), + resource_usage: Py::new(py, ContextResourceUsage::default()) + .expect("failed to allocate ContextResourceUsage"), + previous_context: None, + parent_context: None, + request: None, + scope: None, + }, + ) + .expect("failed to allocate LoggingContext") + } + + #[test] + fn thread_local_defaults_to_sentinel() { + Python::initialize(); + Python::attach(|py| { + // Nothing set on this (fresh) test thread → empty slot (the sentinel). + assert!(current_context(py).is_none()); + }); + } + + #[test] + fn swap_returns_previous_and_updates_thread_local() { + Python::initialize(); + Python::attach(|py| { + let a = test_context(py, "A"); + let b = test_context(py, "B"); + + // Swapping in A returns the previous slot (empty ⇒ sentinel) and + // makes A current. + let prev = swap_current_context(Some(a.clone_ref(py))); + assert!(prev.is_none()); + assert!(current_context(py) + .expect("expected a current context") + .bind(py) + .is(a.bind(py))); + + // Swapping in B returns A. + let prev = swap_current_context(Some(b.clone_ref(py))); + assert!(prev + .expect("expected previous context") + .bind(py) + .is(a.bind(py))); + assert!(current_context(py) + .expect("expected a current context") + .bind(py) + .is(b.bind(py))); + + // Restore the empty slot (sentinel) so we don't leak into any other + // test that happens to reuse this OS thread from the harness pool. + swap_current_context(None); + }); + } + + #[test] + fn task_local_takes_precedence_over_thread_local() { + Python::initialize(); + Python::attach(|py| { + let task_ctx = test_context(py, "TASKCTX"); + + // Outside any scoped task, `current_context` resolves the + // thread-local (here empty: the sentinel). + assert!(current_context(py).is_none()); + assert!(LogContextHandle::current().is_none()); + + let log_context = LogContextHandle { + context: Arc::new(Some(task_ctx.clone_ref(py))), + }; + + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + + rt.block_on(log_context.scope(async { + // Inside the scope, both the Rust handle and the pyfunction (the + // thing the log filter calls) resolve the task-local context — + // even though the thread-local is still empty. + assert!(LogContextHandle::current().is_some()); + Python::attach(|py| { + assert!(current_context(py) + .expect("expected a current context") + .bind(py) + .is(task_ctx.bind(py))); + }); + })); + + // Once the scope ends, we fall back to the thread-local again. + assert!(LogContextHandle::current().is_none()); + assert!(current_context(py).is_none()); + }); + } +} diff --git a/rust/src/logging/mod.rs b/rust/src/logging/mod.rs new file mode 100644 index 00000000000..1d6d898cb46 --- /dev/null +++ b/rust/src/logging/mod.rs @@ -0,0 +1,23 @@ +/* + * 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: + * . + * + */ + +//! Rust counterparts to the `synapse.logging` Python package. +//! +//! Submodules live at the Rust module path that mirrors their Python logging +//! namespace (e.g. `synapse::logging::context` -> `synapse.logging.context`), so +//! that log records emitted from here via the `log` crate land under the matching +//! Python logger without an explicit `target:`. + +pub mod context; diff --git a/synapse/logging/context.py b/synapse/logging/context.py index b6535be3889..5357af98a88 100644 --- a/synapse/logging/context.py +++ b/synapse/logging/context.py @@ -31,7 +31,6 @@ """ import logging -import threading import typing from types import TracebackType from typing import ( @@ -42,6 +41,7 @@ Literal, TypeVar, Union, + cast, overload, ) @@ -52,17 +52,24 @@ from twisted.python.threadpool import ThreadPool from synapse.logging.loggers import ExplicitlyConfiguredLogger +from synapse.synapse_rust.logcontext import ( + DEBUG_LOGGER_NAME, + # Not used in this module, but re-exported: callers import it from here. + ContextResourceUsage, # noqa: F401 + LoggingContext, + current_context as _rust_current_context, + set_current_context as _rust_set_current_context, +) from synapse.util.stringutils import random_string_insecure_fast if TYPE_CHECKING: - from synapse.logging.scopecontextmanager import _LogContextScope from synapse.types import ISynapseReactor logger = logging.getLogger(__name__) original_logger_class = logging.getLoggerClass() logging.setLoggerClass(ExplicitlyConfiguredLogger) -logcontext_debug_logger = logging.getLogger("synapse.logging.context.debug") +logcontext_debug_logger = logging.getLogger(DEBUG_LOGGER_NAME) """ A logger for debugging when the logcontext switches. @@ -74,146 +81,12 @@ # Restore the original logger class logging.setLoggerClass(original_logger_class) -try: - import resource - - # Python doesn't ship with a definition of RUSAGE_THREAD but it's defined - # to be 1 on linux so we hard code it. - RUSAGE_THREAD = 1 - - # If the system doesn't support RUSAGE_THREAD then this should throw an - # exception. - resource.getrusage(RUSAGE_THREAD) - - is_thread_resource_usage_supported = True - - def get_thread_resource_usage() -> "resource.struct_rusage | None": - return resource.getrusage(RUSAGE_THREAD) - -except Exception: - # If the system doesn't support resource.getrusage(RUSAGE_THREAD) then we - # won't track resource usage. - is_thread_resource_usage_supported = False - - def get_thread_resource_usage() -> "resource.struct_rusage | None": - return None - # a hook which can be set during testing to assert that we aren't abusing logcontexts. def logcontext_error(msg: str) -> None: logger.warning(msg) -# get an id for the current thread. -# -# threading.get_ident doesn't actually return an OS-level tid, and annoyingly, -# on Linux it actually returns the same value either side of a fork() call. However -# we only fork in one place, so it's not worth the hoop-jumping to get a real tid. -# -get_thread_id = threading.get_ident - - -class ContextResourceUsage: - """Object for tracking the resources used by a log context - - Attributes: - ru_utime (float): user CPU time (in seconds) - ru_stime (float): system CPU time (in seconds) - db_txn_count (int): number of database transactions done - db_sched_duration_sec (float): amount of time spent waiting for a - database connection - db_txn_duration_sec (float): amount of time spent doing database - transactions (excluding scheduling time) - evt_db_fetch_count (int): number of events requested from the database - """ - - __slots__ = [ - "ru_stime", - "ru_utime", - "db_txn_count", - "db_txn_duration_sec", - "db_sched_duration_sec", - "evt_db_fetch_count", - ] - - def __init__(self, copy_from: "ContextResourceUsage | None" = None) -> None: - """Create a new ContextResourceUsage - - Args: - copy_from: if not None, an object to copy stats from - """ - if copy_from is None: - self.reset() - else: - # FIXME: mypy can't infer the types set via reset() above, so specify explicitly for now - self.ru_utime: float = copy_from.ru_utime - self.ru_stime: float = copy_from.ru_stime - self.db_txn_count: int = copy_from.db_txn_count - - self.db_txn_duration_sec: float = copy_from.db_txn_duration_sec - self.db_sched_duration_sec: float = copy_from.db_sched_duration_sec - self.evt_db_fetch_count: int = copy_from.evt_db_fetch_count - - def copy(self) -> "ContextResourceUsage": - return ContextResourceUsage(copy_from=self) - - def reset(self) -> None: - self.ru_stime = 0.0 - self.ru_utime = 0.0 - self.db_txn_count = 0 - - self.db_txn_duration_sec = 0.0 - self.db_sched_duration_sec = 0.0 - self.evt_db_fetch_count = 0 - - def __repr__(self) -> str: - return ( - "" - ) % ( - self.ru_stime, - self.ru_utime, - self.db_txn_count, - self.db_txn_duration_sec, - self.db_sched_duration_sec, - self.evt_db_fetch_count, - ) - - def __iadd__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": - """Add another ContextResourceUsage's stats to this one's. - - Args: - other: the other resource usage object - """ - self.ru_utime += other.ru_utime - self.ru_stime += other.ru_stime - self.db_txn_count += other.db_txn_count - self.db_txn_duration_sec += other.db_txn_duration_sec - self.db_sched_duration_sec += other.db_sched_duration_sec - self.evt_db_fetch_count += other.evt_db_fetch_count - return self - - def __isub__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": - self.ru_utime -= other.ru_utime - self.ru_stime -= other.ru_stime - self.db_txn_count -= other.db_txn_count - self.db_txn_duration_sec -= other.db_txn_duration_sec - self.db_sched_duration_sec -= other.db_sched_duration_sec - self.evt_db_fetch_count -= other.evt_db_fetch_count - return self - - def __add__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": - res = ContextResourceUsage(copy_from=self) - res += other - return res - - def __sub__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": - res = ContextResourceUsage(copy_from=self) - res -= other - return res - - @attr.s(slots=True, auto_attribs=True) class ContextRequest: """ @@ -237,9 +110,6 @@ class ContextRequest: user_agent: str -LoggingContextOrSentinel = Union["LoggingContext", "_Sentinel"] - - class _Sentinel: """ Sentinel to represent the root context @@ -273,10 +143,10 @@ def __init__(self) -> None: def __str__(self) -> str: return "sentinel" - def start(self, rusage: "resource.struct_rusage | None") -> None: + def start(self, rusage: "tuple[float, float] | None") -> None: pass - def stop(self, rusage: "resource.struct_rusage | None") -> None: + def stop(self, rusage: "tuple[float, float] | None") -> None: pass def add_database_transaction(self, duration_sec: float) -> None: @@ -294,285 +164,40 @@ def __bool__(self) -> Literal[False]: SENTINEL_CONTEXT = _Sentinel() +LoggingContextOrSentinel = Union[LoggingContext, _Sentinel] -class LoggingContext: - """Additional context for log formatting. Contexts are scoped within a - "with" block. - - If a parent is given when creating a new context, then: - - logging fields are copied from the parent to the new context on entry - - when the new context exits, the cpu usage stats are copied from the - child to the parent - Args: - name: Name for the context for logging. - server_name: The name of the server this context is associated with - (`config.server.server_name` or `hs.hostname`) - parent_context (LoggingContext|None): The parent of the new context - request: Synapse Request Context object. Useful to associate all the logs - happening to a given request. +def current_context() -> LoggingContextOrSentinel: + """Get the current logging context. + The storage lives in the Rust extension (see `rust/src/logging/context.rs`), + which represents "no context" as `None` — mapped to `SENTINEL_CONTEXT` here + so callers never see `None`. """ + context = _rust_current_context() + return SENTINEL_CONTEXT if context is None else context - __slots__ = [ - "previous_context", - "name", - "server_name", - "parent_context", - "_resource_usage", - "usage_start", - "main_thread", - "finished", - "request", - "tag", - "scope", - ] - - def __init__( - self, - *, - name: str, - server_name: str, - parent_context: "LoggingContext | None" = None, - request: ContextRequest | None = None, - ) -> None: - self.previous_context = current_context() - - # track the resources used by this context so far - self._resource_usage = ContextResourceUsage() - - # The thread resource usage when the logcontext became active. None - # if the context is not currently active. - self.usage_start: resource.struct_rusage | None = None - - self.name = name - self.server_name = server_name - self.main_thread = get_thread_id() - self.request = None - self.tag = "" - self.scope: "_LogContextScope" | None = None - - # keep track of whether we have hit the __exit__ block for this context - # (suggesting that the the thing that created the context thinks it should - # be finished, and that re-activating it would suggest an error). - self.finished = False - - self.parent_context = parent_context - - # Inherit some fields from the parent context - if self.parent_context is not None: - # which request this corresponds to - self.request = self.parent_context.request - - # we also track the current scope: - self.scope = self.parent_context.scope - - if request is not None: - # the request param overrides the request from the parent context - self.request = request - - def __str__(self) -> str: - return self.name - def __enter__(self) -> "LoggingContext": - """Enters this logging context into thread local storage""" - logcontext_debug_logger.debug("LoggingContext(%s).__enter__", self.name) - old_context = set_current_context(self) - if self.previous_context != old_context: - logcontext_error( - "Expected previous context %r, found %r" - % ( - self.previous_context, - old_context, - ) - ) - return self - - def __exit__( - self, - type: type[BaseException] | None, - value: BaseException | None, - traceback: TracebackType | None, - ) -> None: - """Restore the logging context in thread local storage to the state it - was before this context was entered. - Returns: - None to avoid suppressing any exceptions that were thrown. - """ - logcontext_debug_logger.debug( - "LoggingContext(%s).__exit__ --> %s", self.name, self.previous_context - ) - current = set_current_context(self.previous_context) - if current is not self: - if current is SENTINEL_CONTEXT: - logcontext_error("Expected logging context %s was lost" % (self,)) - else: - logcontext_error( - "Expected logging context %s but found %s" % (self, current) - ) - - # the fact that we are here suggests that the caller thinks that everything - # is done and dusted for this logcontext, and further activity will not get - # recorded against the correct metrics. - self.finished = True - - def start(self, rusage: "resource.struct_rusage | None") -> None: - """ - Record that this logcontext is currently running. - - This should not be called directly: use set_current_context - - Args: - rusage: the resources used by the current thread, at the point of - switching to this logcontext. May be None if this platform doesn't - support getrusuage. - """ - if get_thread_id() != self.main_thread: - logcontext_error("Started logcontext %s on different thread" % (self,)) - return - - if self.finished: - logcontext_error("Re-starting finished log context %s" % (self,)) - - # If we haven't already started record the thread resource usage so - # far - if self.usage_start: - logcontext_error("Re-starting already-active log context %s" % (self,)) - else: - self.usage_start = rusage - - def stop(self, rusage: "resource.struct_rusage | None") -> None: - """ - Record that this logcontext is no longer running. - - This should not be called directly: use set_current_context - - Args: - rusage: the resources used by the current thread, at the point of - switching away from this logcontext. May be None if this platform - doesn't support getrusuage. - """ - - try: - if get_thread_id() != self.main_thread: - logcontext_error("Stopped logcontext %s on different thread" % (self,)) - return - - if not rusage: - return - - # Record the cpu used since we started - if not self.usage_start: - logcontext_error( - "Called stop on logcontext %s without recording a start rusage" - % (self,) - ) - return - - utime_delta, stime_delta = self._get_cputime(rusage) - self.add_cputime(utime_delta, stime_delta) - finally: - self.usage_start = None - - def get_resource_usage(self) -> ContextResourceUsage: - """Get resources used by this logcontext so far. - - Returns: - A *copy* of the object tracking resource usage so far - """ - # we always return a copy, for consistency - res = self._resource_usage.copy() - - # If we are on the correct thread and we're currently running then we - # can include resource usage so far. - is_main_thread = get_thread_id() == self.main_thread - if self.usage_start and is_main_thread: - rusage = get_thread_resource_usage() - assert rusage is not None - utime_delta, stime_delta = self._get_cputime(rusage) - res.ru_utime += utime_delta - res.ru_stime += stime_delta - - return res - - def _get_cputime(self, current: "resource.struct_rusage") -> tuple[float, float]: - """Get the cpu usage time between start() and the given rusage - - Args: - rusage: the current resource usage - - Returns: tuple[float, float]: seconds in user mode, seconds in system mode - """ - assert self.usage_start is not None - - utime_delta = current.ru_utime - self.usage_start.ru_utime - stime_delta = current.ru_stime - self.usage_start.ru_stime - - # sanity check - if utime_delta < 0: - logger.error( - "utime went backwards! %f < %f", - current.ru_utime, - self.usage_start.ru_utime, - ) - utime_delta = 0 - - if stime_delta < 0: - logger.error( - "stime went backwards! %f < %f", - current.ru_stime, - self.usage_start.ru_stime, - ) - stime_delta = 0 - - return utime_delta, stime_delta - - def add_cputime(self, utime_delta: float, stime_delta: float) -> None: - """Update the CPU time usage of this context (and any parents, recursively). - - Args: - utime_delta: additional user time, in seconds, spent in this context. - stime_delta: additional system time, in seconds, spent in this context. - """ - self._resource_usage.ru_utime += utime_delta - self._resource_usage.ru_stime += stime_delta - if self.parent_context: - self.parent_context.add_cputime(utime_delta, stime_delta) - - def add_database_transaction(self, duration_sec: float) -> None: - """Record the use of a database transaction and the length of time it took. - - Args: - duration_sec: The number of seconds the database transaction took. - """ - if duration_sec < 0: - raise ValueError("DB txn time can only be non-negative") - self._resource_usage.db_txn_count += 1 - self._resource_usage.db_txn_duration_sec += duration_sec - if self.parent_context: - self.parent_context.add_database_transaction(duration_sec) - - def add_database_scheduled(self, sched_sec: float) -> None: - """Record a use of the database pool - - Args: - sched_sec: number of seconds it took us to get a connection - """ - if sched_sec < 0: - raise ValueError("DB scheduling time can only be non-negative") - self._resource_usage.db_sched_duration_sec += sched_sec - if self.parent_context: - self.parent_context.add_database_scheduled(sched_sec) +def set_current_context(context: LoggingContextOrSentinel) -> LoggingContextOrSentinel: + """Set the current logging context, returning the context that was + previously current. - def record_event_fetch(self, event_count: int) -> None: - """Record a number of events being fetched from the db + The switch itself lives in the Rust extension: it reads the thread CPU usage + once (`getrusage(RUSAGE_THREAD)`) and does the `stop`/`start` accounting + natively. Rust represents the sentinel as `None`, converted in both + directions here so callers only ever see `LoggingContextOrSentinel`. + """ + # everything blows up if we allow current_context to be set to None, so + # sanity-check that now. + if context is None: + raise TypeError("'context' argument may not be None") - Args: - event_count: number of events being fetched - """ - self._resource_usage.evt_db_fetch_count += event_count - if self.parent_context: - self.parent_context.record_event_fetch(event_count) + # The cast is needed because mypy cannot narrow the Union via the + # `is SENTINEL_CONTEXT` identity check; Rust enforces the type at runtime. + previous = _rust_set_current_context( + None if context is SENTINEL_CONTEXT else cast(LoggingContext, context) + ) + return SENTINEL_CONTEXT if previous is None else previous class LoggingContextFilter(logging.Filter): @@ -733,39 +358,6 @@ def __exit__( ) -_thread_local = threading.local() -_thread_local.current_context = SENTINEL_CONTEXT - - -def current_context() -> LoggingContextOrSentinel: - """Get the current logging context from thread local storage""" - return getattr(_thread_local, "current_context", SENTINEL_CONTEXT) - - -def set_current_context(context: LoggingContextOrSentinel) -> LoggingContextOrSentinel: - """Set the current logging context in thread local storage - Args: - context: The context to activate. - - Returns: - The context that was previously active - """ - # everything blows up if we allow current_context to be set to None, so sanity-check - # that now. - if context is None: - raise TypeError("'context' argument may not be None") - - current = current_context() - - if current is not context: - rusage = get_thread_resource_usage() - current.stop(rusage) - _thread_local.current_context = context - context.start(rusage) - - return current - - def nested_logging_context(suffix: str) -> LoggingContext: """Creates a new logging context as a child of another. diff --git a/synapse/metrics/background_process_metrics.py b/synapse/metrics/background_process_metrics.py index 8ff28034556..d51e08aae51 100644 --- a/synapse/metrics/background_process_metrics.py +++ b/synapse/metrics/background_process_metrics.py @@ -57,8 +57,6 @@ from synapse.metrics._types import Collector if TYPE_CHECKING: - import resource - # Old versions don't have `LiteralString` from typing_extensions import LiteralString @@ -506,7 +504,7 @@ def __init__( desc=name, server_name=server_name, ctx=self ) - def start(self, rusage: "resource.struct_rusage | None") -> None: + def start(self, rusage: "tuple[float, float] | None") -> None: """Log context has started running (again).""" super().start(rusage) diff --git a/synapse/synapse_rust/logcontext.pyi b/synapse/synapse_rust/logcontext.pyi new file mode 100644 index 00000000000..e0a1922eea6 --- /dev/null +++ b/synapse/synapse_rust/logcontext.pyi @@ -0,0 +1,157 @@ +# 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: +# . + +from types import TracebackType +from typing import TYPE_CHECKING, Optional + +from synapse.logging.context import ContextRequest + +if TYPE_CHECKING: + from synapse.logging.scopecontextmanager import _LogContextScope + +DEBUG_LOGGER_NAME: str +"""Name of the opt-in logger for logcontext switch tracing +(`synapse.logging.context.debug`). Shared with the Rust `debug!` target so the +names cannot drift.""" + +class ContextResourceUsage: + """Tracks the resources used by a log context.""" + + ru_stime: float + ru_utime: float + db_txn_count: int + db_txn_duration_sec: float + db_sched_duration_sec: float + evt_db_fetch_count: int + + def __init__(self, copy_from: "Optional[ContextResourceUsage]" = None) -> None: ... + def copy(self) -> "ContextResourceUsage": ... + def reset(self) -> None: ... + def __iadd__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": ... + def __isub__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": ... + def __add__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": ... + def __sub__(self, other: "ContextResourceUsage") -> "ContextResourceUsage": ... + +class LoggingContext: + """Additional context for log formatting. Contexts are scoped within a + "with" block. + + If a parent is given when creating a new context, then: + - logging fields are copied from the parent to the new context on entry + - when the new context exits, the cpu usage stats are copied from the + child to the parent + """ + + # The context that was current when this one was created; None means the + # sentinel (the Rust storage cannot hold the pure-Python `SENTINEL_CONTEXT` + # object, so this is narrower than the pure-Python attribute used to be). + previous_context: Optional[LoggingContext] + name: str + server_name: str + parent_context: "Optional[LoggingContext]" + main_thread: int + finished: bool + request: Optional[ContextRequest] + # Narrower than the runtime: the setter also accepts None (see the Rust + # field docs for why), but all in-tree code treats this as str. + tag: str + scope: "Optional[_LogContextScope]" + + def __init__( + self, + *, + name: str, + server_name: str, + parent_context: "Optional[LoggingContext]" = None, + request: Optional[ContextRequest] = None, + ) -> None: + """ + Args: + name: Name for the context for logging. + server_name: The name of the server this context is associated with + (`config.server.server_name` or `hs.hostname`). + parent_context: The parent of the new context. + request: Synapse Request Context object. Useful to associate all the + logs happening to a given request. + """ + + def __str__(self) -> str: ... + def __enter__(self) -> "LoggingContext": ... + def __exit__( + self, + type: Optional[type[BaseException]], + value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: ... + def start(self, rusage: "Optional[tuple[float, float]]") -> None: + """Record that this logcontext is currently running. + + Should not be called directly: use `set_current_context`. + + Args: + rusage: The thread CPU usage `(ru_utime, ru_stime)` at the point of + switching to this context, or None if the platform doesn't + track it. + """ + + def stop(self, rusage: "Optional[tuple[float, float]]") -> None: + """Record that this logcontext is no longer running. + + Should not be called directly: use `set_current_context`. + + Args: + rusage: The thread CPU usage `(ru_utime, ru_stime)` at the point of + switching away from this context, or None if the platform + doesn't track it. + """ + + def get_resource_usage(self) -> ContextResourceUsage: + """Get the resources used by this logcontext so far. + + Returns: + A *copy* of the object tracking resource usage so far. + """ + + def add_cputime(self, utime_delta: float, stime_delta: float) -> None: + """Update the CPU time usage of this context (and any parents, recursively).""" + + def add_database_transaction(self, duration_sec: float) -> None: + """Record the use of a database transaction and how long it took.""" + + def add_database_scheduled(self, sched_sec: float) -> None: + """Record a use of the database pool (the time taken to get a connection).""" + + def record_event_fetch(self, event_count: int) -> None: + """Record a number of events being fetched from the db.""" + +def current_context() -> Optional[LoggingContext]: + """Get the current logging context, or None for the sentinel. + + Resolves the tokio task-local first (so logging emitted while a Rust task is + being polled is attributed to the context that was current when the task was + spawned), then this OS thread's slot. This is not the Python-facing API: + `synapse.logging.context.current_context` wraps this and returns + `SENTINEL_CONTEXT` instead of `None`. + """ + +def set_current_context( + context: Optional[LoggingContext], +) -> Optional[LoggingContext]: + """Set the current logging context, returning the context that was previously + current. `None` means the sentinel, in both directions. + + Reads the thread CPU usage once via `getrusage(RUSAGE_THREAD)` and does the + `stop`/`start` accounting natively. The annotated type is enforced: raises + `TypeError` unless `context` is a `LoggingContext` (or subclass) or `None`. + This is not the Python-facing API: `synapse.logging.context.set_current_context` + wraps this with the `SENTINEL_CONTEXT` <-> `None` mapping. + """ diff --git a/tests/handlers/test_appservice.py b/tests/handlers/test_appservice.py index 08a4ab624ae..a649f512e64 100644 --- a/tests/handlers/test_appservice.py +++ b/tests/handlers/test_appservice.py @@ -77,6 +77,7 @@ def setUp(self) -> None: self.reactor, self.clock = get_clock() hs = Mock() + hs.hostname = "test_server" def test_run_as_background_process( desc: "LiteralString", diff --git a/tests/storage/test_base.py b/tests/storage/test_base.py index 577229c1199..e67d36cf630 100644 --- a/tests/storage/test_base.py +++ b/tests/storage/test_base.py @@ -105,7 +105,7 @@ def runWithConnection(func, *args, **kwargs): # type: ignore[no-untyped-def] # To fix isinstance(...) checks. fake_engine.__class__ = engine.__class__ # type: ignore[assignment] - db = DatabasePool(Mock(), Mock(config=db_config), fake_engine) + db = DatabasePool(Mock(hostname="test"), Mock(config=db_config), fake_engine) db._db_pool = conn_pool self.datastore = SQLBaseStore(db, None, hs) # type: ignore[arg-type] diff --git a/tests/synapse_rust/test_logcontext.py b/tests/synapse_rust/test_logcontext.py new file mode 100644 index 00000000000..188ee40420e --- /dev/null +++ b/tests/synapse_rust/test_logcontext.py @@ -0,0 +1,292 @@ +# 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: +# . + +"""Cross-language logcontext attribution for Rust. + +The current logcontext lives in the Rust slot (`synapse.synapse_rust.logcontext` +/ `rust/src/logging/context.rs`), visible from both Python (reactor/threadpool +threads) and Rust (tokio tasks). These tests exercise the two guarantees that +gives us, through real production code paths: + +1. Log records emitted from Rust while a task is being polled (e.g. reqwest + connecting) are attributed to the logcontext that was current when Python + called into Rust — not the sentinel. +2. When Rust calls back into Python (`run_python_awaitable`, as the Rust + `/versions` handler does for its per-user feature DB lookup), the Python code + runs in that same logcontext, so its DB-transaction accounting lands on the + right request. +""" + +import logging +import time +from typing import Callable + +from twisted.internet.testing import MemoryReactor + +from synapse.logging.context import ( + LoggingContext, + LoggingContextFilter, + PreserveLoggingContext, + _Sentinel, + current_context, + run_in_background, +) +from synapse.rest import admin +from synapse.rest.client import login +from synapse.server import HomeServer +from synapse.synapse_rust import reset_logging_config +from synapse.synapse_rust.http_client import HttpClient +from synapse.util.clock import Clock + +from tests.unittest import HomeserverTestCase + +logger = logging.getLogger(__name__) + +# Log-target roots that Rust code emits under while running on tokio worker +# threads: the reqwest dependency stack, plus "synapse"/"synapse_rust" because +# the Rust crate is itself named `synapse` (see rust/Cargo.toml). Anything +# emitted under these while a task is being polled should be attributed to the +# caller's logcontext, never the sentinel. +# +# NB: bare "synapse" also matches every *Python* `synapse.*` record, so the +# attribution assertion below implicitly relies on nothing else logging during +# the pump (MemoryReactor with `advance(0)`, so no timed background work fires). +# If this test starts flaking on unrelated records, tighten this filter rather +# than weakening the assertion. +_RUST_LOGGER_ROOTS = frozenset( + {"reqwest", "hyper", "hyper_util", "h2", "rustls", "synapse_rust", "synapse"} +) + + +class RustLogContextTestCase(HomeserverTestCase): + servlets = [ + admin.register_servlets, + login.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. + self._http_client = hs.get_proxied_http_client() + self._rust_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. + shutdown_triggers = self.reactor.triggers.get("shutdown", {}) + for phase in ["before", "during", "after"]: + triggers = shutdown_triggers.get(phase, []) + for callbable, args, kwargs in triggers: + callbable(*args, **kwargs) + + def prepare( + self, reactor: MemoryReactor, clock: Clock, homeserver: HomeServer + ) -> None: + self.user_id = self.register_user("user1", "pass") + + def _check_current_logcontext(self, expected: str) -> None: + context = current_context() + assert isinstance(context, (LoggingContext, _Sentinel)), context + self.assertEqual(str(context), expected, f"expected {expected}, saw {context}") + + def _run_in_logcontext_and_pump( + self, name: str, body: Callable[[dict[str, object]], None] + ) -> dict[str, object]: + """Run `body` fired off inside a fresh `LoggingContext(name)`, pumping the + reactor (and yielding to the Tokio pool) until it sets `result["done"]`. + + Returns the `result` dict `body` populated. Asserts the caller logcontext + is intact afterwards and that we end back in the sentinel. + """ + self._check_current_logcontext("sentinel") + result: dict[str, object] = {} + + with LoggingContext(name=name, server_name="test_server"): + body(result) + + with PreserveLoggingContext(): + # Generous upper bound (the work is a real HTTP round-trip or DB + # hop on a possibly-loaded CI box); the loop exits early via + # `result["done"]`, and we fail below if it never gets set. + for _ in range(50000): + if result.get("done"): + break + # Let the Tokio worker threads make progress... + time.sleep(0) + # ...and run anything they scheduled back on the reactor. + self.reactor.advance(0) + + # The caller's logcontext must be intact after firing off the work. + self._check_current_logcontext(name) + + # ...and we must not have leaked it into the reactor. + self._check_current_logcontext("sentinel") + + self.assertTrue( + result.get("done"), + "work never finished; the test probably didn't pump long enough", + ) + return result + + def test_rust_log_records_attributed_to_caller_logcontext(self) -> None: + """A log record emitted from Rust on a tokio thread (reqwest connecting) + is attributed to the caller's logcontext via `LoggingContextFilter`, not + the sentinel.""" + records: list[tuple[str, object]] = [] + + class CapturingHandler(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + records.append((record.name, getattr(record, "request", ""))) + + handler = CapturingHandler() + # The global filter is what copies `str(current_context())` onto the + # record as `record.request`; attach it so we observe what Synapse would. + handler.addFilter(LoggingContextFilter()) + + root = logging.getLogger() + root.addHandler(handler) + + # Turn up the Rust-side loggers so reqwest actually emits, and refresh + # pyo3-log's cached levels so it forwards them. + saved_levels = { + name: logging.getLogger(name).level for name in _RUST_LOGGER_ROOTS + } + for name in _RUST_LOGGER_ROOTS: + logging.getLogger(name).setLevel(logging.DEBUG) + reset_logging_config() + + try: + server = _StubServer() + self.addCleanup(server.shutdown) + + def body(result: dict[str, object]) -> None: + async def do() -> None: + try: + await self._rust_http_client.get( + url=server.endpoint, + response_limit=1 * 1024 * 1024, + ) + finally: + result["done"] = True + + run_in_background(do) + + self._run_in_logcontext_and_pump("http-caller", body) + finally: + root.removeHandler(handler) + for name, level in saved_levels.items(): + logging.getLogger(name).setLevel(level) + reset_logging_config() + + rust_records = [ + (name, req) + for (name, req) in records + if name.split(".", 1)[0] in _RUST_LOGGER_ROOTS + ] + self.assertTrue( + rust_records, + "expected at least one Rust-origin log record (e.g. reqwest connecting); " + f"captured loggers: {sorted({name for name, _ in records})}", + ) + for name, req in rust_records: + self.assertEqual( + req, + "http-caller", + f"Rust log record from {name!r} was attributed to {req!r}, " + "not the caller's logcontext", + ) + + def test_db_callback_runs_in_caller_logcontext(self) -> None: + """The Rust `/versions` handler's per-user feature lookup calls back into + Python via `run_python_awaitable`; the DB transaction it runs must be + accounted against the caller's logcontext. The failure mode is the + awaitable running in the sentinel instead, silently losing the + accounting.""" + versions_handler = self.hs.get_rust_handlers().versions + + def body(result: dict[str, object]) -> None: + async def do() -> None: + try: + context = current_context() + assert isinstance(context, LoggingContext) + before = context.get_resource_usage().db_txn_count + + # Passing a user id makes the Rust handler do a per-user + # feature DB lookup (msc3881/msc3575 default to off), which + # goes Rust -> run_python_awaitable -> runInteraction. + await versions_handler.get_versions(self.user_id) + + after = context.get_resource_usage().db_txn_count + result["db_txn_delta"] = after - before + finally: + result["done"] = True + + run_in_background(do) + + result = self._run_in_logcontext_and_pump("db-caller", body) + + db_txn_delta = result.get("db_txn_delta", 0) + assert isinstance(db_txn_delta, int) + self.assertGreaterEqual( + db_txn_delta, + 1, + "the Rust handler's DB work was not accounted against the caller's " + "logcontext — run_python_awaitable is not restoring it (ran in the " + "sentinel instead)", + ) + + +class _StubServer: + """A real HTTP server on a random port, served from a background thread.""" + + def __init__(self) -> None: + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + + class Handler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.end_headers() + self.wfile.write(b'{"ok": true}') + + def log_message(self, format: str, *args: object) -> None: + pass + + self._server = HTTPServer(("127.0.0.1", 0), Handler) + self._thread = threading.Thread( + target=self._server.serve_forever, + name="StubServer", + kwargs={"poll_interval": 0.01}, + daemon=True, + ) + self._thread.start() + + @property + def endpoint(self) -> str: + return f"http://127.0.0.1:{self._server.server_port}/" + + def shutdown(self) -> None: + self._server.shutdown() + self._thread.join() diff --git a/tests/test_state.py b/tests/test_state.py index b69a09dab5a..d2e50cf94be 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -245,6 +245,7 @@ def setUp(self) -> None: ) reactor, clock = get_clock() hs.config = default_config(server_name="tesths", parse=True) + hs.hostname = "tesths" hs.get_datastores.return_value = Mock( main=self.dummy_store, state_deletion=dummy_deletion_store, diff --git a/tests/util/test_logcontext.py b/tests/util/test_logcontext.py index a4114cdfcc5..853cb8cf400 100644 --- a/tests/util/test_logcontext.py +++ b/tests/util/test_logcontext.py @@ -20,13 +20,18 @@ # import logging +from contextlib import contextmanager from typing import Callable, Generator, cast +from unittest.mock import patch from twisted.internet import defer, reactor as _reactor +import synapse.logging.context as context_module from synapse.logging.context import ( SENTINEL_CONTEXT, + ContextRequest, LoggingContext, + LoggingContextFilter, PreserveLoggingContext, _Sentinel, current_context, @@ -34,6 +39,7 @@ nested_logging_context, run_coroutine_in_background, run_in_background, + set_current_context, ) from synapse.types import ISynapseReactor from synapse.util.clock import Clock @@ -700,6 +706,226 @@ def test_nested_logging_context(self) -> None: self.assertEqual(nested_context.name, "foo-bar") +# A stand-in rusage `(ru_utime, ru_stime)` for exercising the `start()`/`stop()` +# code paths that only check whether an rusage is present, without depending on +# `RUSAGE_THREAD` support or the real value of the thread's CPU clock. +_TRUTHY_RUSAGE = (0.0, 0.0) + + +@contextmanager +def _capture_logcontext_errors() -> Generator[list[str], None, None]: + """Capture the messages passed to `logcontext_error` while inside the block. + + `logcontext_error` is looked up as a module global at call time, so patching + the module attribute intercepts every call site (`LoggingContext`, + `PreserveLoggingContext`, `run_in_background`, ...). + """ + messages: list[str] = [] + with patch.object(context_module, "logcontext_error", messages.append): + yield messages + + +class LogContextErrorMessageTestCase(unittest.TestCase): + """Characterization tests pinning the exact `logcontext_error` message shapes + and the abuse-detection code paths. + + These exist to guard against accidental drift in the switch machinery (which + lives in Rust: `rust/src/logging/context.rs`): downstream log scraping + depends on the wording, argument order and the conditions that trigger each + warning. Messages that interpolate a context via `%r` embed the object's + `repr()` (id/address), so we reconstruct the expected string from the *same* + live objects rather than hard-coding an address. + """ + + def setUp(self) -> None: + # These tests poke the switch primitives directly; make sure we always + # leave the slot back at the sentinel regardless of what a test does. + self.assertIs(current_context(), SENTINEL_CONTEXT) + self.addCleanup(set_current_context, SENTINEL_CONTEXT) + + def test_enter_wrong_previous_context(self) -> None: + with _capture_logcontext_errors() as messages: + # `ctx` records the sentinel as its `previous_context` (we are at the + # sentinel now), but we enter it while a *different* context is active. + ctx = LoggingContext(name="ctx", server_name="s") + other = LoggingContext(name="other", server_name="s") + other.__enter__() + + ctx.__enter__() + + self.assertEqual( + messages, + [ + "Expected previous context %r, found %r" + % (ctx.previous_context, other) + ], + ) + + def test_exit_context_was_lost(self) -> None: + with _capture_logcontext_errors() as messages: + ctx = LoggingContext(name="ctx", server_name="s") + ctx.__enter__() + # Simulate the context being lost from under us before exit. + set_current_context(SENTINEL_CONTEXT) + + ctx.__exit__(None, None, None) + + self.assertEqual(messages, ["Expected logging context ctx was lost"]) + self.assertTrue(ctx.finished) + + def test_exit_found_different_context(self) -> None: + with _capture_logcontext_errors() as messages: + ctx = LoggingContext(name="ctx", server_name="s") + ctx.__enter__() + other = LoggingContext(name="other", server_name="s") + other.__enter__() + + ctx.__exit__(None, None, None) + + self.assertEqual(messages, ["Expected logging context ctx but found other"]) + + def test_preserve_logging_context_was_lost(self) -> None: + with _capture_logcontext_errors() as messages: + target = LoggingContext(name="target", server_name="s") + preserve = PreserveLoggingContext(target) + preserve.__enter__() + set_current_context(SENTINEL_CONTEXT) + + preserve.__exit__(None, None, None) + + self.assertEqual(messages, ["Expected logging context target was lost"]) + + def test_preserve_logging_context_found_different_context(self) -> None: + with _capture_logcontext_errors() as messages: + target = LoggingContext(name="target", server_name="s") + preserve = PreserveLoggingContext(target) + preserve.__enter__() + other = LoggingContext(name="other", server_name="s") + other.__enter__() + + preserve.__exit__(None, None, None) + + self.assertEqual( + messages, ["Expected logging context target but found other"] + ) + + def test_start_on_different_thread(self) -> None: + with _capture_logcontext_errors() as messages: + ctx = LoggingContext(name="ctx", server_name="s") + # Pretend the context was created on another OS thread. + ctx.main_thread += 1 + + ctx.start(None) + + self.assertEqual(messages, ["Started logcontext ctx on different thread"]) + + def test_stop_on_different_thread(self) -> None: + with _capture_logcontext_errors() as messages: + ctx = LoggingContext(name="ctx", server_name="s") + ctx.main_thread += 1 + + ctx.stop(None) + + self.assertEqual(messages, ["Stopped logcontext ctx on different thread"]) + + def test_restart_finished_context(self) -> None: + with _capture_logcontext_errors() as messages: + ctx = LoggingContext(name="ctx", server_name="s") + ctx.finished = True + + ctx.start(None) + + self.assertEqual(messages, ["Re-starting finished log context ctx"]) + + def test_restart_already_active_context(self) -> None: + with _capture_logcontext_errors() as messages: + ctx = LoggingContext(name="ctx", server_name="s") + ctx.start(_TRUTHY_RUSAGE) + + ctx.start(_TRUTHY_RUSAGE) + + self.assertEqual(messages, ["Re-starting already-active log context ctx"]) + + def test_stop_without_start(self) -> None: + with _capture_logcontext_errors() as messages: + ctx = LoggingContext(name="ctx", server_name="s") + + ctx.stop(_TRUTHY_RUSAGE) + + self.assertEqual( + messages, + ["Called stop on logcontext ctx without recording a start rusage"], + ) + + +class LoggingContextFilterTestCase(unittest.TestCase): + """Characterization tests for `LoggingContextFilter`, the entire observable + surface for logging. Pins that it fills record attributes from a real + logcontext, and that under the sentinel it only sets defaults for attributes + that aren't already present (`safe_set`), which 3rd-party code relies on.""" + + def setUp(self) -> None: + self.assertIs(current_context(), SENTINEL_CONTEXT) + self.filter = LoggingContextFilter() + + def _make_record(self, **attrs: object) -> logging.LogRecord: + record = logging.LogRecord("test", logging.INFO, "path", 1, "msg", None, None) + for key, value in attrs.items(): + setattr(record, key, value) + return record + + def test_fills_record_from_real_context(self) -> None: + request = ContextRequest( + request_id="req-1", + ip_address="1.2.3.4", + site_tag="site", + requester="@u:test", + authenticated_entity="@u:test", + method="GET", + url="/path", + protocol="1.1", + user_agent="UA", + ) + with LoggingContext(name="ctx", server_name="myserver", request=request): + record = self._make_record() + self.assertIs(self.filter.filter(record), True) + + # Read via `__dict__` since these attributes are only ever set + # dynamically (they aren't declared on `logging.LogRecord`). + attrs = record.__dict__ + self.assertEqual(attrs["server_name"], "myserver") + # For backwards compatibility, `str(context)` is stored as `request`. + self.assertEqual(attrs["request"], "ctx") + self.assertEqual(attrs["ip_address"], "1.2.3.4") + self.assertEqual(attrs["site_tag"], "site") + self.assertEqual(attrs["requester"], "@u:test") + self.assertEqual(attrs["authenticated_entity"], "@u:test") + self.assertEqual(attrs["method"], "GET") + self.assertEqual(attrs["url"], "/path") + self.assertEqual(attrs["protocol"], "1.1") + self.assertEqual(attrs["user_agent"], "UA") + + def test_sentinel_sets_defaults_when_absent(self) -> None: + record = self._make_record() + self.assertIs(self.filter.filter(record), True) + + self.assertEqual( + record.__dict__["server_name"], "unknown_server_from_sentinel_context" + ) + self.assertEqual(record.__dict__["request"], "sentinel") + + def test_sentinel_does_not_clobber_preset_attributes(self) -> None: + # Simulate a 3rd-party log record that already carries its own attributes; + # under the sentinel the filter must leave them untouched. + record = self._make_record( + server_name="third-party-server", request="third-party-request" + ) + self.assertIs(self.filter.filter(record), True) + + self.assertEqual(record.__dict__["server_name"], "third-party-server") + self.assertEqual(record.__dict__["request"], "third-party-request") + + # a function which returns a deferred which has been "called", but # which had a function which returned another incomplete deferred on # its callback list, so won't yet call any other new callbacks.