YJIT: Crash with rb_bug() when panicking

Helps with getting good bug reports in the wild. Intended to be
backported to the 3.2.x series.

3 files changed, 31 insertions, 2 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 35b09e1abc..8d29d42bf7 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -88,6 +88,9 @@ fn main() {
         // This function prints info about a value and is useful for debugging
         .allowlist_function("rb_obj_info_dump")
 
+        // For crashing
+        .allowlist_function("rb_bug")
+
         // From shape.h
         .allowlist_function("rb_shape_get_shape_id")
         .allowlist_function("rb_shape_get_shape_by_id")
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 128b249e3b..a405d2aca6 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1074,6 +1074,7 @@ extern "C" {
     pub fn rb_singleton_class(obj: VALUE) -> VALUE;
     pub fn rb_get_alloc_func(klass: VALUE) -> rb_alloc_func_t;
     pub fn rb_method_basic_definition_p(klass: VALUE, mid: ID) -> ::std::os::raw::c_int;
+    pub fn rb_bug(fmt: *const ::std::os::raw::c_char, ...) -> !;
     pub fn rb_gc_writebarrier(old: VALUE, young: VALUE);
     pub fn rb_class_get_superclass(klass: VALUE) -> VALUE;
     pub static mut rb_mKernel: VALUE;
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index 4850dca7a8..305778a939 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -50,13 +50,13 @@ pub extern "C" fn rb_yjit_init_rust() {
 
     // Catch panics to avoid UB for unwinding into C frames.
     // See https://doc.rust-lang.org/nomicon/exception-safety.html
-    // TODO: set a panic handler so the we don't print a message
-    //       everytime we panic.
     let result = std::panic::catch_unwind(|| {
         Invariants::init();
         CodegenGlobals::init();
         YjitExitLocations::init();
 
+        rb_bug_panic_hook();
+
         // YJIT enabled and initialized successfully
         YJIT_ENABLED.store(true, Ordering::Release);
     });
@@ -67,6 +67,31 @@ pub extern "C" fn rb_yjit_init_rust() {
     }
 }
 
+/// At the moment, we abort in all cases we panic.
+/// To aid with getting diagnostics in the wild without requiring
+/// people to set RUST_BACKTRACE=1, register a panic hook that crash using rb_bug().
+/// rb_bug() might not be as good at printing a call trace as Rust's stdlib, but
+/// it dumps some other info that might be relevant.
+///
+/// In case we want to do start doing fancier exception handling with panic=unwind,
+/// we can revisit this later. For now, this helps to get us good bug reports.
+fn rb_bug_panic_hook() {
+    use std::panic;
+    use std::io::{stderr, Write};
+
+    // Probably the default hook. We do this very early during process boot.
+    let previous_hook = panic::take_hook();
+
+    panic::set_hook(Box::new(move |panic_info| {
+        // Not using `eprintln` to avoid double panic.
+        let _ = stderr().write_all(b"ruby: YJIT has panicked. More info to follow...\n");
+
+        previous_hook(panic_info);
+
+        unsafe { rb_bug(b"YJIT panicked\0".as_ref().as_ptr() as *const raw::c_char); }
+    }));
+}
+
 /// Called from C code to begin compiling a function
 /// NOTE: this should be wrapped in RB_VM_LOCK_ENTER(), rb_vm_barrier() on the C side
 #[no_mangle]