perf(plugin-nm): cache install state across workspaces

This lowers memory usage of yarn from around 7GB to 600MB on a repo I have with 223 workspaces and 3MB node_modules/.yarn-state.yml.

Author: sophiebits

packages/plugin-nm/sources/NodeModulesLinker.ts

@@ -23,18 +23,28 @@ const INSTALL_STATE_FILE = `.yarn-state.yml` as Filename;
 const MTIME_ACCURANCY = 1000;
 
 type InstallState = {locatorMap: NodeModulesLocatorMap, locationTree: LocationTree, binSymlinks: BinSymlinkMap, nmMode: NodeModulesMode, mtimeMs: number};
+type InstallStateCacheEntry = {statKey: string | null, promise: Promise<InstallState | null>};
 type BinSymlinkMap = Map<PortablePath, Map<Filename, PortablePath>>;
 type LoadManifest = (locator: LocatorKey, installLocation: PortablePath) => Promise<Pick<Manifest, 'bin'>>;
 
+const installStateCache: Map<string, InstallStateCacheEntry> = new Map();
+
+function getInstallStateCacheKey(project: Project, {unrollAliases}: {unrollAliases: boolean}) {
+  return `${project.cwd}\0${unrollAliases ? `unroll` : `plain`}`;
+}
+
+function clearInstallStateCache(project: Project) {
+  installStateCache.delete(getInstallStateCacheKey(project, {unrollAliases: true}));
+  installStateCache.delete(getInstallStateCacheKey(project, {unrollAliases: false}));
+}
+
 export enum NodeModulesMode {
   CLASSIC = `classic`,
   HARDLINKS_LOCAL = `hardlinks-local`,
   HARDLINKS_GLOBAL = `hardlinks-global`,
 }
 
 export class NodeModulesLinker implements Linker {
-  private installStateCache: Map<string, Promise<InstallState | null>> = new Map();
-
   getCustomDataKey() {
     return JSON.stringify({
       name: `NodeModulesLinker`,
@@ -54,9 +64,7 @@ export class NodeModulesLinker implements Linker {
     if (workspace)
       return workspace.cwd;
 
-    const installState = await miscUtils.getFactoryWithDefault(this.installStateCache, opts.project.cwd, async () => {
-      return await findInstallState(opts.project, {unrollAliases: true});
-    });
+    const installState = await this.findInstallState(opts.project, {unrollAliases: true});
 
     if (installState === null)
       throw new UsageError(`Couldn't find the node_modules state file - running an install might help (findPackageLocation)`);
@@ -79,9 +87,7 @@ export class NodeModulesLinker implements Linker {
     if (!this.isEnabled(opts))
       return null;
 
-    const installState = await miscUtils.getFactoryWithDefault(this.installStateCache, opts.project.cwd, async () => {
-      return await findInstallState(opts.project, {unrollAliases: true});
-    });
+    const installState = await this.findInstallState(opts.project, {unrollAliases: true});
 
     if (installState === null)
       return null;
@@ -107,6 +113,26 @@ export class NodeModulesLinker implements Linker {
     return new NodeModulesInstaller(opts);
   }
 
+  private async findInstallState(project: Project, {unrollAliases}: {unrollAliases: boolean}) {
+    const cacheKey = getInstallStateCacheKey(project, {unrollAliases});
+    const statKey = await getInstallStateStatKey(project);
+    const cached = installStateCache.get(cacheKey);
+
+    if (cached?.statKey === statKey)
+      return await cached.promise;
+
+    let promise: Promise<InstallState | null>;
+    promise = findInstallState(project, {unrollAliases}).catch(error => {
+      const cached = installStateCache.get(cacheKey);
+      if (cached?.promise === promise)
+        installStateCache.delete(cacheKey);
+      throw error;
+    });
+
+    installStateCache.set(cacheKey, {statKey, promise});
+    return await promise;
+  }
+
   private isEnabled(opts: MinimalLinkOptions) {
     return opts.project.configuration.get(`nodeLinker`) === `node-modules`;
   }
@@ -447,13 +473,29 @@ async function writeInstallState(project: Project, locatorMap: NodeModulesLocato
   const rootPath = project.cwd;
   const installStatePath = ppath.join(rootPath, NODE_MODULES, INSTALL_STATE_FILE);
 
-  // Force install state file rewrite, so that it has mtime bigger than all node_modules subfolders
-  if (installChangedByUser)
-    await xfs.removePromise(installStatePath);
+  try {
+    // Force install state file rewrite, so that it has mtime bigger than all node_modules subfolders
+    if (installChangedByUser)
+      await xfs.removePromise(installStatePath);
 
-  await xfs.changeFilePromise(installStatePath, locatorState, {
-    automaticNewlines: true,
-  });
+    await xfs.changeFilePromise(installStatePath, locatorState, {
+      automaticNewlines: true,
+    });
+  } finally {
+    clearInstallStateCache(project);
+  }
+}
+
+async function getInstallStateStatKey(project: Project) {
+  const rootPath = project.cwd;
+  const installStatePath = ppath.join(rootPath, NODE_MODULES, INSTALL_STATE_FILE);
+
+  try {
+    const stats = await xfs.statPromise(installStatePath);
+    return `${stats.mtimeMs}:${stats.size}`;
+  } catch (e) {
+    return null;
+  }
 }
 
 async function findInstallState(project: Project, {unrollAliases = false}: {unrollAliases?: boolean} = {}): Promise<InstallState | null> {