diff --git a/tracing-attributes/Cargo.toml b/tracing-attributes/Cargo.toml index a7774c4615..96044535bb 100644 --- a/tracing-attributes/Cargo.toml +++ b/tracing-attributes/Cargo.toml @@ -39,7 +39,7 @@ async-await = [] [dependencies] proc-macro2 = "1" -syn = { version = "1", features = ["full", "extra-traits"] } +syn = { version = "1", features = ["full", "extra-traits", "visit-mut"] } quote = "1" diff --git a/tracing-attributes/src/lib.rs b/tracing-attributes/src/lib.rs index a8d3aadc64..5d53cb96f9 100644 --- a/tracing-attributes/src/lib.rs +++ b/tracing-attributes/src/lib.rs @@ -176,8 +176,8 @@ use syn::{ /// ``` /// /// 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), +/// (a crate that allows defining async functions in traits, +/// something not currently possible in Rust), /// and hopefully most libraries that exhibit similar behaviors: /// /// ``` @@ -186,18 +186,43 @@ use syn::{ /// /// #[async_trait] /// pub trait Foo { -/// async fn foo(&self, v: usize) -> (); +/// async fn foo(&self, arg: usize); /// } /// /// #[derive(Debug)] -/// struct FooImpl; +/// struct FooImpl(usize); /// /// #[async_trait] /// impl Foo for FooImpl { -/// #[instrument(skip(self))] -/// async fn foo(&self, v: usize) {} +/// #[instrument(fields(value = self.0, tmp = std::any::type_name::()))] +/// async fn foo(&self, arg: usize) {} /// } /// ``` +/// +/// An interesting note on this subject is that references to the `Self` +/// type inside the `fields` argument are only allowed when the instrumented +/// function is a method aka. the function receives `self` as an argument. +/// For example, this *will not work* because it doesn't receive `self`: +/// ```compile_fail +/// # use tracing::instrument; +/// use async_trait::async_trait; +/// +/// #[async_trait] +/// pub trait Bar { +/// async fn bar(); +/// } +/// +/// #[derive(Debug)] +/// struct BarImpl(usize); +/// +/// #[async_trait] +/// impl Bar for BarImpl { +/// #[instrument(fields(tmp = std::any::type_name::()))] +/// async fn bar() {} +/// } +/// ``` +/// Instead, you should manually rewrite any `Self` types as the type for +/// which you implement the trait: `#[instrument(fields(tmp = std::any::type_name::()))]`. /// /// [span]: https://docs.rs/tracing/latest/tracing/span/index.html @@ -211,20 +236,25 @@ pub fn instrument( let input: ItemFn = syn::parse_macro_input!(item as ItemFn); let args = syn::parse_macro_input!(args as InstrumentArgs); + let instrumented_function_name = input.sig.ident.to_string(); + // 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()) - { + if let Some(internal_fun) = get_async_trait_info(&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(); + if fun.sig.ident == internal_fun.name { + *stmt = syn::parse2(gen_body( + fun, + args, + instrumented_function_name, + Some(internal_fun), + )) + .unwrap(); break; } } @@ -240,14 +270,15 @@ pub fn instrument( ) .into() } else { - gen_body(&input, args, None).into() + gen_body(&input, args, instrumented_function_name, None).into() } } fn gen_body( input: &ItemFn, - args: InstrumentArgs, - fun_name: Option, + mut args: InstrumentArgs, + instrumented_function_name: String, + async_trait_fun: 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 @@ -286,14 +317,7 @@ fn gen_body( .name .as_ref() .map(|name| quote!(#name)) - // are we overriding the name because the span is inside a function - // generated by `async-trait`? - .or_else(|| fun_name.as_ref().map(|name| quote!(#name))) - // if neither override is present, use the parsed function's name. - .unwrap_or_else(|| { - let name = ident.to_string(); - quote!(#name) - }); + .unwrap_or_else(|| quote!(#instrumented_function_name)); // generate this inside a closure, so we can return early on errors. let span = (|| { @@ -306,11 +330,21 @@ fn gen_body( 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' + // 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) {} + // } + // ``` .map(|x| { - if fun_name.is_some() && x == "_self" { + // if we are inside a function generated by async-trait, we + // should take care to rewrite "_self" as "self" for + // 'user convenience' + if async_trait_fun.is_some() && x == "_self" { (Ident::new("self", x.span()), x) } else { (x.clone(), x) @@ -318,20 +352,6 @@ fn gen_body( }) .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 &args.skips { if !param_names.iter().map(|(user, _)| user).any(|y| y == skip) { return quote_spanned! {skip.span()=> @@ -364,18 +384,26 @@ fn gen_body( }) .map(|(user_name, real_name)| quote!(#user_name = tracing::field::debug(&#real_name))) .collect(); + + // when async-trait is in use, replace instances of "self" with "_self" inside the fields values + if let (Some(ref async_trait_fun), Some(Fields(ref mut fields))) = + (async_trait_fun, &mut args.fields) + { + let mut replacer = SelfReplacer { + ty: async_trait_fun.self_type.clone(), + }; + for e in fields.iter_mut().filter_map(|f| f.value.as_mut()) { + syn::visit_mut::visit_expr_mut(&mut replacer, e); + } + } + let custom_fields = &args.fields; - let custom_fields = if quoted_fields.is_empty() { - quote! { #custom_fields } - } else { - quote! {, #custom_fields } - }; quote!(tracing::span!( target: #target, #level, #span_name, - #(#quoted_fields),* + #(#quoted_fields,)* #custom_fields )) @@ -782,20 +810,20 @@ mod kw { syn::custom_keyword!(err); } -// 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. +// Get the AST of the inner function we need to hook, if it was generated +// by async-trait. +// 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. +// Were we to instrument its parent, we would only collect information +// regarding the allocation of that future, and not its own span of execution. +// So we inspect the block of the function to find if it matches the pattern +// `async fn foo<...>(...) {...}; Box::pin(foo<...>(...))` and we return +// the name `foo` if that is the case. 'gen_body' 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 { +fn get_async_trait_function(block: &Block, block_is_async: bool) -> Option<&ItemFn> { // are we in an async context? If yes, this isn't a async_trait-like pattern if block_is_async { return None; @@ -816,7 +844,7 @@ fn get_async_trait_name(block: &Block, block_is_async: bool) -> Option { // 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()); + inside_funs.push(fun); } } else if let Stmt::Expr(e) = &stmt { last_expr = Some(e); @@ -846,9 +874,11 @@ fn get_async_trait_name(block: &Block, block_is_async: bool) -> Option { // "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); + for fun in inside_funs { + if fun.sig.ident == func_name { + // we must hook this function now + return Some(fun); + } } } } @@ -858,6 +888,48 @@ fn get_async_trait_name(block: &Block, block_is_async: bool) -> Option { None } +struct AsyncTraitInfo { + name: String, + self_type: Option, +} + +// Return the informations necessary to process a function annotated with async-trait. +fn get_async_trait_info(block: &Block, block_is_async: bool) -> Option { + let fun = get_async_trait_function(block, block_is_async)?; + + // if "_self" is present as an argument, we store its type to be able to rewrite "Self" (the + // parameter type) with the type of "_self" + let self_type = fun + .sig + .inputs + .iter() + .map(|arg| { + if let FnArg::Typed(ty) = arg { + if let Pat::Ident(PatIdent { ident, .. }) = &*ty.pat { + if ident == "_self" { + let mut ty = &*ty.ty; + // extract the inner type if the argument is "&self" or "&mut self" + if let syn::Type::Reference(syn::TypeReference { elem, .. }) = ty { + ty = &*elem; + } + if let syn::Type::Path(tp) = ty { + return Some(tp.clone()); + } + } + } + } + + None + }) + .next() + .flatten(); + + Some(AsyncTraitInfo { + name: fun.sig.ident.to_string(), + self_type, + }) +} + // Return a path as a String fn path_to_string(path: &Path) -> String { use std::fmt::Write; @@ -872,3 +944,27 @@ fn path_to_string(path: &Path) -> String { } res } + +// A visitor struct replacing the "self" and "Self" tokens in user-supplied fields expressions when +// the function is generated by async-trait. +struct SelfReplacer { + ty: Option, +} + +impl syn::visit_mut::VisitMut for SelfReplacer { + fn visit_ident_mut(&mut self, id: &mut Ident) { + if id == "self" { + *id = Ident::new("_self", id.span()) + } + } + + fn visit_type_mut(&mut self, ty: &mut syn::Type) { + if let syn::Type::Path(syn::TypePath { ref mut path, .. }) = ty { + if path_to_string(path) == "Self" { + if let Some(ref true_type) = self.ty { + *path = true_type.path.clone(); + } + } + } + } +} diff --git a/tracing-attributes/tests/async_fn.rs b/tracing-attributes/tests/async_fn.rs index 7e2465e1cc..ac95e94d2e 100644 --- a/tracing-attributes/tests/async_fn.rs +++ b/tracing-attributes/tests/async_fn.rs @@ -150,3 +150,118 @@ fn async_fn_with_async_trait() { handle.assert_finished(); } + +#[test] +fn async_fn_with_async_trait_and_fields_expressions() { + use async_trait::async_trait; + + #[async_trait] + pub trait Test { + async fn call(&mut self, v: usize); + } + + #[derive(Clone, Debug)] + struct TestImpl; + + impl TestImpl { + fn foo(&self) -> usize { + 42 + } + } + + #[async_trait] + impl Test for TestImpl { + // check that self is correctly handled, even when using async_trait + #[instrument(fields(val=self.foo(), test=%v+5))] + async fn call(&mut self, v: usize) {} + } + + let span = span::mock().named("call"); + let (subscriber, handle) = subscriber::mock() + .new_span( + span.clone().with_field( + field::mock("v") + .with_value(&tracing::field::debug(5)) + .and(field::mock("test").with_value(&tracing::field::debug(10))) + .and(field::mock("val").with_value(&42u64)), + ), + ) + .enter(span.clone()) + .exit(span.clone()) + .drop_span(span) + .done() + .run_with_handle(); + + with_default(subscriber, || { + block_on_future(async { TestImpl.call(5).await }); + }); + + handle.assert_finished(); +} + +#[test] +fn async_fn_with_async_trait_and_fields_expressions_with_generic_parameter() { + use async_trait::async_trait; + + #[async_trait] + pub trait Test { + async fn call(); + async fn call_with_self(&self); + async fn call_with_mut_self(&mut self); + } + + #[derive(Clone, Debug)] + struct TestImpl; + + #[async_trait] + impl Test for TestImpl { + // instrumenting this is currently not possible, see https://github.com/tokio-rs/tracing/issues/864#issuecomment-667508801 + //#[instrument(fields(Self=std::any::type_name::()))] + async fn call() {} + + #[instrument(fields(Self=std::any::type_name::()))] + async fn call_with_self(&self) {} + + #[instrument(fields(Self=std::any::type_name::()))] + async fn call_with_mut_self(self: &mut Self) {} + } + + //let span = span::mock().named("call"); + let span2 = span::mock().named("call_with_self"); + let span3 = span::mock().named("call_with_mut_self"); + let (subscriber, handle) = subscriber::mock() + /*.new_span(span.clone() + .with_field( + field::mock("Self").with_value(&"TestImpler"))) + .enter(span.clone()) + .exit(span.clone()) + .drop_span(span)*/ + .new_span( + span2 + .clone() + .with_field(field::mock("Self").with_value(&std::any::type_name::())), + ) + .enter(span2.clone()) + .exit(span2.clone()) + .drop_span(span2) + .new_span( + span3 + .clone() + .with_field(field::mock("Self").with_value(&std::any::type_name::())), + ) + .enter(span3.clone()) + .exit(span3.clone()) + .drop_span(span3) + .done() + .run_with_handle(); + + with_default(subscriber, || { + block_on_future(async { + TestImpl::call().await; + TestImpl.call_with_self().await; + TestImpl.call_with_mut_self().await + }); + }); + + handle.assert_finished(); +}