feat(utils/error): make catching panics work correctly in a multi-threaded setting

Author: kaimast

utilities/src/errors.rs

@@ -15,7 +15,14 @@
 
 use colored::Colorize;
 
-use std::{borrow::Borrow, fmt::Display};
+use std::{any::Any, backtrace::Backtrace, borrow::Borrow, cell::Cell, fmt::Display, panic};
+
+thread_local! {
+/// The message backtrace of the last panic on this thread (if any).
+///
+/// We store this information here instead of directly processing it in a panic hook, because panic hooks are global whereas this can be processed on a per-thread basis.
+/// For example, one thread may execute a program where panics should *not* cause the entire process to terminate, while in another thread there is a panic due to a bug.
+static PANIC_INFO: Cell<Option<(String, Backtrace)>> = const { Cell::new(None) };}
 
 /// Generates an `io::Error` from the given string.
 pub fn io_error<S: ToString>(err: S) -> std::io::Error {
@@ -127,6 +134,17 @@ pub trait PrettyUnwrap {
     fn pretty_unwrap(self) -> Self::Inner;
 }
 
+/// Set the global panic hook for process. Should be called exactly once.
+pub fn set_panic_hook() {
+    std::panic::set_hook(Box::new(|err| {
+        let msg = err.to_string();
+
+        println!("Hook: {msg}");
+        let trace = Backtrace::force_capture();
+        PANIC_INFO.with(move |info| info.set(Some((msg, trace))));
+    }));
+}
+
 /// Helper for `PrettyUnwrap`:
 /// Creates a panic with the `anyhow::Error` nicely formatted.
 #[track_caller]
@@ -155,9 +173,58 @@ impl<T> PrettyUnwrap for anyhow::Result<T> {
     }
 }
 
+/// `try_vm_runtime` executes the given closure in an environment which will safely halt
+/// without producing logs that look like unexpected behavior.
+/// In debug mode, it prints to stderr using the format: "VM safely halted at {location}: {halt message}".
+///
+/// Note: For this to work as expected, panics must be set to `unwind` during compilation (default), and the closure cannot invoke any async code that may potentially execute in a different OS thread.
+#[inline]
+#[track_caller]
+pub fn try_vm_runtime<R, F: FnMut() -> R>(f: F) -> Result<R, Box<dyn Any + Send>> {
+    // Perform the operation that may panic.
+    let result = std::panic::catch_unwind(panic::AssertUnwindSafe(f));
+
+    if result.is_err() {
+        // Get the stored panic and backtrace from the thread-local variable.
+        let (msg, _) = PANIC_INFO.with(|info| info.take()).expect("No panic information stored?");
+
+        #[cfg(debug_assertions)]
+        {
+            // Remove all words up to "panicked".
+            // And prepend with "VM Safely halted"
+            let msg = msg
+                .to_string()
+                .split_ascii_whitespace()
+                .skip_while(|&word| word != "panicked")
+                .collect::<Vec<&str>>()
+                .join(" ")
+                .replacen("panicked", "VM safely halted", 1);
+
+            eprintln!("{msg}");
+        }
+        #[cfg(not(debug_assertions))]
+        {
+            // Discard message
+            let _ = msg;
+        }
+    }
+
+    // Return the result, allowing regular error-handling.
+    result
+}
+
+/// `catch_unwind` calls the given closure `f` and, if `f` panics, returns the panic message and backtrace.
+pub fn catch_unwind<R, F: FnMut() -> R>(f: F) -> Result<R, (String, Backtrace)> {
+    // Perform the operation that may panic.
+    std::panic::catch_unwind(panic::AssertUnwindSafe(f)).map_err(|_| {
+        // Get the stored panic and backtrace from the thread-local variable.
+        PANIC_INFO.with(|info| info.take()).expect("No panic information stored?")
+    })
+}
+
 #[cfg(test)]
 mod tests {
-    use super::{PrettyUnwrap, flatten_anyhow_error, pretty_panic};
+    use super::{PrettyUnwrap, catch_unwind, flatten_anyhow_error, pretty_panic, set_panic_hook, try_vm_runtime};
 
     use anyhow::{Context, Result, anyhow, bail};
     use colored::Colorize;
@@ -215,14 +282,31 @@ mod tests {
         assert_eq!(*result.downcast::<String>().expect("Error was not a string"), expected);
     }
 
+    // Ensure catch_unwind stores the panic message as expected.
+    #[test]
+    fn test_catch_unwind() {
+        set_panic_hook();
+        let result = catch_unwind(move || {
+            panic!("This is my message");
+        });
+        // Remove hook so test asserts work normally again.
+        let _ = std::panic::take_hook();
+
+        let (msg, bt) = result.expect_err("No panic caught");
+        assert!(msg.ends_with("This is my message"));
+
+        // This function should be in the panics backtrace
+        assert!(bt.to_string().contains("test_catch_unwind"));
+    }
+
     /// Ensure catch_unwind does not break `try_vm_runtime`.
     #[test]
     fn test_nested_with_try_vm_runtime() {
-        use crate::try_vm_runtime;
+        set_panic_hook();
 
         let result = std::panic::catch_unwind(|| {
             // try_vm_runtime uses catch_unwind internally
-            let vm_result = try_vm_runtime!(|| {
+            let vm_result = try_vm_runtime(|| {
                 panic!("VM operation failed!");
             });
 

utilities/src/lib.rs

@@ -34,9 +34,6 @@ pub use bytes::*;
 pub mod defer;
 pub use defer::*;
 
-mod vm_error;
-pub use vm_error::*;
-
 pub mod iterator;
 pub use iterator::*;
 
@@ -63,3 +60,16 @@ pub use task::*;
 
 /// Use old name for backward-compatibility.
 pub use errors::io_error as error;
+
+/// This macro provides a VM runtime environment which will safely halt
+/// without producing logs that look like unexpected behavior.
+/// In debug mode, it prints to stderr using the format: "VM safely halted at {location}: {halt message}".
+///
+/// It is more efficient to set the panic hook once and directly use `errors::try_vm_runtime`.
+#[macro_export]
+macro_rules! try_vm_runtime {
+    ($e:expr) => {{
+        $crate::errors::set_panic_hook();
+        $crate::errors::try_vm_runtime($e)
+    }};
+}