Support async methods in scoped_trace in googletest Rust.

This change adds support for async methods in `scoped_trace` by using a runtime-agnostic `InstrumentedFuture` that manages the thread-local trace stack across yield points.

PiperOrigin-RevId: 914068550

Author: th0br0

googletest/src/internal/scoped_trace.rs

@@ -14,6 +14,7 @@
 
 use std::cell::RefCell;
 use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
 
 static NEXT_TRACE_ID: AtomicUsize = AtomicUsize::new(0);
 
@@ -26,19 +27,22 @@ pub struct TraceInfo {
     pub message: String,
 }
 
+struct TraceNode {
+    info: TraceInfo,
+    parent: Option<Arc<TraceNode>>,
+}
+
 thread_local! {
-    static TRACE_STACK: RefCell<Vec<TraceInfo>> = const { RefCell::new(Vec::new()) };
+    static TRACE_STACK: RefCell<Option<Arc<TraceNode>>> = const { RefCell::new(None) };
 }
 
 /// RAII guard to manage the push and pop of trace information.
 ///
-/// This struct is `!Send` and `!Sync` to prevent it from being held across
-/// `.await` points in async tests, which would cause incorrect trace tracking
-/// if the task moves between threads.
+/// This struct is `Send` to allow it to be used in async tests, but it should
+/// not be held across `.await` points directly unless the future is instrumented.
 #[doc(hidden)]
 pub struct ScopedTraceGuard {
     id: usize,
-    _phantom: std::marker::PhantomData<*mut ()>,
 }
 
 impl ScopedTraceGuard {
@@ -47,32 +51,119 @@ impl ScopedTraceGuard {
     pub fn new(message: String) -> Self {
         let caller = std::panic::Location::caller();
         let id = NEXT_TRACE_ID.fetch_add(1, Ordering::Relaxed);
+        let info = TraceInfo { id, file: caller.file(), line: caller.line(), message };
+
         TRACE_STACK.with(|stack| {
-            // Use try_borrow_mut to avoid double panic if called during unwinding.
             if let Ok(mut s) = stack.try_borrow_mut() {
-                s.push(TraceInfo { id, file: caller.file(), line: caller.line(), message });
+                let prev = s.clone();
+                *s = Some(Arc::new(TraceNode { info, parent: prev }));
             }
         });
-        Self { id, _phantom: std::marker::PhantomData }
+        Self { id }
     }
 }
 
 impl Drop for ScopedTraceGuard {
     fn drop(&mut self) {
         TRACE_STACK.with(|stack| {
-            // Use try_borrow_mut to avoid double panic if called during unwinding.
             if let Ok(mut s) = stack.try_borrow_mut() {
-                if let Some(pos) = s.iter().rposition(|t| t.id == self.id) {
-                    s.remove(pos);
-                }
+                *s = remove_from_list(s.clone(), self.id);
             }
         });
     }
 }
 
+/// Removes a node with the given `id` from the trace stack list.
+///
+/// Because the trace stack is structured as a persistent, reverse-linked list
+/// pointing to parent nodes, removing a node from the middle involves rebuilding
+/// the chain from the removed node up to the current head.
+fn remove_from_list(head: Option<Arc<TraceNode>>, id: usize) -> Option<Arc<TraceNode>> {
+    let head = head?;
+    if head.info.id == id {
+        return head.parent.clone();
+    }
+
+    let mut current = head.parent.clone();
+    let mut nodes_to_recreate = vec![head.info.clone()];
+
+    while let Some(node) = current {
+        if node.info.id == id {
+            let mut new_head = node.parent.clone();
+            for info in nodes_to_recreate.into_iter().rev() {
+                new_head = Some(Arc::new(TraceNode { info, parent: new_head }));
+            }
+            return new_head;
+        }
+        nodes_to_recreate.push(node.info.clone());
+        current = node.parent.clone();
+    }
+    Some(head)
+}
+
 /// Retrieves a clone of the current thread's trace stack.
 pub fn get_scoped_traces() -> Vec<TraceInfo> {
-    TRACE_STACK.with(|stack| stack.try_borrow().map(|s| s.clone()).unwrap_or_default())
+    let mut traces = Vec::new();
+    let mut current = TRACE_STACK.with(|stack| stack.try_borrow().ok().and_then(|s| s.clone()));
+    while let Some(node) = current {
+        traces.push(node.info.clone());
+        current = node.parent.clone();
+    }
+    traces.reverse();
+    traces
+}
+
+/// A future that instruments another future with a set of traces.
+pub struct InstrumentedFuture<F> {
+    inner: F,
+    traces: Option<Arc<TraceNode>>,
+}
+
+impl<F> InstrumentedFuture<F> {
+    pub fn new(inner: F) -> Self {
+        Self { inner, traces: None }
+    }
+}
+
+impl<F: std::future::Future> std::future::Future for InstrumentedFuture<F> {
+    type Output = F::Output;
+
+    fn poll(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        // SAFETY: `InstrumentedFuture` provides structural pinning. If `self` is
+        // pinned, it is safe to pin the `inner` future.
+        let this = unsafe { self.get_unchecked_mut() };
+
+        struct TraceSwapGuard<'a>(&'a mut Option<Arc<TraceNode>>);
+        impl<'a> TraceSwapGuard<'a> {
+            fn new(traces: &'a mut Option<Arc<TraceNode>>) -> Self {
+                TRACE_STACK.with(|stack| {
+                    if let Ok(mut s) = stack.try_borrow_mut() {
+                        std::mem::swap(&mut *s, traces);
+                    }
+                });
+                Self(traces)
+            }
+        }
+        impl<'a> Drop for TraceSwapGuard<'a> {
+            fn drop(&mut self) {
+                TRACE_STACK.with(|stack| {
+                    if let Ok(mut s) = stack.try_borrow_mut() {
+                        std::mem::swap(&mut *s, self.0);
+                    }
+                });
+            }
+        }
+
+        // Swap traces into thread-local, ensuring they are swapped back on return/panic.
+        let _guard = TraceSwapGuard::new(&mut this.traces);
+
+        // Poll inner future
+        // SAFETY: As explained above, `this.inner` is properly pinned.
+        unsafe { std::pin::Pin::new_unchecked(&mut this.inner) }.poll(cx)
+    }
 }
 
 // Test-only state and helpers, hidden from production API.

