From fef53847283071e542a59419a4f9e06b8053fd73 Mon Sep 17 00:00:00 2001 From: cynecx Date: Tue, 12 Oct 2021 04:58:36 +0200 Subject: [PATCH] attributes: speculative expand when parsing fails --- tracing-appender/src/non_blocking.rs | 4 +- tracing-attributes/src/lib.rs | 141 +++++++++++++++++++++++---- 2 files changed, 122 insertions(+), 23 deletions(-) diff --git a/tracing-appender/src/non_blocking.rs b/tracing-appender/src/non_blocking.rs index 2448f178c3..009e4f5b61 100644 --- a/tracing-appender/src/non_blocking.rs +++ b/tracing-appender/src/non_blocking.rs @@ -103,7 +103,7 @@ pub const DEFAULT_BUFFERED_LINES_LIMIT: usize = 128_000; #[must_use] #[derive(Debug)] pub struct WorkerGuard { - guard: Option>, + _guard: Option>, sender: Sender, shutdown: Sender<()>, } @@ -250,7 +250,7 @@ impl<'a> MakeWriter<'a> for NonBlocking { impl WorkerGuard { fn new(handle: JoinHandle<()>, sender: Sender, shutdown: Sender<()>) -> Self { WorkerGuard { - guard: Some(handle), + _guard: Some(handle), sender, shutdown, } diff --git a/tracing-attributes/src/lib.rs b/tracing-attributes/src/lib.rs index ea3817a029..42be26c1ff 100644 --- a/tracing-attributes/src/lib.rs +++ b/tracing-attributes/src/lib.rs @@ -88,9 +88,9 @@ use quote::{quote, quote_spanned, ToTokens}; use syn::ext::IdentExt as _; use syn::parse::{Parse, ParseStream}; use syn::{ - punctuated::Punctuated, spanned::Spanned, Block, Expr, ExprAsync, ExprCall, FieldPat, FnArg, - Ident, Item, ItemFn, LitInt, LitStr, Pat, PatIdent, PatReference, PatStruct, PatTuple, - PatTupleStruct, PatType, Path, Signature, Stmt, Token, TypePath, + punctuated::Punctuated, spanned::Spanned, Attribute, Block, Expr, ExprAsync, ExprCall, + FieldPat, FnArg, Ident, Item, ItemFn, LitInt, LitStr, Pat, PatIdent, PatReference, PatStruct, + PatTuple, PatTupleStruct, PatType, Path, Signature, Stmt, Token, TypePath, Visibility, }; /// Instruments a function to create and enter a `tracing` [span] every time /// the function is called. @@ -274,14 +274,45 @@ pub fn instrument( args: proc_macro::TokenStream, item: proc_macro::TokenStream, ) -> proc_macro::TokenStream { - let input = syn::parse_macro_input!(item as ItemFn); let args = syn::parse_macro_input!(args as InstrumentArgs); + instrument_precise(args.clone(), item.clone()).unwrap_or_else(|_err| { + // FIXME: Ideally, we'd like to emit a warning but we can't, + // so we ignore the parsing error for now. + instrument_speculative(args, item) + }) +} + +/// Instrument the function, without parsing the function body (instead using the raw tokens). +fn instrument_speculative( + args: InstrumentArgs, + item: proc_macro::TokenStream, +) -> proc_macro::TokenStream { + let input = syn::parse_macro_input!(item as MaybeItemFn); + let instrumented_function_name = input.sig.ident.to_string(); + gen_function( + &input.as_ref(), + args, + instrumented_function_name.as_str(), + None, + ) + .into() +} + +/// Instrument the function, by fully parsing the function body, +/// which allows us to rewrite some statements related to async_trait-like patterns. +fn instrument_precise( + args: InstrumentArgs, + item: proc_macro::TokenStream, +) -> Result { + let input = syn::parse::(item)?; let instrumented_function_name = input.sig.ident.to_string(); // check for async_trait-like patterns in the block, and instrument // the future instead of the wrapper - if let Some(internal_fun) = get_async_trait_info(&input.block, input.sig.asyncness.is_some()) { + let res = if let Some(internal_fun) = + get_async_trait_info(&input.block, input.sig.asyncness.is_some()) + { // let's rewrite some statements! let mut out_stmts: Vec = input .block @@ -301,7 +332,7 @@ pub fn instrument( out_stmts[iter] = match internal_fun.kind { // async-trait <= 0.1.43 AsyncTraitKind::Function(fun) => gen_function( - fun, + &MaybeItemFnRef::from(fun), args, instrumented_function_name.as_str(), internal_fun.self_type.as_ref(), @@ -335,13 +366,21 @@ pub fn instrument( ) .into() } else { - gen_function(&input, args, instrumented_function_name.as_str(), None).into() - } + gen_function( + &(&input).into(), + args, + instrumented_function_name.as_str(), + None, + ) + .into() + }; + + Ok(res) } /// Given an existing function, generate an instrumented version of that function -fn gen_function( - input: &ItemFn, +fn gen_function<'a, B: ToTokens + 'a>( + input: &MaybeItemFnRef<'a, B>, args: InstrumentArgs, instrumented_function_name: &str, self_type: Option<&syn::TypePath>, @@ -349,12 +388,11 @@ fn gen_function( // these are needed ahead of time, as ItemFn contains the function body _and_ // isn't representable inside a quote!/quote_spanned! macro // (Syn's ToTokens isn't implemented for ItemFn) - let ItemFn { + let MaybeItemFnRef { attrs, vis, - block, sig, - .. + block, } = input; let Signature { @@ -397,8 +435,8 @@ fn gen_function( } /// Instrument a block -fn gen_block( - block: &Block, +fn gen_block( + block: &B, params: &Punctuated, async_context: bool, mut args: InstrumentArgs, @@ -611,7 +649,68 @@ fn gen_block( ) } -#[derive(Default, Debug)] +/// This is a more flexible/imprecise `ItemFn` type, +/// which's block may be anything that implements `ToTokens`. +#[derive(Debug, Clone)] +struct MaybeItemFn { + attrs: Vec, + vis: Visibility, + sig: Signature, + block: B, +} + +impl MaybeItemFn { + fn new(attrs: Vec, vis: Visibility, sig: Signature, block: B) -> Self { + Self { + attrs, + vis, + sig, + block, + } + } + + fn as_ref(&self) -> MaybeItemFnRef<'_, B> { + MaybeItemFnRef { + attrs: &self.attrs, + vis: &self.vis, + sig: &self.sig, + block: &self.block, + } + } +} + +/// This parses a TokenStream as ItemFn/MaybeItemFn but skips the body. +impl Parse for MaybeItemFn { + fn parse(input: ParseStream<'_>) -> syn::Result { + let attrs = input.call(syn::Attribute::parse_outer)?; + let vis: Visibility = input.parse()?; + let sig: Signature = input.parse()?; + let block: TokenStream = input.parse()?; + Ok(Self::new(attrs, vis, sig, block)) + } +} + +/// A generic reference type for `MaybeItemFn`. +#[derive(Debug, Clone)] +struct MaybeItemFnRef<'a, B: ToTokens> { + attrs: &'a Vec, + vis: &'a Visibility, + sig: &'a Signature, + block: &'a B, +} + +impl<'a> From<&'a ItemFn> for MaybeItemFnRef<'a, Box> { + fn from(val: &'a ItemFn) -> Self { + MaybeItemFnRef { + attrs: &val.attrs, + vis: &val.vis, + sig: &val.sig, + block: &val.block, + } + } +} + +#[derive(Clone, Default, Debug)] struct InstrumentArgs { level: Option, name: Option, @@ -803,7 +902,7 @@ impl Parse for Skips { } } -#[derive(Debug, Hash, PartialEq, Eq)] +#[derive(Clone, Debug, Hash, PartialEq, Eq)] enum ErrorMode { Display, Debug, @@ -836,17 +935,17 @@ impl Parse for ErrorMode { } } -#[derive(Debug)] +#[derive(Clone, Debug)] struct Fields(Punctuated); -#[derive(Debug)] +#[derive(Clone, Debug)] struct Field { name: Punctuated, value: Option, kind: FieldKind, } -#[derive(Debug, Eq, PartialEq)] +#[derive(Clone, Debug, Eq, PartialEq)] enum FieldKind { Debug, Display, @@ -930,7 +1029,7 @@ impl ToTokens for FieldKind { } } -#[derive(Debug)] +#[derive(Clone, Debug)] enum Level { Str(LitStr), Int(LitInt),