feat: allow custom pools (#4417)

docs/.vitepress/config.ts

@@ -131,9 +131,13 @@ export default withPwa(defineConfig({
               link: '/advanced/metadata',
             },
             {
-              text: 'Extending default reporters',
+              text: 'Extending Reporters',
               link: '/advanced/reporters',
             },
+            {
+              text: 'Custom Pool',
+              link: '/advanced/pool',
+            },
           ],
         },
       ],

docs/advanced/pool.md

@@ -0,0 +1,90 @@
+# Custom Pool
+
+::: warning
+This is advanced API. If you are just running tests, you probably don't need this. It is primarily used by library authors.
+:::
+
+Vitest runs tests in pools. By default, there are several pools:
+
+- `threads` to run tests using `node:worker_threads` (isolation is provided with a new worker context)
+- `forks` to run tests using `node:child_process` (isolation is provided with a new `child_process.fork` process)
+- `vmThreads` to run tests using `node:worker_threads` (but isolation is provided with `vm` module instead of a new worker context)
+- `browser` to run tests using browser providers
+- `typescript` to run typechecking on tests
+
+You can provide your own pool by specifying a file path:
+
+```ts
+export default defineConfig({
+  test: {
+    // will run every file with a custom pool by default
+    pool: './my-custom-pool.ts',
+    // you can provide options using `poolOptions` object
+    poolOptions: {
+      myCustomPool: {
+        customProperty: true,
+      },
+    },
+    // you can also specify pool for a subset of files
+    poolMatchGlobs: [
+      ['**/*.custom.test.ts', './my-custom-pool.ts'],
+    ],
+  },
+})
+```
+
+## API
+
+The file specified in `pool` option should export a function (can be async) that accepts `Vitest` interface as its first option. This function needs to return an object matching `ProcessPool` interface:
+
+```ts
+import { ProcessPool, WorkspaceProject } from 'vitest/node'
+
+export interface ProcessPool {
+  name: string
+  runTests: (files: [project: WorkspaceProject, testFile: string][], invalidates?: string[]) => Promise<void>
+  close?: () => Promise<void>
+}
+```
+
+The function is called only once (unless the server config was updated), and it's generally a good idea to initialize everything you need for tests inside that function and reuse it when `runTests` is called.
+
+Vitest calls `runTest` when new tests are scheduled to run. It will not call it if `files` is empty. The first argument is an array of tuples: the first element is a reference to a workspace project and the second one is an absolute path to a test file. Files are sorted using [`sequencer`](/config/#sequence.sequencer) before `runTests` is called. It's possible (but unlikely) to have the same file twice, but it will always have a different project - this is implemented via [`vitest.workspace.ts`](/guide/workspace) configuration.
+
+Vitest will wait until `runTests` is executed before finishing a run (i.e., it will emit [`onFinished`](/guide/reporters) only after `runTests` is resolved).
+
+If you are using a custom pool, you will have to provide test files and their results yourself - you can reference [`vitest.state`](https://github.com/vitest-dev/vitest/blob/feat/custom-pool/packages/vitest/src/node/state.ts) for that (most important are `collectFiles` and `updateTasks`). Vitest uses `startTests` function from `@vitest/runner` package to do that.
+
+To communicate between different processes, you can create methods object using `createMethodsRPC` from `vitest/node`, and use any form of communication that you prefer. For example, to use websockets with `birpc` you can write something like this:
+
+```ts
+import { createBirpc } from 'birpc'
+import { parse, stringify } from 'flatted'
+import { WorkspaceProject, createMethodsRPC } from 'vitest/node'
+
+function createRpc(project: WorkspaceProject, wss: WebSocketServer) {
+  return createBirpc(
+    createMethodsRPC(project),
+    {
+      post: msg => wss.send(msg),
+      on: fn => wss.on('message', fn),
+      serialize: stringify,
+      deserialize: parse,
+    },
+  )
+}
+```
+
+To make sure every test is collected, you would call `ctx.state.collectFiles` and report it to Vitest reporters:
+
+```ts
+async function runTests(project: WorkspaceProject, tests: string[]) {
+  // ... running tests, put into "files" and "tasks"
+  const methods = createMethodsRPC(project)
+  await methods.onCollected(files)
+  // most reporters rely on results being updated in "onTaskUpdate"
+  await methods.onTaskUpdate(tasks)
+}
+```
+
+You can see a simple example in [pool/custom-pool.ts](https://github.com/vitest-dev/vitest/blob/feat/custom-pool/test/run/pool-custom-fixtures/pool/custom-pool.ts).

docs/advanced/reporters.md

@@ -1,8 +1,8 @@
-# Extending default reporters
+# Extending Reporters
 
 You can import reporters from `vitest/reporters` and extend them to create your custom reporters.
 
-## Extending built-in reporters
+## Extending Built-in Reporters
 
 In general, you don't need to create your reporter from scratch. `vitest` comes with several default reporting programs that you can extend.
 
@@ -56,7 +56,7 @@ export default defineConfig({
 })
 ```
 
-## Exported reporters
+## Exported Reporters
 
 `vitest` comes with a few [built-in reporters](/guide/reporters) that you can use out of the box.
 

packages/runner/src/suite.ts

@@ -154,6 +154,7 @@ function createSuiteCollector(name: string, factory: SuiteFactory = () => { }, m
       shuffle,
       tasks: [],
       meta: Object.create(null),
+      projectName: '',
     }
 
     setHooks(suite, createSuiteHooks())

packages/runner/src/types/tasks.ts

@@ -52,7 +52,7 @@ export interface Suite extends TaskBase {
   type: 'suite'
   tasks: Task[]
   filepath?: string
-  projectName?: string
+  projectName: string
 }
 
 export interface File extends Suite {

packages/vitest/src/node/config.ts

@@ -6,10 +6,12 @@ import type { ApiConfig, ResolvedConfig, UserConfig, VitestRunMode } from '../ty
 import { defaultBrowserPort, defaultPort } from '../constants'
 import { benchmarkConfigDefaults, configDefaults } from '../defaults'
 import { isCI, stdProvider, toArray } from '../utils'
+import type { BuiltinPool } from '../types/pool-options'
 import { VitestCache } from './cache'
 import { BaseSequencer } from './sequencers/BaseSequencer'
 import { RandomSequencer } from './sequencers/RandomSequencer'
 import type { BenchmarkBuiltinReporters } from './reporters'
+import { builtinPools } from './pool'
 
 const extraInlineDeps = [
   /^(?!.*(?:node_modules)).*\.mjs$/,
@@ -222,6 +224,8 @@ export function resolveConfig(
   if (options.resolveSnapshotPath)
     delete (resolved as UserConfig).resolveSnapshotPath
 
+  resolved.pool ??= 'threads'
+
   if (process.env.VITEST_MAX_THREADS) {
     resolved.poolOptions = {
       ...resolved.poolOptions,
@@ -270,6 +274,22 @@ export function resolveConfig(
     }
   }
 
+  if (!builtinPools.includes(resolved.pool as BuiltinPool)) {
+    resolved.pool = normalize(
+      resolveModule(resolved.pool, { paths: [resolved.root] })
+        ?? resolve(resolved.root, resolved.pool),
+    )
+  }
+  resolved.poolMatchGlobs = (resolved.poolMatchGlobs || []).map(([glob, pool]) => {
+    if (!builtinPools.includes(pool as BuiltinPool)) {
+      pool = normalize(
+        resolveModule(pool, { paths: [resolved.root] })
+          ?? resolve(resolved.root, pool),
+      )
+    }
+    return [glob, pool]
+  })
+
   if (mode === 'benchmark') {
     resolved.benchmark = {
       ...benchmarkConfigDefaults,

packages/vitest/src/node/core.ts

@@ -80,7 +80,7 @@ export class Vitest {
     this.unregisterWatcher?.()
     clearTimeout(this._rerunTimer)
     this.restartsCount += 1
-    this.pool?.close()
+    this.pool?.close?.()
     this.pool = undefined
     this.coverageProvider = undefined
     this.runningPromise = undefined
@@ -761,8 +761,12 @@ export class Vitest {
         if (!this.projects.includes(this.coreWorkspaceProject))
           closePromises.push(this.coreWorkspaceProject.close().then(() => this.server = undefined as any))
 
-        if (this.pool)
-          closePromises.push(this.pool.close().then(() => this.pool = undefined))
+        if (this.pool) {
+          closePromises.push((async () => {
+            await this.pool?.close?.()
+            this.pool = undefined
+          })())
+        }
 
         closePromises.push(...this._onClose.map(fn => fn()))
 

packages/vitest/src/node/index.ts

@@ -4,8 +4,9 @@ export { createVitest } from './create'
 export { VitestPlugin } from './plugins'
 export { startVitest } from './cli-api'
 export { registerConsoleShortcuts } from './stdin'
-export type { WorkspaceSpec } from './pool'
 export type { GlobalSetupContext } from './globalSetup'
+export type { WorkspaceSpec, ProcessPool } from './pool'
+export { createMethodsRPC } from './pools/rpc'
 
 export type { TestSequencer, TestSequencerConstructor } from './sequencers/types'
 export { BaseSequencer } from './sequencers/BaseSequencer'

packages/vitest/src/node/pool.ts

@@ -1,5 +1,6 @@
 import mm from 'micromatch'
-import type { Pool } from '../types'
+import type { Awaitable } from '@vitest/utils'
+import type { BuiltinPool, Pool } from '../types/pool-options'
 import type { Vitest } from './core'
 import { createChildProcessPool } from './pools/child'
 import { createThreadsPool } from './pools/threads'
@@ -9,11 +10,12 @@ import type { WorkspaceProject } from './workspace'
 import { createTypecheckPool } from './pools/typecheck'
 
 export type WorkspaceSpec = [project: WorkspaceProject, testFile: string]
-export type RunWithFiles = (files: WorkspaceSpec[], invalidates?: string[]) => Promise<void>
+export type RunWithFiles = (files: WorkspaceSpec[], invalidates?: string[]) => Awaitable<void>
 
 export interface ProcessPool {
+  name: string
   runTests: RunWithFiles
-  close: () => Promise<void>
+  close?: () => Awaitable<void>
 }
 
 export interface PoolProcessOptions {
@@ -24,6 +26,8 @@ export interface PoolProcessOptions {
   env: Record<string, string>
 }
 
+export const builtinPools: BuiltinPool[] = ['forks', 'threads', 'browser', 'vmThreads', 'typescript']
+
 export function createPool(ctx: Vitest): ProcessPool {
   const pools: Record<Pool, ProcessPool | null> = {
     forks: null,
@@ -48,7 +52,7 @@ export function createPool(ctx: Vitest): ProcessPool {
   }
 
   function getPoolName([project, file]: WorkspaceSpec) {
-    for (const [glob, pool] of project.config.poolMatchGlobs || []) {
+    for (const [glob, pool] of project.config.poolMatchGlobs) {
       if ((pool as Pool) === 'browser')
         throw new Error('Since Vitest 0.31.0 "browser" pool is not supported in "poolMatchGlobs". You can create a workspace to run some of your tests in browser in parallel. Read more: https://vitest.dev/guide/workspace')
       if (mm.isMatch(file, glob, { cwd: project.config.root }))
@@ -82,6 +86,22 @@ export function createPool(ctx: Vitest): ProcessPool {
       },
     }
 
+    const customPools = new Map<string, ProcessPool>()
+    async function resolveCustomPool(filepath: string) {
+      if (customPools.has(filepath))
+        return customPools.get(filepath)!
+      const pool = await ctx.runner.executeId(filepath)
+      if (typeof pool.default !== 'function')
+        throw new Error(`Custom pool "${filepath}" must export a function as default export`)
+      const poolInstance = await pool.default(ctx, options)
+      if (typeof poolInstance?.name !== 'string')
+        throw new Error(`Custom pool "${filepath}" should return an object with "name" property`)
+      if (typeof poolInstance?.runTests !== 'function')
+        throw new Error(`Custom pool "${filepath}" should return an object with "runTests" method`)
+      customPools.set(filepath, poolInstance)
+      return poolInstance as ProcessPool
+    }
+
     const filesByPool: Record<Pool, WorkspaceSpec[]> = {
       forks: [],
       threads: [],
@@ -92,46 +112,63 @@ export function createPool(ctx: Vitest): ProcessPool {
 
     for (const spec of files) {
       const pool = getPoolName(spec)
-      if (!(pool in filesByPool))
-        throw new Error(`Unknown pool name "${pool}" for ${spec[1]}. Available pools: ${Object.keys(filesByPool).join(', ')}`)
+      filesByPool[pool] ??= []
       filesByPool[pool].push(spec)
     }
 
-    await Promise.all(Object.entries(filesByPool).map((entry) => {
+    const Sequencer = ctx.config.sequence.sequencer
+    const sequencer = new Sequencer(ctx)
+
+    async function sortSpecs(specs: WorkspaceSpec[]) {
+      if (ctx.config.shard)
+        specs = await sequencer.shard(specs)
+      return sequencer.sort(specs)
+    }
+
+    await Promise.all(Object.entries(filesByPool).map(async (entry) => {
       const [pool, files] = entry as [Pool, WorkspaceSpec[]]
 
       if (!files.length)
         return null
 
+      const specs = await sortSpecs(files)
+
       if (pool === 'browser') {
         pools.browser ??= createBrowserPool(ctx)
-        return pools.browser.runTests(files, invalidate)
+        return pools.browser.runTests(specs, invalidate)
       }
 
       if (pool === 'vmThreads') {
         pools.vmThreads ??= createVmThreadsPool(ctx, options)
-        return pools.vmThreads.runTests(files, invalidate)
+        return pools.vmThreads.runTests(specs, invalidate)
       }
 
       if (pool === 'threads') {
         pools.threads ??= createThreadsPool(ctx, options)
-        return pools.threads.runTests(files, invalidate)
+        return pools.threads.runTests(specs, invalidate)
       }
 
       if (pool === 'typescript') {
         pools.typescript ??= createTypecheckPool(ctx)
-        return pools.typescript.runTests(files)
+        return pools.typescript.runTests(specs)
+      }
+
+      if (pool === 'forks') {
+        pools.forks ??= createChildProcessPool(ctx, options)
+        return pools.forks.runTests(specs, invalidate)
       }
 
-      pools.forks ??= createChildProcessPool(ctx, options)
-      return pools.forks.runTests(files, invalidate)
+      const poolHandler = await resolveCustomPool(pool)
+      pools[poolHandler.name] ??= poolHandler
+      return poolHandler.runTests(specs, invalidate)
     }))
   }
 
   return {
+    name: 'default',
     runTests,
     async close() {
-      await Promise.all(Object.values(pools).map(p => p?.close()))
+      await Promise.all(Object.values(pools).map(p => p?.close?.()))
     },
   }
 }

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

@@ -44,7 +44,7 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
     if (project.config.browser.isolate) {
       for (const path of paths) {
         if (isCancelled) {
-          ctx.state.cancelFiles(files.slice(paths.indexOf(path)), ctx.config.root)
+          ctx.state.cancelFiles(files.slice(paths.indexOf(path)), ctx.config.root, project.getName())
           break
         }
 
@@ -77,6 +77,7 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
   }
 
   return {
+    name: 'browser',
     async close() {
       ctx.state.browserTestPromises.clear()
       await Promise.all([...providers].map(provider => provider.close()))