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 = info.payload_as_str().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
75fn rust_backtrace_full() -> bool {
76    static RUST_BACKTRACE_FULL: LazyLock<bool> =
77        LazyLock::new(|| matches!(env::var("RUST_BACKTRACE").as_deref(), Ok("full")));
78    *RUST_BACKTRACE_FULL
79}
80
81/// On stable, short backtraces are only available to the default panic hook,
82/// so if we want something similar we have to resort to string processing.
83fn trim_backtrace(full_backtrace: String) -> String {
84    if rust_backtrace_full() {
85        return full_backtrace;
86    }
87
88    let mut buf = String::with_capacity(full_backtrace.len());
89    // Don't print any frames until after the first `__rust_end_short_backtrace`.
90    let mut on = false;
91    // After the short-backtrace state is toggled, skip its associated "at" if present.
92    let mut skip_next_at = false;
93
94    let mut lines = full_backtrace.lines();
95    while let Some(line) = lines.next() {
96        if mem::replace(&mut skip_next_at, false) && line.trim_start().starts_with("at ") {
97            continue;
98        }
99
100        if line.contains("__rust_end_short_backtrace") {
101            on = true;
102            skip_next_at = true;
103            continue;
104        }
105        if line.contains("__rust_begin_short_backtrace") {
106            on = false;
107            skip_next_at = true;
108            continue;
109        }
110
111        if on {
112            writeln!(buf, "{line}").unwrap();
113        }
114    }
115
116    writeln!(
117        buf,
118        "note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace."
119    )
120    .unwrap();
121
122    buf
123}
compiletest/panic_hook.rs

compiletest/
panic_hook.rs
