From b74a4f29f20c2435ea7c6d192353c1e99760edb0 Mon Sep 17 00:00:00 2001
From: Nathan Whitaker <17734409+nathanwhit@users.noreply.github.com>
Date: Sat, 6 Apr 2024 19:22:37 -0700
Subject: [PATCH] perf(lsp): Don't retain `SourceFileObject`s in
 `sourceFileCache` longer than necessary (#23258)

The TS language service requests source files via
[getSourceFile](https://github.com/nathanwhit/deno/blob/7a25fd5ef0a82c2aac76594ccd467e9210e92b80/cli/tsc/99_main_compiler.js#L560).
In that function, we [unconditionally
add](https://github.com/nathanwhit/deno/blob/7a25fd5ef0a82c2aac76594ccd467e9210e92b80/cli/tsc/99_main_compiler.js#L613-L614)
the source file to our sourceFileCache. The issue is that we only remove
things from that cache if the source file [becomes out of
date](https://github.com/nathanwhit/deno/blob/7a25fd5ef0a82c2aac76594ccd467e9210e92b80/cli/tsc/99_main_compiler.js#L777-L783).
For files that don't get changed, we keep them in the cache
indefinitely. So sometimes we keep SourceFile objects from being GC'ed
because they're retained in our cache, even though TS doesn't refer to
them any more. I see this in pretty much all of the heap snapshots I've
taken.

---

The fix here is pretty direct - just store weak references to the
sourcefiles in the cache. It doesn't really change our caching behavior,
it just prevents us from being the only retainer of a `SourceFile`. I
also split the `sourceFileCache` into a separate cache just for assets,
as we rely on those being alive.

The simpler fix is to only cache assets, but presumably that has a perf
impact.

---
In local testing, this PR reduced the size of the JS heap by about 1 GB
when using `deno lsp` in the Typescript repo.
---
 cli/tsc/99_main_compiler.js | 43 +++++++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 14 deletions(-)

diff --git a/cli/tsc/99_main_compiler.js b/cli/tsc/99_main_compiler.js
index bfac52ccea..b606682189 100644
--- a/cli/tsc/99_main_compiler.js
+++ b/cli/tsc/99_main_compiler.js
@@ -147,8 +147,15 @@ delete Object.prototype.__proto__;
     }
   }
 
-  // In the case of the LSP, this will only ever contain the assets.
+  // Cache of asset source files.
   /** @type {Map<string, ts.SourceFile>} */
+  const assetSourceFileCache = new Map();
+
+  // Cache of source files, keyed by specifier.
+  // Stores weak references to ensure we aren't
+  // retaining `ts.SourceFile` objects longer than needed.
+  // This should not include assets, which have a separate cache.
+  /** @type {Map<string, WeakRef<ts.SourceFile>>} */
   const sourceFileCache = new Map();
 
   /** @type {string[]=} */
@@ -576,7 +583,11 @@ delete Object.prototype.__proto__;
       // Needs the original specifier
       specifier = normalizedToOriginalMap.get(specifier) ?? specifier;
 
-      let sourceFile = sourceFileCache.get(specifier);
+      const isAsset = specifier.startsWith(ASSETS_URL_PREFIX);
+
+      let sourceFile = isAsset
+        ? assetSourceFileCache.get(specifier)
+        : sourceFileCache.get(specifier)?.deref();
       if (sourceFile) {
         return sourceFile;
       }
@@ -607,10 +618,12 @@ delete Object.prototype.__proto__;
       );
       sourceFile.moduleName = specifier;
       sourceFile.version = version;
-      if (specifier.startsWith(ASSETS_URL_PREFIX)) {
+      if (isAsset) {
         sourceFile.version = "1";
+        assetSourceFileCache.set(specifier, sourceFile);
+      } else {
+        sourceFileCache.set(specifier, new WeakRef(sourceFile));
       }
-      sourceFileCache.set(specifier, sourceFile);
       scriptVersionCache.set(specifier, version);
       return sourceFile;
     },
@@ -773,9 +786,12 @@ delete Object.prototype.__proto__;
       if (logDebug) {
         debug(`host.getScriptSnapshot("${specifier}")`);
       }
-      let sourceFile = sourceFileCache.get(specifier);
+      const isAsset = specifier.startsWith(ASSETS_URL_PREFIX);
+      let sourceFile = isAsset
+        ? assetSourceFileCache.get(specifier)
+        : sourceFileCache.get(specifier)?.deref();
       if (
-        !specifier.startsWith(ASSETS_URL_PREFIX) &&
+        !isAsset &&
         sourceFile?.version != this.getScriptVersion(specifier)
       ) {
         sourceFileCache.delete(specifier);
@@ -994,13 +1010,11 @@ delete Object.prototype.__proto__;
   function getAssets() {
     /** @type {{ specifier: string; text: string; }[]} */
     const assets = [];
-    for (const sourceFile of sourceFileCache.values()) {
-      if (sourceFile.fileName.startsWith(ASSETS_URL_PREFIX)) {
-        assets.push({
-          specifier: sourceFile.fileName,
-          text: sourceFile.text,
-        });
-      }
+    for (const sourceFile of assetSourceFileCache.values()) {
+      assets.push({
+        specifier: sourceFile.fileName,
+        text: sourceFile.text,
+      });
     }
     return assets;
   }
@@ -1177,8 +1191,9 @@ delete Object.prototype.__proto__;
     "lib.d.ts files have errors",
   );
 
+  assert(buildSpecifier.startsWith(ASSETS_URL_PREFIX));
   // remove this now that we don't need it anymore for warming up tsc
-  sourceFileCache.delete(buildSpecifier);
+  assetSourceFileCache.delete(buildSpecifier);
 
   // exposes the functions that are called by `tsc::exec()` when type
   // checking TypeScript.