From cfa1505068c135785fb5d8a01baaba5823e1f6e9 Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Sat, 26 Apr 2025 07:28:45 +0100
Subject: [PATCH] [red-knot] Fix CLI hang when a dependent query panics
 (#17631)

---
 crates/red_knot/src/main.rs          | 15 +++++---
 crates/red_knot_project/src/lib.rs   | 54 ++++++++++++++++++++++++++--
 crates/ruff_db/src/diagnostic/mod.rs |  3 ++
 3 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/crates/red_knot/src/main.rs b/crates/red_knot/src/main.rs
index 021a777611..a0e47b2d71 100644
--- a/crates/red_knot/src/main.rs
+++ b/crates/red_knot/src/main.rs
@@ -246,11 +246,16 @@ impl MainLoop {
                     // Spawn a new task that checks the project. This needs to be done in a separate thread
                     // to prevent blocking the main loop here.
                     rayon::spawn(move || {
-                        if let Ok(result) = db.check() {
-                            // Send the result back to the main loop for printing.
-                            sender
-                                .send(MainLoopMessage::CheckCompleted { result, revision })
-                                .unwrap();
+                        match db.check() {
+                            Ok(result) => {
+                                // Send the result back to the main loop for printing.
+                                sender
+                                    .send(MainLoopMessage::CheckCompleted { result, revision })
+                                    .unwrap();
+                            }
+                            Err(cancelled) => {
+                                tracing::debug!("Check has been cancelled: {cancelled:?}");
+                            }
                         }
                     });
                 }
diff --git a/crates/red_knot_project/src/lib.rs b/crates/red_knot_project/src/lib.rs
index de3207216f..48a5de4d48 100644
--- a/crates/red_knot_project/src/lib.rs
+++ b/crates/red_knot_project/src/lib.rs
@@ -11,7 +11,7 @@ use red_knot_python_semantic::register_lints;
 use red_knot_python_semantic::types::check_types;
 use ruff_db::diagnostic::{
     create_parse_diagnostic, create_unsupported_syntax_diagnostic, Annotation, Diagnostic,
-    DiagnosticId, Severity, Span,
+    DiagnosticId, Severity, Span, SubDiagnostic,
 };
 use ruff_db::files::File;
 use ruff_db::parsed::parsed_module;
@@ -20,8 +20,10 @@ 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::sync::Arc;
 use thiserror::Error;
+use tracing::error;
 
 pub mod combine;
 
@@ -471,7 +473,16 @@ fn check_file_impl(db: &dyn Db, file: File) -> Vec<Diagnostic> {
             .map(|error| create_unsupported_syntax_diagnostic(file, error)),
     );
 
-    diagnostics.extend(check_types(db.upcast(), file).into_iter().cloned());
+    {
+        let db = AssertUnwindSafe(db);
+        match catch(&**db, file, || check_types(db.upcast(), file)) {
+            Ok(Some(type_check_diagnostics)) => {
+                diagnostics.extend(type_check_diagnostics.into_iter().cloned());
+            }
+            Ok(None) => {}
+            Err(diagnostic) => diagnostics.push(diagnostic),
+        }
+    }
 
     diagnostics.sort_unstable_by_key(|diagnostic| {
         diagnostic
@@ -562,6 +573,45 @@ enum IOErrorKind {
     SourceText(#[from] SourceTextError),
 }
 
+fn catch<F, R>(db: &dyn Db, file: File, f: F) -> Result<Option<R>, Diagnostic>
+where
+    F: FnOnce() -> R + UnwindSafe,
+{
+    match 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>() {
+                Some((*s).to_string())
+            } else {
+                error.downcast_ref::<String>().cloned()
+            };
+
+            let message = if let Some(payload) = payload {
+                format!(
+                    "Panicked while checking `{file}`: `{payload}`",
+                    file = file.path(db)
+                )
+            } else {
+                format!("Panicked while checking `{file}`", file = { file.path(db) })
+            };
+
+            let mut diagnostic = Diagnostic::new(DiagnosticId::Panic, Severity::Fatal, message);
+            diagnostic.sub(SubDiagnostic::new(
+                Severity::Info,
+                "This indicates a bug in Red Knot.",
+            ));
+
+            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));
+
+            Err(diagnostic)
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::db::tests::TestDb;
diff --git a/crates/ruff_db/src/diagnostic/mod.rs b/crates/ruff_db/src/diagnostic/mod.rs
index 02b298afbb..ffcdf2a026 100644
--- a/crates/ruff_db/src/diagnostic/mod.rs
+++ b/crates/ruff_db/src/diagnostic/mod.rs
@@ -461,6 +461,8 @@ impl PartialEq<&str> for LintName {
 /// Uniquely identifies the kind of a diagnostic.
 #[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord, Hash)]
 pub enum DiagnosticId {
+    Panic,
+
     /// Some I/O operation failed
     Io,
 
@@ -521,6 +523,7 @@ impl DiagnosticId {
 
     pub fn as_str(&self) -> Result<&str, DiagnosticAsStrError> {
         Ok(match self {
+            DiagnosticId::Panic => "panic",
             DiagnosticId::Io => "io",
             DiagnosticId::InvalidSyntax => "invalid-syntax",
             DiagnosticId::Lint(name) => {