diff --git a/README.md b/README.md index fdadaee96e..3941a8469b 100644 --- a/README.md +++ b/README.md @@ -242,8 +242,6 @@ async fn write(stream: &mut TcpStream) -> io::Result { Under the hood, the `#[instrument]` macro performs same the explicit span attachment that `Future::instrument` does. -Note: the [`#[tracing::instrument]`](https://github.com/tokio-rs/tracing/issues/399)` macro does not work correctly with the [async-trait](https://github.com/dtolnay/async-trait) crate. This bug is tracked in [#399](https://github.com/tokio-rs/tracing/issues/399). - ## Getting Help First, see if the answer to your question can be found in the API documentation. diff --git a/tracing-attributes/Cargo.toml b/tracing-attributes/Cargo.toml index 0939e6d12a..4e95d6292f 100644 --- a/tracing-attributes/Cargo.toml +++ b/tracing-attributes/Cargo.toml @@ -40,6 +40,7 @@ async-await = [] [dependencies] syn = { version = "1", features = ["full", "extra-traits"] } quote = "1" +proc-macro2 = "1" [dev-dependencies] @@ -47,6 +48,7 @@ tracing = { path = "../tracing", version = "0.1" } tracing-futures = { path = "../tracing-futures", version = "0.2" } tokio-test = { version = "0.2.0" } tracing-core = { path = "../tracing-core", version = "0.1"} +async-trait = "0.1" [badges] maintenance = { status = "experimental" } diff --git a/tracing-attributes/src/lib.rs b/tracing-attributes/src/lib.rs index 01404cadc0..5c769281f3 100644 --- a/tracing-attributes/src/lib.rs +++ b/tracing-attributes/src/lib.rs @@ -68,9 +68,9 @@ use std::iter; use proc_macro::TokenStream; use quote::{quote, quote_spanned, ToTokens}; use syn::{ - spanned::Spanned, AttributeArgs, FieldPat, FnArg, Ident, ItemFn, Lit, LitInt, Meta, MetaList, - MetaNameValue, NestedMeta, Pat, PatIdent, PatReference, PatStruct, PatTuple, PatTupleStruct, - PatType, Path, Signature, + spanned::Spanned, AttributeArgs, Block, Expr, ExprCall, FieldPat, FnArg, Ident, Item, ItemFn, + Lit, LitInt, Meta, MetaList, MetaNameValue, NestedMeta, Pat, PatIdent, PatReference, PatStruct, + PatTuple, PatTupleStruct, PatType, Path, Signature, Stmt, }; /// Instruments a function to create and enter a `tracing` [span] every time @@ -168,6 +168,31 @@ use syn::{ /// # Ok(()) /// } /// ``` +/// +/// It also works with [async-trait](https://crates.io/crates/async-trait) +/// (a crate that allows async functions on traits, +/// something not currently possible with rustc alone), +/// and hopefully most libraries that exhibit similar behaviors: +/// +/// ``` +/// # use tracing::instrument; +/// use async_trait::async_trait; +/// +/// #[async_trait] +/// pub trait Foo { +/// async fn foo(&self, v: usize) -> (); +/// } +/// +/// #[derive(Debug)] +/// struct FooImpl; +/// +/// #[async_trait] +/// impl Foo for FooImpl { +/// #[instrument(skip(self))] +/// async fn foo(&self, v: usize) {} +/// } +/// ``` + /// /// [span]: https://docs.rs/tracing/latest/tracing/span/index.html /// [`tracing`]: https://github.com/tokio-rs/tracing @@ -177,6 +202,44 @@ pub fn instrument(args: TokenStream, item: TokenStream) -> TokenStream { let input: ItemFn = syn::parse_macro_input!(item as ItemFn); let args = syn::parse_macro_input!(args as AttributeArgs); + // check for async_trait-like patterns in the block and wrap the + // internal function with Instrument instead of wrapping the + // async_trait generated wrapper + if let Some(internal_fun_name) = + get_async_trait_name(&input.block, input.sig.asyncness.is_some()) + { + // let's rewrite some statements! + let mut stmts: Vec = input.block.stmts.to_vec(); + for stmt in &mut stmts { + if let Stmt::Item(Item::Fn(fun)) = stmt { + // instrument the function if we considered it as the one we truly want to trace + if fun.sig.ident == internal_fun_name { + *stmt = syn::parse2(gen_body(fun, args, Some(input.sig.ident.to_string()))) + .unwrap(); + break; + } + } + } + + let sig = &input.sig; + let attrs = &input.attrs; + quote!( + #(#attrs) * + #sig { + #(#stmts) * + } + ) + .into() + } else { + gen_body(&input, args, None).into() + } +} + +fn gen_body( + input: &ItemFn, + args: AttributeArgs, + fun_name: Option, +) -> proc_macro2::TokenStream { // these are needed ahead of time, as ItemFn contains the function body _and_ // isn't representable inside a quote!/quote_spanned! macro // (Syn's ToTokens isn't implemented for ItemFn) @@ -206,7 +269,11 @@ pub fn instrument(args: TokenStream, item: TokenStream) -> TokenStream { } = sig; // function name - let ident_str = ident.to_string(); + let ident_str = if let Some(x) = &fun_name { + x.clone() + } else { + ident.to_string() + }; // generate this inside a closure, so we can return early on errors. let span = (|| { @@ -216,42 +283,67 @@ pub fn instrument(args: TokenStream, item: TokenStream) -> TokenStream { Err(err) => return quote!(#err), }; - let param_names: Vec = params + let param_names: Vec<(Ident, Ident)> = params .clone() .into_iter() .flat_map(|param| match param { FnArg::Typed(PatType { pat, .. }) => param_names(*pat), FnArg::Receiver(_) => Box::new(iter::once(Ident::new("self", param.span()))), }) + // if we are inside a function generated by async-trait, we + // should take care to rewrite "_self" as "self" for + // 'user convenience' + .map(|x| { + if fun_name.is_some() && x == "_self" { + (Ident::new("self", x.span()), x) + } else { + (x.clone(), x) + } + }) .collect(); + // TODO: allow the user to rename fields at will (all the + // machinery should be here) + + // Little dance with new (user-exposed) names and old (internal) + // names of identifiers. That way, you can do the following + // even though async_trait rewrite "self" as "_self": + // ``` + // #[async_trait] + // impl Foo for FooImpl { + // #[instrument(skip(self))] + // async fn foo(&self, v: usize) {} + // } + // ``` + for skip in &skips { - if !param_names.contains(skip) { + if !param_names.iter().map(|(user, _)| user).any(|y| y == skip) { return quote_spanned! {skip.span()=> compile_error!("attempting to skip non-existent parameter") }; } } - let param_names: Vec = param_names + // filter out skipped fields + let param_names: Vec<(Ident, Ident)> = param_names .into_iter() - .filter(|ident| !skips.contains(ident)) + .filter(|(ident, _)| !skips.contains(ident)) .collect(); - let fields = match fields(&args, ¶m_names) { + let new_param_names: Vec<&Ident> = param_names.iter().map(|x| &x.0).collect(); + + let fields = match fields(&args, &new_param_names) { Ok(fields) => fields, Err(err) => return quote!(#err), }; - let param_names_clone = param_names.clone(); - let level = level(&args); let target = target(&args); let span_name = name(&args, ident_str); let mut quoted_fields: Vec<_> = param_names - .into_iter() - .map(|i| quote!(#i = tracing::field::debug(&#i))) + .iter() + .map(|(user_name, real_name)| quote!(#user_name = tracing::field::debug(&#real_name))) .collect(); quoted_fields.extend(fields.into_iter().map(|(key, value)| { let value = match value { @@ -276,7 +368,8 @@ pub fn instrument(args: TokenStream, item: TokenStream) -> TokenStream { // If `err` is in args, instrument any resulting `Err`s. let body = if asyncness.is_some() { if instrument_err(&args) { - quote_spanned! {block.span()=> + quote_spanned!(block.span()=> + let __tracing_attr_span = #span; tracing_futures::Instrument::instrument(async move { match async move { #block }.await { Ok(x) => Ok(x), @@ -286,18 +379,20 @@ pub fn instrument(args: TokenStream, item: TokenStream) -> TokenStream { } } }, __tracing_attr_span).await - } + ) } else { - quote_spanned! {block.span()=> - tracing_futures::Instrument::instrument( - async move { #block }, - __tracing_attr_span - ) + quote_spanned!(block.span()=> + let __tracing_attr_span = #span; + tracing_futures::Instrument::instrument( + async move { #block }, + __tracing_attr_span + ) .await - } + ) } } else if instrument_err(&args) { quote_spanned!(block.span()=> + let __tracing_attr_span = #span; let __tracing_attr_guard = __tracing_attr_span.enter(); match { #block } { Ok(x) => Ok(x), @@ -309,6 +404,7 @@ pub fn instrument(args: TokenStream, item: TokenStream) -> TokenStream { ) } else { quote_spanned!(block.span()=> + let __tracing_attr_span = #span; let __tracing_attr_guard = __tracing_attr_span.enter(); #block ) @@ -319,11 +415,9 @@ pub fn instrument(args: TokenStream, item: TokenStream) -> TokenStream { #vis #constness #unsafety #asyncness #abi fn #ident<#gen_params>(#params) #return_type #where_clause { - let __tracing_attr_span = #span; #body } ) - .into() } fn param_names(pat: Pat) -> Box> { @@ -462,7 +556,7 @@ fn target(args: &[NestedMeta]) -> impl ToTokens { fn fields( args: &[NestedMeta], - param_names: &[Ident], + param_names: &[&Ident], ) -> Result<(Vec<(Ident, Option)>), impl ToTokens> { let mut fields = args.iter().filter_map(|arg| match arg { NestedMeta::Meta(Meta::List(MetaList { @@ -579,3 +673,94 @@ fn instrument_err(args: &[NestedMeta]) -> bool { _ => false, }) } + +// Get the name of the inner function we need to hook, if the function +// was generated by async-trait. +// When we are given a function generated by async-trait, that function +// is only a "temporary" one that returns a pinned future, and it is +// that pinned future that needs to be instrumented, otherwise we will +// only collect information on the moment the future was "built", +// and not its true span of execution. +// So we inspect the block of the function to find if we can find the +// pattern `async fn foo<...>(...) {...}; Box::pin(foo<...>(...))` and +// return the name `foo` if that is the case. Our caller will then be +// able to use that information to instrument the proper function. +// (this follows the approach suggested in +// https://github.com/dtolnay/async-trait/issues/45#issuecomment-571245673) +fn get_async_trait_name(block: &Block, block_is_async: bool) -> Option { + // are we in an async context? If yes, this isn't a async_trait-like pattern + if block_is_async { + return None; + } + + // list of async functions declared inside the block + let mut inside_funs = Vec::new(); + // last expression declared in 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) + let mut last_expr = None; + + // obtain the list of direct internal functions and the last + // expression of the block + for stmt in &block.stmts { + if let Stmt::Item(Item::Fn(fun)) = &stmt { + // is the function declared as async? If so, this is a good + // candidate, let's keep it in hand + if fun.sig.asyncness.is_some() { + inside_funs.push(fun.sig.ident.to_string()); + } + } else if let Stmt::Expr(e) = &stmt { + last_expr = Some(e); + } + } + + // let's play with (too much) pattern matching + // is the last expression a function call? + if let Some(Expr::Call(ExprCall { + func: outside_func, + args: outside_args, + .. + })) = last_expr + { + if let Expr::Path(path) = outside_func.as_ref() { + // is it a call to `Box::pin()`? + if "Box::pin" == path_to_string(&path.path) { + // does it takes at least an argument? (if it doesn't, + // it's not gonna compile anyway, but that's no reason + // to (try to) perform an out of bounds access) + if outside_args.is_empty() { + return None; + } + // is the argument to Box::pin a function call itself? + if let Expr::Call(ExprCall { func, args, .. }) = &outside_args[0] { + if let Expr::Path(inside_path) = func.as_ref() { + // "stringify" the path of the function called + let func_name = path_to_string(&inside_path.path); + // is this function directly defined insided the current block? + if inside_funs.contains(&func_name) { + // we must hook this function now + return Some(func_name); + } + } + } + } + } + } + None +} + +// Return a path as a String +fn path_to_string(path: &Path) -> String { + use std::fmt::Write; + // some heuristic to prevent too many allocations + let mut res = String::with_capacity(path.segments.len() * 5); + for i in 0..path.segments.len() { + write!(&mut res, "{}", path.segments[i].ident) + .expect("writing to a String should never fail"); + if i < path.segments.len() - 1 { + res.push_str("::"); + } + } + res +} diff --git a/tracing-attributes/tests/async_fn.rs b/tracing-attributes/tests/async_fn.rs index a73502fb0f..7e2465e1cc 100644 --- a/tracing-attributes/tests/async_fn.rs +++ b/tracing-attributes/tests/async_fn.rs @@ -5,8 +5,9 @@ mod support; use support::*; use tracing::subscriber::with_default; +use tracing_attributes::instrument; -#[tracing_attributes::instrument] +#[instrument] async fn test_async_fn(polls: usize) -> Result<(), ()> { let future = PollN::new_ok(polls); tracing::trace!(awaiting = true); @@ -30,3 +31,122 @@ fn async_fn_only_enters_for_polls() { }); handle.assert_finished(); } + +#[test] +fn async_fn_nested() { + #[instrument] + async fn test_async_fns_nested() { + test_async_fns_nested_other().await + } + + #[instrument] + async fn test_async_fns_nested_other() { + tracing::trace!(nested = true); + } + + let span = span::mock().named("test_async_fns_nested"); + let span2 = span::mock().named("test_async_fns_nested_other"); + let (subscriber, handle) = subscriber::mock() + .new_span(span.clone()) + .enter(span.clone()) + .new_span(span2.clone()) + .enter(span2.clone()) + .event(event::mock().with_fields(field::mock("nested").with_value(&true))) + .exit(span2.clone()) + .drop_span(span2) + .exit(span.clone()) + .drop_span(span) + .done() + .run_with_handle(); + + with_default(subscriber, || { + block_on_future(async { test_async_fns_nested().await }); + }); + + handle.assert_finished(); +} + +#[test] +fn async_fn_with_async_trait() { + use async_trait::async_trait; + + // test the correctness of the metadata obtained by #[instrument] + // (function name, functions parameters) when async-trait is used + #[async_trait] + pub trait TestA { + async fn foo(&mut self, v: usize); + } + + // test nesting of async fns with aync-trait + #[async_trait] + pub trait TestB { + async fn bar(&self); + } + + // test skip(self) with async-await + #[async_trait] + pub trait TestC { + async fn baz(&self); + } + + #[derive(Debug)] + struct TestImpl(usize); + + #[async_trait] + impl TestA for TestImpl { + #[instrument] + async fn foo(&mut self, v: usize) { + self.baz().await; + self.0 = v; + self.bar().await + } + } + + #[async_trait] + impl TestB for TestImpl { + #[instrument] + async fn bar(&self) { + tracing::trace!(val = self.0); + } + } + + #[async_trait] + impl TestC for TestImpl { + #[instrument(skip(self))] + async fn baz(&self) { + tracing::trace!(val = self.0); + } + } + + let span = span::mock().named("foo"); + let span2 = span::mock().named("bar"); + let span3 = span::mock().named("baz"); + let (subscriber, handle) = subscriber::mock() + .new_span( + span.clone() + .with_field(field::mock("self")) + .with_field(field::mock("v")), + ) + .enter(span.clone()) + .new_span(span3.clone()) + .enter(span3.clone()) + .event(event::mock().with_fields(field::mock("val").with_value(&2u64))) + .exit(span3.clone()) + .drop_span(span3) + .new_span(span2.clone().with_field(field::mock("self"))) + .enter(span2.clone()) + .event(event::mock().with_fields(field::mock("val").with_value(&5u64))) + .exit(span2.clone()) + .drop_span(span2) + .exit(span.clone()) + .drop_span(span) + .done() + .run_with_handle(); + + with_default(subscriber, || { + let mut test = TestImpl(2); + block_on_future(async { test.foo(5).await }); + }); + + handle.assert_finished(); +} diff --git a/tracing/README.md b/tracing/README.md index 7d483112f7..548c479c8a 100644 --- a/tracing/README.md +++ b/tracing/README.md @@ -260,8 +260,6 @@ async fn write(stream: &mut TcpStream) -> io::Result { Under the hood, the `#[instrument]` macro performs same the explicit span attachment that `Future::instrument` does. -Note: the [`#[tracing::instrument]`](https://github.com/tokio-rs/tracing/issues/399)` macro does work correctly with the [async-trait](https://github.com/dtolnay/async-trait) crate. This bug is tracked in [#399](https://github.com/tokio-rs/tracing/issues/399). - ### Concepts This crate provides macros for creating `Span`s and `Event`s, which represent