From 62e4b39cc1c14fdea9ea9b84d1b5e0cb10ac262f Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 16 Jul 2026 15:05:27 +0000 Subject: [PATCH] Fix the logcontext restore protocol in run_python_awaitable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two fixes to the `starter` closure that drives a Python awaitable on the reactor thread on behalf of a tokio task: - Run the fallible `run_in_background`/`addCallbacks` section inside a closure so an error can no longer skip the restore: previously the `?` returned early and permanently leaked the restored context onto the reactor thread, misattributing all subsequent reactor work. - Don't restore a context that has already finished (the request completed or was cancelled while the task was still running — `create_deferred` does not propagate cancellation). Doing so tripped the 'Re-starting finished log context' abuse check and accounted work against a context whose metrics were already finalised; run in the reactor's current context (normally the sentinel) instead, as the pre-port code did. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01JFbRtswu7rsHrttJFauUUb --- rust/src/deferred.rs | 75 ++++++++++++++++++++++++++++--------- rust/src/logging/context.rs | 6 +++ 2 files changed, 63 insertions(+), 18 deletions(-) diff --git a/rust/src/deferred.rs b/rust/src/deferred.rs index f7a1b18e99..e0fb4caa92 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,7 +26,7 @@ use pyo3::{ }; use tokio::sync::oneshot; -use crate::logging::context::set_current_context; +use crate::logging::context::{set_current_context, LoggingContext, DEBUG_LOGGER_NAME}; use crate::tokio_runtime::runtime; create_exception!( @@ -238,31 +239,69 @@ where // the awaitable is driven in it. `run_in_background` then starts the // awaitable in the current logcontext and follows the logcontext rules // from there. + // + // 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 run in the reactor's current context (normally the + // sentinel) instead. Both `__exit__` (which sets `finished`) and this + // check run on the reactor thread, so the check cannot race. let previous = match &logcontext { - Some(ctx) => Some(set_current_context(py, ctx.as_py(py))?), + Some(ctx) => { + let ctx = ctx.as_py(py); + let finished = ctx + .cast::() + .map(|c| c.borrow().is_finished()) + .unwrap_or(false); + if finished { + if log_enabled!(target: DEBUG_LOGGER_NAME, Level::Debug) { + debug!( + target: DEBUG_LOGGER_NAME, + "run_python_awaitable: captured logcontext {} has \ + finished; running in the sentinel", + ctx.str()? + ); + } + None + } else { + Some(set_current_context(py, ctx)?) + } + } None => None, }; // 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. - let deferred = logging_context.call_method1( - intern!(py, "run_in_background"), - (awaitable_factory.bind(py),), - )?; - deferred.call_method1( - intern!(py, "addCallbacks"), - (on_success.bind(py), on_error.bind(py)), - )?; + // current logcontext while following the logcontext rules. Run inside + // a closure so that an error cannot skip the restore below — that + // would leak the restored context onto the reactor thread permanently, + // misattributing everything the reactor does next. + let result = (|| -> PyResult<()> { + let deferred = logging_context.call_method1( + intern!(py, "run_in_background"), + (awaitable_factory.bind(py),), + )?; + deferred.call_method1( + intern!(py, "addCallbacks"), + (on_success.bind(py), on_error.bind(py)), + )?; + Ok(()) + })(); // Return the reactor thread to its previous logcontext (normally the - // sentinel). `run_in_background` has already arranged for the - // awaitable's completion to reset to the sentinel and has restored the - // calling logcontext synchronously, so this leaves the reactor as we - // found it rather than leaking `ctx` into it. - if let Some(previous) = previous { - set_current_context(py, previous.into_bound(py))?; - } + // sentinel) whether or not the above succeeded. `run_in_background` + // has already arranged for the awaitable's completion to reset to the + // sentinel and has restored the calling logcontext synchronously, so + // this leaves the reactor as we found it rather than leaking `ctx` + // into it. If both the drive and the restore fail, report the first + // error. + let restore = match previous { + Some(previous) => set_current_context(py, previous.into_bound(py)).map(drop), + None => Ok(()), + }; + result.and(restore)?; Ok(py.None()) }, diff --git a/rust/src/logging/context.rs b/rust/src/logging/context.rs index 28861fe858..aded12406d 100644 --- a/rust/src/logging/context.rs +++ b/rust/src/logging/context.rs @@ -766,6 +766,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 + } + /// 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