Merge pull request #20419 from ShoyuVanilla/flyck-gen

internal: Make flycheck generational

4 files changed, 148 insertions, 38 deletions

diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/diagnostics.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/diagnostics.rs
index 438a2a0ba1e..cd7c632d105 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/diagnostics.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/diagnostics.rs
@@ -26,6 +26,19 @@ pub struct DiagnosticsMapConfig {
 
 pub(crate) type DiagnosticsGeneration = usize;
 
+#[derive(Debug, Clone)]
+pub(crate) struct WorkspaceFlycheckDiagnostic {
+    pub(crate) generation: DiagnosticsGeneration,
+    pub(crate) per_package:
+        FxHashMap<Option<Arc<PackageId>>, FxHashMap<FileId, Vec<lsp_types::Diagnostic>>>,
+}
+
+impl WorkspaceFlycheckDiagnostic {
+    fn new(generation: DiagnosticsGeneration) -> Self {
+        WorkspaceFlycheckDiagnostic { generation, per_package: Default::default() }
+    }
+}
+
 #[derive(Debug, Default, Clone)]
 pub(crate) struct DiagnosticCollection {
     // FIXME: should be FxHashMap<FileId, Vec<ra_id::Diagnostic>>
@@ -33,9 +46,7 @@ pub(crate) struct DiagnosticCollection {
         FxHashMap<FileId, (DiagnosticsGeneration, Vec<lsp_types::Diagnostic>)>,
     pub(crate) native_semantic:
         FxHashMap<FileId, (DiagnosticsGeneration, Vec<lsp_types::Diagnostic>)>,
-    // FIXME: should be Vec<flycheck::Diagnostic>
-    pub(crate) check:
-        Vec<FxHashMap<Option<Arc<PackageId>>, FxHashMap<FileId, Vec<lsp_types::Diagnostic>>>>,
+    pub(crate) check: Vec<WorkspaceFlycheckDiagnostic>,
     pub(crate) check_fixes: CheckFixes,
     changes: FxHashSet<FileId>,
     /// Counter for supplying a new generation number for diagnostics.
@@ -57,7 +68,7 @@ impl DiagnosticCollection {
         let Some(check) = self.check.get_mut(flycheck_id) else {
             return;
         };
-        self.changes.extend(check.drain().flat_map(|(_, v)| v.into_keys()));
+        self.changes.extend(check.per_package.drain().flat_map(|(_, v)| v.into_keys()));
         if let Some(fixes) = Arc::make_mut(&mut self.check_fixes).get_mut(flycheck_id) {
             fixes.clear();
         }
@@ -66,7 +77,9 @@ impl DiagnosticCollection {
     pub(crate) fn clear_check_all(&mut self) {
         Arc::make_mut(&mut self.check_fixes).clear();
         self.changes.extend(
-            self.check.iter_mut().flat_map(|it| it.drain().flat_map(|(_, v)| v.into_keys())),
+            self.check
+                .iter_mut()
+                .flat_map(|it| it.per_package.drain().flat_map(|(_, v)| v.into_keys())),
         )
     }
 
@@ -79,7 +92,7 @@ impl DiagnosticCollection {
             return;
         };
         let package_id = Some(package_id);
-        if let Some(checks) = check.remove(&package_id) {
+        if let Some(checks) = check.per_package.remove(&package_id) {
             self.changes.extend(checks.into_keys());
         }
         if let Some(fixes) = Arc::make_mut(&mut self.check_fixes).get_mut(flycheck_id) {
@@ -87,6 +100,16 @@ impl DiagnosticCollection {
         }
     }
 
+    pub(crate) fn clear_check_older_than(
+        &mut self,
+        flycheck_id: usize,
+        generation: DiagnosticsGeneration,
+    ) {
+        if self.check[flycheck_id].generation < generation {
+            self.clear_check(flycheck_id);
+        }
+    }
+
     pub(crate) fn clear_native_for(&mut self, file_id: FileId) {
         self.native_syntax.remove(&file_id);
         self.native_semantic.remove(&file_id);
@@ -96,15 +119,23 @@ impl DiagnosticCollection {
     pub(crate) fn add_check_diagnostic(
         &mut self,
         flycheck_id: usize,
+        generation: DiagnosticsGeneration,
         package_id: &Option<Arc<PackageId>>,
         file_id: FileId,
         diagnostic: lsp_types::Diagnostic,
         fix: Option<Box<Fix>>,
     ) {
         if self.check.len() <= flycheck_id {
-            self.check.resize_with(flycheck_id + 1, Default::default);
+            self.check
+                .resize_with(flycheck_id + 1, || WorkspaceFlycheckDiagnostic::new(generation));
+        }
+
+        // Getting message from old generation. Might happen in restarting checks.
+        if self.check[flycheck_id].generation > generation {
+            return;
         }
         let diagnostics = self.check[flycheck_id]
+            .per_package
             .entry(package_id.clone())
             .or_default()
             .entry(file_id)
@@ -177,7 +208,7 @@ impl DiagnosticCollection {
         let check = self
             .check
             .iter()
-            .flat_map(|it| it.values())
+            .flat_map(|it| it.per_package.values())
             .filter_map(move |it| it.get(&file_id))
             .flatten();
         native_syntax.chain(native_semantic).chain(check)
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/flycheck.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/flycheck.rs
index e4e0bcdc1cd..233bedec690 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/flycheck.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/flycheck.rs
@@ -1,7 +1,12 @@
 //! Flycheck provides the functionality needed to run `cargo check` to provide
 //! LSP diagnostics based on the output of the command.
 
-use std::{fmt, io, process::Command, time::Duration};
+use std::{
+    fmt, io,
+    process::Command,
+    sync::atomic::{AtomicUsize, Ordering},
+    time::Duration,
+};
 
 use cargo_metadata::PackageId;
 use crossbeam_channel::{Receiver, Sender, select_biased, unbounded};
@@ -18,7 +23,10 @@ pub(crate) use cargo_metadata::diagnostic::{
 use toolchain::Tool;
 use triomphe::Arc;
 
-use crate::command::{CargoParser, CommandHandle};
+use crate::{
+    command::{CargoParser, CommandHandle},
+    diagnostics::DiagnosticsGeneration,
+};
 
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub(crate) enum InvocationStrategy {
@@ -138,36 +146,54 @@ pub(crate) struct FlycheckHandle {
     sender: Sender<StateChange>,
     _thread: stdx::thread::JoinHandle,
     id: usize,
+    generation: AtomicUsize,
 }
 
 impl FlycheckHandle {
     pub(crate) fn spawn(
         id: usize,
+        generation: DiagnosticsGeneration,
         sender: Sender<FlycheckMessage>,
         config: FlycheckConfig,
         sysroot_root: Option<AbsPathBuf>,
         workspace_root: AbsPathBuf,
         manifest_path: Option<AbsPathBuf>,
     ) -> FlycheckHandle {
-        let actor =
-            FlycheckActor::new(id, sender, config, sysroot_root, workspace_root, manifest_path);
+        let actor = FlycheckActor::new(
+            id,
+            generation,
+            sender,
+            config,
+            sysroot_root,
+            workspace_root,
+            manifest_path,
+        );
         let (sender, receiver) = unbounded::<StateChange>();
         let thread =
             stdx::thread::Builder::new(stdx::thread::ThreadIntent::Worker, format!("Flycheck{id}"))
                 .spawn(move || actor.run(receiver))
                 .expect("failed to spawn thread");
-        FlycheckHandle { id, sender, _thread: thread }
+        FlycheckHandle { id, generation: generation.into(), sender, _thread: thread }
     }
 
     /// Schedule a re-start of the cargo check worker to do a workspace wide check.
     pub(crate) fn restart_workspace(&self, saved_file: Option<AbsPathBuf>) {
-        self.sender.send(StateChange::Restart { package: None, saved_file, target: None }).unwrap();
+        let generation = self.generation.fetch_add(1, Ordering::Relaxed) + 1;
+        self.sender
+            .send(StateChange::Restart { generation, package: None, saved_file, target: None })
+            .unwrap();
     }
 
     /// Schedule a re-start of the cargo check worker to do a package wide check.
     pub(crate) fn restart_for_package(&self, package: String, target: Option<Target>) {
+        let generation = self.generation.fetch_add(1, Ordering::Relaxed) + 1;
         self.sender
-            .send(StateChange::Restart { package: Some(package), saved_file: None, target })
+            .send(StateChange::Restart {
+                generation,
+                package: Some(package),
+                saved_file: None,
+                target,
+            })
             .unwrap();
     }
 
@@ -179,23 +205,31 @@ impl FlycheckHandle {
     pub(crate) fn id(&self) -> usize {
         self.id
     }
+
+    pub(crate) fn generation(&self) -> DiagnosticsGeneration {
+        self.generation.load(Ordering::Relaxed)
+    }
+}
+
+#[derive(Debug)]
+pub(crate) enum ClearDiagnosticsKind {
+    All,
+    OlderThan(DiagnosticsGeneration),
+    Package(Arc<PackageId>),
 }
 
 pub(crate) enum FlycheckMessage {
     /// Request adding a diagnostic with fixes included to a file
     AddDiagnostic {
         id: usize,
+        generation: DiagnosticsGeneration,
         workspace_root: Arc<AbsPathBuf>,
         diagnostic: Diagnostic,
         package_id: Option<Arc<PackageId>>,
     },
 
     /// Request clearing all outdated diagnostics.
-    ClearDiagnostics {
-        id: usize,
-        /// The package whose diagnostics to clear, or if unspecified, all diagnostics.
-        package_id: Option<Arc<PackageId>>,
-    },
+    ClearDiagnostics { id: usize, kind: ClearDiagnosticsKind },
 
     /// Request check progress notification to client
     Progress {
@@ -208,18 +242,23 @@ pub(crate) enum FlycheckMessage {
 impl fmt::Debug for FlycheckMessage {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
-            FlycheckMessage::AddDiagnostic { id, workspace_root, diagnostic, package_id } => f
+            FlycheckMessage::AddDiagnostic {
+                id,
+                generation,
+                workspace_root,
+                diagnostic,
+                package_id,
+            } => f
                 .debug_struct("AddDiagnostic")
                 .field("id", id)
+                .field("generation", generation)
                 .field("workspace_root", workspace_root)
                 .field("package_id", package_id)
                 .field("diagnostic_code", &diagnostic.code.as_ref().map(|it| &it.code))
                 .finish(),
-            FlycheckMessage::ClearDiagnostics { id, package_id } => f
-                .debug_struct("ClearDiagnostics")
-                .field("id", id)
-                .field("package_id", package_id)
-                .finish(),
+            FlycheckMessage::ClearDiagnostics { id, kind } => {
+                f.debug_struct("ClearDiagnostics").field("id", id).field("kind", kind).finish()
+            }
             FlycheckMessage::Progress { id, progress } => {
                 f.debug_struct("Progress").field("id", id).field("progress", progress).finish()
             }
@@ -237,7 +276,12 @@ pub(crate) enum Progress {
 }
 
 enum StateChange {
-    Restart { package: Option<String>, saved_file: Option<AbsPathBuf>, target: Option<Target> },
+    Restart {
+        generation: DiagnosticsGeneration,
+        package: Option<String>,
+        saved_file: Option<AbsPathBuf>,
+        target: Option<Target>,
+    },
     Cancel,
 }
 
@@ -246,6 +290,7 @@ struct FlycheckActor {
     /// The workspace id of this flycheck instance.
     id: usize,
 
+    generation: DiagnosticsGeneration,
     sender: Sender<FlycheckMessage>,
     config: FlycheckConfig,
     manifest_path: Option<AbsPathBuf>,
@@ -283,6 +328,7 @@ pub(crate) const SAVED_FILE_PLACEHOLDER: &str = "$saved_file";
 impl FlycheckActor {
     fn new(
         id: usize,
+        generation: DiagnosticsGeneration,
         sender: Sender<FlycheckMessage>,
         config: FlycheckConfig,
         sysroot_root: Option<AbsPathBuf>,
@@ -292,6 +338,7 @@ impl FlycheckActor {
         tracing::info!(%id, ?workspace_root, "Spawning flycheck");
         FlycheckActor {
             id,
+            generation,
             sender,
             config,
             sysroot_root,
@@ -327,7 +374,12 @@ impl FlycheckActor {
                     tracing::debug!(flycheck_id = self.id, "flycheck cancelled");
                     self.cancel_check_process();
                 }
-                Event::RequestStateChange(StateChange::Restart { package, saved_file, target }) => {
+                Event::RequestStateChange(StateChange::Restart {
+                    generation,
+                    package,
+                    saved_file,
+                    target,
+                }) => {
                     // Cancel the previously spawned process
                     self.cancel_check_process();
                     while let Ok(restart) = inbox.recv_timeout(Duration::from_millis(50)) {
@@ -337,6 +389,8 @@ impl FlycheckActor {
                         }
                     }
 
+                    self.generation = generation;
+
                     let Some(command) =
                         self.check_command(package.as_deref(), saved_file.as_deref(), target)
                     else {
@@ -383,7 +437,16 @@ impl FlycheckActor {
                         // Clear everything for good measure
                         self.send(FlycheckMessage::ClearDiagnostics {
                             id: self.id,
-                            package_id: None,
+                            kind: ClearDiagnosticsKind::All,
+                        });
+                    } else if res.is_ok() {
+                        // We clear diagnostics for packages on
+                        // `[CargoCheckMessage::CompilerArtifact]` but there seem to be setups where
+                        // cargo may not report an artifact to our runner at all. To handle such
+                        // cases, clear stale diagnostics when flycheck completes successfully.
+                        self.send(FlycheckMessage::ClearDiagnostics {
+                            id: self.id,
+                            kind: ClearDiagnosticsKind::OlderThan(self.generation),
                         });
                     }
                     self.clear_diagnostics_state();
@@ -412,7 +475,7 @@ impl FlycheckActor {
                             );
                             self.send(FlycheckMessage::ClearDiagnostics {
                                 id: self.id,
-                                package_id: Some(package_id),
+                                kind: ClearDiagnosticsKind::Package(package_id),
                             });
                         }
                     }
@@ -435,7 +498,7 @@ impl FlycheckActor {
                                 );
                                 self.send(FlycheckMessage::ClearDiagnostics {
                                     id: self.id,
-                                    package_id: Some(package_id.clone()),
+                                    kind: ClearDiagnosticsKind::Package(package_id.clone()),
                                 });
                             }
                         } else if self.diagnostics_received
@@ -444,11 +507,12 @@ impl FlycheckActor {
                             self.diagnostics_received = DiagnosticsReceived::YesAndClearedForAll;
                             self.send(FlycheckMessage::ClearDiagnostics {
                                 id: self.id,
-                                package_id: None,
+                                kind: ClearDiagnosticsKind::All,
                             });
                         }
                         self.send(FlycheckMessage::AddDiagnostic {
                             id: self.id,
+                            generation: self.generation,
                             package_id,
                             workspace_root: self.root.clone(),
                             diagnostic,
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs
index f5932d8cff0..c6762f31832 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs
@@ -20,7 +20,7 @@ use crate::{
     config::Config,
     diagnostics::{DiagnosticsGeneration, NativeDiagnosticsFetchKind, fetch_native_diagnostics},
     discover::{DiscoverArgument, DiscoverCommand, DiscoverProjectMessage},
-    flycheck::{self, FlycheckMessage},
+    flycheck::{self, ClearDiagnosticsKind, FlycheckMessage},
     global_state::{
         FetchBuildDataResponse, FetchWorkspaceRequest, FetchWorkspaceResponse, GlobalState,
         file_id_to_url, url_to_file_id,
@@ -1008,7 +1008,13 @@ impl GlobalState {
 
     fn handle_flycheck_msg(&mut self, message: FlycheckMessage) {
         match message {
-            FlycheckMessage::AddDiagnostic { id, workspace_root, diagnostic, package_id } => {
+            FlycheckMessage::AddDiagnostic {
+                id,
+                generation,
+                workspace_root,
+                diagnostic,
+                package_id,
+            } => {
                 let snap = self.snapshot();
                 let diagnostics = crate::diagnostics::to_proto::map_rust_diagnostic_to_lsp(
                     &self.config.diagnostics_map(None),
@@ -1020,6 +1026,7 @@ impl GlobalState {
                     match url_to_file_id(&self.vfs.read().0, &diag.url) {
                         Ok(Some(file_id)) => self.diagnostics.add_check_diagnostic(
                             id,
+                            generation,
                             &package_id,
                             file_id,
                             diag.diagnostic,
@@ -1035,12 +1042,17 @@ impl GlobalState {
                     };
                 }
             }
-            FlycheckMessage::ClearDiagnostics { id, package_id: None } => {
+            FlycheckMessage::ClearDiagnostics { id, kind: ClearDiagnosticsKind::All } => {
                 self.diagnostics.clear_check(id)
             }
-            FlycheckMessage::ClearDiagnostics { id, package_id: Some(package_id) } => {
-                self.diagnostics.clear_check_for_package(id, package_id)
-            }
+            FlycheckMessage::ClearDiagnostics {
+                id,
+                kind: ClearDiagnosticsKind::OlderThan(generation),
+            } => self.diagnostics.clear_check_older_than(id, generation),
+            FlycheckMessage::ClearDiagnostics {
+                id,
+                kind: ClearDiagnosticsKind::Package(package_id),
+            } => self.diagnostics.clear_check_for_package(id, package_id),
             FlycheckMessage::Progress { id, progress } => {
                 let (state, message) = match progress {
                     flycheck::Progress::DidStart => (Progress::Begin, None),
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/reload.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/reload.rs
index aa38aa72d44..f8f29eee126 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/reload.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/reload.rs
@@ -855,11 +855,13 @@ impl GlobalState {
                 invocation_strategy.clone()
             }
         };
+        let next_gen = self.flycheck.iter().map(|f| f.generation() + 1).max().unwrap_or_default();
 
         self.flycheck = match invocation_strategy {
             crate::flycheck::InvocationStrategy::Once => {
                 vec![FlycheckHandle::spawn(
                     0,
+                    next_gen,
                     sender,
                     config,
                     None,
@@ -898,6 +900,7 @@ impl GlobalState {
                     .map(|(id, (root, manifest_path), sysroot_root)| {
                         FlycheckHandle::spawn(
                             id,
+                            next_gen,
                             sender.clone(),
                             config.clone(),
                             sysroot_root,