perf(lsp): only store parsed sources for open documents (#23454)

5 files changed, 66 insertions, 89 deletions

diff --git a/cli/lsp/testing/collectors.rs b/cli/lsp/testing/collectors.rs
index 8579ccc7d..508e50f9b 100644
--- a/cli/lsp/testing/collectors.rs
+++ b/cli/lsp/testing/collectors.rs
@@ -451,13 +451,9 @@ pub struct TestCollector {
 }
 
 impl TestCollector {
-  pub fn new(
-    specifier: ModuleSpecifier,
-    script_version: String,
-    text_info: SourceTextInfo,
-  ) -> Self {
+  pub fn new(specifier: ModuleSpecifier, text_info: SourceTextInfo) -> Self {
     Self {
-      test_module: TestModule::new(specifier, script_version),
+      test_module: TestModule::new(specifier),
       vars: HashSet::new(),
       fns: HashMap::new(),
       text_info,
@@ -653,8 +649,7 @@ pub mod tests {
     })
     .unwrap();
     let text_info = parsed_module.text_info().clone();
-    let mut collector =
-      TestCollector::new(specifier, "1".to_string(), text_info);
+    let mut collector = TestCollector::new(specifier, text_info);
     parsed_module.module().visit_with(&mut collector);
     collector.take()
   }
@@ -671,7 +666,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9"
             .to_string(),
@@ -704,7 +698,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9"
             .to_string(),
@@ -747,7 +740,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![
           (
             "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9".to_string(),
@@ -809,7 +801,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![
           (
             "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9".to_string(),
@@ -862,7 +853,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9"
             .to_string(),
@@ -897,7 +887,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![
           (
             "86b4c821900e38fc89f24bceb0e45193608ab3f9d2a6019c7b6a5aceff5d7df2".to_string(),
@@ -939,7 +928,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9"
             .to_string(),
@@ -973,7 +961,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9"
             .to_string(),
@@ -1008,7 +995,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![
           (
             "87f28e06f5ddadd90a74a93b84df2e31b9edced8301b0ad4c8fbab8d806ec99d".to_string(),
@@ -1022,7 +1008,7 @@ pub mod tests {
             },
           ),
           (
-            "e0f6a73647b763f82176c98a019e54200b799a32007f9859fb782aaa9e308568".to_string(), 
+            "e0f6a73647b763f82176c98a019e54200b799a32007f9859fb782aaa9e308568".to_string(),
             TestDefinition {
               id: "e0f6a73647b763f82176c98a019e54200b799a32007f9859fb782aaa9e308568".to_string(),
               name: "someFunction".to_string(),
@@ -1063,7 +1049,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "e0f6a73647b763f82176c98a019e54200b799a32007f9859fb782aaa9e308568"
             .to_string(),
@@ -1097,7 +1082,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "6d05d6dc35548b86a1e70acaf24a5bc2dd35db686b35b685ad5931d201b4a918"
             .to_string(),
@@ -1138,7 +1122,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![
           (
             "3799fc549a32532145ffc8532b0cd943e025bbc19a02e2cde9be94f87bceb829".to_string(),
diff --git a/cli/lsp/testing/definitions.rs b/cli/lsp/testing/definitions.rs
index ab47beec9..43a07c2e3 100644
--- a/cli/lsp/testing/definitions.rs
+++ b/cli/lsp/testing/definitions.rs
@@ -28,16 +28,13 @@ pub struct TestDefinition {
 #[derive(Debug, Clone, PartialEq)]
 pub struct TestModule {
   pub specifier: ModuleSpecifier,
-  /// The version of the document that the discovered tests relate to.
-  pub script_version: String,
   pub defs: HashMap<String, TestDefinition>,
 }
 
 impl TestModule {
-  pub fn new(specifier: ModuleSpecifier, script_version: String) -> Self {
+  pub fn new(specifier: ModuleSpecifier) -> Self {
     Self {
       specifier,
-      script_version,
       defs: Default::default(),
     }
   }
diff --git a/cli/lsp/testing/execution.rs b/cli/lsp/testing/execution.rs
index e1189ec2c..73916d0c2 100644
--- a/cli/lsp/testing/execution.rs
+++ b/cli/lsp/testing/execution.rs
@@ -3,6 +3,7 @@
 use super::definitions::TestDefinition;
 use super::definitions::TestModule;
 use super::lsp_custom;
+use super::server::TestServerTests;
 
 use crate::args::flags_from_vec;
 use crate::args::DenoSubcommand;
@@ -21,7 +22,6 @@ use deno_core::error::JsError;
 use deno_core::futures::future;
 use deno_core::futures::stream;
 use deno_core::futures::StreamExt;
-use deno_core::parking_lot::Mutex;
 use deno_core::parking_lot::RwLock;
 use deno_core::unsync::spawn;
 use deno_core::unsync::spawn_blocking;
@@ -42,7 +42,7 @@ use tower_lsp::lsp_types as lsp;
 /// any filters to be applied to those tests
 fn as_queue_and_filters(
   params: &lsp_custom::TestRunRequestParams,
-  tests: &HashMap<ModuleSpecifier, TestModule>,
+  tests: &HashMap<ModuleSpecifier, (TestModule, String)>,
 ) -> (
   HashSet<ModuleSpecifier>,
   HashMap<ModuleSpecifier, LspTestFilter>,
@@ -52,7 +52,7 @@ fn as_queue_and_filters(
 
   if let Some(include) = &params.include {
     for item in include {
-      if let Some(test_definitions) = tests.get(&item.text_document.uri) {
+      if let Some((test_definitions, _)) = tests.get(&item.text_document.uri) {
         queue.insert(item.text_document.uri.clone());
         if let Some(id) = &item.id {
           if let Some(test) = test_definitions.get(id) {
@@ -74,7 +74,7 @@ fn as_queue_and_filters(
   }
 
   for item in &params.exclude {
-    if let Some(test_definitions) = tests.get(&item.text_document.uri) {
+    if let Some((test_definitions, _)) = tests.get(&item.text_document.uri) {
       if let Some(id) = &item.id {
         // there is no way to exclude a test step
         if item.step_id.is_none() {
@@ -91,7 +91,7 @@ fn as_queue_and_filters(
     }
   }
 
-  queue.retain(|s| !tests.get(s).unwrap().is_empty());
+  queue.retain(|s| !tests.get(s).unwrap().0.is_empty());
 
   (queue, filters)
 }
@@ -147,19 +147,19 @@ pub struct TestRun {
   kind: lsp_custom::TestRunKind,
   filters: HashMap<ModuleSpecifier, LspTestFilter>,
   queue: HashSet<ModuleSpecifier>,
-  tests: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>>,
+  tests: TestServerTests,
   token: CancellationToken,
   workspace_settings: config::WorkspaceSettings,
 }
 
 impl TestRun {
-  pub fn new(
+  pub async fn init(
     params: &lsp_custom::TestRunRequestParams,
-    tests: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>>,
+    tests: TestServerTests,
     workspace_settings: config::WorkspaceSettings,
   ) -> Self {
     let (queue, filters) = {
-      let tests = tests.lock();
+      let tests = tests.lock().await;
       as_queue_and_filters(params, &tests)
     };
 
@@ -176,13 +176,13 @@ impl TestRun {
 
   /// Provide the tests of a test run as an enqueued module which can be sent
   /// to the client to indicate tests are enqueued for testing.
-  pub fn as_enqueued(&self) -> Vec<lsp_custom::EnqueuedTestModule> {
-    let tests = self.tests.lock();
+  pub async fn as_enqueued(&self) -> Vec<lsp_custom::EnqueuedTestModule> {
+    let tests = self.tests.lock().await;
     self
       .queue
       .iter()
       .map(|s| {
-        let ids = if let Some(test_module) = tests.get(s) {
+        let ids = if let Some((test_module, _)) = tests.get(s) {
           if let Some(filter) = self.filters.get(s) {
             filter.as_ids(test_module)
           } else {
@@ -332,7 +332,7 @@ impl TestRun {
           match event {
             test::TestEvent::Register(description) => {
               for (_, description) in description.into_iter() {
-                reporter.report_register(description);
+                reporter.report_register(description).await;
                 // TODO(mmastrac): we shouldn't need to clone here - we can re-use the descriptions
                 tests.write().insert(description.id, description.clone());
               }
@@ -378,7 +378,7 @@ impl TestRun {
               summary.uncaught_errors.push((origin, error));
             }
             test::TestEvent::StepRegister(description) => {
-              reporter.report_step_register(&description);
+              reporter.report_step_register(&description).await;
               test_steps.insert(description.id, description);
             }
             test::TestEvent::StepWait(id) => {
@@ -541,7 +541,7 @@ struct LspTestReporter {
   client: Client,
   id: u32,
   maybe_root_uri: Option<ModuleSpecifier>,
-  files: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>>,
+  files: TestServerTests,
   tests: IndexMap<usize, LspTestDescription>,
   current_test: Option<usize>,
 }
@@ -551,7 +551,7 @@ impl LspTestReporter {
     run: &TestRun,
     client: Client,
     maybe_root_uri: Option<&ModuleSpecifier>,
-    files: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>>,
+    files: TestServerTests,
   ) -> Self {
     Self {
       client,
@@ -576,12 +576,12 @@ impl LspTestReporter {
 
   fn report_plan(&mut self, _plan: &test::TestPlan) {}
 
-  fn report_register(&mut self, desc: &test::TestDescription) {
-    let mut files = self.files.lock();
+  async fn report_register(&mut self, desc: &test::TestDescription) {
+    let mut files = self.files.lock().await;
     let specifier = ModuleSpecifier::parse(&desc.location.file_name).unwrap();
-    let test_module = files
+    let (test_module, _) = files
       .entry(specifier.clone())
-      .or_insert_with(|| TestModule::new(specifier, "1".to_string()));
+      .or_insert_with(|| (TestModule::new(specifier), "1".to_string()));
     let (static_id, is_new) = test_module.register_dynamic(desc);
     self.tests.insert(
       desc.id,
@@ -681,12 +681,12 @@ impl LspTestReporter {
     }
   }
 
-  fn report_step_register(&mut self, desc: &test::TestStepDescription) {
-    let mut files = self.files.lock();
+  async fn report_step_register(&mut self, desc: &test::TestStepDescription) {
+    let mut files = self.files.lock().await;
     let specifier = ModuleSpecifier::parse(&desc.location.file_name).unwrap();
-    let test_module = files
+    let (test_module, _) = files
       .entry(specifier.clone())
-      .or_insert_with(|| TestModule::new(specifier, "1".to_string()));
+      .or_insert_with(|| (TestModule::new(specifier), "1".to_string()));
     let (static_id, is_new) = test_module.register_step_dynamic(
       desc,
       self.tests.get(&desc.parent_id).unwrap().static_id(),
@@ -828,7 +828,6 @@ mod tests {
     };
     let test_module = TestModule {
       specifier: specifier.clone(),
-      script_version: "1".to_string(),
       defs: vec![
         (test_def_a.id.clone(), test_def_a.clone()),
         (test_def_b.id.clone(), test_def_b.clone()),
@@ -836,10 +835,10 @@ mod tests {
       .into_iter()
       .collect(),
     };
-    tests.insert(specifier.clone(), test_module.clone());
+    tests.insert(specifier.clone(), (test_module.clone(), "1".to_string()));
     tests.insert(
       non_test_specifier.clone(),
-      TestModule::new(non_test_specifier, "1".to_string()),
+      (TestModule::new(non_test_specifier), "1".to_string()),
     );
     let (queue, filters) = as_queue_and_filters(&params, &tests);
     assert_eq!(json!(queue), json!([specifier]));
diff --git a/cli/lsp/testing/mod.rs b/cli/lsp/testing/mod.rs
index a11d3a8cc..d285481f0 100644
--- a/cli/lsp/testing/mod.rs
+++ b/cli/lsp/testing/mod.rs
@@ -6,6 +6,8 @@ mod execution;
 pub mod lsp_custom;
 mod server;
 
+pub use collectors::TestCollector;
+pub use definitions::TestModule;
 pub use lsp_custom::TEST_RUN_CANCEL_REQUEST;
 pub use lsp_custom::TEST_RUN_REQUEST;
 pub use server::TestServer;
diff --git a/cli/lsp/testing/server.rs b/cli/lsp/testing/server.rs
index bdf238078..ff59e1033 100644
--- a/cli/lsp/testing/server.rs
+++ b/cli/lsp/testing/server.rs
@@ -1,6 +1,5 @@
 // Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
 
-use super::collectors::TestCollector;
 use super::definitions::TestModule;
 use super::execution::TestRun;
 use super::lsp_custom;
@@ -12,7 +11,6 @@ use crate::lsp::documents::DocumentsFilter;
 use crate::lsp::language_server::StateSnapshot;
 use crate::lsp::performance::Performance;
 
-use deno_ast::swc::visit::VisitWith;
 use deno_core::error::AnyError;
 use deno_core::parking_lot::Mutex;
 use deno_core::serde_json::json;
@@ -36,6 +34,9 @@ fn as_delete_notification(uri: ModuleSpecifier) -> TestingNotification {
   )
 }
 
+pub type TestServerTests =
+  Arc<tokio::sync::Mutex<HashMap<ModuleSpecifier, (TestModule, String)>>>;
+
 /// The main structure which handles requests and sends notifications related
 /// to the Testing API.
 #[derive(Debug)]
@@ -47,7 +48,7 @@ pub struct TestServer {
   /// A map of run ids to test runs
   runs: Arc<Mutex<HashMap<u32, TestRun>>>,
   /// Tests that are discovered from a versioned document
-  tests: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>>,
+  tests: TestServerTests,
   /// A channel for requesting that changes to documents be statically analyzed
   /// for tests
   update_channel: mpsc::UnboundedSender<Arc<StateSnapshot>>,
@@ -59,8 +60,7 @@ impl TestServer {
     performance: Arc<Performance>,
     maybe_root_uri: Option<ModuleSpecifier>,
   ) -> Self {
-    let tests: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>> =
-      Arc::new(Mutex::new(HashMap::new()));
+    let tests = Default::default();
 
     let (update_channel, mut update_rx) =
       mpsc::unbounded_channel::<Arc<StateSnapshot>>();
@@ -88,7 +88,7 @@ impl TestServer {
             None => break,
             Some(snapshot) => {
               let mark = performance.mark("lsp.testing_update");
-              let mut tests = tests.lock();
+              let mut tests = tests.lock().await;
               // we create a list of test modules we currently are tracking
               // eliminating any we go over when iterating over the document
               let mut keys: HashSet<ModuleSpecifier> =
@@ -106,37 +106,33 @@ impl TestServer {
                 }
                 keys.remove(specifier);
                 let script_version = document.script_version();
-                let valid = if let Some(test) = tests.get(specifier) {
-                  test.script_version == script_version
-                } else {
-                  false
-                };
+                let valid =
+                  if let Some((_, old_script_version)) = tests.get(specifier) {
+                    old_script_version == &script_version
+                  } else {
+                    false
+                  };
                 if !valid {
-                  if let Some(Ok(parsed_source)) =
-                    document.maybe_parsed_source()
-                  {
-                    let was_empty = tests
-                      .remove(specifier)
-                      .map(|tm| tm.is_empty())
-                      .unwrap_or(true);
-                    let mut collector = TestCollector::new(
-                      specifier.clone(),
-                      script_version,
-                      parsed_source.text_info().clone(),
+                  let was_empty = tests
+                    .remove(specifier)
+                    .map(|(tm, _)| tm.is_empty())
+                    .unwrap_or(true);
+                  let test_module = document
+                    .maybe_test_module()
+                    .await
+                    .map(|tm| tm.as_ref().clone())
+                    .unwrap_or_else(|| TestModule::new(specifier.clone()));
+                  if !test_module.is_empty() {
+                    client.send_test_notification(
+                      test_module.as_replace_notification(mru.as_ref()),
                     );
-                    parsed_source.module().visit_with(&mut collector);
-                    let test_module = collector.take();
-                    if !test_module.is_empty() {
-                      client.send_test_notification(
-                        test_module.as_replace_notification(mru.as_ref()),
-                      );
-                    } else if !was_empty {
-                      client.send_test_notification(as_delete_notification(
-                        specifier.clone(),
-                      ));
-                    }
-                    tests.insert(specifier.clone(), test_module);
+                  } else if !was_empty {
+                    client.send_test_notification(as_delete_notification(
+                      specifier.clone(),
+                    ));
                   }
+                  tests
+                    .insert(specifier.clone(), (test_module, script_version));
                 }
               }
               for key in keys {
@@ -205,14 +201,14 @@ impl TestServer {
   }
 
   /// A request from the client to start a test run.
-  pub fn run_request(
+  pub async fn run_request(
     &self,
     params: lsp_custom::TestRunRequestParams,
     workspace_settings: config::WorkspaceSettings,
   ) -> LspResult<Option<Value>> {
     let test_run =
-      { TestRun::new(&params, self.tests.clone(), workspace_settings) };
-    let enqueued = test_run.as_enqueued();
+      { TestRun::init(&params, self.tests.clone(), workspace_settings).await };
+    let enqueued = test_run.as_enqueued().await;
     {
       let mut runs = self.runs.lock();
       runs.insert(params.id, test_run);