perf: improve performance of forks pool (#5592)

Author: sheremet-va

packages/vite-node/src/server.ts

@@ -31,6 +31,11 @@ export class ViteNodeServer {
     web: new Map<string, Promise<TransformResult | null | undefined>>(),
   }
 
+  private durations = {
+    ssr: new Map<string, number[]>(),
+    web: new Map<string, number[]>(),
+  }
+
   private existingOptimizedDeps = new Set<string>()
 
   fetchCaches = {
@@ -102,6 +107,12 @@ export class ViteNodeServer {
     return shouldExternalize(id, this.options.deps, this.externalizeCache)
   }
 
+  public getTotalDuration() {
+    const ssrDurations = [...this.durations.ssr.values()].flat()
+    const webDurations = [...this.durations.web.values()].flat()
+    return [...ssrDurations, ...webDurations].reduce((a, b) => a + b, 0)
+  }
+
   private async ensureExists(id: string): Promise<boolean> {
     if (this.existingOptimizedDeps.has(id))
       return true
@@ -138,16 +149,20 @@ export class ViteNodeServer {
   }
 
   async fetchModule(id: string, transformMode?: 'web' | 'ssr'): Promise<FetchResult> {
-    const moduleId = normalizeModuleId(id)
     const mode = transformMode || this.getTransformMode(id)
+    return this.fetchResult(id, mode)
+      .then((r) => {
+        return this.options.sourcemap !== true ? { ...r, map: undefined } : r
+      })
+  }
+
+  async fetchResult(id: string, mode: 'web' | 'ssr') {
+    const moduleId = normalizeModuleId(id)
     this.assertMode(mode)
     const promiseMap = this.fetchPromiseMap[mode]
     // reuse transform for concurrent requests
     if (!promiseMap.has(moduleId)) {
       promiseMap.set(moduleId, this._fetchModule(moduleId, mode)
-        .then((r) => {
-          return this.options.sourcemap !== true ? { ...r, map: undefined } : r
-        })
         .finally(() => {
           promiseMap.delete(moduleId)
         }))
@@ -268,6 +283,12 @@ export class ViteNodeServer {
       result,
     }
 
+    const durations = this.durations[transformMode].get(filePath) || []
+    this.durations[transformMode].set(
+      filePath,
+      [...durations, duration ?? 0],
+    )
+
     this.fetchCaches[transformMode].set(filePath, cacheEntry)
     this.fetchCache.set(filePath, cacheEntry)
 

packages/vitest/src/integrations/env/loader.ts

@@ -1,3 +1,4 @@
+import { readFileSync } from 'node:fs'
 import { normalize, resolve } from 'pathe'
 import { ViteNodeRunner } from 'vite-node/client'
 import type { ViteNodeRunnerOptions } from 'vite-node'
@@ -26,7 +27,12 @@ export async function loadEnvironment(ctx: ContextRPC, rpc: WorkerRPC): Promise<
     return environments[name]
   const loader = await createEnvironmentLoader({
     root: ctx.config.root,
-    fetchModule: id => rpc.fetch(id, 'ssr'),
+    fetchModule: async (id) => {
+      const result = await rpc.fetch(id, 'ssr')
+      if (result.id)
+        return { code: readFileSync(result.id, 'utf-8') }
+      return result
+    },
     resolveId: (id, importer) => rpc.resolveId(id, importer, 'ssr'),
   })
   const root = loader.root

packages/vitest/src/node/pools/rpc.ts

@@ -1,7 +1,12 @@
+import { mkdir, writeFile } from 'node:fs/promises'
 import type { RawSourceMap } from 'vite-node'
+import { join } from 'pathe'
 import type { RuntimeRPC } from '../../types'
 import type { WorkspaceProject } from '../workspace'
 
+const created = new Set()
+const promises = new Map<string, Promise<void>>()
+
 export function createMethodsRPC(project: WorkspaceProject): RuntimeRPC {
   const ctx = project.ctx
   return {
@@ -20,8 +25,31 @@ export function createMethodsRPC(project: WorkspaceProject): RuntimeRPC {
       const r = await project.vitenode.transformRequest(id)
       return r?.map as RawSourceMap | undefined
     },
-    fetch(id, transformMode) {
-      return project.vitenode.fetchModule(id, transformMode)
+    async fetch(id, transformMode) {
+      const result = await project.vitenode.fetchResult(id, transformMode)
+      const code = result.code
+      if (result.externalize)
+        return result
+      if ('id' in result)
+        return { id: result.id as string }
+
+      if (!code)
+        throw new Error(`Failed to fetch module ${id}`)
+
+      const dir = join(project.tmpDir, transformMode)
+      const tmp = join(dir, id.replace(/[/\\?%*:|"<>]/g, '_').replace('\0', '__x00__'))
+      if (promises.has(tmp)) {
+        await promises.get(tmp)
+        return { id: tmp }
+      }
+      if (!created.has(dir)) {
+        await mkdir(dir, { recursive: true })
+        created.add(dir)
+      }
+      promises.set(tmp, writeFile(tmp, code, 'utf-8').finally(() => promises.delete(tmp)))
+      await promises.get(tmp)
+      Object.assign(result, { id: tmp })
+      return { id: tmp }
     },
     resolveId(id, importer, transformMode) {
       return project.vitenode.resolveId(id, importer, transformMode)

packages/vitest/src/node/reporters/base.ts

@@ -232,7 +232,7 @@ export abstract class BaseReporter implements Reporter {
     const setupTime = files.reduce((acc, test) => acc + Math.max(0, test.setupDuration || 0), 0)
     const testsTime = files.reduce((acc, test) => acc + Math.max(0, test.result?.duration || 0), 0)
     const transformTime = this.ctx.projects
-      .flatMap(w => Array.from(w.vitenode.fetchCache.values()).map(i => i.duration || 0))
+      .flatMap(w => w.vitenode.getTotalDuration())
       .reduce((a, b) => a + b, 0)
     const environmentTime = files.reduce((acc, file) => acc + Math.max(0, file.environmentLoad || 0), 0)
     const prepareTime = files.reduce((acc, file) => acc + Math.max(0, file.prepareDuration || 0), 0)

packages/vitest/src/node/workspace.ts

@@ -1,4 +1,6 @@
 import { promises as fs } from 'node:fs'
+import { rm } from 'node:fs/promises'
+import { tmpdir } from 'node:os'
 import fg from 'fast-glob'
 import mm from 'micromatch'
 import { dirname, isAbsolute, join, relative, resolve, toNamespacedPath } from 'pathe'
@@ -8,10 +10,10 @@ import { ViteNodeServer } from 'vite-node/server'
 import c from 'picocolors'
 import { createBrowserServer } from '../integrations/browser/server'
 import type { ProvidedContext, ResolvedConfig, UserConfig, UserWorkspaceConfig, Vitest } from '../types'
-import { deepMerge } from '../utils'
 import type { Typechecker } from '../typecheck/typechecker'
 import type { BrowserProvider } from '../types/browser'
 import { getBrowserProvider } from '../integrations/browser'
+import { deepMerge, nanoid } from '../utils/base'
 import { isBrowserEnabled, resolveConfig } from './config'
 import { WorkspaceVitestPlugin } from './plugins/workspace'
 import { createViteServer } from './vite'
@@ -78,6 +80,9 @@ export class WorkspaceProject {
 
   testFilesList: string[] | null = null
 
+  public readonly id = nanoid()
+  public readonly tmpDir = join(tmpdir(), this.id)
+
   private _globalSetups: GlobalSetupFile[] | undefined
   private _provided: ProvidedContext = {} as any
 
@@ -402,11 +407,19 @@ export class WorkspaceProject {
         this.server.close(),
         this.typechecker?.stop(),
         this.browser?.close(),
+        this.clearTmpDir(),
       ].filter(Boolean)).then(() => this._provided = {} as any)
     }
     return this.closingPromise
   }
 
+  private async clearTmpDir() {
+    try {
+      await rm(this.tmpDir, { force: true, recursive: true })
+    }
+    catch {}
+  }
+
   async initBrowserProvider() {
     if (!this.isBrowserEnabled())
       return

packages/vitest/src/runtime/execute.ts

@@ -1,5 +1,6 @@
 import vm from 'node:vm'
 import { pathToFileURL } from 'node:url'
+import { readFileSync } from 'node:fs'
 import type { ModuleCacheMap } from 'vite-node/client'
 import { DEFAULT_REQUEST_STUBS, ViteNodeRunner } from 'vite-node/client'
 import { isInternalRequest, isNodeBuiltin, isPrimitive, toFilePath } from 'vite-node/utils'
@@ -104,7 +105,12 @@ export async function startVitestExecutor(options: ContextExecutorOptions) {
         return { externalize: id }
       }
 
-      return rpc().fetch(id, getTransformMode())
+      const result = await rpc().fetch(id, getTransformMode())
+      if (result.id && !result.externalize) {
+        const code = readFileSync(result.id, 'utf-8')
+        return { code }
+      }
+      return result
     },
     resolveId(id, importer) {
       return rpc().resolveId(id, importer, getTransformMode())

packages/vitest/src/types/rpc.ts

@@ -9,7 +9,10 @@ import type { AfterSuiteRunMeta } from './worker'
 type TransformMode = 'web' | 'ssr'
 
 export interface RuntimeRPC {
-  fetch: (id: string, environment: TransformMode) => Promise<FetchResult>
+  fetch: (id: string, environment: TransformMode) => Promise<{
+    externalize?: string
+    id?: string
+  }>
   transform: (id: string, environment: TransformMode) => Promise<FetchResult>
   resolveId: (id: string, importer: string | undefined, environment: TransformMode) => Promise<ViteNodeResolveId | null>
   getSourceMap: (id: string, force?: boolean) => Promise<RawSourceMap | undefined>

packages/vitest/src/utils/base.ts

@@ -169,3 +169,14 @@ export function escapeRegExp(s: string) {
 export function wildcardPatternToRegExp(pattern: string): RegExp {
   return new RegExp(`^${pattern.split('*').map(escapeRegExp).join('.*')}$`, 'i')
 }
+
+// port from nanoid
+// https://github.com/ai/nanoid
+const urlAlphabet
+  = 'useandom-26T198340PX75pxJACKVERYMINDBUSHWOLF_GQZbfghjklqvwyzrict'
+export function nanoid(size = 21) {
+  let id = ''
+  let i = size
+  while (i--) id += urlAlphabet[(Math.random() * 64) | 0]
+  return id
+}

packages/web-worker/src/utils.ts

@@ -1,3 +1,4 @@
+import { readFileSync } from 'node:fs'
 import type { WorkerGlobalState } from 'vitest'
 import ponyfillStructuredClone from '@ungap/structured-clone'
 import createDebug from 'debug'
@@ -65,8 +66,13 @@ export function getRunnerOptions(): any {
   const { config, rpc, mockMap, moduleCache } = state
 
   return {
-    fetchModule(id: string) {
-      return rpc.fetch(id, 'web')
+    async fetchModule(id: string) {
+      const result = await rpc.fetch(id, 'web')
+      if (result.id && !result.externalize) {
+        const code = readFileSync(result.id, 'utf-8')
+        return { code }
+      }
+      return result
     },
     resolveId(id: string, importer?: string) {
       return rpc.resolveId(id, importer, 'web')