#[instrument]: add ret(err, ok)

tracing-attributes/src/attr.rs

@@ -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.
@@ -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>,
@@ -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>,
 }
@@ -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) {
@@ -127,12 +136,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![,]>()?;
@@ -147,6 +156,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)
  }
 }
@@ -224,6 +249,64 @@ impl Parse for LitStrOrIdent {
  }
 }
 
+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: LitStrOrIdent,
  _p: std::marker::PhantomData<T>,
@@ -333,7 +416,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![%]>()?;
@@ -459,6 +542,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);
 }

tracing-attributes/src/expand.rs

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