don't expand given block twice

Author: hkmatsumoto

tracing-attributes/src/expand.rs

@@ -197,8 +197,8 @@ fn gen_block<B: ToTokens>(
     };
 
     let ret_event = match args.ret_mode {
-        Some(FormatMode::Display) => Some(quote!(tracing::trace!(return = %#block))),
-        Some(FormatMode::Debug) => Some(quote!(tracing::trace!(return = ?#block))),
+        Some(FormatMode::Display) => Some(quote!(tracing::trace!(return = %x))),
+        Some(FormatMode::Debug) => Some(quote!(tracing::trace!(return = ?x))),
         _ => None,
     };
 
@@ -207,14 +207,19 @@ 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, ret_event) {
             (Some(err_event), Some(ret_event)) => quote_spanned!(block.span()=>
                 async move {
                     #ret_event;
                     match async move { #block }.await {
                         #[allow(clippy::unit_arg)]
-                        Ok(x) => Ok(x),
+                        Ok(x) => {
+                            #ret_event;
+                            Ok(x)
+                        },
                         Err(e) => {
                             #err_event;
                             Err(e)
@@ -236,8 +241,9 @@ fn gen_block<B: ToTokens>(
             ),
             (None, Some(ret_event)) => quote_spanned!(block.span()=>
                 async move {
+                    let x = async move { #block }.await;
                     #ret_event;
-                    #block
+                    x
                 }
             ),
             (None, None) => quote_spanned!(block.span()=>
@@ -282,11 +288,13 @@ fn gen_block<B: ToTokens>(
     match (err_event, ret_event) {
         (Some(err_event), Some(ret_event)) => quote_spanned! {block.span()=>
             #span
-            #ret_event;
             #[allow(clippy::redundant_closure_call)]
             match (move || #block)() {
                 #[allow(clippy::unit_arg)]
-                Ok(x) => Ok(x),
+                Ok(x) => {
+                    #ret_event;
+                    Ok(x)
+                },
                 Err(e) => {
                     #err_event;
                     Err(e)
@@ -307,8 +315,9 @@ fn gen_block<B: ToTokens>(
         ),
         (None, Some(ret_event)) => quote_spanned!(block.span()=>
             #span
+            let x = (move || #block)();
             #ret_event;
-            #block
+            x
         ),
         (None, None) => quote_spanned!(block.span() =>
             // Because `quote` produces a stream of tokens _without_ whitespace, the

tracing-attributes/tests/ret.rs

@@ -35,6 +35,38 @@ fn test() {
     handle.assert_finished();
 }
 
+#[instrument(ret)]
+fn ret_mut(a: &mut i32) -> i32 {
+    *a *= 2;
+    tracing::info!(?a);
+    *a
+}
+
+#[test]
+fn test_mut() {
+    let span = span::mock().named("ret_mut");
+    let (collector, handle) = collector::mock()
+        .new_span(span.clone())
+        .enter(span.clone())
+        .event(
+            event::mock()
+                .with_fields(field::mock("a").with_value(&tracing::field::display(2)))
+                .at_level(Level::INFO),
+        )
+        .event(
+            event::mock()
+                .with_fields(field::mock("return").with_value(&tracing::field::display(2)))
+                .at_level(Level::TRACE),
+        )
+        .exit(span.clone())
+        .drop_span(span)
+        .done()
+        .run_with_handle();
+
+    with_default(collector, || ret_mut(&mut 1));
+    handle.assert_finished();
+}
+
 #[instrument(ret)]
 async fn ret_async() -> i32 {
     42
@@ -121,14 +153,14 @@ fn test_ret_and_err() {
     let (collector, handle) = collector::mock()
         .new_span(span.clone())
         .enter(span.clone())
-        .event(event::mock().with_fields(
-            field::mock("return").with_value(&tracing::field::debug(u8::try_from(1234))),
-        ))
         .event(
-            event::mock().with_fields(
-                field::mock("error")
-                    .with_value(&tracing::field::display(u8::try_from(1234).unwrap_err())),
-            ),
+            event::mock()
+                .with_fields(
+                    field::mock("error")
+                        .with_value(&tracing::field::display(u8::try_from(1234).unwrap_err()))
+                        .only(),
+                )
+                .at_level(Level::ERROR),
         )
         .exit(span.clone())
         .drop_span(span)
@@ -138,3 +170,32 @@ fn test_ret_and_err() {
     with_default(collector, || ret_and_err().ok());
     handle.assert_finished();
 }
+
+#[instrument(err, ret(Debug))]
+fn ret_and_ok() -> Result<u8, TryFromIntError> {
+    u8::try_from(123)
+}
+
+#[test]
+fn test_ret_and_ok() {
+    let span = span::mock().named("ret_and_ok");
+    let (collector, handle) = collector::mock()
+        .new_span(span.clone())
+        .enter(span.clone())
+        .event(
+            event::mock()
+                .with_fields(
+                    field::mock("return")
+                        .with_value(&tracing::field::display(u8::try_from(123).unwrap()))
+                        .only(),
+                )
+                .at_level(Level::TRACE),
+        )
+        .exit(span.clone())
+        .drop_span(span)
+        .done()
+        .run_with_handle();
+
+    with_default(collector, || ret_and_ok().ok());
+    handle.assert_finished();
+}