attributes: implement #[instrument(ret)] (#1716)

* attributes: implement `#[instrument(ret)]`

## Motivation

Currently, users have to explicitly call logging functions at the end of
functions if they wanted to record values being returned. For example,
```rust
fn increment(a: i32) -> i32 {
    let succ = a + 1;
    tracing::debug!(?succ);
    succ
}
```
The code above will be significantly simpler if we provide the
functionality to record returned values.

## Solution

This PR adds the exact feature we're talking here, which enables users
to write the code below instead.
```rust
#[instrument(ret)]
fn increment(a: i32) -> i32 {
    a + 1
}
```

Co-authored-by: Eliza Weisman <[email protected]>

Author: hkmatsumoto

tracing-attributes/src/attr.rs

@@ -14,7 +14,8 @@ pub(crate) struct InstrumentArgs {
     pub(crate) skips: HashSet<Ident>,
     pub(crate) skip_all: bool,
     pub(crate) fields: Option<Fields>,
-    pub(crate) err_mode: Option<ErrorMode>,
+    pub(crate) err_mode: Option<FormatMode>,
+    pub(crate) ret_mode: Option<FormatMode>,
     /// Errors describing any unrecognized parse inputs that we skipped.
     parse_warnings: Vec<syn::Error>,
 }
@@ -152,8 +153,12 @@ impl Parse for InstrumentArgs {
                 args.fields = Some(input.parse()?);
             } else if lookahead.peek(kw::err) {
                 let _ = input.parse::<kw::err>();
-                let mode = ErrorMode::parse(input)?;
+                let mode = FormatMode::parse(input)?;
                 args.err_mode = Some(mode);
+            } else if lookahead.peek(kw::ret) {
+                let _ = input.parse::<kw::ret>()?;
+                let mode = FormatMode::parse(input)?;
+                args.ret_mode = Some(mode);
             } else if lookahead.peek(Token![,]) {
                 let _ = input.parse::<Token![,]>()?;
             } else {
@@ -212,29 +217,30 @@ impl Parse for Skips {
 }
 
 #[derive(Clone, Debug, Hash, PartialEq, Eq)]
-pub(crate) enum ErrorMode {
+pub(crate) enum FormatMode {
+    Default,
     Display,
     Debug,
 }
 
-impl Default for ErrorMode {
+impl Default for FormatMode {
     fn default() -> Self {
-        ErrorMode::Display
+        FormatMode::Default
     }
 }
 
-impl Parse for ErrorMode {
+impl Parse for FormatMode {
     fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
         if !input.peek(syn::token::Paren) {
-            return Ok(ErrorMode::default());
+            return Ok(FormatMode::default());
         }
         let content;
         let _ = syn::parenthesized!(content in input);
         let maybe_mode: Option<Ident> = content.parse()?;
-        maybe_mode.map_or(Ok(ErrorMode::default()), |ident| {
+        maybe_mode.map_or(Ok(FormatMode::default()), |ident| {
             match ident.to_string().as_str() {
-                "Debug" => Ok(ErrorMode::Debug),
-                "Display" => Ok(ErrorMode::Display),
+                "Debug" => Ok(FormatMode::Debug),
+                "Display" => Ok(FormatMode::Display),
                 _ => Err(syn::Error::new(
                     ident.span(),
                     "unknown error mode, must be Debug or Display",
@@ -370,4 +376,5 @@ mod kw {
     syn::custom_keyword!(target);
     syn::custom_keyword!(name);
     syn::custom_keyword!(err);
+    syn::custom_keyword!(ret);
 }

tracing-attributes/src/expand.rs

@@ -9,7 +9,7 @@ use syn::{
 };
 
 use crate::{
-    attr::{ErrorMode, Field, Fields, InstrumentArgs},
+    attr::{Field, Fields, FormatMode, InstrumentArgs},
     MaybeItemFnRef,
 };
 
@@ -191,8 +191,18 @@ fn gen_block<B: ToTokens>(
     })();
 
     let err_event = match args.err_mode {
-        Some(ErrorMode::Display) => Some(quote!(tracing::error!(error = %e))),
-        Some(ErrorMode::Debug) => Some(quote!(tracing::error!(error = ?e))),
+        Some(FormatMode::Default) | Some(FormatMode::Display) => {
+            Some(quote!(tracing::error!(error = %e)))
+        }
+        Some(FormatMode::Debug) => Some(quote!(tracing::error!(error = ?e))),
+        _ => None,
+    };
+
+    let ret_event = match args.ret_mode {
+        Some(FormatMode::Display) => Some(quote!(tracing::event!(#level, return = %x))),
+        Some(FormatMode::Default) | Some(FormatMode::Debug) => {
+            Some(quote!(tracing::event!(#level, return = ?x)))
+        }
         _ => None,
     };
 
@@ -201,9 +211,26 @@ fn gen_block<B: ToTokens>(
     // which is `instrument`ed using `tracing-futures`. Otherwise, this will
     // enter the span and then perform the rest of the body.
     // If `err` is in args, instrument any resulting `Err`s.
+    // 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 {
-            Some(err_event) => quote_spanned!(block.span()=>
+        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)
+                        }
+                    }
+                }
+            ),
+            (Some(err_event), None) => quote_spanned!(block.span()=>
                 async move {
                     match async move { #block }.await {
                         #[allow(clippy::unit_arg)]
@@ -215,7 +242,14 @@ fn gen_block<B: ToTokens>(
                     }
                 }
             ),
-            None => quote_spanned!(block.span()=>
+            (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()=>
                 async move { #block }
             ),
         };
@@ -254,8 +288,23 @@ fn gen_block<B: ToTokens>(
         }
     );
 
-    if let Some(err_event) = err_event {
-        return quote_spanned!(block.span()=>
+    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)() {
@@ -266,22 +315,28 @@ fn gen_block<B: ToTokens>(
                     Err(e)
                 }
             }
-        );
-    }
-
-    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.
-        // Therefore, suppress the lint inside the generated code...
-        #[allow(clippy::suspicious_else_formatting)]
-        {
+        ),
+        (None, Some(ret_event)) => quote_spanned!(block.span()=>
             #span
-            // ...but turn the lint back on inside the function body.
-            #[warn(clippy::suspicious_else_formatting)]
-            #block
-        }
-    )
+            #[allow(clippy::redundant_closure_call)]
+            let x = (move || #block)();
+            #ret_event;
+            x
+        ),
+        (None, 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.
+            // Therefore, suppress the lint inside the generated code...
+            #[allow(clippy::suspicious_else_formatting)]
+            {
+                #span
+                // ...but turn the lint back on inside the function body.
+                #[warn(clippy::suspicious_else_formatting)]
+                #block
+            }
+        ),
+    }
 }
 
 /// Indicates whether a field should be recorded as `Value` or `Debug`.

tracing-attributes/src/lib.rs

@@ -368,6 +368,35 @@ mod expand;
 /// }
 /// ```
 ///
+/// Adding the `ret` argument to `#[instrument]` will emit an event with the function's
+/// return value when the function returns:
+///
+/// ```
+/// # use tracing_attributes::instrument;
+/// #[instrument(ret)]
+/// fn my_function() -> i32 {
+///     42
+/// }
+/// ```
+/// The return value event will have the same level as the span generated by `#[instrument]`.
+/// By default, this will be `TRACE`, but if the level is overridden, the event will be at the same
+/// level.
+///
+/// **Note**:  if the function returns a `Result<T, E>`, `ret` will record returned values if and
+/// only if the function returns [`Result::Ok`].
+///
+/// By default, returned values will be recorded using their [`std::fmt::Debug`] implementations.
+/// If a returned value implements [`std::fmt::Display`], it can be recorded using its `Display`
+/// implementation instead, by writing `ret(Display)`:
+///
+/// ```
+/// # use tracing_attributes::instrument;
+/// #[instrument(ret(Display))]
+/// fn my_function() -> i32 {
+///     42
+/// }
+/// ```
+///
 /// If the function returns a `Result<T, E>` and `E` implements `std::fmt::Display`, you can add
 /// `err` or `err(Display)` to emit error events when the function returns `Err`:
 ///
@@ -391,6 +420,17 @@ mod expand;
 /// }
 /// ```
 ///
+/// The `ret` and `err` arguments can be combined in order to record an event if a
+/// function returns [`Result::Ok`] or [`Result::Err`]:
+///
+/// ```
+/// # use tracing_attributes::instrument;
+/// #[instrument(err, ret)]
+/// fn my_function(arg: usize) -> Result<(), std::io::Error> {
+///     Ok(())
+/// }
+/// ```
+///
 /// `async fn`s may also be instrumented:
 ///
 /// ```