perf(lsp): release unused documents (#23398)

cli/lsp/documents.rs

@@ -1015,6 +1015,11 @@ impl Documents {
     Ok(())
   }
 
+  pub fn release(&self, specifier: &ModuleSpecifier) {
+    self.file_system_docs.remove_document(specifier);
+    self.file_system_docs.set_dirty(true);
+  }
+
   /// Return `true` if the provided specifier can be resolved to a document,
   /// otherwise `false`.
   pub fn contains_import(

cli/lsp/language_server.rs

@@ -364,14 +364,11 @@ impl LanguageServer {
           .client
           .show_message(MessageType::WARNING, err);
       }
-      {
-        let mut inner = self.0.write().await;
-        let lockfile = inner.config.tree.root_lockfile().cloned();
-        inner.documents.refresh_lockfile(lockfile);
-        inner.refresh_npm_specifiers().await;
-      }
-      // now refresh the data in a read
-      self.0.read().await.post_cache(result.mark).await;
+      let mut inner = self.0.write().await;
+      let lockfile = inner.config.tree.root_lockfile().cloned();
+      inner.documents.refresh_lockfile(lockfile);
+      inner.refresh_npm_specifiers().await;
+      inner.post_cache(result.mark).await;
     }
     Ok(Some(json!(true)))
   }
@@ -1421,7 +1418,16 @@ impl Inner {
       self.recreate_npm_services_if_necessary().await;
       self.refresh_documents_config().await;
       self.diagnostics_server.invalidate_all();
-      self.ts_server.restart(self.snapshot()).await;
+      self
+        .project_changed(
+          &changes
+            .iter()
+            .map(|(s, _)| (s, ChangeKind::Modified))
+            .collect::<Vec<_>>(),
+          false,
+        )
+        .await;
+      self.ts_server.cleanup_semantic_cache(self.snapshot()).await;
       self.send_diagnostics_update();
       self.send_testing_update();
     }
@@ -3544,13 +3550,14 @@ impl Inner {
     }))
   }
 
