core: Make asyncify preserve the current span (#2790)

Author: kim

crates/core/src/util/mod.rs

@@ -2,6 +2,7 @@ use futures::{Future, FutureExt};
 use std::borrow::Cow;
 use std::pin::pin;
 use tokio::sync::oneshot;
+use tracing::Span;
 
 pub mod prometheus_handle;
 
@@ -40,13 +41,19 @@ where
     F: FnOnce() -> R + Send + 'static,
     R: Send + 'static,
 {
-    tokio::task::spawn_blocking(f)
-        .await
-        .unwrap_or_else(|e| match e.try_into_panic() {
-            Ok(panic_payload) => std::panic::resume_unwind(panic_payload),
-            // the only other variant is cancelled, which shouldn't happen because we don't cancel it.
-            Err(e) => panic!("Unexpected JoinError: {e}"),
-        })
+    // Ensure that `f` executes in the current span context.
+    // If there is no current span, or it is disabled, `span` is disabled.
+    let span = Span::current();
+    tokio::task::spawn_blocking(move || {
+        let _enter = span.enter();
+        f()
+    })
+    .await
+    .unwrap_or_else(|e| match e.try_into_panic() {
+        Ok(panic_payload) => std::panic::resume_unwind(panic_payload),
+        // the only other variant is cancelled, which shouldn't happen because we don't cancel it.
+        Err(e) => panic!("Unexpected JoinError: {e}"),
+    })
 }
 
 /// Await `fut`, while also polling `also`.