From 863ff326b79e2b1a9cbd08511a572a4189dd2e9f Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 16 Jul 2026 15:06:07 +0000 Subject: [PATCH] Enforce the logcontext storage invariant in swap_current_context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The current-context storage is asymmetric: reads give the tokio task-local precedence, but the switch primitive only ever writes the OS-thread slot. A switch performed while a scoped task is being polled would therefore be write-only — invisible to current_context() and never restored — silently misattributing everything that follows. That invariant was previously enforced only by comments; check it at the single write chokepoint and log an error if it is ever violated. Also update the module doc, which still described the earlier phase's architecture where Python kept the accounting policy — set_current_context is fully native as of this branch. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JFbRtswu7rsHrttJFauUUb --- rust/src/logging/context.rs | 29 +++++++++++++++++++++-------- 1 file changed, 21 insertions(+), 8 deletions(-) diff --git a/rust/src/logging/context.rs b/rust/src/logging/context.rs index aded12406d..9a7dd431d5 100644 --- a/rust/src/logging/context.rs +++ b/rust/src/logging/context.rs @@ -37,12 +37,13 @@ //! task is being polled are attributed to the task's captured context with no //! per-record stamping machinery. //! -//! Python keeps the accounting policy: `set_current_context` still does the -//! `getrusage` start/stop bookkeeping and merely uses [`swap_current_context`] for -//! the raw slot write. The switch primitive is only ever driven on the reactor -//! (or threadpool) threads — never on tokio worker threads — so it always writes -//! the thread-local, and the task-local (populated only by [`LogContext::scope`] -//! at spawn time) takes read precedence during a poll. +//! The accounting policy is native too: [`set_current_context`] reads the thread +//! rusage via libc, runs the `stop`/`start` bookkeeping, and uses +//! [`swap_current_context`] for the raw slot write. The switch primitive is only +//! ever driven on the reactor (or threadpool) threads — never on tokio worker +//! threads — so it always writes the thread-local, and the task-local (populated +//! only by [`LogContext::scope`] at spawn time) takes read precedence during a +//! poll. [`swap_current_context`] checks that invariant rather than trusting it. use std::{cell::RefCell, future::Future}; @@ -932,8 +933,8 @@ pub fn current_context(py: Python<'_>) -> Py { /// previously current *on this thread*. /// /// This is the raw slot write only — it does **not** do any resource-usage -/// accounting or thread-affinity checks; `synapse.logging.context.set_current_context` -/// wraps this with the `getrusage` start/stop bookkeeping. +/// accounting or thread-affinity checks; [`set_current_context`] wraps this with +/// the `getrusage` start/stop bookkeeping. /// /// Note this deliberately only touches the thread-local slot, never the tokio /// task-local: the switch primitive is only ever driven on reactor/threadpool @@ -941,6 +942,18 @@ pub fn current_context(py: Python<'_>) -> Py { /// [`LogContext::scope`]. #[pyfunction] pub fn swap_current_context(py: Python<'_>, context: Py) -> Py { + // 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, so the check costs nothing 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" + ); + } + let previous = THREAD_LOCAL_CONTEXT.with(|slot| slot.borrow_mut().replace(context)); match previous { Some(ctx) => ctx,