capture panic message inside panic_handler (#1890)

Summary:

The main motivation is by making this change, we can log panic message here too:

https://www.internalfb.com/code/fbsource/[4a662228cd8bdf2bdf9b760e705cc9958f85e55c]/fbcode/monarch/hyperactor/src/panic_handler.rs?lines=47

it currently does not and caused quite some confusion for me, since I thought there is a bug in the panic catching logic.

Differential Revision: D87086712

Author: pzhan9

hyperactor/src/panic_handler.rs

@@ -15,10 +15,19 @@ use std::future::Future;
 use std::ops::Deref;
 use std::panic;
 
+/// A struct to store the message and backtrace from a panic.
+#[derive(Clone)]
+pub struct PanicInfo {
+    /// The message from the panic.
+    pub message: String,
+    /// The backtrace from the panic.
+    pub backtrace: String,
+}
+
 tokio::task_local! {
     /// A task_local variable to store the backtrace from a panic, so it can be
     /// retrieved later.
-    static BACKTRACE: RefCell<Option<String>>;
+    static BACKTRACE: RefCell<Option<PanicInfo>>;
 }
 
 /// Call this from the main method of your application, and use it in conjunction
@@ -33,9 +42,23 @@ pub fn set_panic_hook() {
             || "unavailable".to_owned(),
             |loc: &panic::Location<'_>| format!("{}:{}:{}", loc.file(), loc.line(), loc.column()),
         );
+        // Extract the panic message from the payload
+        let panic_msg = if let Some(&s) = info.payload().downcast_ref::<&str>() {
+            s.to_string()
+        } else if let Some(s) = info.payload().downcast_ref::<String>() {
+            s.clone()
+        } else {
+            "panic message cannot be downcasted".to_string()
+        };
+
+        tracing::error!("stacktrace"=%backtrace, "panic at {loc}: {panic_msg}");
+
         let _result = BACKTRACE.try_with(|entry| match entry.try_borrow_mut() {
             Ok(mut entry_ref) => {
-                *entry_ref = Some(format!("panicked at {loc}\n{backtrace}"));
+                *entry_ref = Some(PanicInfo {
+                    message: panic_msg,
+                    backtrace: format!("panicked at {loc}\n{backtrace}"),
+                });
             }
             Err(borrow_mut_error) => {
                 eprintln!(
@@ -44,7 +67,6 @@ pub fn set_panic_hook() {
                 );
             }
         });
-        tracing::error!("stacktrace"=%backtrace, "panic at {loc}");
 
         // Execute the default hood to preserve the default behavior.
         prev(info);
@@ -62,11 +84,11 @@ where
 
 /// Take the backtrace from the task_local variable, and reset the task_local to
 /// None. Return error if the backtrace is not stored, or cannot be retrieved.
-pub fn take_panic_backtrace() -> Result<String, anyhow::Error> {
+pub fn take_panic_backtrace() -> Result<PanicInfo, anyhow::Error> {
     BACKTRACE.try_with(|entry| {
         entry.try_borrow_mut().map(|mut entry_ref| {
             let result = match entry_ref.deref() {
-                Some(bt) => Ok(bt.to_string()),
+                Some(bt) => Ok(bt.clone()),
                 None => Err(anyhow::anyhow!("nothing is stored in task_local")),
             };
             // Clear the task_local because the backtrace has been retrieve.
@@ -143,11 +165,9 @@ mod tests {
             .await;
             assert!(result.is_err());
             if backtrace_captured {
-                assert!(
-                    take_panic_backtrace()
-                        .unwrap()
-                        .contains("verify_inner_panic")
-                );
+                let info = take_panic_backtrace().unwrap();
+                assert_eq!(info.message, "wow!");
+                assert!(info.backtrace.contains("verify_inner_panic"));
             } else {
                 assert!(take_panic_backtrace().is_err());
             }
@@ -169,11 +189,9 @@ mod tests {
             assert!(result.is_ok());
 
             // Verify the outer task can get its own backtrace.
-            assert!(
-                take_panic_backtrace()
-                    .unwrap()
-                    .contains("test_nested_tasks")
-            );
+            let info = take_panic_backtrace().unwrap();
+            assert_eq!(info.message, "boom!");
+            assert!(info.backtrace.contains("test_nested_tasks"));
         })
         .await;
     }

hyperactor/src/proc.rs

@@ -1247,23 +1247,14 @@ impl<A: Actor> Instance<A> {
             .await
         {
             Ok(result) => result,
-            Err(err) => {
-                // This is only the error message. Backtrace is not included.
+            Err(_) => {
                 did_panic = true;
-                let err_msg = err
-                    .downcast_ref::<&str>()
-                    .copied()
-                    .or_else(|| err.downcast_ref::<String>().map(|s| s.as_str()))
-                    .unwrap_or("panic cannot be downcasted");
-
-                let backtrace = panic_handler::take_panic_backtrace()
+                let panic_info = panic_handler::take_panic_backtrace()
+                    .map(|info| format!("{}\n{}", info.message, info.backtrace))
                     .unwrap_or_else(|e| format!("Cannot take backtrace due to: {:?}", e));
                 Err(ActorError::new(
                     self.self_id(),
-                    ActorErrorKind::panic(anyhow::anyhow!(
-                        "{}
-{}", err_msg, backtrace
-                    )),
+                    ActorErrorKind::panic(anyhow::anyhow!(panic_info)),
                 ))
             }
         };