diff --git a/crates/ruff_db/src/cancellation.rs b/crates/ruff_db/src/cancellation.rs new file mode 100644 index 0000000000..172f1d6d2f --- /dev/null +++ b/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, +} + +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, +} + +impl CancellationToken { + pub fn is_cancelled(&self) -> bool { + self.cancelled.load(std::sync::atomic::Ordering::Relaxed) + } +} diff --git a/crates/ruff_db/src/diagnostic/mod.rs b/crates/ruff_db/src/diagnostic/mod.rs index e966cdd208..328bf70fd6 100644 --- a/crates/ruff_db/src/diagnostic/mod.rs +++ b/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, } 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, + ) -> 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, } } } diff --git a/crates/ruff_db/src/diagnostic/render/concise.rs b/crates/ruff_db/src/diagnostic/render/concise.rs index 0d8c477f65..95a9e114e6 100644 --- a/crates/ruff_db/src/diagnostic/render/concise.rs +++ b/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, diff --git a/crates/ruff_db/src/diagnostic/render/full.rs b/crates/ruff_db/src/diagnostic/render/full.rs index bbc71f93d0..986077fac1 100644 --- a/crates/ruff_db/src/diagnostic/render/full.rs +++ b/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() { diff --git a/crates/ruff_db/src/lib.rs b/crates/ruff_db/src/lib.rs index efc5b2c974..d9a661a518 100644 --- a/crates/ruff_db/src/lib.rs +++ b/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; diff --git a/crates/ty/src/lib.rs b/crates/ty/src/lib.rs index 2934eb00a5..ae13949712 100644 --- a/crates/ty/src/lib.rs +++ b/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, + source: CancellationTokenSource, } impl MainLoopCancellationToken { fn stop(self) { + self.source.cancel(); self.sender.send(MainLoopMessage::Exit).unwrap(); } }