diff --git a/docs/log_contexts.md b/docs/log_contexts.md index f2ed81683f..ea60fde20e 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 diff --git a/rust/src/deferred.rs b/rust/src/deferred.rs index 62a1ce6b90..bd40c57518 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),), - ); + // 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(()) + })?; - let deferred = deferred?; - deferred.call_method1( - intern!(py, "addCallbacks"), - (on_success.bind(py), on_error.bind(py)), - )?; 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/logging/context.rs b/rust/src/logging/context.rs index 9648748615..6618729e7b 100644 --- a/rust/src/logging/context.rs +++ b/rust/src/logging/context.rs @@ -21,29 +21,36 @@ //! including from spawned tokio tasks — could not be attributed to the request //! that caused it. //! -//! This module holds that storage — a per-OS-thread slot -//! ([`THREAD_LOCAL_CONTEXT`]) used by the reactor thread and any -//! reactor-managed threadpool threads — along with the logcontext classes -//! themselves. `LoggingContextFilter` (and therefore `pyo3-log`) resolves the -//! context by calling [`current_context`] at log-record time. +//! This module holds that storage and unifies two sources of truth so that a +//! single [`current_context`] answer is correct from *both* worlds: //! -//! The slot holds an `Option>`: `None` means "no context" — +//! 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. -//! -//! TODO: tokio tasks do not yet see a logcontext — a worker thread's slot is -//! always empty, so Rust-emitted log records still land in the sentinel. A -//! task-scoped capture (carried with the task across `.await` points and -//! consulted by [`current_context`] ahead of the thread slot) follows in the -//! next change. +//! [`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; +use std::{cell::RefCell, future::Future}; use log::{debug, error, log_enabled, Level}; use pyo3::call::PyCallArgs; @@ -72,6 +79,64 @@ thread_local! { 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 @@ -686,6 +751,12 @@ impl LoggingContext { } 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 — @@ -863,13 +934,41 @@ pub fn set_current_context( 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 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`. +/// 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))) } @@ -880,10 +979,27 @@ pub fn current_context(py: Python<'_>) -> Option> { /// 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)) } @@ -909,6 +1025,8 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> #[cfg(test)] mod tests { + use std::sync::Arc; + use pyo3::types::PyString; use super::*; @@ -977,4 +1095,42 @@ mod tests { 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/synapse/synapse_rust/logcontext.pyi b/synapse/synapse_rust/logcontext.pyi index 92b303c9ff..e0a1922eea 100644 --- a/synapse/synapse_rust/logcontext.pyi +++ b/synapse/synapse_rust/logcontext.pyi @@ -136,7 +136,9 @@ class LoggingContext: def current_context() -> Optional[LoggingContext]: """Get the current logging context, or None for the sentinel. - Resolves this OS thread's slot. This is not the Python-facing API: + 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`. """ diff --git a/tests/synapse_rust/test_logcontext.py b/tests/synapse_rust/test_logcontext.py new file mode 100644 index 0000000000..188ee40420 --- /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()