itsscb/tracing
1
0
Fork 0
mirror of https://github.com/tokio-rs/tracing.git synced 2025-10-02 15:24:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
##  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).
This commit is contained in:
Simon THOBY 2020-08-07 20:46:26 +02:00 committed by GitHub
parent e486bdf81c
commit 9ae4676054
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 275 additions and 64 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tracing-attributes/Cargo.toml
View File

@ -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"

212
tracing-attributes/src/lib.rs
View File

@ -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,
/// 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:
///
/// ```
@ -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))]
/// 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
@ -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) =
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())))
if fun.sig.ident == internal_fun.name {
*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,
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
@ -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
// 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 = (|| {
@ -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() {
quote! { #custom_fields }
} else {
quote! {, #custom_fields }
// 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;
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
// 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> {
// 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();
}
}
}
}
}

115
tracing-attributes/tests/async_fn.rs
View File

@ -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();
}