fix(coverage): improve memory usage by writing temporary files on file system (#4603)

Author: AriPerkkio

Diff for: docs/config/index.md

@@ -1323,6 +1323,15 @@ See [istanbul documentation](https://github.com/istanbuljs/nyc#ignoring-methods)
 
 Watermarks for statements, lines, branches and functions. See [istanbul documentation](https://github.com/istanbuljs/nyc#high-and-low-watermarks) for more information.
 
+#### coverage.processingConcurrency
+
+- **Type:** `boolean`
+- **Default:** `Math.min(20, os.cpu().length)`
+- **Available for providers:** `'v8' | 'istanbul'`
+- **CLI:** `--coverage.processingConcurrency=<number>`
+
+Concurrency limit used when processing the coverage results.
+
 #### coverage.customProviderModule
 
 - **Type:** `string`

Diff for: packages/coverage-istanbul/package.json

@@ -45,6 +45,7 @@
     "vitest": "^1.0.0-0"
   },
   "dependencies": {
+    "debug": "^4.3.4",
     "istanbul-lib-coverage": "^3.2.2",
     "istanbul-lib-instrument": "^6.0.1",
     "istanbul-lib-report": "^3.0.1",
@@ -55,6 +56,7 @@
     "test-exclude": "^6.0.0"
   },
   "devDependencies": {
+    "@types/debug": "^4.1.12",
     "@types/istanbul-lib-coverage": "^2.0.6",
     "@types/istanbul-lib-instrument": "^1.7.7",
     "@types/istanbul-lib-report": "^3.0.3",

Diff for: packages/coverage-istanbul/src/provider.ts

@@ -5,9 +5,10 @@ import { coverageConfigDefaults, defaultExclude, defaultInclude } from 'vitest/c
 import { BaseCoverageProvider } from 'vitest/coverage'
 import c from 'picocolors'
 import { parseModule } from 'magicast'
+import createDebug from 'debug'
 import libReport from 'istanbul-lib-report'
 import reports from 'istanbul-reports'
-import type { CoverageMap, CoverageMapData } from 'istanbul-lib-coverage'
+import type { CoverageMap } from 'istanbul-lib-coverage'
 import libCoverage from 'istanbul-lib-coverage'
 import libSourceMaps from 'istanbul-lib-source-maps'
 import { type Instrumenter, createInstrumenter } from 'istanbul-lib-instrument'
@@ -17,7 +18,8 @@ import _TestExclude from 'test-exclude'
 import { COVERAGE_STORE_KEY } from './constants'
 
 type Options = ResolvedCoverageOptions<'istanbul'>
-type CoverageByTransformMode = Record<AfterSuiteRunMeta['transformMode'], CoverageMapData[]>
+type Filename = string
+type CoverageFilesByTransformMode = Record<AfterSuiteRunMeta['transformMode'], Filename[]>
 type ProjectName = NonNullable<AfterSuiteRunMeta['projectName']> | typeof DEFAULT_PROJECT
 
 interface TestExclude {
@@ -35,6 +37,8 @@ interface TestExclude {
 }
 
 const DEFAULT_PROJECT = Symbol.for('default-project')
+const debug = createDebug('vitest:coverage')
+let uniqueId = 0
 
 export class IstanbulCoverageProvider extends BaseCoverageProvider implements CoverageProvider {
   name = 'istanbul'
@@ -44,13 +48,9 @@ export class IstanbulCoverageProvider extends BaseCoverageProvider implements Co
   instrumenter!: Instrumenter
   testExclude!: InstanceType<TestExclude>
 
-  /**
-   * Coverage objects collected from workers.
-   * Some istanbul utilizers write these into file system instead of storing in memory.
-   * If storing in memory causes issues, we can simply write these into fs in `onAfterSuiteRun`
-   * and read them back when merging coverage objects in `onAfterAllFilesRun`.
-   */
-  coverages = new Map<ProjectName, CoverageByTransformMode>()
+  coverageFiles = new Map<ProjectName, CoverageFilesByTransformMode>()
+  coverageFilesDirectory!: string
+  pendingPromises: Promise<void>[] = []
 
   initialize(ctx: Vitest) {
     const config: CoverageIstanbulOptions = ctx.config.coverage
@@ -96,6 +96,8 @@ export class IstanbulCoverageProvider extends BaseCoverageProvider implements Co
       extension: this.options.extension,
       relativePath: !this.options.allowExternal,
     })
+
+    this.coverageFilesDirectory = resolve(this.options.reportsDirectory, '.tmp')
   }
 
   resolveOptions() {
@@ -121,43 +123,79 @@ export class IstanbulCoverageProvider extends BaseCoverageProvider implements Co
    * backwards compatibility is a breaking change.
    */
   onAfterSuiteRun({ coverage, transformMode, projectName }: AfterSuiteRunMeta) {
+    if (!coverage)
+      return
+
     if (transformMode !== 'web' && transformMode !== 'ssr')
       throw new Error(`Invalid transform mode: ${transformMode}`)
 
-    let entry = this.coverages.get(projectName || DEFAULT_PROJECT)
+    let entry = this.coverageFiles.get(projectName || DEFAULT_PROJECT)
 
     if (!entry) {
       entry = { web: [], ssr: [] }
-      this.coverages.set(projectName || DEFAULT_PROJECT, entry)
+      this.coverageFiles.set(projectName || DEFAULT_PROJECT, entry)
     }
 
-    entry[transformMode].push(coverage as CoverageMapData)
+    const filename = resolve(this.coverageFilesDirectory, `coverage-${uniqueId++}.json`)
+    entry[transformMode].push(filename)
+
+    const promise = fs.writeFile(filename, JSON.stringify(coverage), 'utf-8')
+    this.pendingPromises.push(promise)
   }
 
   async clean(clean = true) {
     if (clean && existsSync(this.options.reportsDirectory))
       await fs.rm(this.options.reportsDirectory, { recursive: true, force: true, maxRetries: 10 })
 
-    this.coverages = new Map()
+    if (existsSync(this.coverageFilesDirectory))
+      await fs.rm(this.coverageFilesDirectory, { recursive: true, force: true, maxRetries: 10 })
+
+    await fs.mkdir(this.coverageFilesDirectory, { recursive: true })
+
+    this.coverageFiles = new Map()
+    this.pendingPromises = []
   }
 
   async reportCoverage({ allTestsRun }: ReportContext = {}) {
-    const coverageMaps = await Promise.all(
-      Array.from(this.coverages.values()).map(coverages => [
-        mergeAndTransformCoverage(coverages.ssr),
-        mergeAndTransformCoverage(coverages.web),
-      ]).flat(),
-    )
+    const coverageMap = libCoverage.createCoverageMap({})
+    let index = 0
+    const total = this.pendingPromises.length
+
+    await Promise.all(this.pendingPromises)
+    this.pendingPromises = []
+
+    for (const coveragePerProject of this.coverageFiles.values()) {
+      for (const filenames of [coveragePerProject.ssr, coveragePerProject.web]) {
+        const coverageMapByTransformMode = libCoverage.createCoverageMap({})
+
+        for (const chunk of toSlices(filenames, this.options.processingConcurrency)) {
+          if (debug.enabled) {
+            index += chunk.length
+            debug('Covered files %d/%d', index, total)
+          }
+
+          await Promise.all(chunk.map(async (filename) => {
+            const contents = await fs.readFile(filename, 'utf-8')
+            const coverage = JSON.parse(contents) as CoverageMap
+
+            coverageMapByTransformMode.merge(coverage)
+          }))
+        }
+
+        // Source maps can change based on projectName and transform mode.
+        // Coverage transform re-uses source maps so we need to separate transforms from each other.
+        const transformedCoverage = await transformCoverage(coverageMapByTransformMode)
+        coverageMap.merge(transformedCoverage)
+      }
+    }
 
     if (this.options.all && allTestsRun) {
-      const coveredFiles = coverageMaps.map(map => map.files()).flat()
+      const coveredFiles = coverageMap.files()
       const uncoveredCoverage = await this.getCoverageMapForUncoveredFiles(coveredFiles)
 
-      coverageMaps.push(await mergeAndTransformCoverage([uncoveredCoverage]))
+      coverageMap.merge(await transformCoverage(uncoveredCoverage))
     }
 
-    const coverageMap = mergeCoverageMaps(...coverageMaps)
-
     const context = libReport.createContext({
       dir: this.options.reportsDirectory,
       coverageMap,
@@ -206,6 +244,9 @@ export class IstanbulCoverageProvider extends BaseCoverageProvider implements Co
         })
       }
     }
+
+    await fs.rm(this.coverageFilesDirectory, { recursive: true })
+    this.coverageFiles = new Map()
   }
 
   async getCoverageMapForUncoveredFiles(coveredFiles: string[]) {
@@ -218,31 +259,31 @@ export class IstanbulCoverageProvider extends BaseCoverageProvider implements Co
 
     const coverageMap = libCoverage.createCoverageMap({})
 
-    for (const filename of uncoveredFiles) {
+    // Note that these cannot be run parallel as synchronous instrumenter.lastFileCoverage
+    // returns the coverage of the last transformed file
+    for (const [index, filename] of uncoveredFiles.entries()) {
+      debug('Uncovered file %s %d/%d', filename, index, uncoveredFiles.length)
+
+      // Make sure file is not served from cache
+      // so that instrumenter loads up requested file coverage
+      if (this.ctx.vitenode.fetchCache.has(filename))
+        this.ctx.vitenode.fetchCache.delete(filename)
+
       await this.ctx.vitenode.transformRequest(filename)
 
       const lastCoverage = this.instrumenter.lastFileCoverage()
       coverageMap.addFileCoverage(lastCoverage)
     }
 
-    return coverageMap.data
+    return coverageMap
   }
 }
 
-async function mergeAndTransformCoverage(coverages: CoverageMapData[]) {
-  const mergedCoverage = mergeCoverageMaps(...coverages)
-  includeImplicitElseBranches(mergedCoverage)
+async function transformCoverage(coverageMap: CoverageMap) {
+  includeImplicitElseBranches(coverageMap)
 
   const sourceMapStore = libSourceMaps.createSourceMapStore()
-  return await sourceMapStore.transformCoverage(mergedCoverage)
-}
-
-function mergeCoverageMaps(...coverageMaps: (CoverageMap | CoverageMapData)[]) {
-  return coverageMaps.reduce<CoverageMap>((coverage, previousCoverageMap) => {
-    const map = libCoverage.createCoverageMap(coverage)
-    map.merge(previousCoverageMap)
-    return map
-  }, libCoverage.createCoverageMap({}))
+  return await sourceMapStore.transformCoverage(coverageMap)
 }
 
 /**
@@ -302,3 +343,19 @@ function hasTerminalReporter(reporters: Options['reporter']) {
     || reporter === 'text-lcov'
     || reporter === 'teamcity')
 }
+
+function toSlices<T>(array: T[], size: number): T[][] {
+  return array.reduce<T[][]>((chunks, item) => {
+    const index = Math.max(0, chunks.length - 1)
+    const lastChunk = chunks[index] || []
+    chunks[index] = lastChunk
+
+    if (lastChunk.length >= size)
+      chunks.push([item])
+
+    else
+      lastChunk.push(item)
+
+    return chunks
+  }, [])
+}

Diff for: packages/coverage-v8/package.json

@@ -47,6 +47,7 @@
   "dependencies": {
     "@ampproject/remapping": "^2.2.1",
     "@bcoe/v8-coverage": "^0.2.3",
+    "debug": "^4.3.4",
     "istanbul-lib-coverage": "^3.2.2",
     "istanbul-lib-report": "^3.0.1",
     "istanbul-lib-source-maps": "^4.0.1",
@@ -59,6 +60,7 @@
     "v8-to-istanbul": "^9.2.0"
   },
   "devDependencies": {
+    "@types/debug": "^4.1.12",
     "@types/istanbul-lib-coverage": "^2.0.6",
     "@types/istanbul-lib-report": "^3.0.3",
     "@types/istanbul-lib-source-maps": "^4.0.4",