perf(lsp): independent diagnostic source publishes (#13427)

Cargo.lock

@@ -764,6 +764,7 @@ dependencies = [
  "text-size",
  "text_lines",
  "tokio",
+ "tokio-util",
  "trust-dns-client",
  "trust-dns-server",
  "typed-arena",

cli/Cargo.toml

@@ -91,6 +91,7 @@ tempfile = "=3.2.0"
 text-size = "=1.1.0"
 text_lines = "=0.4.1"
 tokio = { version = "=1.14", features = ["full"] }
+tokio-util = "=0.6.9"
 typed-arena = "2.0.1"
 uuid = { version = "=0.8.2", features = ["v4", "serde"] }
 walkdir = "=2.3.2"

cli/bench/lsp_bench_standalone.rs

@@ -78,12 +78,13 @@ fn wait_for_deno_lint_diagnostic(
       let version = msg.get("version").unwrap().as_u64().unwrap();
       if document_version == version {
         let diagnostics = msg.get("diagnostics").unwrap().as_array().unwrap();
-        let first = &diagnostics[0];
+        for diagnostic in diagnostics {
-        let source = first.get("source").unwrap().as_str().unwrap();
+          let source = diagnostic.get("source").unwrap().as_str().unwrap();
           if source == "deno-lint" {
             return;
           }
         }
+      }
     } else {
       todo!() // handle_misc_message
     }

cli/lsp/diagnostics.rs

@@ -32,73 +32,19 @@ use tokio::sync::Mutex;
 use tokio::time::sleep;
 use tokio::time::Duration;
 use tokio::time::Instant;
+use tokio_util::sync::CancellationToken;
 
 pub type DiagnosticRecord =
   (ModuleSpecifier, Option<i32>, Vec<lsp::Diagnostic>);
 pub type DiagnosticVec = Vec<DiagnosticRecord>;
+type DiagnosticMap =
+  HashMap<ModuleSpecifier, (Option<i32>, Vec<lsp::Diagnostic>)>;
 type TsDiagnosticsMap = HashMap<String, Vec<diagnostics::Diagnostic>>;
 
-#[derive(Debug, Hash, Clone, PartialEq, Eq)]
-pub(crate) enum DiagnosticSource {
-  Deno,
-  DenoLint,
-  TypeScript,
-}
-
-#[derive(Debug, Default)]
-struct DiagnosticCollection {
-  map: HashMap<(ModuleSpecifier, DiagnosticSource), Vec<lsp::Diagnostic>>,
-  versions: HashMap<ModuleSpecifier, HashMap<DiagnosticSource, i32>>,
-  changes: HashSet<ModuleSpecifier>,
-}
-
-impl DiagnosticCollection {
-  pub fn get(
-    &self,
-    specifier: &ModuleSpecifier,
-    source: DiagnosticSource,
-  ) -> impl Iterator<Item = &lsp::Diagnostic> {
-    self
-      .map
-      .get(&(specifier.clone(), source))
-      .into_iter()
-      .flatten()
-  }
-
-  pub fn get_version(
-    &self,
-    specifier: &ModuleSpecifier,
-    source: &DiagnosticSource,
-  ) -> Option<i32> {
-    let source_version = self.versions.get(specifier)?;
-    source_version.get(source).cloned()
-  }
-
-  pub fn set(&mut self, source: DiagnosticSource, record: DiagnosticRecord) {
-    let (specifier, maybe_version, diagnostics) = record;
-    self
-      .map
-      .insert((specifier.clone(), source.clone()), diagnostics);
-    if let Some(version) = maybe_version {
-      let source_version = self.versions.entry(specifier.clone()).or_default();
-      source_version.insert(source, version);
-    }
-    self.changes.insert(specifier);
-  }
-
-  pub fn take_changes(&mut self) -> Option<HashSet<ModuleSpecifier>> {
-    if self.changes.is_empty() {
-      None
-    } else {
-      Some(mem::take(&mut self.changes))
-    }
-  }
-}
-
 #[derive(Debug)]
 pub(crate) struct DiagnosticsServer {
   channel: Option<mpsc::UnboundedSender<()>>,
-  collection: Arc<Mutex<DiagnosticCollection>>,
+  ts_diagnostics: Arc<Mutex<DiagnosticMap>>,
   client: Client,
   performance: Arc<Performance>,
   ts_server: Arc<TsServer>,
@@ -112,37 +58,40 @@ impl DiagnosticsServer {
   ) -> Self {
     DiagnosticsServer {
       channel: Default::default(),
-      collection: Default::default(),
+      ts_diagnostics: Default::default(),
       client,
       performance,
       ts_server,
     }
   }
 
-  pub(crate) async fn get(
+  pub(crate) async fn get_ts_diagnostics(
     &self,
     specifier: &ModuleSpecifier,
-    source: DiagnosticSource,
+    document_version: Option<i32>,
   ) -> Vec<lsp::Diagnostic> {
-    self
+    let ts_diagnostics = self.ts_diagnostics.lock().await;
-      .collection
+    if let Some((diagnostics_doc_version, diagnostics)) =
-      .lock()
+      ts_diagnostics.get(specifier)
-      .await
+    {
-      .get(specifier, source)
+      // only get the diagnostics if they're up to date
-      .cloned()
+      if document_version == *diagnostics_doc_version {
-      .collect()
+        return diagnostics.clone();
+      }
+    }
+    Vec::new()
   }
 
   pub(crate) async fn invalidate(&self, specifiers: Vec<ModuleSpecifier>) {
-    let mut collection = self.collection.lock().await;
+    let mut ts_diagnostics = self.ts_diagnostics.lock().await;
     for specifier in &specifiers {
-      collection.versions.remove(specifier);
+      ts_diagnostics.remove(specifier);
     }
   }
 
   pub(crate) async fn invalidate_all(&self) {
-    let mut collection = self.collection.lock().await;
+    let mut ts_diagnostics = self.ts_diagnostics.lock().await;
-    collection.versions.clear();
+    ts_diagnostics.clear();
   }
 
   pub(crate) fn start(
@@ -151,49 +100,28 @@ impl DiagnosticsServer {
   ) {
     let (tx, mut rx) = mpsc::unbounded_channel::<()>();
     self.channel = Some(tx);
-    let collection = self.collection.clone();
     let client = self.client.clone();
     let performance = self.performance.clone();
+    let stored_ts_diagnostics = self.ts_diagnostics.clone();
     let ts_server = self.ts_server.clone();
 
     let _join_handle = thread::spawn(move || {
       let runtime = create_basic_runtime();
 
       runtime.block_on(async {
-        // Debounce timer delay. 150ms between keystrokes is about 45 WPM, so we
+        let mut token = CancellationToken::new();
-        // want something that is longer than that, but not too long to
+        let mut ts_handle: Option<tokio::task::JoinHandle<()>> = None;
-        // introduce detectable UI delay; 200ms is a decent compromise.
+        let mut lint_handle: Option<tokio::task::JoinHandle<()>> = None;
-        const DELAY: Duration = Duration::from_millis(200);
+        let mut deps_handle: Option<tokio::task::JoinHandle<()>> = None;
-        // If the debounce timer isn't active, it will be set to expire "never",
-        // which is actually just 1 year in the future.
-        const NEVER: Duration = Duration::from_secs(365 * 24 * 60 * 60);
-
-        // A flag that is set whenever something has changed that requires the
-        // diagnostics collection to be updated.
-        let mut dirty = false;
-
-        let debounce_timer = sleep(NEVER);
-        tokio::pin!(debounce_timer);
 
         loop {
-          // "race" the next message off the rx queue or the debounce timer.
+          match rx.recv().await {
-          // The debounce timer gets reset every time a message comes off the
-          // queue. When the debounce timer expires, a snapshot of the most
-          // up-to-date state is used to produce diagnostics.
-          tokio::select! {
-            maybe_request = rx.recv() => {
-              match maybe_request {
             // channel has closed
             None => break,
-                Some(_) => {
+            Some(()) => {
-                  dirty = true;
+              // cancel the previous run
-                  debounce_timer.as_mut().reset(Instant::now() + DELAY);
+              token.cancel();
-                }
+              token = CancellationToken::new();
-              }
-            }
-            _ = debounce_timer.as_mut(), if dirty => {
-              dirty = false;
-              debounce_timer.as_mut().reset(Instant::now() + NEVER);
 
               let (snapshot, config, maybe_lint_config) = {
                 let language_server = language_server.lock().await;
@@ -203,15 +131,131 @@ impl DiagnosticsServer {
                   language_server.maybe_lint_config.clone(),
                 )
               };
-              update_diagnostics(
+
-                &client,
+              let previous_ts_handle = ts_handle.take();
-                collection.clone(),
+              ts_handle = Some(tokio::spawn({
-                snapshot,
+                let performance = performance.clone();
-                config,
+                let ts_server = ts_server.clone();
+                let client = client.clone();
+                let token = token.clone();
+                let stored_ts_diagnostics = stored_ts_diagnostics.clone();
+                let snapshot = snapshot.clone();
+                let config = config.clone();
+                async move {
+                  if let Some(previous_handle) = previous_ts_handle {
+                    // Wait on the previous run to complete in order to prevent
+                    // multiple threads queueing up a lot of tsc requests.
+                    // Do not race this with cancellation because we want a
+                    // chain of events to wait for all the previous diagnostics to complete
+                    previous_handle.await;
+                  }
+
+                  // Debounce timer delay. 150ms between keystrokes is about 45 WPM, so we
+                  // want something that is longer than that, but not too long to
+                  // introduce detectable UI delay; 200ms is a decent compromise.
+                  const DELAY: Duration = Duration::from_millis(200);
+                  tokio::select! {
+                    _ = token.cancelled() => { return; }
+                    _ = tokio::time::sleep(DELAY) => {}
+                  };
+
+                  let mark =
+                    performance.mark("update_diagnostics_ts", None::<()>);
+                  let diagnostics =
+                    generate_ts_diagnostics(snapshot.clone(), &ts_server)
+                      .await
+                      .map_err(|err| {
+                        error!(
+                          "Error generating TypeScript diagnostics: {}",
+                          err
+                        );
+                      })
+                      .unwrap_or_default();
+
+                  if !token.is_cancelled() {
+                    {
+                      let mut stored_ts_diagnostics =
+                        stored_ts_diagnostics.lock().await;
+                      *stored_ts_diagnostics = diagnostics
+                        .iter()
+                        .map(|(specifier, version, diagnostics)| {
+                          (specifier.clone(), (*version, diagnostics.clone()))
+                        })
+                        .collect();
+                    }
+
+                    for (specifier, version, diagnostics) in diagnostics {
+                      client
+                        .publish_diagnostics(specifier, diagnostics, version)
+                        .await;
+                    }
+                    performance.measure(mark);
+                  }
+                }
+              }));
+
+              let previous_deps_handle = deps_handle.take();
+              deps_handle = Some(tokio::spawn({
+                let performance = performance.clone();
+                let client = client.clone();
+                let token = token.clone();
+                let snapshot = snapshot.clone();
+                let config = config.clone();
+                async move {
+                  if let Some(previous_handle) = previous_deps_handle {
+                    previous_handle.await;
+                  }
+                  let mark =
+                    performance.mark("update_diagnostics_deps", None::<()>);
+                  let diagnostics = generate_deps_diagnostics(
+                    snapshot.clone(),
+                    config.clone(),
+                    token.clone(),
+                  )
+                  .await;
+
+                  if !token.is_cancelled() {
+                    for (specifier, version, diagnostics) in diagnostics {
+                      client
+                        .publish_diagnostics(specifier, diagnostics, version)
+                        .await;
+                    }
+                    performance.measure(mark);
+                  }
+                }
+              }));
+
+              let previous_lint_handle = lint_handle.take();
+              lint_handle = Some(tokio::spawn({
+                let performance = performance.clone();
+                let client = client.clone();
+                let token = token.clone();
+                let snapshot = snapshot.clone();
+                let config = config.clone();
+                async move {
+                  if let Some(previous_handle) = previous_lint_handle {
+                    previous_handle.await;
+                  }
+                  let mark =
+                    performance.mark("update_diagnostics_lint", None::<()>);
+                  let diagnostics = generate_lint_diagnostics(
+                    &snapshot,
+                    &config,
                     maybe_lint_config,
-                &ts_server,
+                    token.clone(),
-                performance.clone(),
+                  )
-              ).await;
+                  .await;
+
+                  if !token.is_cancelled() {
+                    for (specifier, version, diagnostics) in diagnostics {
+                      client
+                        .publish_diagnostics(specifier, diagnostics, version)
+                        .await;
+                    }
+                    performance.measure(mark);
+                  }
+                }
+              }));
             }
           }
         }
@@ -336,23 +380,22 @@ fn ts_json_to_diagnostics(
 
 async fn generate_lint_diagnostics(
   snapshot: &language_server::StateSnapshot,
-  collection: Arc<Mutex<DiagnosticCollection>>,
   config: &ConfigSnapshot,
   maybe_lint_config: Option<LintConfig>,
-) -> Result<DiagnosticVec, AnyError> {
+  token: CancellationToken,
+) -> DiagnosticVec {
   let documents = snapshot.documents.documents(true, true);
   let workspace_settings = config.settings.workspace.clone();
 
-  tokio::task::spawn(async move {
   let mut diagnostics_vec = Vec::new();
   if workspace_settings.lint {
     for document in documents {
+      // exit early if cancelled
+      if token.is_cancelled() {
+        break;
+      }
+
       let version = document.maybe_lsp_version();
-        let current_version = collection
-          .lock()
-          .await
-          .get_version(document.specifier(), &DiagnosticSource::DenoLint);
-        if version != current_version {
       let is_allowed = match &maybe_lint_config {
         Some(lint_config) => {
           lint_config.files.matches_specifier(document.specifier())
@@ -390,37 +433,20 @@ async fn generate_lint_diagnostics(
       ));
     }
   }
-    }
+  diagnostics_vec
-    Ok(diagnostics_vec)
-  })
-  .await
-  .unwrap()
 }
 
 async fn generate_ts_diagnostics(
   snapshot: Arc<language_server::StateSnapshot>,
-  collection: Arc<Mutex<DiagnosticCollection>>,
   ts_server: &tsc::TsServer,
 ) -> Result<DiagnosticVec, AnyError> {
   let mut diagnostics_vec = Vec::new();
-  let specifiers: Vec<ModuleSpecifier> = {
+  let specifiers = snapshot
-    let collection = collection.lock().await;
-    snapshot
     .documents
     .documents(true, true)
     .iter()
-      .filter_map(|d| {
+    .map(|d| d.specifier().clone())
-        let version = d.maybe_lsp_version();
+    .collect::<Vec<_>>();
-        let current_version =
-          collection.get_version(d.specifier(), &DiagnosticSource::TypeScript);
-        if version != current_version {
-          Some(d.specifier().clone())
-        } else {
-          None
-        }
-      })
-      .collect()
-  };
   if !specifiers.is_empty() {
     let req = tsc::RequestMethod::GetDiagnostics(specifiers);
     let ts_diagnostics_map: TsDiagnosticsMap =
@@ -552,21 +578,17 @@ fn diagnose_dependency(
 async fn generate_deps_diagnostics(
   snapshot: Arc<language_server::StateSnapshot>,
   config: Arc<ConfigSnapshot>,
-  collection: Arc<Mutex<DiagnosticCollection>>,
+  token: CancellationToken,
-) -> Result<DiagnosticVec, AnyError> {
+) -> DiagnosticVec {
-  tokio::task::spawn(async move {
   let mut diagnostics_vec = Vec::new();
 
   for document in snapshot.documents.documents(true, true) {
+    if token.is_cancelled() {
+      break;
+    }
     if !config.specifier_enabled(document.specifier()) {
       continue;
     }
-      let version = document.maybe_lsp_version();
-      let current_version = collection
-        .lock()
-        .await
-        .get_version(document.specifier(), &DiagnosticSource::Deno);
-      if version != current_version {
     let mut diagnostics = Vec::new();
     for (_, dependency) in document.dependencies() {
       diagnose_dependency(
@@ -586,134 +608,12 @@ async fn generate_deps_diagnostics(
     }
     diagnostics_vec.push((
       document.specifier().clone(),
-          version,
+      document.maybe_lsp_version(),
       diagnostics,
     ));
   }
-    }
 
-    Ok(diagnostics_vec)
+  diagnostics_vec
-  })
-  .await
-  .unwrap()
-}
-
-/// Publishes diagnostics to the client.
-async fn publish_diagnostics(
-  client: &Client,
-  collection: &mut DiagnosticCollection,
-  snapshot: &language_server::StateSnapshot,
-  config: &ConfigSnapshot,
-) {
-  if let Some(changes) = collection.take_changes() {
-    for specifier in changes {
-      let mut diagnostics: Vec<lsp::Diagnostic> =
-        if config.settings.workspace.lint {
-          collection
-            .get(&specifier, DiagnosticSource::DenoLint)
-            .cloned()
-            .collect()
-        } else {
-          Vec::new()
-        };
-      if config.specifier_enabled(&specifier) {
-        diagnostics.extend(
-          collection
-            .get(&specifier, DiagnosticSource::TypeScript)
-            .cloned(),
-        );
-        diagnostics
-          .extend(collection.get(&specifier, DiagnosticSource::Deno).cloned());
-      }
-      let version = snapshot
-        .documents
-        .get(&specifier)
-        .map(|d| d.maybe_lsp_version())
-        .flatten();
-      client
-        .publish_diagnostics(specifier.clone(), diagnostics, version)
-        .await;
-    }
-  }
-}
-
-/// Updates diagnostics for any specifiers that don't have the correct version
-/// generated and publishes the diagnostics to the client.
-async fn update_diagnostics(
-  client: &Client,
-  collection: Arc<Mutex<DiagnosticCollection>>,
-  snapshot: Arc<language_server::StateSnapshot>,
-  config: Arc<ConfigSnapshot>,
-  maybe_lint_config: Option<LintConfig>,
-  ts_server: &tsc::TsServer,
-  performance: Arc<Performance>,
-) {
-  let mark = performance.mark("update_diagnostics", None::<()>);
-
-  let lint = async {
-    let mark = performance.mark("update_diagnostics_lint", None::<()>);
-    let collection = collection.clone();
-    let diagnostics = generate_lint_diagnostics(
-      &snapshot,
-      collection.clone(),
-      &config,
-      maybe_lint_config,
-    )
-    .await
-    .map_err(|err| {
-      error!("Error generating lint diagnostics: {}", err);
-    })
-    .unwrap_or_default();
-
-    let mut collection = collection.lock().await;
-    for diagnostic_record in diagnostics {
-      collection.set(DiagnosticSource::DenoLint, diagnostic_record);
-    }
-    publish_diagnostics(client, &mut collection, &snapshot, &config).await;
-    performance.measure(mark);
-  };
-
-  let ts = async {
-    let mark = performance.mark("update_diagnostics_ts", None::<()>);
-    let collection = collection.clone();
-    let diagnostics =
-      generate_ts_diagnostics(snapshot.clone(), collection.clone(), ts_server)
-        .await
-        .map_err(|err| {
-          error!("Error generating TypeScript diagnostics: {}", err);
-        })
-        .unwrap_or_default();
-    let mut collection = collection.lock().await;
-    for diagnostic_record in diagnostics {
-      collection.set(DiagnosticSource::TypeScript, diagnostic_record);
-    }
-    publish_diagnostics(client, &mut collection, &snapshot, &config).await;
-    performance.measure(mark);
-  };
-
-  let deps = async {
-    let mark = performance.mark("update_diagnostics_deps", None::<()>);
-    let collection = collection.clone();
-    let diagnostics = generate_deps_diagnostics(
-      snapshot.clone(),
-      config.clone(),
-      collection.clone(),
-    )
-    .await
-    .map_err(|err| {
-      error!("Error generating Deno diagnostics: {}", err);
-    })
-    .unwrap_or_default();
-    let mut collection = collection.lock().await;
-    for diagnostic_record in diagnostics {
-      collection.set(DiagnosticSource::Deno, diagnostic_record);
-    }
-    publish_diagnostics(client, &mut collection, &snapshot, &config).await;
-    performance.measure(mark);
-  };
-
-  tokio::join!(lint, ts, deps);
-  performance.measure(mark);
 }
 
 #[cfg(test)]
@@ -765,23 +665,17 @@ mod tests {
 
   fn setup(
     sources: &[(&str, &str, i32, LanguageId)],
-  ) -> (
+  ) -> (StateSnapshot, PathBuf, ConfigSnapshot) {
-    StateSnapshot,
-    Arc<Mutex<DiagnosticCollection>>,
-    PathBuf,
-    ConfigSnapshot,
-  ) {
     let temp_dir = TempDir::new().expect("could not create temp dir");
     let location = temp_dir.path().join("deps");
     let state_snapshot = mock_state_snapshot(sources, &location);
-    let collection = Arc::new(Mutex::new(DiagnosticCollection::default()));
     let config = mock_config();
-    (state_snapshot, collection, location, config)
+    (state_snapshot, location, config)
   }
 
   #[tokio::test]
   async fn test_generate_lint_diagnostics() {
-    let (snapshot, collection, _, config) = setup(&[(
+    let (snapshot, _, config) = setup(&[(
       "file:///a.ts",
       r#"import * as b from "./b.ts";
 
@@ -791,10 +685,9 @@ console.log(a);
       1,
       LanguageId::TypeScript,
     )]);
-    let result =
+    let diagnostics =
-      generate_lint_diagnostics(&snapshot, collection, &config, None).await;
+      generate_lint_diagnostics(&snapshot, &config, None, Default::default())
-    assert!(result.is_ok());
+        .await;
-    let diagnostics = result.unwrap();
     assert_eq!(diagnostics.len(), 1);
     let (_, _, diagnostics) = &diagnostics[0];
     assert_eq!(diagnostics.len(), 2);

cli/lsp/language_server.rs

@@ -36,7 +36,6 @@ use super::config::Config;
 use super::config::ConfigSnapshot;
 use super::config::SETTINGS_SECTION;
 use super::diagnostics;
-use super::diagnostics::DiagnosticSource;
 use super::diagnostics::DiagnosticsServer;
 use super::documents::to_hover_text;
 use super::documents::to_lsp_range;
@@ -1198,7 +1197,7 @@ impl Inner {
       let mut code_actions = CodeActionCollection::default();
       let file_diagnostics = self
         .diagnostics_server
-        .get(&specifier, DiagnosticSource::TypeScript)
+        .get_ts_diagnostics(&specifier, asset_or_doc.document_lsp_version())
         .await;
       for diagnostic in &fixable_diagnostics {
         match diagnostic.source.as_deref() {

cli/tests/integration/lsp_tests.rs

@@ -1,6 +1,7 @@
 // Copyright 2018-2022 the Deno authors. All rights reserved. MIT license.
 
 use deno_ast::ModuleSpecifier;
+use deno_core::serde::de::DeserializeOwned;
 use deno_core::serde::Deserialize;
 use deno_core::serde::Serialize;
 use deno_core::serde_json;
@@ -17,10 +18,15 @@ use test_util::lsp::LspClient;
 use test_util::testdata_path;
 
 fn load_fixture(path: &str) -> Value {
-  let fixtures_path = testdata_path().join("lsp");
+  load_fixture_as(path)
-  let path = fixtures_path.join(path);
+}
-  let fixture_str = fs::read_to_string(path).unwrap();
+
-  serde_json::from_str(&fixture_str).unwrap()
+fn load_fixture_as<T>(path: &str) -> T
+where
+  T: DeserializeOwned,
+{
+  let fixture_str = load_fixture_str(path);
+  serde_json::from_str::<T>(&fixture_str).unwrap()
 }
 
 fn load_fixture_str(path: &str) -> String {
@@ -64,6 +70,11 @@ where
     )
     .unwrap();
 
+  read_diagnostics(client).0
+}
+
+fn read_diagnostics(client: &mut LspClient) -> CollectedDiagnostics {
+  // diagnostics come in batches of three unless they're cancelled
   let mut diagnostics = vec![];
   for _ in 0..3 {
     let (method, response) = client
@@ -72,8 +83,7 @@ where
     assert_eq!(method, "textDocument/publishDiagnostics");
     diagnostics.push(response.unwrap());
   }
-
+  CollectedDiagnostics(diagnostics)
-  diagnostics
 }
 
 fn shutdown(client: &mut LspClient) {
@@ -83,6 +93,110 @@ fn shutdown(client: &mut LspClient) {
   client.write_notification("exit", json!(null)).unwrap();
 }
 
+struct TestSession {
+  client: LspClient,
+  open_file_count: usize,
+}
+
+impl TestSession {
+  pub fn from_file(init_path: &str) -> Self {
+    Self::from_client(init(init_path))
+  }
+
+  pub fn from_client(client: LspClient) -> Self {
+    Self {
+      client,
+      open_file_count: 0,
+    }
+  }
+
+  pub fn did_open<V>(&mut self, params: V) -> CollectedDiagnostics
+  where
+    V: Serialize,
+  {
+    self
+      .client
+      .write_notification("textDocument/didOpen", params)
+      .unwrap();
+
+    let (id, method, _) = self.client.read_request::<Value>().unwrap();
+    assert_eq!(method, "workspace/configuration");
+    self
+      .client
+      .write_response(
+        id,
+        json!([{
+          "enable": true,
+          "codeLens": {
+            "test": true
+          }
+        }]),
+      )
+      .unwrap();
+
+    self.open_file_count += 1;
+    self.read_diagnostics()
+  }
+
+  pub fn read_diagnostics(&mut self) -> CollectedDiagnostics {
+    let mut all_diagnostics = Vec::new();
+    for _ in 0..self.open_file_count {
+      all_diagnostics.extend(read_diagnostics(&mut self.client).0);
+    }
+    CollectedDiagnostics(all_diagnostics)
+  }
+
+  pub fn shutdown_and_exit(&mut self) {
+    shutdown(&mut self.client);
+  }
+}
+
+#[derive(Debug, Clone)]
+struct CollectedDiagnostics(Vec<lsp::PublishDiagnosticsParams>);
+
+impl CollectedDiagnostics {
+  pub fn all(&self) -> Vec<lsp::Diagnostic> {
+    self
+      .0
+      .iter()
+      .flat_map(|p| p.diagnostics.clone().into_iter())
+      .collect()
+  }
+
+  pub fn with_source(&self, source: &str) -> lsp::PublishDiagnosticsParams {
+    self
+      .0
+      .iter()
+      .find(|p| {
+        p.diagnostics
+          .iter()
+          .any(|d| d.source == Some(source.to_string()))
+      })
+      .map(ToOwned::to_owned)
+      .unwrap()
+  }
+
+  pub fn with_file_and_source(
+    &self,
+    specifier: &str,
+    source: &str,
+  ) -> lsp::PublishDiagnosticsParams {
+    let specifier = ModuleSpecifier::parse(specifier).unwrap();
+    self
+      .0
+      .iter()
+      .find(|p| {
+        p.uri == specifier
+          && p
+            .diagnostics
+            .iter()
+            .any(|d| d.source == Some(source.to_string()))
+      })
+      .map(ToOwned::to_owned)
+      .unwrap()
+  }
+}
+
 #[test]
 fn lsp_startup_shutdown() {
   let mut client = init("initialize_params.json");
@@ -361,7 +475,7 @@ fn lsp_import_assertions() {
     )
     .unwrap();
 
-  let mut diagnostics = did_open(
+  let diagnostics = CollectedDiagnostics(did_open(
     &mut client,
     json!({
       "textDocument": {
@@ -371,11 +485,14 @@ fn lsp_import_assertions() {
         "text": "import a from \"./test.json\";\n\nconsole.log(a);\n"
       }
     }),
-  );
+  ));
 
-  let last = diagnostics.pop().unwrap();
   assert_eq!(
-    json!(last.diagnostics),
+    json!(
+      diagnostics
+        .with_file_and_source("file:///a/a.ts", "deno")
+        .diagnostics
+    ),
     json!([
       {
         "range": {
@@ -2521,13 +2638,11 @@ fn lsp_code_actions_deno_cache() {
   client
     .write_response(id, json!([{ "enable": true }]))
     .unwrap();
-  let (method, _) = client.read_notification::<Value>().unwrap();
+  let diagnostics = read_diagnostics(&mut client);
-  assert_eq!(method, "textDocument/publishDiagnostics");
+  assert_eq!(
-  let (method, _) = client.read_notification::<Value>().unwrap();
+    diagnostics.with_source("deno"),
-  assert_eq!(method, "textDocument/publishDiagnostics");
+    load_fixture_as("diagnostics_deno_deps.json")
-  let (method, params) = client.read_notification().unwrap();
+  );
-  assert_eq!(method, "textDocument/publishDiagnostics");
-  assert_eq!(params, Some(load_fixture("diagnostics_deno_deps.json")));
 
   let (maybe_res, maybe_err) = client
     .write_request(
@@ -2545,31 +2660,26 @@ fn lsp_code_actions_deno_cache() {
 
 #[test]
 fn lsp_code_actions_imports() {
-  let mut client = init("initialize_params.json");
+  let mut session = TestSession::from_file("initialize_params.json");
-  did_open(
+  session.did_open(json!({
-    &mut client,
-    json!({
     "textDocument": {
       "uri": "file:///a/file00.ts",
       "languageId": "typescript",
       "version": 1,
       "text": "export const abc = \"abc\";\nexport const def = \"def\";\n"
     }
-    }),
+  }));
-  );
+  session.did_open(json!({
-  did_open(
-    &mut client,
-    json!({
     "textDocument": {
       "uri": "file:///a/file01.ts",
       "languageId": "typescript",
       "version": 1,
       "text": "\nconsole.log(abc);\nconsole.log(def)\n"
     }
-    }),
+  }));
-  );
 
-  let (maybe_res, maybe_err) = client
+  let (maybe_res, maybe_err) = session
+    .client
     .write_request(
       "textDocument/codeAction",
       load_fixture("code_action_params_imports.json"),
@@ -2580,7 +2690,8 @@ fn lsp_code_actions_imports() {
     maybe_res,
     Some(load_fixture("code_action_response_imports.json"))
   );
-  let (maybe_res, maybe_err) = client
+  let (maybe_res, maybe_err) = session
+    .client
     .write_request(
       "codeAction/resolve",
       load_fixture("code_action_resolve_params_imports.json"),
@@ -2591,7 +2702,8 @@ fn lsp_code_actions_imports() {
     maybe_res,
     Some(load_fixture("code_action_resolve_response_imports.json"))
   );
-  shutdown(&mut client);
+
+  session.shutdown_and_exit();
 }
 
 #[test]
@@ -2707,10 +2819,7 @@ fn lsp_code_actions_deadlock() {
     .unwrap();
   assert!(maybe_err.is_none());
   assert!(maybe_res.is_some());
-  for _ in 0..3 {
+  read_diagnostics(&mut client);
-    let (method, _) = client.read_notification::<Value>().unwrap();
-    assert_eq!(method, "textDocument/publishDiagnostics");
-  }
   client
     .write_notification(
       "textDocument/didChange",
@@ -2818,12 +2927,8 @@ fn lsp_code_actions_deadlock() {
   assert!(maybe_err.is_none());
   assert!(maybe_res.is_some());
 
-  for _ in 0..3 {
+  read_diagnostics(&mut client);
-    let (method, _) = client.read_notification::<Value>().unwrap();
-    assert_eq!(method, "textDocument/publishDiagnostics");
-  }
 
-  assert!(client.queue_is_empty());
   shutdown(&mut client);
 }
 
@@ -3128,7 +3233,7 @@ fn lsp_cache_location() {
     load_fixture("did_open_params_import_hover.json"),
   );
   let diagnostics = diagnostics.into_iter().flat_map(|x| x.diagnostics);
-  assert_eq!(diagnostics.count(), 14);
+  assert_eq!(diagnostics.count(), 7);
   let (maybe_res, maybe_err) = client
     .write_request::<_, _, Value>(
       "deno/cache",
@@ -3233,24 +3338,23 @@ fn lsp_tls_cert() {
   client
     .write_request::<_, _, Value>("initialize", params)
     .unwrap();
-
   client.write_notification("initialized", json!({})).unwrap();
-  did_open(
+  let mut session = TestSession::from_client(client);
-    &mut client,
+
-    json!({
+  session.did_open(json!({
     "textDocument": {
       "uri": "file:///a/file_01.ts",
       "languageId": "typescript",
       "version": 1,
       "text": "export const a = \"a\";\n",
     }
-    }),
+  }));
-  );
   let diagnostics =
-    did_open(&mut client, load_fixture("did_open_params_tls_cert.json"));
+    session.did_open(load_fixture("did_open_params_tls_cert.json"));
-  let diagnostics = diagnostics.into_iter().flat_map(|x| x.diagnostics);
+  let diagnostics = diagnostics.all();
-  assert_eq!(diagnostics.count(), 14);
+  assert_eq!(diagnostics.len(), 7);
-  let (maybe_res, maybe_err) = client
+  let (maybe_res, maybe_err) = session
+    .client
     .write_request::<_, _, Value>(
       "deno/cache",
       json!({
@@ -3263,7 +3367,8 @@ fn lsp_tls_cert() {
     .unwrap();
   assert!(maybe_err.is_none());
   assert!(maybe_res.is_some());
-  let (maybe_res, maybe_err) = client
+  let (maybe_res, maybe_err) = session
+    .client
     .write_request(
       "textDocument/hover",
       json!({
@@ -3297,7 +3402,8 @@ fn lsp_tls_cert() {
       }
     }))
   );
-  let (maybe_res, maybe_err) = client
+  let (maybe_res, maybe_err) = session
+    .client
     .write_request::<_, _, Value>(
       "textDocument/hover",
       json!({
@@ -3331,7 +3437,7 @@ fn lsp_tls_cert() {
       }
     }))
   );
-  shutdown(&mut client);
+  session.shutdown_and_exit();
 }
 
 #[test]
@@ -3366,17 +3472,10 @@ fn lsp_diagnostics_warn() {
     .unwrap();
   assert!(maybe_err.is_none());
   assert!(maybe_res.is_some());
-  let (method, _) = client.read_notification::<Value>().unwrap();
+  let diagnostics = read_diagnostics(&mut client);
-  assert_eq!(method, "textDocument/publishDiagnostics");
-  let (method, _) = client.read_notification::<Value>().unwrap();
-  assert_eq!(method, "textDocument/publishDiagnostics");
-  let (method, maybe_params) = client
-    .read_notification::<lsp::PublishDiagnosticsParams>()
-    .unwrap();
-  assert_eq!(method, "textDocument/publishDiagnostics");
   assert_eq!(
-    maybe_params,
+    diagnostics.with_source("deno"),
-    Some(lsp::PublishDiagnosticsParams {
+    lsp::PublishDiagnosticsParams {
       uri: Url::parse("file:///a/file.ts").unwrap(),
       diagnostics: vec![lsp::Diagnostic {
         range: lsp::Range {
@@ -3396,7 +3495,7 @@ fn lsp_diagnostics_warn() {
         ..Default::default()
       }],
       version: Some(1),
-    })
+    }
   );
   shutdown(&mut client);
 }
@@ -3485,58 +3584,40 @@ fn lsp_diagnostics_deno_types() {
     .unwrap();
   assert!(maybe_res.is_some());
   assert!(maybe_err.is_none());
-  let (method, _) = client.read_notification::<Value>().unwrap();
+  let diagnostics = read_diagnostics(&mut client);
-  assert_eq!(method, "textDocument/publishDiagnostics");
+  assert_eq!(diagnostics.all().len(), 5);
-  let (method, _) = client.read_notification::<Value>().unwrap();
-  assert_eq!(method, "textDocument/publishDiagnostics");
-  let (method, maybe_params) = client
-    .read_notification::<lsp::PublishDiagnosticsParams>()
-    .unwrap();
-  assert_eq!(method, "textDocument/publishDiagnostics");
-  assert!(maybe_params.is_some());
-  let params = maybe_params.unwrap();
-  assert_eq!(params.diagnostics.len(), 5);
   shutdown(&mut client);
 }
 
 #[test]
 fn lsp_diagnostics_refresh_dependents() {
-  let mut client = init("initialize_params.json");
+  let mut session = TestSession::from_file("initialize_params.json");
-  did_open(
+  session.did_open(json!({
-    &mut client,
-    json!({
     "textDocument": {
       "uri": "file:///a/file_00.ts",
       "languageId": "typescript",
       "version": 1,
       "text": "export const a = \"a\";\n",
     },
-    }),
+  }));
-  );
+  session.did_open(json!({
-  did_open(
-    &mut client,
-    json!({
     "textDocument": {
       "uri": "file:///a/file_01.ts",
       "languageId": "typescript",
       "version": 1,
       "text": "export * from \"./file_00.ts\";\n",
     },
-    }),
+  }));
-  );
+  let diagnostics = session.did_open(json!({
-  let diagnostics = did_open(
-    &mut client,
-    json!({
     "textDocument": {
       "uri": "file:///a/file_02.ts",
       "languageId": "typescript",
       "version": 1,
       "text": "import { a, b } from \"./file_01.ts\";\n\nconsole.log(a, b);\n"
     }
-    }),
+  }));
-  );
   assert_eq!(
-    json!(diagnostics[2]),
+    json!(diagnostics.with_file_and_source("file:///a/file_02.ts", "deno-ts")),
     json!({
       "uri": "file:///a/file_02.ts",
       "diagnostics": [
@@ -3560,7 +3641,10 @@ fn lsp_diagnostics_refresh_dependents() {
       "version": 1
     })
   );
-  client
+
+  // fix the code causing the diagnostic
+  session
+    .client
     .write_notification(
       "textDocument/didChange",
       json!({
@@ -3586,34 +3670,11 @@ fn lsp_diagnostics_refresh_dependents() {
       }),
     )
     .unwrap();
-  let (method, _) = client.read_notification::<Value>().unwrap();
+  let diagnostics = session.read_diagnostics();
-  assert_eq!(method, "textDocument/publishDiagnostics");
+  assert_eq!(diagnostics.all().len(), 0); // no diagnostics now
-  let (method, _) = client.read_notification::<Value>().unwrap();
-  assert_eq!(method, "textDocument/publishDiagnostics");
-  let (method, _) = client.read_notification::<Value>().unwrap();
-  assert_eq!(method, "textDocument/publishDiagnostics");
-  // ensure that the server publishes any inflight diagnostics
-  std::thread::sleep(std::time::Duration::from_millis(250));
-  client
-    .write_request::<_, _, Value>("shutdown", json!(null))
-    .unwrap();
-  client.write_notification("exit", json!(null)).unwrap();
 
-  let queue_len = client.queue_len();
+  session.shutdown_and_exit();
-  assert!(!client.queue_is_empty());
+  assert_eq!(session.client.queue_len(), 0);
-  for i in 0..queue_len {
-    let (method, maybe_params) = client
-      .read_notification::<lsp::PublishDiagnosticsParams>()
-      .unwrap();
-    assert_eq!(method, "textDocument/publishDiagnostics");
-    // the last 3 diagnostic publishes should be the clear of any diagnostics
-    if queue_len - i <= 3 {
-      assert!(maybe_params.is_some());
-      let params = maybe_params.unwrap();
-      assert_eq!(params.diagnostics, Vec::with_capacity(0));
-    }
-  }
-  assert!(client.queue_is_empty());
 }
 
 #[derive(Deserialize)]
@@ -3665,7 +3726,7 @@ fn lsp_performance() {
     .unwrap();
   assert!(maybe_err.is_none());
   if let Some(res) = maybe_res {
-    assert!(res.averages.len() >= 6);
+    assert_eq!(res.averages.len(), 13);
   } else {
     panic!("unexpected result");
   }
@@ -4350,13 +4411,11 @@ fn lsp_lint_with_config() {
     .into_iter()
     .flat_map(|x| x.diagnostics)
     .collect::<Vec<_>>();
-  assert_eq!(diagnostics.len(), 3);
+  assert_eq!(diagnostics.len(), 1);
-  for diagnostic in diagnostics {
   assert_eq!(
-      diagnostic.code,
+    diagnostics[0].code,
     Some(lsp::NumberOrString::String("ban-untagged-todo".to_string()))
   );
-  }
   shutdown(&mut client);
 }