Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 5 commits into from
Aug 7, 2020
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1be63e3
attributes - #[instrument] - support use of the 'self' variable in as…
nightmared Jul 30, 2020
3ce8cc3
fix style issues
nightmared Aug 3, 2020
65a141d
some more style improvements
nightmared Aug 4, 2020
4e2a102
attributes: document the shortcomings of #[instrument] with async-trait
nightmared Aug 5, 2020
e3c9259
Apply docs suggestions from code review
hawkw Aug 7, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion tracing-attributes/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"


Expand Down
222 changes: 159 additions & 63 deletions tracing-attributes/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
///
/// ```
Expand All @@ -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::<Self>()))]
/// 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
Expand All @@ -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<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,
instrumented_function_name,
Some(internal_fun),
))
.unwrap();
break;
}
}
Expand All @@ -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<String>,
mut args: InstrumentArgs,
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
Expand Down Expand Up @@ -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 = (|| {
Expand All @@ -306,32 +330,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()=>
Expand Down Expand Up @@ -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

))
Expand Down Expand Up @@ -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<String> {
fn get_async_trait_function(block: &Block, block_is_async: bool) -> Option<&ItemFn> {
hawkw marked this conversation as resolved.
Show resolved Hide resolved
// are we in an async context? If yes, this isn't a async_trait-like pattern
if block_is_async {
return None;
Expand All @@ -816,7 +844,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);
Expand Down Expand Up @@ -846,9 +874,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);
}
}
}
}
Expand All @@ -858,6 +888,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;
Expand All @@ -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<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();
}
}
}
}
}
Loading