googletest/src/lib.rs

@@ -81,6 +81,7 @@ macro_rules! scoped_trace {
 pub use googletest_macro::gtest;
 pub use googletest_macro::test;
 
+pub use internal::scoped_trace::InstrumentedFuture;
 use internal::test_outcome::{TestAssertionFailure, TestOutcome};
 
 /// A `Result` whose `Err` variant indicates a test failure.

googletest_macro/src/lib.rs

@@ -136,10 +136,14 @@ pub fn gtest(
         let mut invocation = quote! {
             (#maybe_async move || -> #result_type {
                 #block
-            })() #maybe_await
+            })()
         };
         if output_type.is_none() {
-            invocation = quote! { { let () = #invocation; googletest::Result::<()>::Ok(()) } };
+            if is_async {
+                invocation = quote! { async move { let () = #invocation.await; googletest::Result::<()>::Ok(()) } };
+            } else {
+                invocation = quote! { { let () = #invocation; googletest::Result::<()>::Ok(()) } };
+            }
         }
         invocation
     } else {
@@ -155,9 +159,19 @@ pub fn gtest(
             (#maybe_async move || -> #result_type {
                 #sig { #block }
                 #closure_body
-            })() #maybe_await
+            })()
         }
     };
+
+    let invocation = if is_async {
+        quote! {
+            ::googletest::internal::scoped_trace::InstrumentedFuture::new(
+                #invocation
+            ).await
+        }
+    } else {
+        quote! { #invocation }
+    };
     if !attrs.iter().any(is_test_attribute) && !is_rstest_enabled {
         let test_attr: Attribute = parse_quote! {
             #[::core::prelude::v1::test]

integration_tests/async_scoped_trace_panic_test.golden

@@ -0,0 +1,3 @@
+Google Test trace:
+  async_scoped_trace_panic_test.rs:4: Sync trace message
+  async_scoped_trace_panic_test.rs:14: Outer async trace message

integration_tests/async_scoped_trace_panic_test.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+BINARY_PATH="$1"
+
+if [ -z "$BINARY_PATH" ]; then
+  echo "Usage: $0 <binary_path>"
+  exit 1
+fi
+
+# Run the binary and capture output.
+set +e
+OUTPUT=$("${BINARY_PATH}" 2>&1)
+STATUS=$?
+set -e
+
+# 101 is the standard Rust panic exit code.
+if [[ ${STATUS} -ne 101 ]]; then
+  echo "Expected panic (exit code 101), but script exited with ${STATUS}."
+  echo "Output:"
+  echo "${OUTPUT}"
+  exit 1
+fi
+
+echo "${OUTPUT}" | \
+sed -n -e '/^Google Test trace:/p' -e '/^  .*trace/p' | \
+sed 's#^.*async_scoped_trace_panic_test.rs#  async_scoped_trace_panic_test.rs#g'

integration_tests/src/async_scoped_trace_panic_test.rs

@@ -0,0 +1,27 @@
+use googletest::prelude::*;
+use googletest::InstrumentedFuture;
+fn sync_function_with_trace() {
+    scoped_trace!("Sync trace message");
+    panic!("Intentional panic in sync function");
+}
+
+#[tokio::main]
+async fn main() {
+    // Initialize Google Test to install panic hook
+    googletest::internal::test_outcome::TestOutcome::init_current_test_outcome();
+
+    let fut = async {
+        scoped_trace!("Outer async trace message");
+
+        // Yield to verify async traces are preserved across yields
+        tokio::task::yield_now().await;
+
+        // Call sync function that adds a trace and panics
+        sync_function_with_trace();
+    };
+
+    // Wrap in InstrumentedFuture to preserve traces across yields
+    let fut = InstrumentedFuture::new(fut);
+
+    fut.await;
+}

integration_tests/src/async_test_with_expect_that.rs

@@ -75,4 +75,54 @@ mod tests {
         f.0 = 6;
         Ok(())
     }
+
+    #[gtest]
+    #[tokio::test]
+    async fn async_test_with_scoped_trace() -> Result<()> {
+        scoped_trace!("Outer trace");
+        brief_sleep().await;
+        {
+            scoped_trace!("Inner trace");
+            brief_sleep().await;
+            let traces = googletest::internal::scoped_trace::get_scoped_traces();
+            verify_eq!(traces.len(), 2)?;
+            verify_eq!(traces[0].message, "Outer trace")?;
+            verify_eq!(traces[1].message, "Inner trace")?;
+        }
+        brief_sleep().await;
+        let traces = googletest::internal::scoped_trace::get_scoped_traces();
+        verify_eq!(traces.len(), 1)?;
+        verify_eq!(traces[0].message, "Outer trace")?;
+        Ok(())
+    }
+
+    fn sync_subroutine_with_trace() {
+        scoped_trace!("Sync subroutine trace");
+        let traces = googletest::internal::scoped_trace::get_scoped_traces();
+        assert!(traces.iter().any(|t| t.message == "Sync subroutine trace"));
+    }
+
+    #[gtest]
+    #[tokio::test]
+    async fn async_test_mixed_scoped_trace() -> Result<()> {
+        scoped_trace!("Outer async trace");
+        brief_sleep().await;
+
+        sync_subroutine_with_trace();
+
+        {
+            scoped_trace!("Inner async trace");
+            brief_sleep().await;
+            let traces = googletest::internal::scoped_trace::get_scoped_traces();
+            verify_eq!(traces.len(), 2)?;
+            verify_eq!(traces[0].message, "Outer async trace")?;
+            verify_eq!(traces[1].message, "Inner async trace")?;
+        }
+
+        brief_sleep().await;
+        let traces = googletest::internal::scoped_trace::get_scoped_traces();
+        verify_eq!(traces.len(), 1)?;
+        verify_eq!(traces[0].message, "Outer async trace")?;
+        Ok(())
+    }
 }