attributes: correctly instrument fns returning an async move block (#1866)

This should take care of half of #1863, meaning it should instrument
functions returning an unboxed `async move` block.

This unfortunately does not cover generic `impl Future` returns, such as
`std::future::ready()` and similar :-(

Author: Swatinem

tracing-attributes/src/expand.rs

@@ -437,43 +437,56 @@ fn param_names(pat: Pat, record_type: RecordType) -> Box<dyn Iterator<Item = (Id
     }
 }
 
-enum AsyncTraitKind<'a> {
-    // old construction. Contains the function
+/// The specific async code pattern that was detected
+enum AsyncKind<'a> {
+    /// Immediately-invoked async fn, as generated by `async-trait <= 0.1.43`:
+    /// `async fn foo<...>(...) {...}; Box::pin(foo<...>(...))`
     Function(&'a ItemFn),
-    // new construction. Contains a reference to the async block
-    Async(&'a ExprAsync),
+    /// A function returning an async (move) block, optionally `Box::pin`-ed,
+    /// as generated by `async-trait >= 0.1.44`:
+    /// `Box::pin(async move { ... })`
+    Async {
+        async_expr: &'a ExprAsync,
+        pinned_box: bool,
+    },
 }
 
-pub(crate) struct AsyncTraitInfo<'block> {
+pub(crate) struct AsyncInfo<'block> {
     // statement that must be patched
     source_stmt: &'block Stmt,
-    kind: AsyncTraitKind<'block>,
+    kind: AsyncKind<'block>,
     self_type: Option<syn::TypePath>,
     input: &'block ItemFn,
 }
 
-impl<'block> AsyncTraitInfo<'block> {
-    /// Get the AST of the inner function we need to hook, if it was generated
-    /// by async-trait.
+impl<'block> AsyncInfo<'block> {
+    /// Get the AST of the inner function we need to hook, if it looks like a
+    /// manual future implementation.
     ///
-    /// When we are given a function annotated by async-trait, that function
-    /// is only a placeholder that returns a pinned future containing the
-    /// user logic, and it is that pinned future that needs to be instrumented.
+    /// When we are given a function that returns a (pinned) future containing the
+    /// user logic, it is that (pinned) future that needs to be instrumented.
     /// Were we to instrument its parent, we would only collect information
     /// regarding the allocation of that future, and not its own span of execution.
-    /// Depending on the version of async-trait, we inspect the block of the function
-    /// to find if it matches the pattern
     ///
-    /// `async fn foo<...>(...) {...}; Box::pin(foo<...>(...))` (<=0.1.43), or if
-    /// it matches `Box::pin(async move { ... }) (>=0.1.44). We the return the
-    /// statement that must be instrumented, along with some other informations.
+    /// We inspect the block of the function to find if it matches any of the
+    /// following patterns:
+    ///
+    /// - Immediately-invoked async fn, as generated by `async-trait <= 0.1.43`:
+    ///   `async fn foo<...>(...) {...}; Box::pin(foo<...>(...))`
+    ///
+    /// - A function returning an async (move) block, optionally `Box::pin`-ed,
+    ///   as generated by `async-trait >= 0.1.44`:
+    ///   `Box::pin(async move { ... })`
+    ///
+    /// We the return the statement that must be instrumented, along with some
+    /// other information.
     /// 'gen_body' will then be able to use that information to instrument the
     /// proper function/future.
     ///
     /// (this follows the approach suggested in
     /// https://github.com/dtolnay/async-trait/issues/45#issuecomment-571245673)
     pub(crate) fn from_fn(input: &'block ItemFn) -> Option<Self> {
-        // are we in an async context? If yes, this isn't a async_trait-like pattern
+        // are we in an async context? If yes, this isn't a manual async-like pattern
         if input.sig.asyncness.is_some() {
             return None;
         }
@@ -491,10 +504,8 @@ impl<'block> AsyncTraitInfo<'block> {
             None
         });
 
-        // last expression of the block (it determines the return value
-        // of the block, so that if we are working on a function whose
-        // `trait` or `impl` declaration is annotated by async_trait,
-        // this is quite likely the point where the future is pinned)
+        // last expression of the block: it determines the return value of the
+        // block, this is quite likely a `Box::pin` statement or an async block
         let (last_expr_stmt, last_expr) = block.stmts.iter().rev().find_map(|stmt| {
             if let Stmt::Expr(expr) = stmt {
                 Some((stmt, expr))
@@ -503,6 +514,19 @@ impl<'block> AsyncTraitInfo<'block> {
             }
         })?;
 
+        // is the last expression an async block?
+        if let Expr::Async(async_expr) = last_expr {
+            return Some(AsyncInfo {
+                source_stmt: last_expr_stmt,
+                kind: AsyncKind::Async {
+                    async_expr,
+                    pinned_box: false,
+                },
+                self_type: None,
+                input,
+            });
+        }
+
         // is the last expression a function call?
         let (outside_func, outside_args) = match last_expr {
             Expr::Call(ExprCall { func, args, .. }) => (func, args),
@@ -528,12 +552,12 @@ impl<'block> AsyncTraitInfo<'block> {
         // Is the argument to Box::pin an async block that
         // captures its arguments?
         if let Expr::Async(async_expr) = &outside_args[0] {
-            // check that the move 'keyword' is present
-            async_expr.capture?;
-
-            return Some(AsyncTraitInfo {
+            return Some(AsyncInfo {
                 source_stmt: last_expr_stmt,
-                kind: AsyncTraitKind::Async(async_expr),
+                kind: AsyncKind::Async {
+                    async_expr,
+                    pinned_box: true,
+                },
                 self_type: None,
                 input,
             });
@@ -579,15 +603,15 @@ impl<'block> AsyncTraitInfo<'block> {
             }
         }
 
-        Some(AsyncTraitInfo {
+        Some(AsyncInfo {
             source_stmt: stmt_func_declaration,
-            kind: AsyncTraitKind::Function(func),
+            kind: AsyncKind::Function(func),
             self_type,
             input,
         })
     }
 
-    pub(crate) fn gen_async_trait(
+    pub(crate) fn gen_async(
         self,
         args: InstrumentArgs,
         instrumented_function_name: &str,
@@ -611,15 +635,18 @@ impl<'block> AsyncTraitInfo<'block> {
         {
             // instrument the future by rewriting the corresponding statement
             out_stmts[iter] = match self.kind {
-                // async-trait <= 0.1.43
-                AsyncTraitKind::Function(fun) => gen_function(
+                // `Box::pin(immediately_invoked_async_fn())`
+                AsyncKind::Function(fun) => gen_function(
                     fun.into(),
                     args,
                     instrumented_function_name,
                     self.self_type.as_ref(),
                 ),
-                // async-trait >= 0.1.44
-                AsyncTraitKind::Async(async_expr) => {
+                // `async move { ... }`, optionally pinned
+                AsyncKind::Async {
+                    async_expr,
+                    pinned_box,
+                } => {
                     let instrumented_block = gen_block(
                         &async_expr.block,
                         &self.input.sig.inputs,
@@ -629,8 +656,14 @@ impl<'block> AsyncTraitInfo<'block> {
                         None,
                     );
                     let async_attrs = &async_expr.attrs;
-                    quote! {
-                        Box::pin(#(#async_attrs) * async move { #instrumented_block })
+                    if pinned_box {
+                        quote! {
+                            Box::pin(#(#async_attrs) * async move { #instrumented_block })
+                        }
+                    } else {
+                        quote! {
+                            #(#async_attrs) * async move { #instrumented_block }
+                        }
                     }
                 }
             };

tracing-attributes/src/lib.rs

@@ -530,7 +530,7 @@ fn instrument_speculative(
 }
 
 /// Instrument the function, by fully parsing the function body,
-/// which allows us to rewrite some statements related to async_trait-like patterns.
+/// which allows us to rewrite some statements related to async-like patterns.
 fn instrument_precise(
     args: attr::InstrumentArgs,
     item: proc_macro::TokenStream,
@@ -540,8 +540,8 @@ fn instrument_precise(
 
     // check for async_trait-like patterns in the block, and instrument
     // the future instead of the wrapper
-    if let Some(async_trait) = expand::AsyncTraitInfo::from_fn(&input) {
-        return Ok(async_trait.gen_async_trait(args, instrumented_function_name.as_str()));
+    if let Some(async_like) = expand::AsyncInfo::from_fn(&input) {
+        return Ok(async_like.gen_async(args, instrumented_function_name.as_str()));
     }
 
     Ok(expand::gen_function(

tracing-attributes/tests/async_fn.rs

@@ -371,3 +371,66 @@ fn out_of_scope_fields() {
 
     handle.assert_finished();
 }
+
+#[test]
+fn manual_impl_future() {
+    #[allow(clippy::manual_async_fn)]
+    #[instrument]
+    fn manual_impl_future() -> impl Future<Output = ()> {
+        async {
+            tracing::trace!(poll = true);
+        }
+    }
+
+    let span = span::mock().named("manual_impl_future");
+    let poll_event = || event::mock().with_fields(field::mock("poll").with_value(&true));
+
+    let (collector, handle) = collector::mock()
+        // await manual_impl_future
+        .new_span(span.clone())
+        .enter(span.clone())
+        .event(poll_event())
+        .exit(span.clone())
+        .drop_span(span)
+        .done()
+        .run_with_handle();
+
+    with_default(collector, || {
+        block_on_future(async {
+            manual_impl_future().await;
+        });
+    });
+
+    handle.assert_finished();
+}
+
+#[test]
+fn manual_box_pin() {
+    #[instrument]
+    fn manual_box_pin() -> Pin<Box<dyn Future<Output = ()>>> {
+        Box::pin(async {
+            tracing::trace!(poll = true);
+        })
+    }
+
+    let span = span::mock().named("manual_box_pin");
+    let poll_event = || event::mock().with_fields(field::mock("poll").with_value(&true));
+
+    let (collector, handle) = collector::mock()
+        // await manual_box_pin
+        .new_span(span.clone())
+        .enter(span.clone())
+        .event(poll_event())
+        .exit(span.clone())
+        .drop_span(span)
+        .done()
+        .run_with_handle();
+
+    with_default(collector, || {
+        block_on_future(async {
+            manual_box_pin().await;
+        });
+    });
+
+    handle.assert_finished();
+}