diff --git a/crates/cli/src/server.rs b/crates/cli/src/server.rs index 969a93e23..114fb809c 100644 --- a/crates/cli/src/server.rs +++ b/crates/cli/src/server.rs @@ -184,13 +184,11 @@ async fn log_response_middleware( let response = next.run(request).await; - let Some(log_context) = LogContext::current() else { + let Some(stats) = LogContext::maybe_with(LogContext::stats) else { tracing::error!("Missing log context for request, this is a bug!"); return response; }; - let stats = log_context.stats(); - let status_code = response.status(); match status_code.as_u16() { 100..=399 => tracing::info!( diff --git a/crates/context/src/fmt.rs b/crates/context/src/fmt.rs index f2e619690..d9ae7baa5 100644 --- a/crates/context/src/fmt.rs +++ b/crates/context/src/fmt.rs @@ -113,13 +113,14 @@ where write!(&mut writer, "{} ", style.apply_to(metadata.name()))?; } - if let Some(log_context) = LogContext::current() { + LogContext::maybe_with(|log_context| { let log_context = Style::new() .bold() .force_styling(ansi) .apply_to(log_context); - write!(&mut writer, "{log_context} - ")?; - } + write!(&mut writer, "{log_context} - ") + }) + .transpose()?; let field_fromatter = DefaultFields::new(); field_fromatter.format_fields(writer.by_ref(), event)?; diff --git a/crates/context/src/lib.rs b/crates/context/src/lib.rs index a0d6b384b..655d407e9 100644 --- a/crates/context/src/lib.rs +++ b/crates/context/src/lib.rs @@ -76,9 +76,12 @@ impl LogContext { } } - /// Get a copy of the current log context, if any - pub fn current() -> Option { - CURRENT_LOG_CONTEXT.try_with(Self::clone).ok() + /// Run a closure with the current log context, if any + pub fn maybe_with(f: F) -> Option + where + F: FnOnce(&Self) -> R, + { + CURRENT_LOG_CONTEXT.try_with(f).ok() } /// Run the async function `f` with the given log context. It will wrap the diff --git a/crates/tasks/src/new_queue.rs b/crates/tasks/src/new_queue.rs index 81fef84d4..ea055e2f8 100644 --- a/crates/tasks/src/new_queue.rs +++ b/crates/tasks/src/new_queue.rs @@ -789,14 +789,14 @@ impl JobTracker { ); let result = job.run(&state, context.clone()).await; - let Some(log_context) = LogContext::current() else { + let Some(context_stats) = + LogContext::maybe_with(mas_context::LogContext::stats) + else { // This should never happen, but if it does it's fine: we're recovering fine // from panics in those tasks panic!("Missing log context, this should never happen"); }; - let context_stats = log_context.stats(); - // We log the result here so that it's attached to the right span & log context match &result { Ok(()) => {