[red-knot] Capture backtrace in "check-failed" diagnostic (#17641)

Co-authored-by: David Peter <sharkdp@users.noreply.github.com>

crates/red_knot_project/src/lib.rs

@@ -20,7 +20,7 @@ use ruff_db::system::{SystemPath, SystemPathBuf};
 use rustc_hash::FxHashSet;
 use salsa::Durability;
 use salsa::Setter;
-use std::panic::{catch_unwind, AssertUnwindSafe, UnwindSafe};
+use std::panic::{AssertUnwindSafe, UnwindSafe};
 use std::sync::Arc;
 use thiserror::Error;
 use tracing::error;
@@ -577,26 +577,29 @@ fn catch<F, R>(db: &dyn Db, file: File, f: F) -> Result<Option<R>, Diagnostic>
 where
     F: FnOnce() -> R + UnwindSafe,
 {
-    match catch_unwind(|| {
+    match ruff_db::panic::catch_unwind(|| {
         // Ignore salsa errors
         salsa::Cancelled::catch(f).ok()
     }) {
         Ok(result) => Ok(result),
         Err(error) => {
-            let payload = if let Some(s) = error.downcast_ref::<&str>() {
+            use std::fmt::Write;
-                Some((*s).to_string())
+            let mut message = String::new();
-            } else {
+            message.push_str("Panicked");
-                error.downcast_ref::<String>().cloned()
-            };
 
-            let message = if let Some(payload) = payload {
+            if let Some(location) = error.location {
-                format!(
+                let _ = write!(&mut message, " at {location}");
-                    "Panicked while checking `{file}`: `{payload}`",
+            }
-                    file = file.path(db)
+
-                )
+            let _ = write!(
-            } else {
+                &mut message,
-                format!("Panicked while checking `{file}`", file = { file.path(db) })
+                " when checking `{file}`",
-            };
+                file = file.path(db)
+            );
+
+            if let Some(payload) = error.payload.as_str() {
+                let _ = write!(&mut message, ": `{payload}`");
+            }
 
             let mut diagnostic = Diagnostic::new(DiagnosticId::Panic, Severity::Fatal, message);
             diagnostic.sub(SubDiagnostic::new(
@@ -606,6 +609,28 @@ where
 
             let report_message = "If you could open an issue at https://github.com/astral-sh/ruff/issues/new?title=%5Bred-knot%5D:%20panic we'd be very appreciative!";
             diagnostic.sub(SubDiagnostic::new(Severity::Info, report_message));
+            diagnostic.sub(SubDiagnostic::new(
+                Severity::Info,
+                format!(
+                    "Platform: {os} {arch}",
+                    os = std::env::consts::OS,
+                    arch = std::env::consts::ARCH
+                ),
+            ));
+            diagnostic.sub(SubDiagnostic::new(
+                Severity::Info,
+                format!(
+                    "Args: {args:?}",
+                    args = std::env::args().collect::<Vec<_>>()
+                ),
+            ));
+
+            if let Some(backtrace) = error.backtrace {
+                diagnostic.sub(SubDiagnostic::new(
+                    Severity::Info,
+                    format!("Backtrace:\n{backtrace}"),
+                ));
+            }
 
             Err(diagnostic)
         }

crates/red_knot_test/src/lib.rs

@@ -324,8 +324,8 @@ fn run_test(
                         Some(location) => messages.push(format!("panicked at {location}")),
                         None => messages.push("panicked at unknown location".to_string()),
                     }
-                    match info.payload {
+                    match info.payload.as_str() {
-                        Some(payload) => messages.push(payload),
+                        Some(message) => messages.push(message.to_string()),
                         // Mimic the default panic hook's rendering of the panic payload if it's
                         // not a string.
                         None => messages.push("Box<dyn Any>".to_string()),

crates/ruff_db/src/panic.rs

@@ -2,20 +2,35 @@ use std::cell::Cell;
 use std::panic::Location;
 use std::sync::OnceLock;
 
-#[derive(Default, Debug)]
+#[derive(Debug)]
 pub struct PanicError {
     pub location: Option<String>,
-    pub payload: Option<String>,
+    pub payload: Payload,
     pub backtrace: Option<std::backtrace::Backtrace>,
 }
 
+#[derive(Debug)]
+pub struct Payload(Box<dyn std::any::Any + Send>);
+
+impl Payload {
+    pub fn as_str(&self) -> Option<&str> {
+        if let Some(s) = self.0.downcast_ref::<String>() {
+            Some(s)
+        } else if let Some(s) = self.0.downcast_ref::<&str>() {
+            Some(s)
+        } else {
+            None
+        }
+    }
+}
+
 impl std::fmt::Display for PanicError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         writeln!(f, "panicked at")?;
         if let Some(location) = &self.location {
             write!(f, " {location}")?;
         }
-        if let Some(payload) = &self.payload {
+        if let Some(payload) = self.payload.as_str() {
             write!(f, ":\n{payload}")?;
         }
         if let Some(backtrace) = &self.backtrace {
@@ -27,8 +42,7 @@ impl std::fmt::Display for PanicError {
 
 thread_local! {
     static CAPTURE_PANIC_INFO: Cell<bool> = const { Cell::new(false) };
-    static OUR_HOOK_RAN: Cell<bool> = const { Cell::new(false) };
+    static LAST_BACKTRACE: Cell<(Option<std::backtrace::Backtrace>, Option<String>)> = const { Cell::new((None, None)) };
-    static LAST_PANIC: Cell<Option<PanicError>> = const { Cell::new(None) };
 }
 
 fn install_hook() {
@@ -36,25 +50,15 @@ fn install_hook() {
     ONCE.get_or_init(|| {
         let prev = std::panic::take_hook();
         std::panic::set_hook(Box::new(move |info| {
-            OUR_HOOK_RAN.with(|cell| cell.set(true));
             let should_capture = CAPTURE_PANIC_INFO.with(Cell::get);
             if !should_capture {
                 return (*prev)(info);
             }
-            let payload = if let Some(s) = info.payload().downcast_ref::<&str>() {
+
-                Some(s.to_string())
-            } else {
-                info.payload().downcast_ref::<String>().cloned()
-            };
             let location = info.location().map(Location::to_string);
-            let backtrace = std::backtrace::Backtrace::force_capture();
+            let backtrace = Some(std::backtrace::Backtrace::force_capture());
-            LAST_PANIC.with(|cell| {
+
-                cell.set(Some(PanicError {
+            LAST_BACKTRACE.set((backtrace, location));
-                    payload,
-                    location,
-                    backtrace: Some(backtrace),
-                }));
-            });
         }));
     });
 }
@@ -70,7 +74,7 @@ fn install_hook() {
 /// stderr).
 ///
 /// We assume that there is nothing else running in this process that needs to install a competing
-/// panic hook.  We are careful to install our custom hook only once, and we do not ever restore
+/// panic hook. We are careful to install our custom hook only once, and we do not ever restore
 /// the previous hook (since you can always retain the previous hook's behavior by not calling this
 /// wrapper).
 pub fn catch_unwind<F, R>(f: F) -> Result<R, PanicError>
@@ -78,15 +82,78 @@ where
     F: FnOnce() -> R + std::panic::UnwindSafe,
 {
     install_hook();
-    OUR_HOOK_RAN.with(|cell| cell.set(false));
+    let prev_should_capture = CAPTURE_PANIC_INFO.replace(true);
-    let prev_should_capture = CAPTURE_PANIC_INFO.with(|cell| cell.replace(true));
+    let result = std::panic::catch_unwind(f).map_err(|payload| {
-    let result = std::panic::catch_unwind(f).map_err(|_| {
+        // Try to get the backtrace and location from our custom panic hook.
-        let our_hook_ran = OUR_HOOK_RAN.with(Cell::get);
+        // The custom panic hook only runs once when `panic!` is called (or similar). It doesn't
-        if !our_hook_ran {
+        // run when the panic is propagated with `std::panic::resume_unwind`. The panic hook
-            panic!("detected a competing panic hook");
+        // is also not called when the panic is raised with `std::panic::resum_unwind` as is the
+        // case for salsa unwinds (see the ignored test below).
+        // Because of that, always take the payload from `catch_unwind` because it may have been transformed
+        // by an inner `std::panic::catch_unwind` handlers and only use the information
+        // from the custom handler to enrich the error with the backtrace and location.
+        let (backtrace, location) = LAST_BACKTRACE.with(Cell::take);
+
+        PanicError {
+            location,
+            payload: Payload(payload),
+            backtrace,
         }
-        LAST_PANIC.with(Cell::take).unwrap_or_default()
     });
-    CAPTURE_PANIC_INFO.with(|cell| cell.set(prev_should_capture));
+    CAPTURE_PANIC_INFO.set(prev_should_capture);
     result
 }
+
+#[cfg(test)]
+mod tests {
+    use salsa::{Database, Durability};
+
+    #[test]
+    #[ignore = "super::catch_unwind installs a custom panic handler, which could effect test isolation"]
+    fn no_backtrace_for_salsa_cancelled() {
+        #[salsa::input]
+        struct Input {
+            value: u32,
+        }
+
+        #[salsa::tracked]
+        fn test_query(db: &dyn Database, input: Input) -> u32 {
+            loop {
+                // This should throw a cancelled error
+                let _ = input.value(db);
+            }
+        }
+
+        let db = salsa::DatabaseImpl::new();
+
+        let input = Input::new(&db, 42);
+
+        let result = std::thread::scope(move |scope| {
+            {
+                let mut db = db.clone();
+                scope.spawn(move || {
+                    // This will cancel the other thread by throwing a `salsa::Cancelled` error.
+                    db.synthetic_write(Durability::MEDIUM);
+                });
+            }
+
+            {
+                scope.spawn(move || {
+                    super::catch_unwind(|| {
+                        test_query(&db, input);
+                    })
+                })
+            }
+            .join()
+            .unwrap()
+        });
+
+        match result {
+            Ok(_) => panic!("Expected query to panic"),
+            Err(err) => {
+                // Panics triggered with `resume_unwind` have no backtrace.
+                assert!(err.backtrace.is_none());
+            }
+        }
+    }
+}