perf: avoid unnecessary clones of the log context

Author: sandhose

crates/cli/src/main.rs

@@ -6,7 +6,14 @@
 
 #![allow(clippy::module_name_repetitions)]
 
-use std::{io::IsTerminal, process::ExitCode, sync::Arc};
+use std::{
+    io::IsTerminal,
+    process::ExitCode,
+    sync::{
+        Arc,
+        atomic::{AtomicU64, Ordering},
+    },
+};
 
 use anyhow::Context;
 use clap::Parser;
@@ -52,6 +59,10 @@ impl sentry::TransportFactory for SentryTransportFactory {
 fn main() -> anyhow::Result<ExitCode> {
     let mut builder = tokio::runtime::Builder::new_multi_thread();
     builder.enable_all();
+    builder.thread_name_fn(|| {
+        static IDX: AtomicU64 = AtomicU64::new(0);
+        format!("tokio-{}", IDX.fetch_add(1, Ordering::Relaxed))
+    });
 
     #[cfg(tokio_unstable)]
     builder

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!(

crates/context/src/fmt.rs

@@ -113,13 +113,13 @@ 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)?;

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<Self> {
-        CURRENT_LOG_CONTEXT.try_with(Self::clone).ok()
+    /// Run a closure with the current log context, if any
+    pub fn maybe_with<F, R>(f: F) -> Option<R>
+    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

crates/tasks/src/new_queue.rs

@@ -789,14 +789,12 @@ 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(()) => {