attributes: support use of the 'self' variable in async-trait contexts (#875)

##  Motivation

Fixes #864.

## Solution

Visit the fields and replace 'self' with '_self', and 'Self' with the
type of the impl (when we can acces it, see
https://github.com/tokio-rs/tracing/issues/864#issuecomment-667508801
for current limitations).

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

@@ -184,8 +184,8 @@ use syn::{
 /// ```
 ///
 /// It also works with [async-trait](https://crates.io/crates/async-trait)
-/// (a crate that allows async functions on traits,
+/// (a crate that allows defining async functions in traits,
-/// something not currently possible with rustc alone),
+/// something not currently possible in Rust),
 /// and hopefully most libraries that exhibit similar behaviors:
 ///
 /// ```
@@ -194,18 +194,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))]
+///     #[instrument(fields(value = self.0, tmp = std::any::type_name::<Self>()))]
-///     async fn foo(&self, v: usize) {}
+///     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::<Self>()))]
+///     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::<Bar>()))]`.
 
 ///
 /// [span]: https://docs.rs/tracing/latest/tracing/span/index.html
@@ -219,19 +244,24 @@ 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) =
+    if let Some(internal_fun) = get_async_trait_info(&input.block, input.sig.asyncness.is_some()) {
-        get_async_trait_name(&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 {
+                if fun.sig.ident == internal_fun.name {
-                    *stmt = syn::parse2(gen_body(fun, args, Some(input.sig.ident.to_string())))
+                    *stmt = syn::parse2(gen_body(
+                        fun,
+                        args,
+                        instrumented_function_name,
+                        Some(internal_fun),
+                    ))
                     .unwrap();
                     break;
                 }
@@ -248,14 +278,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,
+    mut args: InstrumentArgs,
-    fun_name: Option<String>,
+    instrumented_function_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
@@ -294,14 +325,7 @@ fn gen_body(
         .name
         .as_ref()
         .map(|name| quote!(#name))
-        // are we overriding the name because the span is inside a function
+        .unwrap_or_else(|| quote!(#instrumented_function_name));
-        // 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)
-        });
 
     // generate this inside a closure, so we can return early on errors.
     let span = (|| {
@@ -314,21 +338,6 @@ 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'
-            .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":
@@ -339,6 +348,17 @@ fn gen_body(
             //     async fn foo(&self, v: usize) {}
             // }
             // ```
+            .map(|x| {
+                // 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();
 
         for skip in &args.skips {
             if !param_names.iter().map(|(user, _)| user).any(|y| y == skip) {
@@ -372,18 +392,26 @@ fn gen_body(
             })
             .map(|(user_name, real_name)| quote!(#user_name = tracing::field::debug(&#real_name)))
             .collect();
-        let custom_fields = &args.fields;
+
-        let custom_fields = if quoted_fields.is_empty() {
+        // when async-trait is in use, replace instances of "self" with "_self" inside the fields values
-            quote! { #custom_fields }
+        if let (Some(ref async_trait_fun), Some(Fields(ref mut fields))) =
-        } else {
+            (async_trait_fun, &mut args.fields)
-            quote! {, #custom_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;
 
         quote!(tracing::span!(
             target: #target,
             #level,
             #span_name,
-            #(#quoted_fields),*
+            #(#quoted_fields,)*
             #custom_fields
 
         ))
@@ -790,20 +818,20 @@ mod kw {
     syn::custom_keyword!(err);
 }
 
-// Get the name of the inner function we need to hook, if the function
+// Get the AST of the inner function we need to hook, if it was generated
-// was generated by async-trait.
+// by async-trait.
-// When we are given a function generated by async-trait, that function
+// When we are given a function annotated by async-trait, that function
-// is only a "temporary" one that returns a pinned future, and it is
+// is only a placeholder that returns a pinned future containing the
-// that pinned future that needs to be instrumented, otherwise we will
+// user logic, and it is that pinned future that needs to be instrumented.
-// only collect information on the moment the future was "built",
+// Were we to instrument its parent, we would only collect information
-// and not its true span of execution.
+// regarding the allocation of that future, and not its own span of execution.
-// So we inspect the block of the function to find if we can find the
+// So we inspect the block of the function to find if it matches the pattern
-// pattern `async fn foo<...>(...) {...}; Box::pin(foo<...>(...))` and
+// `async fn foo<...>(...) {...}; Box::pin(foo<...>(...))` and we return
-// return the name `foo` if that is the case. Our caller will then be
+// the name `foo` if that is the case. 'gen_body' will then be able
-// able to use that information to instrument the proper function.
+// 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;
@@ -824,7 +852,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);
@@ -854,9 +882,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) {
+                        for fun in inside_funs {
+                            if fun.sig.ident == func_name {
                                 // we must hook this function now
-                            return Some(func_name);
+                                return Some(fun);
+                            }
                         }
                     }
                 }
@@ -866,6 +896,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;
@@ -880,3 +952,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();
+}