|
| 1 | +// Copyright 2026 Google LLC |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +use std::cell::RefCell; |
| 16 | +use std::sync::atomic::{AtomicUsize, Ordering}; |
| 17 | + |
| 18 | +static NEXT_TRACE_ID: AtomicUsize = AtomicUsize::new(0); |
| 19 | + |
| 20 | +/// Information about a scoped trace. |
| 21 | +#[derive(Clone, Debug)] |
| 22 | +pub struct TraceInfo { |
| 23 | + pub id: usize, |
| 24 | + pub file: &'static str, |
| 25 | + pub line: u32, |
| 26 | + pub message: String, |
| 27 | +} |
| 28 | + |
| 29 | +thread_local! { |
| 30 | + static TRACE_STACK: RefCell<Vec<TraceInfo>> = const { RefCell::new(Vec::new()) }; |
| 31 | +} |
| 32 | + |
| 33 | +/// RAII guard to manage the push and pop of trace information. |
| 34 | +/// |
| 35 | +/// This struct is `!Send` and `!Sync` to prevent it from being held across |
| 36 | +/// `.await` points in async tests, which would cause incorrect trace tracking |
| 37 | +/// if the task moves between threads. |
| 38 | +#[doc(hidden)] |
| 39 | +pub struct ScopedTraceGuard { |
| 40 | + id: usize, |
| 41 | + _phantom: std::marker::PhantomData<*mut ()>, |
| 42 | +} |
| 43 | + |
| 44 | +impl ScopedTraceGuard { |
| 45 | + #[doc(hidden)] |
| 46 | + #[track_caller] |
| 47 | + pub fn new(message: String) -> Self { |
| 48 | + let caller = std::panic::Location::caller(); |
| 49 | + let id = NEXT_TRACE_ID.fetch_add(1, Ordering::Relaxed); |
| 50 | + TRACE_STACK.with(|stack| { |
| 51 | + // Use try_borrow_mut to avoid double panic if called during unwinding. |
| 52 | + if let Ok(mut s) = stack.try_borrow_mut() { |
| 53 | + s.push(TraceInfo { id, file: caller.file(), line: caller.line(), message }); |
| 54 | + } |
| 55 | + }); |
| 56 | + Self { id, _phantom: std::marker::PhantomData } |
| 57 | + } |
| 58 | +} |
| 59 | + |
| 60 | +impl Drop for ScopedTraceGuard { |
| 61 | + fn drop(&mut self) { |
| 62 | + TRACE_STACK.with(|stack| { |
| 63 | + // Use try_borrow_mut to avoid double panic if called during unwinding. |
| 64 | + if let Ok(mut s) = stack.try_borrow_mut() { |
| 65 | + if let Some(pos) = s.iter().rposition(|t| t.id == self.id) { |
| 66 | + s.remove(pos); |
| 67 | + } |
| 68 | + } |
| 69 | + }); |
| 70 | + } |
| 71 | +} |
| 72 | + |
| 73 | +/// Retrieves a clone of the current thread's trace stack. |
| 74 | +pub fn get_scoped_traces() -> Vec<TraceInfo> { |
| 75 | + TRACE_STACK.with(|stack| stack.try_borrow().map(|s| s.clone()).unwrap_or_default()) |
| 76 | +} |
| 77 | + |
| 78 | +// Test-only state and helpers, hidden from production API. |
| 79 | +#[cfg(test)] |
| 80 | +pub(crate) mod test_helpers { |
| 81 | + use super::*; |
| 82 | + use std::cell::Cell; |
| 83 | + |
| 84 | + thread_local! { |
| 85 | + pub static CAPTURED_TRACES_IN_HOOK: RefCell<Vec<TraceInfo>> = const { RefCell::new(Vec::new()) }; |
| 86 | + pub static USE_CAPTURE_HOOK: Cell<bool> = const { Cell::new(false) }; |
| 87 | + } |
| 88 | + |
| 89 | + pub fn enable_capture_in_hook(enable: bool) { |
| 90 | + USE_CAPTURE_HOOK.with(|v| v.set(enable)); |
| 91 | + } |
| 92 | + |
| 93 | + pub fn get_captured_traces_in_hook() -> Vec<TraceInfo> { |
| 94 | + CAPTURED_TRACES_IN_HOOK.with(|v| v.borrow().clone()) |
| 95 | + } |
| 96 | + |
| 97 | + pub fn clear_captured_traces_in_hook() { |
| 98 | + CAPTURED_TRACES_IN_HOOK.with(|v| v.borrow_mut().clear()); |
| 99 | + } |
| 100 | +} |
| 101 | + |
| 102 | +#[cfg(test)] |
| 103 | +mod tests { |
| 104 | + use super::test_helpers::*; |
| 105 | + use super::*; |
| 106 | + |
| 107 | + #[test] |
| 108 | + fn test_scoped_trace_fatal() { |
| 109 | + enable_capture_in_hook(true); |
| 110 | + clear_captured_traces_in_hook(); |
| 111 | + |
| 112 | + // Ensure hook is installed |
| 113 | + crate::internal::test_outcome::TestOutcome::init_current_test_outcome(); |
| 114 | + |
| 115 | + let _ = std::panic::catch_unwind(|| { |
| 116 | + let _guard = ScopedTraceGuard::new("Second trace".to_string()); |
| 117 | + panic!("Intentional panic"); |
| 118 | + }); |
| 119 | + |
| 120 | + let captured = get_captured_traces_in_hook(); |
| 121 | + assert_eq!(captured.len(), 1); |
| 122 | + assert_eq!(captured[0].message, "Second trace"); |
| 123 | + |
| 124 | + enable_capture_in_hook(false); |
| 125 | + } |
| 126 | +} |
0 commit comments