fix(pool): properly reuse the vm pool (#8758)

Author: sheremet-va

packages/vitest/src/node/pool.ts

@@ -123,7 +123,7 @@ export function createPool(ctx: Vitest): ProcessPool {
             workerId: workerId++,
           },
           project,
-          env: options.env,
+          env: { ...options.env, ...project.config.env },
           execArgv: [...options.execArgv, ...project.config.execArgv],
           worker: pool,
           isolate: project.config.isolate,

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

@@ -281,17 +281,6 @@ function isEqualRunner(runner: PoolRunner, task: PoolTask) {
     runner.worker.name === task.worker
     && runner.project === task.project
     && runner.environment === task.context.environment.name
-    && runner.worker.execArgv.every((arg, index) => task.execArgv[index] === arg)
-    && isEnvEqual(runner.worker.env, task.env)
+    && (!runner.worker.canReuse || runner.worker.canReuse(task))
   )
 }
-
-function isEnvEqual(a: PoolOptions['env'], b: PoolTask['env']) {
-  const keys = Object.keys(a)
-
-  if (keys.length !== Object.keys(b).length) {
-    return false
-  }
-
-  return keys.every(key => a[key] === b[key])
-}

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

@@ -13,13 +13,11 @@ export interface PoolOptions {
   cacheFs?: boolean
   environment: string
   execArgv: string[]
-  env: Record<string, string>
+  env: Partial<NodeJS.ProcessEnv>
 }
 
 export interface PoolWorker {
   readonly name: string
-  readonly execArgv: string[]
-  readonly env: Record<string, string>
   readonly reportMemory?: boolean
   readonly cacheFs?: boolean
 
@@ -30,13 +28,20 @@ export interface PoolWorker {
 
   start: () => Promise<void>
   stop: () => Promise<void>
+
+  /**
+   * This is called on workers that already satisfy certain constraints:
+   * - The task has the same project
+   * - The task has the same environment
+   */
+  canReuse?: (task: PoolTask) => boolean
 }
 
 export interface PoolTask {
   worker: 'forks' | 'threads' | 'vmForks' | 'vmThreads' | (string & {})
   project: TestProject
   isolate: boolean
-  env: Record<string, string>
+  env: Partial<NodeJS.ProcessEnv>
   execArgv: string[]
   context: ContextRPC
   memoryLimit: number | null

packages/vitest/src/node/pools/workers/forksWorker.ts

@@ -10,11 +10,11 @@ const SIGKILL_TIMEOUT = 500 /** jest does 500ms by default, let's follow it */
 /** @experimental */
 export class ForksPoolWorker implements PoolWorker {
   public readonly name: string = 'forks'
-  public readonly execArgv: string[]
-  public readonly env: Record<string, string>
   public readonly cacheFs: boolean = true
 
   protected readonly entrypoint: string
+  protected execArgv: string[]
+  protected env: Partial<NodeJS.ProcessEnv>
 
   private _fork?: ChildProcess
 

packages/vitest/src/node/pools/workers/threadsWorker.ts

@@ -5,9 +5,10 @@ import { Worker } from 'node:worker_threads'
 /** @experimental */
 export class ThreadsPoolWorker implements PoolWorker {
   public readonly name: string = 'threads'
-  public readonly execArgv: string[]
-  public readonly env: Record<string, string>
+
   protected readonly entrypoint: string
+  protected execArgv: string[]
+  protected env: Partial<NodeJS.ProcessEnv>
 
   private _thread?: Worker
 

packages/vitest/src/node/pools/workers/typecheckWorker.ts

@@ -13,15 +13,11 @@ import { Typechecker } from '../../../typecheck/typechecker'
 /** @experimental */
 export class TypecheckPoolWorker implements PoolWorker {
   public readonly name: string = 'typecheck'
-  public readonly execArgv: string[]
-  public readonly env: Record<string, string>
   private readonly project: TestProject
 
   private _eventEmitter = new EventEmitter()
 
   constructor(options: PoolOptions) {
-    this.execArgv = options.execArgv
-    this.env = options.env
     this.project = options.project
   }
 
@@ -33,6 +29,10 @@ export class TypecheckPoolWorker implements PoolWorker {
     // noop, onMessage handles it
   }
 
+  canReuse(): boolean {
+    return true
+  }
+
   send(message: WorkerRequest): void {
     void onMessage(message, this.project).then((response) => {
       if (response) {

packages/vitest/src/node/pools/workers/vmForksWorker.ts

@@ -9,7 +9,8 @@ export class VmForksPoolWorker extends ForksPoolWorker {
   protected readonly entrypoint: string
 
   constructor(options: PoolOptions) {
-    super({ ...options, execArgv: [...options.execArgv, '--experimental-vm-modules'] })
+    super(options)
+    this.execArgv.push('--experimental-vm-modules')
 
     /** Loads {@link file://./../../../runtime/workers/vmForks.ts} */
     this.entrypoint = resolve(options.distPath, 'workers/vmForks.js')

packages/vitest/src/node/pools/workers/vmThreadsWorker.ts

@@ -9,7 +9,8 @@ export class VmThreadsPoolWorker extends ThreadsPoolWorker {
   protected readonly entrypoint: string
 
   constructor(options: PoolOptions) {
-    super({ ...options, execArgv: [...options.execArgv, '--experimental-vm-modules'] })
+    super(options)
+    this.execArgv.push('--experimental-vm-modules')
 
     /** Loads {@link file://./../../../runtime/workers/vmThreads.ts} */
     this.entrypoint = resolve(options.distPath, 'workers/vmThreads.js')

test/cli/fixtures/custom-pool/pool/custom-pool.ts

@@ -20,13 +20,9 @@ export class CustomRuntimeWorker implements PoolWorker {
   public readonly name = 'custom'
   private vitest: Vitest
   private customEvents = new EventEmitter()
-  readonly execArgv: string[]
-  readonly env: Record<string, string>
   private project: TestProject
 
   constructor(options: PoolOptions, private settings: OptionsCustomPool) {
-    this.execArgv = options.execArgv
-    this.env = options.env
     this.vitest = options.project.vitest
     this.project = options.project
   }