compiletest: Capture panic messages via a custom panic hook

src/tools/compiletest/src/executor.rs

@@ -13,6 +13,7 @@ use std::sync::{Arc, Mutex, mpsc};
 use std::{env, hint, io, mem, panic, thread};
 
 use crate::common::{Config, TestPaths};
+use crate::panic_hook;
 
 mod deadline;
 mod json;
@@ -120,6 +121,11 @@ fn run_test_inner(
     completion_sender: mpsc::Sender<TestCompletion>,
 ) {
     let is_capture = !runnable_test.config.nocapture;
+
+    // Install a panic-capture buffer for use by the custom panic hook.
+    if is_capture {
+        panic_hook::set_capture_buf(Default::default());
+    }
     let capture_buf = is_capture.then(|| Arc::new(Mutex::new(vec![])));
 
     if let Some(capture_buf) = &capture_buf {
@@ -128,6 +134,13 @@ fn run_test_inner(
 
     let panic_payload = panic::catch_unwind(move || runnable_test.run()).err();
 
+    if let Some(panic_buf) = panic_hook::take_capture_buf() {
+        let panic_buf = panic_buf.lock().unwrap_or_else(|e| e.into_inner());
+        // For now, forward any captured panic message to (captured) stderr.
+        // FIXME(Zalathar): Once we have our own output-capture buffer for
+        // non-panic output, append the panic message to that buffer instead.
+        eprint!("{panic_buf}");
+    }
     if is_capture {
         io::set_output_capture(None);
     }

src/tools/compiletest/src/lib.rs

@@ -15,6 +15,7 @@ pub mod directives;
 pub mod errors;
 mod executor;
 mod json;
+mod panic_hook;
 mod raise_fd_limit;
 mod read2;
 pub mod runtest;
@@ -493,6 +494,8 @@ pub fn opt_str2(maybestr: Option<String>) -> String {
 pub fn run_tests(config: Arc<Config>) {
     debug!(?config, "run_tests");
 
+    panic_hook::install_panic_hook();
+
     // If we want to collect rustfix coverage information,
     // we first make sure that the coverage file does not exist.
     // It will be created later on.

src/tools/compiletest/src/panic_hook.rs

@@ -0,0 +1,136 @@
+use std::backtrace::{Backtrace, BacktraceStatus};
+use std::cell::Cell;
+use std::fmt::{Display, Write};
+use std::panic::PanicHookInfo;
+use std::sync::{Arc, LazyLock, Mutex};
+use std::{env, mem, panic, thread};
+
+type PanicHook = Box<dyn Fn(&PanicHookInfo<'_>) + Sync + Send + 'static>;
+type CaptureBuf = Arc<Mutex<String>>;
+
+thread_local!(
+    static CAPTURE_BUF: Cell<Option<CaptureBuf>> = const { Cell::new(None) };
+);
+
+/// Installs a custom panic hook that will divert panic output to a thread-local
+/// capture buffer, but only for threads that have a capture buffer set.
+///
+/// Otherwise, the custom hook delegates to a copy of the default panic hook.
+pub(crate) fn install_panic_hook() {
+    let default_hook = panic::take_hook();
+    panic::set_hook(Box::new(move |info| custom_panic_hook(&default_hook, info)));
+}
+
+pub(crate) fn set_capture_buf(buf: CaptureBuf) {
+    CAPTURE_BUF.set(Some(buf));
+}
+
+pub(crate) fn take_capture_buf() -> Option<CaptureBuf> {
+    CAPTURE_BUF.take()
+}
+
+fn custom_panic_hook(default_hook: &PanicHook, info: &panic::PanicHookInfo<'_>) {
+    // Temporarily taking the capture buffer means that if a panic occurs in
+    // the subsequent code, that panic will fall back to the default hook.
+    let Some(buf) = take_capture_buf() else {
+        // There was no capture buffer, so delegate to the default hook.
+        default_hook(info);
+        return;
+    };
+
+    let mut out = buf.lock().unwrap_or_else(|e| e.into_inner());
+
+    let thread = thread::current().name().unwrap_or("(test runner)").to_owned();
+    let location = get_location(info);
+    let payload = payload_as_str(info).unwrap_or("Box<dyn Any>");
+    let backtrace = Backtrace::capture();
+
+    writeln!(out, "\nthread '{thread}' panicked at {location}:\n{payload}").unwrap();
+    match backtrace.status() {
+        BacktraceStatus::Captured => {
+            let bt = trim_backtrace(backtrace.to_string());
+            write!(out, "stack backtrace:\n{bt}",).unwrap();
+        }
+        BacktraceStatus::Disabled => {
+            writeln!(
+                out,
+                "note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace",
+            )
+            .unwrap();
+        }
+        _ => {}
+    }
+
+    drop(out);
+    set_capture_buf(buf);
+}
+
+fn get_location<'a>(info: &'a PanicHookInfo<'_>) -> &'a dyn Display {
+    match info.location() {
+        Some(location) => location,
+        None => &"(unknown)",
+    }
+}
+
+/// FIXME(Zalathar): Replace with `PanicHookInfo::payload_as_str` when that's
+/// stable in beta.
+fn payload_as_str<'a>(info: &'a PanicHookInfo<'_>) -> Option<&'a str> {
+    let payload = info.payload();
+    if let Some(s) = payload.downcast_ref::<&str>() {
+        Some(s)
+    } else if let Some(s) = payload.downcast_ref::<String>() {
+        Some(s)
+    } else {
+        None
+    }
+}
+
+fn rust_backtrace_full() -> bool {
+    static RUST_BACKTRACE_FULL: LazyLock<bool> =
+        LazyLock::new(|| matches!(env::var("RUST_BACKTRACE").as_deref(), Ok("full")));
+    *RUST_BACKTRACE_FULL
+}
+
+/// On stable, short backtraces are only available to the default panic hook,
+/// so if we want something similar we have to resort to string processing.
+fn trim_backtrace(full_backtrace: String) -> String {
+    if rust_backtrace_full() {
+        return full_backtrace;
+    }
+
+    let mut buf = String::with_capacity(full_backtrace.len());
+    // Don't print any frames until after the first `__rust_end_short_backtrace`.
+    let mut on = false;
+    // After the short-backtrace state is toggled, skip its associated "at" if present.
+    let mut skip_next_at = false;
+
+    let mut lines = full_backtrace.lines();
+    while let Some(line) = lines.next() {
+        if mem::replace(&mut skip_next_at, false) && line.trim_start().starts_with("at ") {
+            continue;
+        }
+
+        if line.contains("__rust_end_short_backtrace") {
+            on = true;
+            skip_next_at = true;
+            continue;
+        }
+        if line.contains("__rust_begin_short_backtrace") {
+            on = false;
+            skip_next_at = true;
+            continue;
+        }
+
+        if on {
+            writeln!(buf, "{line}").unwrap();
+        }
+    }
+
+    writeln!(
+        buf,
+        "note: Some details are omitted, run with `RUST_BACKTRACE=full` for a verbose backtrace."
+    )
+    .unwrap();
+
+    buf
+}