Skip to content

Commit

Permalink
Unify handling of return values in #[instrument]
Browse files Browse the repository at this point in the history
  • Loading branch information
@Xiretza
Xiretza committed May 23, 2024
1 parent 8835356 commit 7f47ae7
Show file tree
Hide file tree
Showing 4 changed files with 268 additions and 108 deletions.
98 changes: 91 additions & 7 deletions tracing-attributes/src/attr.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use syn::{punctuated::Punctuated, Expr, Ident, LitInt, LitStr, Path, Token};
use proc_macro2::TokenStream;
use quote::{quote, quote_spanned, ToTokens};
use syn::ext::IdentExt as _;
use syn::parse::{Parse, ParseStream};
use syn::parse::{Parse, ParseBuffer, ParseStream};

/// Arguments to `#[instrument(err(...))]` and `#[instrument(ret(...))]` which describe how the
/// return value event should be emitted.
Expand All @@ -14,6 +14,15 @@ pub(crate) struct EventArgs {
pub(crate) mode: FormatMode,
}

#[derive(Clone, Debug)]
pub(crate) enum RetArgs {
AnyType(EventArgs),
Result {
ok: Option<EventArgs>,
err: Option<EventArgs>,
},
}

#[derive(Clone, Default, Debug)]
pub(crate) struct InstrumentArgs {
level: Option<Level>,
Expand All @@ -23,8 +32,7 @@ pub(crate) struct InstrumentArgs {
pub(crate) follows_from: Option<Expr>,
pub(crate) skips: HashSet<Ident>,
pub(crate) fields: Option<Fields>,
pub(crate) err_args: Option<EventArgs>,
pub(crate) ret_args: Option<EventArgs>,
pub(crate) ret_args: Option<RetArgs>,
/// Errors describing any unrecognized parse inputs that we skipped.
parse_warnings: Vec<syn::Error>,
}
Expand Down Expand Up @@ -74,6 +82,7 @@ impl InstrumentArgs {
impl Parse for InstrumentArgs {
fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
let mut args = Self::default();
let mut err_args = None;
while !input.is_empty() {
let lookahead = input.lookahead1();
if lookahead.peek(kw::name) {
Expand Down Expand Up @@ -125,12 +134,12 @@ impl Parse for InstrumentArgs {
}
args.fields = Some(input.parse()?);
} else if lookahead.peek(kw::err) {
let error_fork = input.fork();
let _ = input.parse::<kw::err>();
let err_args = EventArgs::parse(input)?;
args.err_args = Some(err_args);
err_args = Some((error_fork, EventArgs::parse(input)?));
} else if lookahead.peek(kw::ret) {
let _ = input.parse::<kw::ret>()?;
let ret_args = EventArgs::parse(input)?;
let ret_args = RetArgs::parse(input)?;
args.ret_args = Some(ret_args);
} else if lookahead.peek(Token![,]) {
let _ = input.parse::<Token![,]>()?;
Expand All @@ -145,6 +154,22 @@ impl Parse for InstrumentArgs {
let _ = input.parse::<proc_macro2::TokenTree>();
}
}

if let Some((error_fork, err)) = err_args {
// emulate old behaviour: err() sets Err formatting, ret() sets Ok formatting
let ok = match args.ret_args.take() {
None => None,
Some(RetArgs::AnyType(ok)) => Some(ok),
Some(RetArgs::Result { .. }) => {
return Err(
error_fork.error("`err()` and `ret(err())`/`ret(ok())` are incompatible")
)
}
};

args.ret_args = Some(RetArgs::Result { ok, err: Some(err) });
}

Ok(args)
}
}
Expand Down Expand Up @@ -198,6 +223,64 @@ impl Parse for EventArgs {
}
}

impl Parse for RetArgs {
fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
if !input.peek(syn::token::Paren) {
return Ok(Self::AnyType(EventArgs::default()));
}

{
let ahead = input.fork();

let content;
let _ = syn::parenthesized!(content in ahead);

if !content.peek(kw::ok) && !content.peek(kw::err) {
return Ok(Self::AnyType(input.parse()?));
}
}

let content;
let _ = syn::parenthesized!(content in input);

fn parse_nested<KW: ToTokens + Parse>(
content: &ParseBuffer<'_>,
args: &mut Option<EventArgs>,
) -> syn::Result<()> {
let kw = content.parse::<KW>()?;
if args.is_some() {
return Err(content.error(format!(
"expected only a single `{}` argument",
kw.into_token_stream()
)));
}
*args = Some(content.parse()?);
if content.peek(Token![,]) {
let _ = content.parse::<Token![,]>()?;
}
Ok(())
}

let mut ok_args = None;
let mut err_args = None;
while !content.is_empty() {
let lookahead = content.lookahead1();
if lookahead.peek(kw::ok) {
parse_nested::<kw::ok>(&content, &mut ok_args)?;
} else if lookahead.peek(kw::err) {
parse_nested::<kw::err>(&content, &mut err_args)?;
} else {
return Err(lookahead.error());
}
}

Ok(Self::Result {
ok: ok_args,
err: err_args,
})
}
}

