From f5ad56a6d56cc0dc74ef113031a4dcc6f2de9e2c Mon Sep 17 00:00:00 2001 From: Arpad Borsos Date: Thu, 27 Jan 2022 21:00:11 +0100 Subject: [PATCH] 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 :-( --- tracing-attributes/src/expand.rs | 105 ++++++++++++++++++--------- tracing-attributes/src/lib.rs | 6 +- tracing-attributes/tests/async_fn.rs | 63 ++++++++++++++++ 3 files changed, 135 insertions(+), 39 deletions(-) diff --git a/tracing-attributes/src/expand.rs b/tracing-attributes/src/expand.rs index e9b4f898c6..3928441e1f 100644 --- a/tracing-attributes/src/expand.rs +++ b/tracing-attributes/src/expand.rs @@ -437,43 +437,56 @@ fn param_names(pat: Pat, record_type: RecordType) -> Box { - // 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, 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 { - // 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 } + } } } }; diff --git a/tracing-attributes/src/lib.rs b/tracing-attributes/src/lib.rs index a321962ca6..b44c77b815 100644 --- a/tracing-attributes/src/lib.rs +++ b/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( diff --git a/tracing-attributes/tests/async_fn.rs b/tracing-attributes/tests/async_fn.rs index 1f5e3d0bf6..e0a6073beb 100644 --- a/tracing-attributes/tests/async_fn.rs +++ b/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 { + 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 (subscriber, handle) = subscriber::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(subscriber, || { + block_on_future(async { + manual_impl_future().await; + }); + }); + + handle.assert_finished(); +} + +#[test] +fn manual_box_pin() { + #[instrument] + fn manual_box_pin() -> Pin>> { + 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 (subscriber, handle) = subscriber::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(subscriber, || { + block_on_future(async { + manual_box_pin().await; + }); + }); + + handle.assert_finished(); +}