Allow disk-based diagnostic progress begin/end events to interleave

When multiple saves occur, we can have multiple start events followed by multiple end events. We don't want to update our project diagnostics view until all pending progress is finished.

Co-Authored-By: Antonio Scandurra <me@as-cii.com>

crates/editor/src/display_map.rs

@@ -446,10 +446,11 @@ impl ToDisplayPoint for Anchor {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::{movement, test::*};
+    use crate::movement;
-    use gpui::{color::Color, elements::*, MutableAppContext};
+    use gpui::{color::Color, elements::*, test::observe, MutableAppContext};
     use language::{Buffer, Language, LanguageConfig, RandomCharIter, SelectionGoal};
     use rand::{prelude::*, Rng};
+    use smol::stream::StreamExt;
     use std::{env, sync::Arc};
     use theme::SyntaxTheme;
     use util::test::sample_text;
@@ -493,7 +494,7 @@ mod tests {
         let map = cx.add_model(|cx| {
             DisplayMap::new(buffer.clone(), tab_size, font_id, font_size, wrap_width, cx)
         });
-        let (_observer, notifications) = Observer::new(&map, &mut cx);
+        let mut notifications = observe(&map, &mut cx);
         let mut fold_count = 0;
         let mut blocks = Vec::new();
 
@@ -589,7 +590,7 @@ mod tests {
             }
 
             if map.read_with(&cx, |map, cx| map.is_rewrapping(cx)) {
-                notifications.recv().await.unwrap();
+                notifications.next().await.unwrap();
             }
 
             let snapshot = map.update(&mut cx, |map, cx| map.snapshot(cx));

crates/editor/src/display_map/wrap_map.rs

@@ -1014,11 +1014,12 @@ mod tests {
     use super::*;
     use crate::{
         display_map::{fold_map::FoldMap, tab_map::TabMap},
-        test::Observer,
         MultiBuffer,
     };
+    use gpui::test::observe;
     use language::RandomCharIter;
     use rand::prelude::*;
+    use smol::stream::StreamExt;
     use std::{cmp, env};
     use text::Rope;
 
@@ -1072,10 +1073,10 @@ mod tests {
 
         let (wrap_map, _) =
             cx.update(|cx| WrapMap::new(tabs_snapshot.clone(), font_id, font_size, wrap_width, cx));
-        let (_observer, notifications) = Observer::new(&wrap_map, &mut cx);
+        let mut notifications = observe(&wrap_map, &mut cx);
 
         if wrap_map.read_with(&cx, |map, _| map.is_rewrapping()) {
-            notifications.recv().await.unwrap();
+            notifications.next().await.unwrap();
         }
 
         let (initial_snapshot, _) = wrap_map.update(&mut cx, |map, cx| {
@@ -1148,7 +1149,7 @@ mod tests {
             if wrap_map.read_with(&cx, |map, _| map.is_rewrapping()) && rng.gen_bool(0.4) {
                 log::info!("Waiting for wrapping to finish");
                 while wrap_map.read_with(&cx, |map, _| map.is_rewrapping()) {
-                    notifications.recv().await.unwrap();
+                    notifications.next().await.unwrap();
                 }
                 wrap_map.read_with(&cx, |map, _| assert!(map.pending_edits.is_empty()));
             }
@@ -1236,7 +1237,7 @@ mod tests {
         if wrap_map.read_with(&cx, |map, _| map.is_rewrapping()) {
             log::info!("Waiting for wrapping to finish");
             while wrap_map.read_with(&cx, |map, _| map.is_rewrapping()) {
-                notifications.recv().await.unwrap();
+                notifications.next().await.unwrap();
             }
         }
         wrap_map.read_with(&cx, |map, _| assert!(map.pending_edits.is_empty()));

crates/editor/src/test.rs

@@ -1,33 +1,6 @@
-use gpui::{Entity, ModelHandle};
-use smol::channel;
-use std::marker::PhantomData;
-
 #[cfg(test)]
 #[ctor::ctor]
 fn init_logger() {
     // std::env::set_var("RUST_LOG", "info");
     env_logger::init();
 }
-
-pub struct Observer<T>(PhantomData<T>);
-
-impl<T: 'static> Entity for Observer<T> {
-    type Event = ();
-}
-
-impl<T: Entity> Observer<T> {
-    pub fn new(
-        handle: &ModelHandle<T>,
-        cx: &mut gpui::TestAppContext,
-    ) -> (ModelHandle<Self>, channel::Receiver<()>) {
-        let (notify_tx, notify_rx) = channel::unbounded();
-        let observer = cx.add_model(|cx| {
-            cx.observe(handle, move |_, _, _| {
-                let _ = notify_tx.try_send(());
-            })
-            .detach();
-            Observer(PhantomData)
-        });
-        (observer, notify_rx)
-    }
-}

crates/gpui/src/app.rs

@@ -992,7 +992,7 @@ impl MutableAppContext {
         })
     }
 
-    fn observe<E, H, F>(&mut self, handle: &H, mut callback: F) -> Subscription
+    pub fn observe<E, H, F>(&mut self, handle: &H, mut callback: F) -> Subscription
     where
         E: Entity,
         E::Event: 'static,

crates/gpui/src/test.rs

@@ -7,7 +7,13 @@ use std::{
     },
 };
 
-use crate::{executor, platform, FontCache, MutableAppContext, Platform, TestAppContext};
+use futures::StreamExt;
+use smol::channel;
+
+use crate::{
+    executor, platform, Entity, FontCache, Handle, MutableAppContext, Platform, Subscription,
+    TestAppContext,
+};
 
 #[cfg(test)]
 #[ctor::ctor]
@@ -87,3 +93,47 @@ pub fn run_test(
         }
     }
 }
+
+pub struct Observation<T> {
+    rx: channel::Receiver<T>,
+    _subscription: Subscription,
+}
+
+impl<T> futures::Stream for Observation<T> {
+    type Item = T;
+
+    fn poll_next(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        self.rx.poll_next_unpin(cx)
+    }
+}
+
+pub fn observe<T: Entity>(entity: &impl Handle<T>, cx: &mut TestAppContext) -> Observation<()> {
+    let (tx, rx) = smol::channel::unbounded();
+    let _subscription = cx.update(|cx| {
+        cx.observe(entity, move |_, _| {
+            let _ = smol::block_on(tx.send(()));
+        })
+    });
+
+    Observation { rx, _subscription }
+}
+
+pub fn subscribe<T: Entity>(
+    entity: &impl Handle<T>,
+    cx: &mut TestAppContext,
+) -> Observation<T::Event>
+where
+    T::Event: Clone,
+{
+    let (tx, rx) = smol::channel::unbounded();
+    let _subscription = cx.update(|cx| {
+        cx.subscribe(entity, move |_, event, _| {
+            let _ = smol::block_on(tx.send(event.clone()));
+        })
+    });
+
+    Observation { rx, _subscription }
+}

crates/language/src/language.rs

@@ -237,6 +237,7 @@ impl LanguageServerConfig {
         (
             Self {
                 fake_server: Some((server, started)),
+                disk_based_diagnostics_progress_token: Some("fakeServer/check".to_string()),
                 ..Default::default()
             },
             fake,

crates/lsp/src/lsp.rs

@@ -514,6 +514,22 @@ impl FakeLanguageServer {
         notification.params
     }
 
+    pub async fn start_progress(&mut self, token: impl Into<String>) {
+        self.notify::<notification::Progress>(ProgressParams {
+            token: NumberOrString::String(token.into()),
+            value: ProgressParamsValue::WorkDone(WorkDoneProgress::Begin(Default::default())),
+        })
+        .await;
+    }
+
+    pub async fn end_progress(&mut self, token: impl Into<String>) {
+        self.notify::<notification::Progress>(ProgressParams {
+            token: NumberOrString::String(token.into()),
+            value: ProgressParamsValue::WorkDone(WorkDoneProgress::End(Default::default())),
+        })
+        .await;
+    }
+
     async fn send(&mut self, message: Vec<u8>) {
         self.stdout
             .write_all(CONTENT_LEN_HEADER.as_bytes())

crates/project/src/worktree.rs

@@ -67,7 +67,7 @@ pub enum Worktree {
     Remote(RemoteWorktree),
 }
 
-#[derive(Debug)]
+#[derive(Clone, Debug, Eq, PartialEq)]
 pub enum Event {
     DiskBasedDiagnosticsUpdated,
     DiagnosticsUpdated(Arc<Path>),
@@ -1120,6 +1120,7 @@ impl LocalWorktree {
             })
             .detach();
 
+            let mut pending_disk_based_diagnostics: i32 = 0;
             language_server
                 .on_notification::<lsp::notification::Progress, _>(move |params| {
                     let token = match params.token {
@@ -1130,8 +1131,15 @@ impl LocalWorktree {
                     if token == disk_based_diagnostics_progress_token {
                         match params.value {
                             lsp::ProgressParamsValue::WorkDone(progress) => match progress {
+                                lsp::WorkDoneProgress::Begin(_) => {
+                                    pending_disk_based_diagnostics += 1;
+                                }
                                 lsp::WorkDoneProgress::End(_) => {
-                                    smol::block_on(disk_based_diagnostics_done_tx.send(())).ok();
+                                    pending_disk_based_diagnostics -= 1;
+                                    if pending_disk_based_diagnostics == 0 {
+                                        smol::block_on(disk_based_diagnostics_done_tx.send(()))
+                                            .ok();
+                                    }
                                 }
                                 _ => {}
                             },
@@ -3107,6 +3115,7 @@ mod tests {
     use anyhow::Result;
     use client::test::{FakeHttpClient, FakeServer};
     use fs::RealFs;
+    use gpui::test::subscribe;
     use language::{tree_sitter_rust, DiagnosticEntry, LanguageServerConfig};
     use language::{Diagnostic, LanguageConfig};
     use lsp::Url;
@@ -3756,6 +3765,10 @@ mod tests {
     async fn test_language_server_diagnostics(mut cx: gpui::TestAppContext) {
         let (language_server_config, mut fake_server) =
             LanguageServerConfig::fake(cx.background()).await;
+        let progress_token = language_server_config
+            .disk_based_diagnostics_progress_token
+            .clone()
+            .unwrap();
         let mut languages = LanguageRegistry::new();
         languages.add(Arc::new(Language::new(
             LanguageConfig {
@@ -3795,6 +3808,13 @@ mod tests {
             .await
             .unwrap();
 
+        let mut events = subscribe(&tree, &mut cx);
+
+        fake_server.start_progress(&progress_token).await;
+        fake_server.start_progress(&progress_token).await;
+        fake_server.end_progress(&progress_token).await;
+        fake_server.start_progress(&progress_token).await;
+
         fake_server
             .notify::<lsp::notification::PublishDiagnostics>(lsp::PublishDiagnosticsParams {
                 uri: Url::from_file_path(dir.path().join("a.rs")).unwrap(),
@@ -3808,6 +3828,18 @@ mod tests {
             })
             .await;
 
+        let event = events.next().await.unwrap();
+        assert_eq!(
+            event,
+            Event::DiagnosticsUpdated(Arc::from(Path::new("a.rs")))
+        );
+
+        fake_server.end_progress(&progress_token).await;
+        fake_server.end_progress(&progress_token).await;
+
+        let event = events.next().await.unwrap();
+        assert_eq!(event, Event::DiskBasedDiagnosticsUpdated);
+
         let buffer = tree
             .update(&mut cx, |tree, cx| tree.open_buffer("a.rs", cx))
             .await