-  async fn post_cache(&self, mark: PerformanceMark) {
+  async fn post_cache(&mut self, mark: PerformanceMark) {
     // Now that we have dependencies loaded, we need to re-analyze all the files.
     // For that we're invalidating all the existing diagnostics and restarting
     // the language server for TypeScript (as it might hold to some stale
     // documents).
     self.diagnostics_server.invalidate_all();
-    self.ts_server.restart(self.snapshot()).await;
+    self.project_changed(&[], false).await;
+    self.ts_server.cleanup_semantic_cache(self.snapshot()).await;
     self.send_diagnostics_update();
     self.send_testing_update();
 

cli/lsp/tsc.rs

@@ -356,6 +356,21 @@ impl TsServer {
     Ok(diagnostics_map)
   }
 
+  pub async fn cleanup_semantic_cache(&self, snapshot: Arc<StateSnapshot>) {
+    let req = TscRequest {
+      method: "cleanupSemanticCache",
+      args: json!([]),
+    };
+    self
+      .request::<()>(snapshot, req)
+      .await
+      .map_err(|err| {
+        log::error!("Failed to request to tsserver {}", err);
+        LspError::invalid_request()
+      })
+      .ok();
+  }
+
   pub async fn find_references(
     &self,
     snapshot: Arc<StateSnapshot>,
@@ -1010,14 +1025,6 @@ impl TsServer {
     })
   }
 
-  pub async fn restart(&self, snapshot: Arc<StateSnapshot>) {
-    let req = TscRequest {
-      method: "$restart",
-      args: json!([]),
-    };
-    self.request::<bool>(snapshot, req).await.unwrap();
-  }
-
   async fn request<R>(
     &self,
     snapshot: Arc<StateSnapshot>,
@@ -4032,6 +4039,21 @@ fn op_load<'s>(
   Ok(serialized)
 }
 
+#[op2(fast)]
+fn op_release(
+  state: &mut OpState,
+  #[string] specifier: &str,
+) -> Result<(), AnyError> {
+  let state = state.borrow_mut::<State>();
+  let mark = state
+    .performance
+    .mark_with_args("tsc.op.op_release", specifier);
+  let specifier = state.specifier_map.normalize(specifier)?;
+  state.state_snapshot.documents.release(&specifier);
+  state.performance.measure(mark);
+  Ok(())
+}
+
 #[op2]
 #[serde]
 fn op_resolve(
@@ -4244,6 +4266,7 @@ deno_core::extension!(deno_tsc,
     op_is_cancelled,
     op_is_node_file,
     op_load,
+    op_release,
     op_resolve,
     op_respond,
     op_script_names,

cli/tsc/99_main_compiler.js

@@ -155,6 +155,12 @@ delete Object.prototype.__proto__;
   /** @type {Map<string, ts.SourceFile>} */
   const sourceFileCache = new Map();
 
+  /** @type {Map<string, string>} */
+  const sourceTextCache = new Map();
+
+  /** @type {Map<string, number>} */
+  const sourceRefCounts = new Map();
+
   /** @type {string[]=} */
   let scriptFileNamesCache;
 
@@ -172,6 +178,8 @@ delete Object.prototype.__proto__;
   /** @type {number | null} */
   let projectVersionCache = null;
 
+  let lastRequestMethod = null;
+
   const ChangeKind = {
     Opened: 0,
     Modified: 1,
@@ -250,6 +258,8 @@ delete Object.prototype.__proto__;
         );
         documentRegistrySourceFileCache.set(mapKey, sourceFile);
       }
+      const sourceRefCount = sourceRefCounts.get(fileName) ?? 0;
+      sourceRefCounts.set(fileName, sourceRefCount + 1);
       return sourceFile;
     },
 
@@ -333,8 +343,20 @@ delete Object.prototype.__proto__;
     },
 
     releaseDocumentWithKey(path, key, _scriptKind, _impliedNodeFormat) {
-      const mapKey = path + key;
-      documentRegistrySourceFileCache.delete(mapKey);
+      const sourceRefCount = sourceRefCounts.get(path) ?? 1;
+      if (sourceRefCount <= 1) {
+        sourceRefCounts.delete(path);
+        // We call `cleanupSemanticCache` for other purposes, don't bust the
+        // source cache in this case.
+        if (lastRequestMethod != "cleanupSemanticCache") {
+          const mapKey = path + key;
+          documentRegistrySourceFileCache.delete(mapKey);
+          sourceTextCache.delete(path);
+          ops.op_release(path);
+        }
+      } else {
+        sourceRefCounts.set(path, sourceRefCount - 1);
+      }
     },
 
     reportStats() {
@@ -807,19 +829,26 @@ delete Object.prototype.__proto__;
       if (logDebug) {
         debug(`host.getScriptSnapshot("${specifier}")`);
       }
-      let sourceFile = sourceFileCache.get(specifier);
-      if (!sourceFile) {
-        sourceFile = this.getSourceFile(
-          specifier,
-          specifier.endsWith(".json")
-            ? ts.ScriptTarget.JSON
-            : ts.ScriptTarget.ESNext,
-        );
-      }
+      const sourceFile = sourceFileCache.get(specifier);
       if (sourceFile) {
+        // This case only occurs for assets.
         return ts.ScriptSnapshot.fromString(sourceFile.text);
       }
-      return undefined;
+      let sourceText = sourceTextCache.get(specifier);
+      if (sourceText == undefined) {
+        /** @type {{ data: string, version: string, isCjs: boolean }} */
+        const fileInfo = ops.op_load(specifier);
+        if (!fileInfo) {
+          return undefined;
+        }
+        if (fileInfo.isCjs) {
+          isCjsCache.add(specifier);
+        }
+        sourceTextCache.set(specifier, fileInfo.data);
+        scriptVersionCache.set(specifier, fileInfo.version);
+        sourceText = fileInfo.data;
+      }
+      return ts.ScriptSnapshot.fromString(sourceText);
     },
   };
 
@@ -1047,6 +1076,7 @@ delete Object.prototype.__proto__;
     if (logDebug) {
       debug(`serverRequest()`, id, method, args);
     }
+    lastRequestMethod = method;
     switch (method) {
       case "$projectChanged": {
         /** @type {[string, number][]} */
@@ -1058,6 +1088,7 @@ delete Object.prototype.__proto__;
 
         if (configChanged) {
           tsConfigCache = null;
+          isNodeSourceFileCache.clear();
         }
 
         projectVersionCache = newProjectVersion;
@@ -1068,7 +1099,7 @@ delete Object.prototype.__proto__;
             opened = true;
           }
           scriptVersionCache.delete(script);
-          sourceFileCache.delete(script);
+          sourceTextCache.delete(script);
         }
 
         if (configChanged || opened) {
@@ -1077,10 +1108,6 @@ delete Object.prototype.__proto__;
 
         return respond(id);
       }
-      case "$restart": {
-        serverRestart();
-        return respond(id, true);
-      }
       case "$getSupportedCodeFixes": {
         return respond(
           id,
@@ -1152,12 +1179,6 @@ delete Object.prototype.__proto__;
     debug("serverInit()");
   }
 
-  function serverRestart() {
-    languageService = ts.createLanguageService(host, documentRegistry);
-    isNodeSourceFileCache.clear();
-    debug("serverRestart()");
-  }
-
   // A build time only op that provides some setup information that is used to
   // ensure the snapshot is setup properly.
   /** @type {{ buildSpecifier: string; libs: string[]; nodeBuiltInModuleNames: string[] }} */