compiletest/
panic_hook.rs

1use std::backtrace::{Backtrace, BacktraceStatus};
2use std::cell::Cell;
3use std::fmt::{Display, Write};
4use std::panic::PanicHookInfo;
5use std::sync::{Arc, LazyLock, Mutex};
6use std::{env, mem, panic, thread};
7
8type PanicHook = Box<dyn Fn(&PanicHookInfo<'_>) + Sync + Send + 'static>;
9type CaptureBuf = Arc<Mutex<String>>;
10
11thread_local!(
12    static CAPTURE_BUF: Cell<Option<CaptureBuf>> = const { Cell::new(None) };
13);
14
15/// Installs a custom panic hook that will divert panic output to a thread-local
16/// capture buffer, but only for threads that have a capture buffer set.
17///
18/// Otherwise, the custom hook delegates to a copy of the default panic hook.
19pub(crate) fn install_panic_hook() {
20    let default_hook = panic::take_hook();
21    panic::set_hook(Box::new(move |info| custom_panic_hook(&default_hook, info)));
22}
23
24pub(crate) fn set_capture_buf(buf: CaptureBuf) {
25    CAPTURE_BUF.set(Some(buf));
26}
27
28pub(crate) fn take_capture_buf() -> Option<CaptureBuf> {
29    CAPTURE_BUF.take()
30}
31
32fn custom_panic_hook(default_hook: &PanicHook, info: &panic::PanicHookInfo<'_>) {
33    // Temporarily taking the capture buffer means that if a panic occurs in
34    // the subsequent code, that panic will fall back to the default hook.
35    let Some(buf) = take_capture_buf() else {
36        // There was no capture buffer, so delegate to the default hook.
37        default_hook(info);
38        return;
39    };
40
41    let mut out = buf.lock().unwrap_or_else(|e| e.into_inner());
42
43    let thread = thread::current().name().unwrap_or("(test runner)").to_owned();
44    let location = get_location(info);
45    let payload = payload_as_str(info).unwrap_or("Box<dyn Any>");
46    let backtrace = Backtrace::capture();
47
48    writeln!(out, "\nthread '{thread}' panicked at {location}:\n{payload}").unwrap();
49    match backtrace.status() {
50        BacktraceStatus::Captured => {
51            let bt = trim_backtrace(backtrace.to_string());
52            write!(out, "stack backtrace:\n{bt}",).unwrap();
53        }
54        BacktraceStatus::Disabled => {
55            writeln!(
56                out,
57                "note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace",
58            )
59            .unwrap();
60        }
61        _ => {}
62    }
63
64    drop(out);
65    set_capture_buf(buf);
66}
67
68fn get_location<'a>(info: &'a PanicHookInfo<'_>) -> &'a dyn Display {
69    match info.location() {
70        Some(location) => location,
71        None => &"(unknown)",
72    }
73}
74
75/// FIXME(Zalathar): Replace with `PanicHookInfo::payload_as_str` when that's
76/// stable in beta.
77fn payload_as_str<'a>(info: &'a PanicHookInfo<'_>) -> Option<&'a str> {
78    let payload = info.payload();
79    if let Some(s) = payload.downcast_ref::<&str>() {
80        Some(s)
81    } else if let Some(s) = payload.downcast_ref::<String>() {
82        Some(s)
83    } else {
84        None
85    }
86}
87
88fn rust_backtrace_full() -> bool {
89    static RUST_BACKTRACE_FULL: LazyLock<bool> =
90        LazyLock::new(|| matches!(env::var("RUST_BACKTRACE").as_deref(), Ok("full")));
91    *RUST_BACKTRACE_FULL
92}
93
94/// On stable, short backtraces are only available to the default panic hook,
95/// so if we want something similar we have to resort to string processing.
96fn trim_backtrace(full_backtrace: String) -> String {
97    if rust_backtrace_full() {
98        return full_backtrace;
99    }
100
101    let mut buf = String::with_capacity(full_backtrace.len());
102    // Don't print any frames until after the first `__rust_end_short_backtrace`.
103    let mut on = false;
104    // After the short-backtrace state is toggled, skip its associated "at" if present.
105    let mut skip_next_at = false;
106
107    let mut lines = full_backtrace.lines();
108    while let Some(line) = lines.next() {
109        if mem::replace(&mut skip_next_at, false) && line.trim_start().starts_with("at ") {
110            continue;
111        }
112
113        if line.contains("__rust_end_short_backtrace") {
114            on = true;
115            skip_next_at = true;
116            continue;
117        }
118        if line.contains("__rust_begin_short_backtrace") {
119            on = false;
120            skip_next_at = true;
121            continue;
122        }
123
124        if on {
125            writeln!(buf, "{line}").unwrap();
126        }
127    }
128
129    writeln!(
130        buf,
131        "note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace."
132    )
133    .unwrap();
134
135    buf
136}
compiletest/panic_hook.rs

compiletest/
panic_hook.rs
