Clear old diagnostics when restarting a language server (#2509)

Max Brunsfeld created 

Previously, restarting a language server did not clear out the
diagnostics that were published by that server. Those diagnostics would
stick around forever.

Release Notes:

* Fixed a bug where restarting a language server would sometimes leave
buffers with stale diagnostics associated with that server.

Change summary

crates/language/src/buffer.rs         | 15 +++-
crates/language/src/diagnostic_set.rs |  4 +
crates/project/src/project.rs         | 17 +++++
crates/project/src/project_tests.rs   | 89 +++++++++++++++++++++++++++++
crates/project/src/worktree.rs        | 39 ++++++++++++
5 files changed, 160 insertions(+), 4 deletions(-)

Detailed changes

crates/language/src/buffer.rs 🔗 

@@ -1644,10 +1644,17 @@ impl Buffer {
         cx: &mut ModelContext<Self>,
     ) {
         if lamport_timestamp > self.diagnostics_timestamp {
-            match self.diagnostics.binary_search_by_key(&server_id, |e| e.0) {
-                Err(ix) => self.diagnostics.insert(ix, (server_id, diagnostics)),
-                Ok(ix) => self.diagnostics[ix].1 = diagnostics,
-            };
+            let ix = self.diagnostics.binary_search_by_key(&server_id, |e| e.0);
+            if diagnostics.len() == 0 {
+                if let Ok(ix) = ix {
+                    self.diagnostics.remove(ix);
+                }
+            } else {
+                match ix {
+                    Err(ix) => self.diagnostics.insert(ix, (server_id, diagnostics)),
+                    Ok(ix) => self.diagnostics[ix].1 = diagnostics,
+                };
+            }
             self.diagnostics_timestamp = lamport_timestamp;
             self.diagnostics_update_count += 1;
             self.text.lamport_clock.observe(lamport_timestamp);

crates/language/src/diagnostic_set.rs 🔗 

@@ -80,6 +80,10 @@ impl DiagnosticSet {
         }
     }
 
+    pub fn len(&self) -> usize {
+        self.diagnostics.summary().count
+    }
+
     pub fn iter(&self) -> impl Iterator<Item = &DiagnosticEntry<Anchor>> {
         self.diagnostics.iter()
     }

crates/project/src/project.rs 🔗 

@@ -2565,6 +2565,23 @@ impl Project {
                 }
             }
 
+            for buffer in self.opened_buffers.values() {
+                if let Some(buffer) = buffer.upgrade(cx) {
+                    buffer.update(cx, |buffer, cx| {
+                        buffer.update_diagnostics(server_id, Default::default(), cx);
+                    });
+                }
+            }
+            for worktree in &self.worktrees {
+                if let Some(worktree) = worktree.upgrade(cx) {
+                    worktree.update(cx, |worktree, cx| {
+                        if let Some(worktree) = worktree.as_local_mut() {
+                            worktree.clear_diagnostics_for_language_server(server_id, cx);
+                        }
+                    });
+                }
+            }
+
             self.language_server_statuses.remove(&server_id);
             cx.notify();

crates/project/src/project_tests.rs 🔗 

@@ -926,6 +926,95 @@ async fn test_restarting_server_with_diagnostics_running(cx: &mut gpui::TestAppC
     });
 }
 
