prefix_logs_with: Ensure the macro works correctly for futures

prdoc/pr_11095.prdoc

@@ -0,0 +1,8 @@
+title: '`prefix_logs_with`: Ensure the macro works correctly for futures'
+doc:
+- audience: Node Dev
+  description: |-
+    When setting up a tracing span in an async future, it may gets invalidated by any `await` point. The problem is that after continuing a future, it may runs on a different thread where the `span` isn't active anymore. The solution for this is to `instrument` the future properly.
+crates:
+- name: sc-tracing-proc-macro
+  bump: patch

substrate/client/tracing/proc-macro/src/lib.rs

@@ -123,19 +123,40 @@ pub fn prefix_logs_with(arg: TokenStream, item: TokenStream) -> TokenStream {
 
 	let syn::ItemFn { attrs, vis, sig, block } = item_fn;
 
-	(quote! {
-		#(#attrs)*
-		#vis #sig {
-			let span = #crate_name::tracing::info_span!(
-				#crate_name::logging::PREFIX_LOG_SPAN,
-				name = #prefix_expr,
-			);
-			let _enter = span.enter();
+	if sig.asyncness.is_some() {
+		// For async functions, use `Instrument::instrument` to properly propagate the span
+		// across `.await` points. Using `span.enter()` in async functions is incorrect because
+		// the guard can be held across `.await` points where the task may migrate between
+		// threads, losing the span from thread-local state.
+		(quote! {
+			#(#attrs)*
+			#vis #sig {
+				let span = #crate_name::tracing::info_span!(
+					#crate_name::logging::PREFIX_LOG_SPAN,
+					name = #prefix_expr,
+				);
 
-			#block
-		}
-	})
-	.into()
+				#crate_name::tracing::Instrument::instrument(async move {
+					#block
+				}, span).await
+			}
+		})
+		.into()
+	} else {
+		(quote! {
+			#(#attrs)*
+			#vis #sig {
+				let span = #crate_name::tracing::info_span!(
+					#crate_name::logging::PREFIX_LOG_SPAN,
+					name = #prefix_expr,
+				);
+				let _enter = span.enter();
+
+				#block
+			}
+		})
+		.into()
+	}
 }
 
 /// Resolve the correct path for sc_tracing: