From 9b9a00178c916438e310090ead99e4d1610d5ad1 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Fri, 17 Jul 2026 09:19:39 +0000 Subject: [PATCH 1/6] Add characterization tests for logcontext error messages and the filter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pin the Python-observable behaviour of the logcontext machinery ahead of porting it to Rust: - LogContextErrorMessageTestCase pins the exact logcontext_error message wording, argument order and the conditions that trigger each abuse warning. These are load-bearing: tests and downstream log scraping match on them. Messages interpolating a context via %r embed the object's repr (id/address), so expected strings are reconstructed from the same live objects rather than hard-coded. - LoggingContextFilterTestCase pins LoggingContextFilter — the entire observable surface for logging: record attributes are filled from a real logcontext, and (crucially for 3rd-party code) under the sentinel only absent attributes get defaults. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JFbRtswu7rsHrttJFauUUb --- tests/util/test_logcontext.py | 227 ++++++++++++++++++++++++++++++++++ 1 file changed, 227 insertions(+) diff --git a/tests/util/test_logcontext.py b/tests/util/test_logcontext.py index a4114cdfcc5..36514c94723 100644 --- a/tests/util/test_logcontext.py +++ b/tests/util/test_logcontext.py @@ -20,13 +20,19 @@ # import logging +import resource +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 +40,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 +707,226 @@ def test_nested_logging_context(self) -> None: self.assertEqual(nested_context.name, "foo-bar") +# A real (truthy) rusage for exercising the `start()`/`stop()` code paths that +# only check whether an rusage is present. `RUSAGE_SELF` (not `RUSAGE_THREAD`) +# so this works on platforms without per-thread rusage support. +_TRUTHY_RUSAGE = resource.getrusage(resource.RUSAGE_SELF) + + +@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: + 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. From 49c163679db040bf352e146604b5f99a4e4c3ff1 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Fri, 17 Jul 2026 09:25:02 +0000 Subject: [PATCH 2/6] Port ContextResourceUsage to a Rust pyclass The first slice of moving the logcontext machinery into the Rust extension: ContextResourceUsage becomes a native class in the new synapse.synapse_rust.logcontext module (rust/src/logging/context.rs), re-exported from synapse.logging.context so callers are unchanged. The public attribute surface, operators and repr are a compatibility contract with the Python callers (Measure, request/background-process metrics, the task scheduler, ...). Keeping the tracker native lets the upcoming switch machinery do its rusage accounting without allocating a Python object per operation. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JFbRtswu7rsHrttJFauUUb --- rust/src/lib.rs | 2 + rust/src/logging/context.rs | 154 ++++++++++++++++++++++++++++ rust/src/logging/mod.rs | 23 +++++ synapse/logging/context.py | 105 +------------------ synapse/synapse_rust/logcontext.pyi | 31 ++++++ 5 files changed, 214 insertions(+), 101 deletions(-) create mode 100644 rust/src/logging/context.rs create mode 100644 rust/src/logging/mod.rs create mode 100644 synapse/synapse_rust/logcontext.pyi 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..26f19ed5d1e --- /dev/null +++ b/rust/src/logging/context.rs @@ -0,0 +1,154 @@ +/* + * 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 counterparts to `synapse.logging.context`. +//! +//! Currently just [`ContextResourceUsage`]; the goal is per-operation resource +//! accounting with no Python allocation on the switch path. +//! +//! TODO: the storage for the "current" logcontext, `LoggingContext` itself and +//! the sentinel follow — see `synapse.logging.context` for the Python +//! implementations being replaced. + +use pyo3::prelude::*; + +/// 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 + } +} + +/// 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::()?; + + 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(()) +} 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..01421828362 100644 --- a/synapse/logging/context.py +++ b/synapse/logging/context.py @@ -52,6 +52,10 @@ from twisted.python.threadpool import ThreadPool from synapse.logging.loggers import ExplicitlyConfiguredLogger +from synapse.synapse_rust.logcontext import ( + # Not used in this module, but re-exported: callers import it from here. + ContextResourceUsage, # noqa: F401 +) from synapse.util.stringutils import random_string_insecure_fast if TYPE_CHECKING: @@ -113,107 +117,6 @@ def logcontext_error(msg: str) -> None: 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: """ diff --git a/synapse/synapse_rust/logcontext.pyi b/synapse/synapse_rust/logcontext.pyi new file mode 100644 index 00000000000..dc0e9555184 --- /dev/null +++ b/synapse/synapse_rust/logcontext.pyi @@ -0,0 +1,31 @@ +# 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 typing import Optional + +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": ... From c3dfbfc6ef7e725819c075afafaca27954c9e40e Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Fri, 17 Jul 2026 11:03:12 +0000 Subject: [PATCH 3/6] Move the logcontext storage and LoggingContext to Rust MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The "current logcontext" slot moves from a Python threading.local into the Rust extension, together with a native port of LoggingContext. A Python thread-local is invisible to Rust: each tokio worker thread would see its own slot, permanently at the sentinel, so logging emitted from Rust could not be attributed to the request that caused it. This lays the storage groundwork; a follow-up change gives tokio tasks a task-scoped capture (see the module-doc TODO). Design notes, for review: - The slot is typed: Option>, with None representing the sentinel. The _Sentinel class and SENTINEL_CONTEXT singleton stay pure Python, unchanged; thin wrappers on current_context and set_current_context in synapse.logging.context convert between the singleton and None at the boundary, so no Rust code ever sees or produces the sentinel object. pyo3's extraction enforces the type: anything that is not a LoggingContext (or subclass) or None raises TypeError. - The accounting policy is native too: set_current_context reads the thread rusage once via libc (no per-switch struct_rusage allocation) and runs the stop/start bookkeeping inline for base LoggingContexts, only dispatching through Python for subclasses (BackgroundProcessLoggingContext) so their overrides run. start()/stop() now take an Optional[tuple[float, float]] instead of a struct_rusage, the get_thread_resource_usage/is_thread_resource_usage_supported/ get_thread_id module helpers are gone, LoggingContext.previous_context is now Optional[LoggingContext] (None where it used to hold SENTINEL_CONTEXT), and the nominally-private _resource_usage attribute is no longer exposed (nothing read it; use get_resource_usage()) — worth an upgrade note when this is released, as out-of-tree code may rely on the old shapes. - The switch path avoids per-operation allocation and Python round-trips: names are stored as Py (LoggingContextFilter reads server_name and str(context) per log record process-wide, now INCREF-only), error messages materialise the context name only in the cold branches, and the thread id is read via PyThread_get_thread_ident (the exact value threading.get_ident() returns) rather than by calling into Python. - The attribute surface, method set and error-message wording are a compatibility contract, pinned by the characterization tests (which now exercise the tuple-based start/stop API). - The opt-in synapse.logging.context.debug switch traces are emitted from Rust via pyo3-log, whose level cache only refreshes on reset_logging_config(); docs/log_contexts.md documents the manhole procedure. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JFbRtswu7rsHrttJFauUUb --- Cargo.lock | 1 + docs/log_contexts.md | 15 + rust/Cargo.toml | 1 + rust/src/logging/context.rs | 838 +++++++++++++++++- synapse/logging/context.py | 379 +------- synapse/metrics/background_process_metrics.py | 4 +- synapse/synapse_rust/logcontext.pyi | 126 ++- tests/util/test_logcontext.py | 21 +- 8 files changed, 1019 insertions(+), 366 deletions(-) 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/docs/log_contexts.md b/docs/log_contexts.md index 76710e10e09..f2ed81683fc 100644 --- a/docs/log_contexts.md +++ b/docs/log_contexts.md @@ -564,3 +564,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/logging/context.rs b/rust/src/logging/context.rs index 26f19ed5d1e..96487486150 100644 --- a/rust/src/logging/context.rs +++ b/rust/src/logging/context.rs @@ -13,16 +13,64 @@ * */ -//! Native counterparts to `synapse.logging.context`. +//! Native storage for the Synapse "current logcontext". //! -//! Currently just [`ContextResourceUsage`]; the goal is per-operation resource -//! accounting with no Python allocation on the switch path. +//! 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. //! -//! TODO: the storage for the "current" logcontext, `LoggingContext` itself and -//! the sentinel follow — see `synapse.logging.context` for the Python -//! implementations being replaced. +//! 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. +//! +//! The slot holds 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. +//! +//! 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. +use std::cell::RefCell; + +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) }; +} /// Tracks the resources used by a log context. /// @@ -137,10 +185,716 @@ impl ContextResourceUsage { } } +/// 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 { + /// 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) +} + +/// 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`. +#[pyfunction] +pub fn current_context(py: Python<'_>) -> Option> { + 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. +/// +/// 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> { + 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)?; @@ -152,3 +906,75 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> Ok(()) } + +#[cfg(test)] +mod tests { + 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); + }); + } +} diff --git a/synapse/logging/context.py b/synapse/logging/context.py index 01421828362..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, ) @@ -53,20 +53,23 @@ 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. @@ -78,45 +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 - - @attr.s(slots=True, auto_attribs=True) class ContextRequest: """ @@ -140,9 +110,6 @@ class ContextRequest: user_agent: str -LoggingContextOrSentinel = Union["LoggingContext", "_Sentinel"] - - class _Sentinel: """ Sentinel to represent the root context @@ -176,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: @@ -197,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): @@ -636,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 index dc0e9555184..92b303c9ff1 100644 --- a/synapse/synapse_rust/logcontext.pyi +++ b/synapse/synapse_rust/logcontext.pyi @@ -10,7 +10,18 @@ # See the GNU Affero General Public License for more details: # . -from typing import Optional +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.""" @@ -29,3 +40,116 @@ class 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 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/util/test_logcontext.py b/tests/util/test_logcontext.py index 36514c94723..853cb8cf400 100644 --- a/tests/util/test_logcontext.py +++ b/tests/util/test_logcontext.py @@ -20,7 +20,6 @@ # import logging -import resource from contextlib import contextmanager from typing import Callable, Generator, cast from unittest.mock import patch @@ -707,10 +706,10 @@ def test_nested_logging_context(self) -> None: self.assertEqual(nested_context.name, "foo-bar") -# A real (truthy) rusage for exercising the `start()`/`stop()` code paths that -# only check whether an rusage is present. `RUSAGE_SELF` (not `RUSAGE_THREAD`) -# so this works on platforms without per-thread rusage support. -_TRUTHY_RUSAGE = resource.getrusage(resource.RUSAGE_SELF) +# 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 @@ -730,12 +729,12 @@ 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: - 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. + 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: From c7d3246550bec554e24875d5cff132e6360b338e Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Fri, 17 Jul 2026 09:41:56 +0000 Subject: [PATCH 4/6] Attribute Rust-spawned work to the caller's logcontext MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Give tokio tasks a captured logcontext, resolving the module-doc TODO: - LogContextHandle is a cheap, clone-able, GIL-free handle in the same Option> representation the storage slots use. create_deferred captures the caller's context at the FFI boundary and scopes it onto the spawned task via a tokio task-local, which rides with the task across .await points. current_context() gives the task-local read precedence, so log records emitted while a task is polled — via LoggingContextFilter and pyo3-log — are attributed to the captured context with no per-record stamping. - The switch primitive is only ever driven on reactor/threadpool threads, never during a tokio-scoped poll (where the write would be invisible to reads); swap_current_context enforces that invariant with an error log rather than trusting it. - run_python_awaitable restores the captured context on the reactor thread before driving the awaitable, so Python called back from Rust (e.g. DatabasePool.runInteraction from the Rust /versions handler) runs in — and accounts its DB usage against — the right request. The restore protocol lives in a new with_logcontext helper (the Rust equivalent of `with PreserveLoggingContext(...)`): an error cannot skip the restore (which would leak the context onto the reactor thread permanently), and a context that has already finished is not re-started (create_deferred does not propagate cancellation, so a task can outlive its request; see the TODO) — such work runs in the sentinel instead. - tests/synapse_rust/test_logcontext.py exercises both guarantees through real production code paths: reqwest's log records carry the caller's request id, and the /versions handler's DB transaction lands on the caller's usage accounting. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JFbRtswu7rsHrttJFauUUb --- docs/log_contexts.md | 35 ++- rust/src/deferred.rs | 117 +++++++++-- rust/src/logging/context.rs | 190 +++++++++++++++-- synapse/synapse_rust/logcontext.pyi | 4 +- tests/synapse_rust/test_logcontext.py | 292 ++++++++++++++++++++++++++ 5 files changed, 600 insertions(+), 38 deletions(-) create mode 100644 tests/synapse_rust/test_logcontext.py diff --git a/docs/log_contexts.md b/docs/log_contexts.md index f2ed81683fc..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 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/logging/context.rs b/rust/src/logging/context.rs index 96487486150..6618729e7bd 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 92b303c9ff1..e0a1922eea6 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 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() From 42f5a32f02071aff07d088465cd7287e7afe7c44 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Fri, 17 Jul 2026 15:22:48 +0000 Subject: [PATCH 5/6] Add changelog Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JFbRtswu7rsHrttJFauUUb --- changelog.d/19979.misc | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/19979.misc 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. From 9e75f288c24caf7462f4d9938dcbf92b97ecc01e Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Mon, 20 Jul 2026 13:08:43 +0000 Subject: [PATCH 6/6] tests: give mock homeservers a real string hostname MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Rust `LoggingContext` requires `server_name` to be a `str`. Three test suites built a mock homeserver whose `hostname` was left as an auto-generated `Mock`, which flowed into `server_name` (via `DatabasePool`, `StateHandler` and `ApplicationServicesHandler`) and raised `TypeError: argument 'server_name': 'Mock' object is not an instance of 'str'` once a `LoggingContext`/`Measure` was constructed — 42 trial errors in CI. Set `hostname` to the server name each suite already uses, so the mocks match what production passes. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01VF56cZ93AqpuGCf8yguCcR --- tests/handlers/test_appservice.py | 1 + tests/storage/test_base.py | 2 +- tests/test_state.py | 1 + 3 files changed, 3 insertions(+), 1 deletion(-) 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/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,