[ty] Abort printing diagnostics when pressing Ctrl+C (#22083)

crates/ruff_db/src/cancellation.rs

@@ -0,0 +1,51 @@
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+
+/// Signals a [`CancellationToken`] that it should be canceled.
+#[derive(Debug, Clone)]
+pub struct CancellationTokenSource {
+    cancelled: Arc<AtomicBool>,
+}
+
+impl Default for CancellationTokenSource {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl CancellationTokenSource {
+    pub fn new() -> Self {
+        Self {
+            cancelled: Arc::new(AtomicBool::new(false)),
+        }
+    }
+
+    pub fn is_cancellation_requested(&self) -> bool {
+        self.cancelled.load(std::sync::atomic::Ordering::Relaxed)
+    }
+
+    /// Creates a new token that uses this source.
+    pub fn token(&self) -> CancellationToken {
+        CancellationToken {
+            cancelled: self.cancelled.clone(),
+        }
+    }
+
+    /// Requests cancellation for operations using this token.
+    pub fn cancel(&self) {
+        self.cancelled
+            .store(true, std::sync::atomic::Ordering::Relaxed);
+    }
+}
+
+/// Token signals whether an operation should be canceled.
+#[derive(Debug, Clone)]
+pub struct CancellationToken {
+    cancelled: Arc<AtomicBool>,
+}
+
+impl CancellationToken {
+    pub fn is_cancelled(&self) -> bool {
+        self.cancelled.load(std::sync::atomic::Ordering::Relaxed)
+    }
+}

crates/ruff_db/src/diagnostic/mod.rs

@@ -11,6 +11,7 @@ pub use self::render::{
     ceil_char_boundary,
     github::{DisplayGithubDiagnostics, GithubRenderer},
 };
+use crate::cancellation::CancellationToken;
 use crate::{Db, files::File};
 
 mod render;
@@ -1312,6 +1313,8 @@ pub struct DisplayDiagnosticConfig {
     show_fix_diff: bool,
     /// The lowest applicability that should be shown when reporting diagnostics.
     fix_applicability: Applicability,
+
+    cancellation_token: Option<CancellationToken>,
 }
 
 impl DisplayDiagnosticConfig {
@@ -1385,6 +1388,20 @@ impl DisplayDiagnosticConfig {
     pub fn fix_applicability(&self) -> Applicability {
         self.fix_applicability
     }
+
+    pub fn with_cancellation_token(
+        mut self,
+        token: Option<CancellationToken>,
+    ) -> DisplayDiagnosticConfig {
+        self.cancellation_token = token;
+        self
+    }
+
+    pub fn is_canceled(&self) -> bool {
+        self.cancellation_token
+            .as_ref()
+            .is_some_and(|token| token.is_cancelled())
+    }
 }
 
 impl Default for DisplayDiagnosticConfig {
@@ -1398,6 +1415,7 @@ impl Default for DisplayDiagnosticConfig {
             show_fix_status: false,
             show_fix_diff: false,
             fix_applicability: Applicability::Safe,
+            cancellation_token: None,
         }
     }
 }

crates/ruff_db/src/diagnostic/render/concise.rs

@@ -28,6 +28,10 @@ impl<'a> ConciseRenderer<'a> {
 
         let sep = fmt_styled(":", stylesheet.separator);
         for diag in diagnostics {
+            if self.config.is_canceled() {
+                return Ok(());
+            }
+
             if let Some(span) = diag.primary_span() {
                 write!(
                     f,

crates/ruff_db/src/diagnostic/render/full.rs

@@ -53,6 +53,10 @@ impl<'a> FullRenderer<'a> {
             .hyperlink(stylesheet.hyperlink);
 
         for diag in diagnostics {
+            if self.config.is_canceled() {
+                return Ok(());
+            }
+
             let resolved = Resolved::new(self.resolver, diag, self.config);
             let renderable = resolved.to_renderable(self.config.context);
             for diag in renderable.diagnostics.iter() {

crates/ruff_db/src/lib.rs

@@ -12,6 +12,7 @@ use std::hash::BuildHasherDefault;
 use std::num::NonZeroUsize;
 use ty_static::EnvVars;
 
+pub mod cancellation;
 pub mod diagnostic;
 pub mod display;
 pub mod file_revision;

crates/ty/src/lib.rs

@@ -10,9 +10,9 @@ use ty_static::EnvVars;
 
 use std::fmt::Write;
 use std::process::{ExitCode, Termination};
+use std::sync::Mutex;
 
 use anyhow::Result;
-use std::sync::Mutex;
 
 use crate::args::{CheckCommand, Command, TerminalColor};
 use crate::logging::{VerbosityLevel, setup_tracing};
@@ -22,6 +22,7 @@ use clap::{CommandFactory, Parser};
 use colored::Colorize;
 use crossbeam::channel as crossbeam_channel;
 use rayon::ThreadPoolBuilder;
+use ruff_db::cancellation::{CancellationToken, CancellationTokenSource};
 use ruff_db::diagnostic::{
     Diagnostic, DiagnosticId, DisplayDiagnosticConfig, DisplayDiagnostics, Severity,
 };
@@ -227,6 +228,11 @@ struct MainLoop {
     printer: Printer,
 
     project_options_overrides: ProjectOptionsOverrides,
+
+    /// Cancellation token that gets set by Ctrl+C.
+    /// Used for long-running operations on the main thread. Operations on background threads
+    /// use Salsa's cancellation mechanism.
+    cancellation_token: CancellationToken,
 }
 
 impl MainLoop {
@@ -236,6 +242,9 @@ impl MainLoop {
     ) -> (Self, MainLoopCancellationToken) {
         let (sender, receiver) = crossbeam_channel::bounded(10);
 
+        let cancellation_token_source = CancellationTokenSource::new();
+        let cancellation_token = cancellation_token_source.token();
+
         (
             Self {
                 sender: sender.clone(),
@@ -243,8 +252,12 @@ impl MainLoop {
                 watcher: None,
                 project_options_overrides,
                 printer,
+                cancellation_token,
+            },
+            MainLoopCancellationToken {
+                sender,
+                source: cancellation_token_source,
             },
-            MainLoopCancellationToken { sender },
         )
     }
 
@@ -316,6 +329,7 @@ impl MainLoop {
                     let display_config = DisplayDiagnosticConfig::default()
                         .format(terminal_settings.output_format.into())
                         .color(colored::control::SHOULD_COLORIZE.should_colorize())
+                        .with_cancellation_token(Some(self.cancellation_token.clone()))
                         .show_fix_diff(true);
 
                     if check_revision == revision {
@@ -359,19 +373,21 @@ impl MainLoop {
                                 )?;
                             }
 
-                            if is_human_readable {
-                                writeln!(
-                                    self.printer.stream_for_failure_summary(),
-                                    "Found {} diagnostic{}",
-                                    diagnostics_count,
-                                    if diagnostics_count > 1 { "s" } else { "" }
-                                )?;
-                            }
+                            if !self.cancellation_token.is_cancelled() {
+                                if is_human_readable {
+                                    writeln!(
+                                        self.printer.stream_for_failure_summary(),
+                                        "Found {} diagnostic{}",
+                                        diagnostics_count,
+                                        if diagnostics_count > 1 { "s" } else { "" }
+                                    )?;
+                                }
 
-                            if exit_status.is_internal_error() {
-                                tracing::warn!(
-                                    "A fatal error occurred while checking some files. Not all project files were analyzed. See the diagnostics list above for details."
-                                );
+                                if exit_status.is_internal_error() {
+                                    tracing::warn!(
+                                        "A fatal error occurred while checking some files. Not all project files were analyzed. See the diagnostics list above for details."
+                                    );
+                                }
                             }
 
                             if self.watcher.is_none() {
@@ -498,10 +514,12 @@ impl ty_project::ProgressReporter for IndicatifReporter {
 #[derive(Debug)]
 struct MainLoopCancellationToken {
     sender: crossbeam_channel::Sender<MainLoopMessage>,
+    source: CancellationTokenSource,
 }
 
 impl MainLoopCancellationToken {
     fn stop(self) {
+        self.source.cancel();
         self.sender.send(MainLoopMessage::Exit).unwrap();
     }
 }