perf: lowers peak memory usage when using --changed on a large graph (#10866)

Author: jszumski

packages/vitest/src/node/specifications.ts

@@ -2,6 +2,7 @@ import type { Vitest } from './core'
 import type { TestProject } from './project'
 import type { TestSpecification } from './test-specification'
 import { existsSync } from 'node:fs'
+import os from 'node:os'
 import { join, relative, resolve } from 'pathe'
 import pm from 'picomatch'
 import { isWindows } from '../utils/env'
@@ -145,34 +146,81 @@ export class VitestSpecifications {
       return []
     }
 
-    const testGraphs = await Promise.all(
-      specs.map(async (spec) => {
-        const deps = await this.getTestDependencies(spec)
-        return [spec, deps] as const
-      }),
-    )
+    // The module graph, and so the dependency edges, are per project.
+    const specsByProject = new Map<TestProject, TestSpecification[]>()
+    for (const spec of specs) {
+      let projectSpecs = specsByProject.get(spec.project)
+      if (!projectSpecs) {
+        specsByProject.set(spec.project, projectSpecs = [])
+      }
+      projectSpecs.push(spec)
+    }
+
+    const affectedByProject = new Map<TestProject, Set<string>>()
+    for (const [project, projectSpecs] of specsByProject) {
+      affectedByProject.set(
+        project,
+        await this.getAffectedModules(project, projectSpecs, related),
+      )
+    }
+
+    return specs.filter(spec => affectedByProject.get(spec.project)!.has(spec.moduleId))
+  }
 
-    const runningTests: TestSpecification[] = []
+  /**
+   * Returns every module in `project` that transitively imports one of `related`.
+   *
+   * Expands each module's imports at most once into a shared reverse-edge map,
+   * then walks that map backwards from the changed files.
+   */
+  private async getAffectedModules(
+    project: TestProject,
+    specs: TestSpecification[],
+    related: string[],
+  ): Promise<Set<string>> {
+    const importers = new Map<string, Set<string>>()
+    const visited = new Set<string>()
+    const existsCache = new Map<string, boolean>()
 
-    for (const [specification, deps] of testGraphs) {
-      // if deps or the test itself were changed
-      if (related.some(path => path === specification.moduleId || deps.has(path))) {
-        runningTests.push(specification)
+    // limit concurrency to lower peak memory usage on large graphs
+    const TRANSFORM_CONCURRENCY = os.availableParallelism?.() ?? os.cpus().length
+    let active = 0
+    const waiters: Array<() => void> = []
+    const withLimit = async <T>(fn: () => Promise<T>): Promise<T> => {
+      if (active >= TRANSFORM_CONCURRENCY) {
+        await new Promise<void>(resolve => waiters.push(resolve))
+      }
+      active++
+      try {
+        return await fn()
+      }
+      finally {
+        active--
+        waiters.shift()?.()
       }
     }
 
-    return runningTests
-  }
+    const cachedExists = (filepath: string): boolean => {
+      const cached = existsCache.get(filepath)
+      if (cached !== undefined) {
+        return cached
+      }
+      const result = existsSync(filepath)
+      existsCache.set(filepath, result)
+      return result
+    }
 
-  private async getTestDependencies(spec: TestSpecification, deps = new Set<string>()): Promise<Set<string>> {
-    const addImports = async (project: TestProject, filepath: string) => {
-      if (deps.has(filepath)) {
+    const addImports = async (filepath: string) => {
+      // `visited` is shared by every spec in the project, so a module is
+      // expanded once per run instead of once per test file that reaches it.
+      if (visited.has(filepath)) {
         return
       }
-      deps.add(filepath)
+      visited.add(filepath)
 
-      const mod = project.vite.environments.ssr.moduleGraph.getModuleById(filepath)
-      const transformed = mod?.transformResult || await project.vite.environments.ssr.transformRequest(filepath)
+      const environment = project.vite.environments.ssr
+      const mod = environment.moduleGraph.getModuleById(filepath)
+      const transformed = mod?.transformResult || await withLimit(() => environment.transformRequest(filepath))
       if (!transformed) {
         return
       }
@@ -181,15 +229,35 @@ export class VitestSpecifications {
         const fsPath = dep.startsWith('/@fs/')
           ? dep.slice(isWindows ? 5 : 4)
           : join(project.config.root, dep)
-        if (!fsPath.includes('node_modules') && !deps.has(fsPath) && existsSync(fsPath)) {
-          await addImports(project, fsPath)
+        if (fsPath.includes('node_modules') || !cachedExists(fsPath)) {
+          return
         }
+        let importedBy = importers.get(fsPath)
+        if (!importedBy) {
+          importers.set(fsPath, importedBy = new Set())
+        }
+        importedBy.add(filepath)
+        await addImports(fsPath)
       }))
     }
 
-    await addImports(spec.project, spec.moduleId)
-    deps.delete(spec.moduleId)
+    await Promise.all(specs.map(spec => addImports(spec.moduleId)))
+
+    const affected = new Set<string>(related)
+    const queue = [...related]
+    while (queue.length) {
+      const importedBy = importers.get(queue.pop()!)
+      if (!importedBy) {
+        continue
+      }
+      for (const importer of importedBy) {
+        if (!affected.has(importer)) {
+          affected.add(importer)
+          queue.push(importer)
+        }
+      }
+    }
 
-    return deps
+    return affected
   }
 }