struct StrArg<T> {
value: LitStr,
_p: std::marker::PhantomData<T>,
Expand Down Expand Up @@ -307,7 +390,7 @@ impl Parse for Field {
kind = FieldKind::Debug;
};
let name = Punctuated::parse_separated_nonempty_with(input, Ident::parse_any)?;
let value = if input.peek(Token![=]) {
let value = if kind == FieldKind::Value && input.peek(Token![=]) {
input.parse::<Token![=]>()?;
if input.peek(Token![%]) {
input.parse::<Token![%]>()?;
Expand Down Expand Up @@ -433,6 +516,7 @@ mod kw {
syn::custom_keyword!(parent);
syn::custom_keyword!(follows_from);
syn::custom_keyword!(name);
syn::custom_keyword!(ok);
syn::custom_keyword!(err);
syn::custom_keyword!(ret);
}
171 changes: 80 additions & 91 deletions tracing-attributes/src/expand.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ use syn::{
Path, ReturnType, Signature, Stmt, Token, Type, TypePath,
};

use crate::attr::{EventArgs, RetArgs};
use crate::{
attr::{Field, Fields, FormatMode, InstrumentArgs, Level},
MaybeItemFn, MaybeItemFnRef,
Expand Down Expand Up @@ -236,34 +237,28 @@ fn gen_block<B: ToTokens>(

let target = args.target();

let err_event = match args.err_args {
Some(event_args) => {
let level_tokens = event_args.level(Level::Error);
match event_args.mode {
FormatMode::Default | FormatMode::Display => Some(quote!(
tracing::event!(target: #target, #level_tokens, error = %e)
)),
FormatMode::Debug => Some(quote!(
tracing::event!(target: #target, #level_tokens, error = ?e)
)),
}
let make_err_event = |event_args: &EventArgs| {
let level_tokens = event_args.level(Level::Error);
match event_args.mode {
FormatMode::Default | FormatMode::Display => quote!(
tracing::event!(target: #target, #level_tokens, error = %e);
),
FormatMode::Debug => quote!(
tracing::event!(target: #target, #level_tokens, error = ?e);
),
}
_ => None,
};

let ret_event = match args.ret_args {
Some(event_args) => {
let level_tokens = event_args.level(args_level);
match event_args.mode {
FormatMode::Display => Some(quote!(
tracing::event!(target: #target, #level_tokens, return = %x)
)),
FormatMode::Default | FormatMode::Debug => Some(quote!(
tracing::event!(target: #target, #level_tokens, return = ?x)
)),
}
let make_ret_event = |event_args: &EventArgs| {
let level_tokens = event_args.level(args_level);
match event_args.mode {
FormatMode::Display => quote!(
tracing::event!(target: #target, #level_tokens, return = %x);
),
FormatMode::Default | FormatMode::Debug => quote!(
tracing::event!(target: #target, #level_tokens, return = ?x);
),
}
_ => None,
};

// Generate the instrumented function body.
Expand All @@ -274,42 +269,39 @@ fn gen_block<B: ToTokens>(
// If `ret` is in args, instrument any resulting `Ok`s when the function
// returns `Result`s, otherwise instrument any resulting values.
if async_context {
let mk_fut = match (err_event, ret_event) {
(Some(err_event), Some(ret_event)) => quote_spanned!(block.span()=>
async move {
match async move #block.await {
#[allow(clippy::unit_arg)]
Ok(x) => {
#ret_event;
Ok(x)
},
Err(e) => {
#err_event;
Err(e)
let mk_fut = match args.ret_args {
Some(RetArgs::Result { ok, err }) => {
let ok_event = ok.as_ref().map(make_ret_event).unwrap_or(quote! {});
let err_event = err.as_ref().map(make_err_event).unwrap_or(quote! {});

quote_spanned!(block.span()=>
async move {
match async move #block.await {
#[allow(clippy::unit_arg)]
Ok(x) => {
#ok_event
Ok(x)
},
Err(e) => {
#err_event
Err(e)
}
}
}
}
),
(Some(err_event), None) => quote_spanned!(block.span()=>
async move {
match async move #block.await {
#[allow(clippy::unit_arg)]
Ok(x) => Ok(x),
Err(e) => {
#err_event;
Err(e)
}
)
}
Some(RetArgs::AnyType(args)) => {
let ret_event = make_ret_event(&args);

quote_spanned!(block.span()=>
async move {
let x = async move #block.await;
#ret_event
x
}
}
),
(None, Some(ret_event)) => quote_spanned!(block.span()=>
async move {
let x = async move #block.await;
#ret_event;
x
}
),
(None, None) => quote_spanned!(block.span()=>
)
}
None => quote_spanned!(block.span()=>
async move #block
),
};
Expand Down Expand Up @@ -350,42 +342,39 @@ fn gen_block<B: ToTokens>(
}
);

match (err_event, ret_event) {
(Some(err_event), Some(ret_event)) => quote_spanned! {block.span()=>
#span
#[allow(clippy::redundant_closure_call)]
match (move || #block)() {
#[allow(clippy::unit_arg)]
Ok(x) => {
#ret_event;
Ok(x)
},
Err(e) => {
#err_event;
Err(e)
}
}
},
(Some(err_event), None) => quote_spanned!(block.span()=>
#span
#[allow(clippy::redundant_closure_call)]
match (move || #block)() {
#[allow(clippy::unit_arg)]
Ok(x) => Ok(x),
Err(e) => {
#err_event;
Err(e)
match args.ret_args {
Some(RetArgs::Result { ok, err }) => {
let ret_event = ok.as_ref().map(make_ret_event).unwrap_or(quote! {});
let err_event = err.as_ref().map(make_err_event).unwrap_or(quote! {});

quote_spanned! {block.span()=>
#span
#[allow(clippy::redundant_closure_call)]
match (move || #block)() {
#[allow(clippy::unit_arg)]
Ok(x) => {
#ret_event;
Ok(x)
},
Err(e) => {
#err_event;
Err(e)
}
}
}
),
(None, Some(ret_event)) => quote_spanned!(block.span()=>
#span
#[allow(clippy::redundant_closure_call)]
let x = (move || #block)();
#ret_event;
x
),
(None, None) => quote_spanned!(block.span() =>
}
Some(RetArgs::AnyType(args)) => {
let ret_event = make_ret_event(&args);

quote_spanned!(block.span()=>
#span
#[allow(clippy::redundant_closure_call)]
let x = (move || #block)();
#ret_event;
x
)
}
None => quote_spanned!(block.span() =>
// Because `quote` produces a stream of tokens _without_ whitespace, the
// `if` and the block will appear directly next to each other. This
// generates a clippy lint about suspicious `if/else` formatting.
Expand Down

0 comments on commit 7f47ae7

Please sign in to comment.