perf(pool): resolve all environments first (#8759)

Author: sheremet-va

packages/vitest/src/node/pool.ts

@@ -11,7 +11,7 @@ import { version as viteVersion } from 'vite'
 import { rootDir } from '../paths'
 import { isWindows } from '../utils/env'
 import { getWorkerMemoryLimit, stringToBytes } from '../utils/memory-limit'
-import { groupFilesByEnv } from '../utils/test-helpers'
+import { getSpecificationsEnvironments } from '../utils/test-helpers'
 import { createBrowserPool } from './pools/browser'
 import { Pool } from './pools/pool'
 
@@ -86,6 +86,7 @@ export function createPool(ctx: Vitest): ProcessPool {
     let workerId = 0
 
     const sorted = await sequencer.sort(specs)
+    const environments = await getSpecificationsEnvironments(specs)
     const groups = groupSpecs(sorted)
 
     for (const group of groups) {
@@ -108,16 +109,18 @@ export function createPool(ctx: Vitest): ProcessPool {
           continue
         }
 
-        const byEnv = await groupFilesByEnv(specs)
-        const env = Object.values(byEnv)[0][0]
+        const environment = environments.get(specs[0])!
+        if (!environment) {
+          throw new Error(`Cannot find the environment. This is a bug in Vitest.`)
+        }
 
         taskGroup.push({
           context: {
             pool,
             config: project.serializedConfig,
             files: specs.map(spec => ({ filepath: spec.moduleId, testLocations: spec.testLines })),
             invalidates,
-            environment: env.environment,
+            environment,
             projectName: project.name,
             providedContext: project.getProvidedContext(),
             workerId: workerId++,
@@ -313,7 +316,7 @@ function groupSpecs(specs: TestSpecification[]) {
 
   // Tests in a single group are executed with `maxWorkers` parallelism.
   // Next group starts running after previous finishes - allows real sequential tests.
-  interface Groups { specs: SpecsForRunner[]; maxWorkers: number }
+  interface Groups { specs: SpecsForRunner[]; maxWorkers: number; typecheck?: boolean }
   const groups: Groups[] = []
 
   // Files without file parallelism but without explicit sequence.groupOrder
@@ -349,11 +352,17 @@ function groupSpecs(specs: TestSpecification[]) {
     groups[order].specs.push([spec])
   })
 
-  for (const project in typechecks) {
-    const order = Math.max(0, ...groups.keys()) + 1
+  let order = Math.max(0, ...groups.keys()) + 1
+
+  for (const projectName in typechecks) {
+    const maxWorkers = resolveMaxWorkers(typechecks[projectName][0].project)
+    const previous = groups[order - 1]
+    if (previous && previous.typecheck && maxWorkers !== previous.maxWorkers) {
+      order += 1
+    }
 
-    groups[order] ||= { specs: [], maxWorkers: resolveMaxWorkers(typechecks[project][0].project) }
-    groups[order].specs.push(typechecks[project])
+    groups[order] ||= { specs: [], maxWorkers, typecheck: true }
+    groups[order].specs.push(typechecks[projectName])
   }
 
   if (sequential.specs.length) {

packages/vitest/src/utils/test-helpers.ts

@@ -1,28 +1,22 @@
-import type { TestProject } from '../node/project'
 import type { TestSpecification } from '../node/spec'
 import type { EnvironmentOptions, VitestEnvironment } from '../node/types/config'
 import type { ContextTestEnvironment } from '../types/worker'
 import { promises as fs } from 'node:fs'
-import { groupBy } from './base'
 
-export const envsOrder: string[] = ['node', 'jsdom', 'happy-dom', 'edge-runtime']
-
-export interface FileByEnv {
-  file: string
-  env: VitestEnvironment
-  envOptions: EnvironmentOptions | null
-}
-
-export async function groupFilesByEnv(
-  files: Array<TestSpecification>,
-): Promise<Record<string, {
-  file: { filepath: string; testLocations: number[] | undefined }
-  project: TestProject
-  environment: ContextTestEnvironment
-}[]>> {
-  const filesWithEnv = await Promise.all(
-    files.map(async ({ moduleId: filepath, project, testLines }) => {
-      const code = await fs.readFile(filepath, 'utf-8')
+export async function getSpecificationsEnvironments(
+  specifications: Array<TestSpecification>,
+): Promise<WeakMap<TestSpecification, ContextTestEnvironment>> {
+  const environments = new WeakMap<TestSpecification, ContextTestEnvironment>()
+  const cache = new Map<string, string>()
+  await Promise.all(
+    specifications.map(async (spec) => {
+      const { moduleId: filepath, project } = spec
+      // reuse if projects have the same test files
+      let code = cache.get(filepath)
+      if (!code) {
+        code = await fs.readFile(filepath, 'utf-8')
+        cache.set(filepath, code)
+      }
 
       // 1. Check for control comments in the file
       let env = code.match(/@(?:vitest|jest)-environment\s+([\w-]+)\b/)?.[1]
@@ -43,16 +37,8 @@ export async function groupFilesByEnv(
           ? ({ [envKey]: envOptions } as EnvironmentOptions)
           : null,
       }
-      return {
-        file: {
-          filepath,
-          testLocations: testLines,
-        },
-        project,
-        environment,
-      }
+      environments.set(spec, environment)
     }),
   )
-
-  return groupBy(filesWithEnv, ({ environment }) => environment.name)
+  return environments
 }