+#[gpui::test]
+async fn test_restarting_server_with_diagnostics_published(cx: &mut gpui::TestAppContext) {
+    init_test(cx);
+
+    let mut language = Language::new(
+        LanguageConfig {
+            path_suffixes: vec!["rs".to_string()],
+            ..Default::default()
+        },
+        None,
+    );
+    let mut fake_servers = language
+        .set_fake_lsp_adapter(Arc::new(FakeLspAdapter {
+            ..Default::default()
+        }))
+        .await;
+
+    let fs = FakeFs::new(cx.background());
+    fs.insert_tree("/dir", json!({ "a.rs": "x" })).await;
+
+    let project = Project::test(fs, ["/dir".as_ref()], cx).await;
+    project.update(cx, |project, _| project.languages.add(Arc::new(language)));
+
+    let buffer = project
+        .update(cx, |project, cx| project.open_local_buffer("/dir/a.rs", cx))
+        .await
+        .unwrap();
+
+    // Publish diagnostics
+    let fake_server = fake_servers.next().await.unwrap();
+    fake_server.notify::<lsp::notification::PublishDiagnostics>(lsp::PublishDiagnosticsParams {
+        uri: Url::from_file_path("/dir/a.rs").unwrap(),
+        version: None,
+        diagnostics: vec![lsp::Diagnostic {
+            range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 0)),
+            severity: Some(lsp::DiagnosticSeverity::ERROR),
+            message: "the message".to_string(),
+            ..Default::default()
+        }],
+    });
+
+    cx.foreground().run_until_parked();
+    buffer.read_with(cx, |buffer, _| {
+        assert_eq!(
+            buffer
+                .snapshot()
+                .diagnostics_in_range::<_, usize>(0..1, false)
+                .map(|entry| entry.diagnostic.message.clone())
+                .collect::<Vec<_>>(),
+            ["the message".to_string()]
+        );
+    });
+    project.read_with(cx, |project, cx| {
+        assert_eq!(
+            project.diagnostic_summary(cx),
+            DiagnosticSummary {
+                error_count: 1,
+                warning_count: 0,
+            }
+        );
+    });
+
+    project.update(cx, |project, cx| {
+        project.restart_language_servers_for_buffers([buffer.clone()], cx);
+    });
+
+    // The diagnostics are cleared.
+    cx.foreground().run_until_parked();
+    buffer.read_with(cx, |buffer, _| {
+        assert_eq!(
+            buffer
+                .snapshot()
+                .diagnostics_in_range::<_, usize>(0..1, false)
+                .map(|entry| entry.diagnostic.message.clone())
+                .collect::<Vec<_>>(),
+            Vec::<String>::new(),
+        );
+    });
+    project.read_with(cx, |project, cx| {
+        assert_eq!(
+            project.diagnostic_summary(cx),
+            DiagnosticSummary {
+                error_count: 0,
+                warning_count: 0,
+            }
+        );
+    });
+}
+
 #[gpui::test]
 async fn test_restarted_server_reporting_invalid_buffer_version(cx: &mut gpui::TestAppContext) {
     init_test(cx);

crates/project/src/worktree.rs 🔗 

@@ -737,6 +737,45 @@ impl LocalWorktree {
         self.diagnostics.get(path).cloned().unwrap_or_default()
     }
 
+    pub fn clear_diagnostics_for_language_server(
+        &mut self,
+        server_id: LanguageServerId,
+        _: &mut ModelContext<Worktree>,
+    ) {
+        let worktree_id = self.id().to_proto();
+        self.diagnostic_summaries
+            .retain(|path, summaries_by_server_id| {
+                if summaries_by_server_id.remove(&server_id).is_some() {
+                    if let Some(share) = self.share.as_ref() {
+                        self.client
+                            .send(proto::UpdateDiagnosticSummary {
+                                project_id: share.project_id,
+                                worktree_id,
+                                summary: Some(proto::DiagnosticSummary {
+                                    path: path.to_string_lossy().to_string(),
+                                    language_server_id: server_id.0 as u64,
+                                    error_count: 0,
+                                    warning_count: 0,
+                                }),
+                            })
+                            .log_err();
+                    }
+                    !summaries_by_server_id.is_empty()
+                } else {
+                    true
+                }
+            });
+
+        self.diagnostics.retain(|_, diagnostics_by_server_id| {
+            if let Ok(ix) = diagnostics_by_server_id.binary_search_by_key(&server_id, |e| e.0) {
+                diagnostics_by_server_id.remove(ix);
+                !diagnostics_by_server_id.is_empty()
+            } else {
+                true
+            }
+        });
+    }
+
     pub fn update_diagnostics(
         &mut self,
         server_id: LanguageServerId,