#[instrument] - support use of the 'self' variable in async-trait contexts

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"
 
 

tracing-attributes/src/lib.rs

@@ -211,20 +211,21 @@ pub fn instrument(
     let input: ItemFn = syn::parse_macro_input!(item as ItemFn);
     let args = syn::parse_macro_input!(args as InstrumentArgs);
 
+    let span_declared_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<Stmt> = 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, span_declared_name, Some(internal_fun)))
+                            .unwrap();
                     break;
                 }
             }
@@ -240,14 +241,15 @@ pub fn instrument(
         )
         .into()
     } else {
-        gen_body(&input, args, None).into()
+        gen_body(&input, args, span_declared_name, None).into()
     }
 }
 
 fn gen_body(
     input: &ItemFn,
-    args: InstrumentArgs,
-    fun_name: Option<String>,
+    mut args: InstrumentArgs,
+    span_declared_name: String,
+    async_trait_fun: Option<AsyncTraitInfo>,
 ) -> 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 +288,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!(#span_declared_name));
 
     // generate this inside a closure, so we can return early on errors.
     let span = (|| {
@@ -306,32 +301,28 @@ 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)
                 }
             })
             .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 +355,30 @@ 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(async_trait_fun) = async_trait_fun {
+            if let Some(Fields(ref mut fields)) = args.fields {
+                for mut field in fields.iter_mut() {
+                    if let Some(ref mut e) = field.value {
+                        syn::visit_mut::visit_expr_mut(
+                            &mut SelfReplacer {
+                                ty: async_trait_fun.self_type.clone(),
+                            },
+                            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
 
         ))
@@ -795,7 +798,7 @@ mod kw {
 // 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<String> {
+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 +819,7 @@ fn get_async_trait_name(block: &Block, block_is_async: bool) -> Option<String> {
             // 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 +849,11 @@ fn get_async_trait_name(block: &Block, block_is_async: bool) -> Option<String> {
                         // "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 +863,48 @@ fn get_async_trait_name(block: &Block, block_is_async: bool) -> Option<String> {
     None
 }
 
+struct AsyncTraitInfo {
+    name: String,
+    self_type: Option<syn::TypePath>,
+}
+
+// Return the informations necessary to process a function annotated with async-trait.
+fn get_async_trait_info(block: &Block, block_is_async: bool) -> Option<AsyncTraitInfo> {
+    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 +919,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<syn::TypePath>,
+}
+
+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();
+                }
+            }
+        }
+    }
+}

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::<Self>()))]
+        async fn call() {}
+
+        #[instrument(fields(Self=std::any::type_name::<Self>()))]
+        async fn call_with_self(&self) {}
+
+        #[instrument(fields(Self=std::any::type_name::<Self>()))]
+        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::<TestImpl>())),
+        )
+        .enter(span2.clone())
+        .exit(span2.clone())
+        .drop_span(span2)
+        .new_span(
+            span3
+                .clone()
+                .with_field(field::mock("Self").with_value(&std::any::type_name::<TestImpl>())),
+        )
+        .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();
+}