perf(browser): improve the parallelisation logic

packages/browser/src/node/pool.ts

@@ -1,7 +1,8 @@
+import type { DeferPromise } from '@vitest/utils'
 import type { BrowserProvider, ProcessPool, Vitest, WorkspaceProject, WorkspaceSpec } from 'vitest/node'
 import crypto from 'node:crypto'
 import * as nodeos from 'node:os'
-import { relative } from 'pathe'
+import { createDefer } from '@vitest/utils'
 import { createDebugger } from 'vitest/node'
 
 const debug = createDebugger('vitest:browser:pool')
@@ -19,7 +20,12 @@ async function waitForTests(
 export function createBrowserPool(ctx: Vitest): ProcessPool {
   const providers = new Set<BrowserProvider>()
 
-  const executeTests = async (method: 'run' | 'collect', project: WorkspaceProject, files: string[]) => {
+  const executeTests = async (
+    defer: DeferPromise<void>,
+    method: 'run' | 'collect',
+    project: WorkspaceProject,
+    files: string[],
+  ) => {
     ctx.state.clearFiles(project, files)
     const browser = project.browser!
 
@@ -54,58 +60,49 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
       })
     }
 
-    const filesPerThread = Math.ceil(files.length / threadsCount)
-
-    // TODO: make it smarter,
-    // Currently if we run 4/4/4/4 tests, and one of the chunks ends,
-    // but there are pending tests in another chunks, we can't redistribute them
-    const chunks: string[][] = []
-    for (let i = 0; i < files.length; i += filesPerThread) {
-      const chunk = files.slice(i, i + filesPerThread)
-      chunks.push(chunk)
-    }
+    const orchestrators = [...browser.state.orchestrators.entries()]
 
-    debug?.(
-      `[%s] Running %s tests in %s chunks (%s threads)`,
-      project.getName() || 'core',
-      files.length,
-      chunks.length,
-      threadsCount,
-    )
+    browser.state.onReady(async (contextId, orchestrator) => {
+      const file = files.shift()
+      if (!file) {
+        browser.state.cleanListeners()
+        defer.resolve()
+        // No more files to run
+        // resolve the context
+        return
+      }
+      waitForTests(method, contextId, project, [file])
+      orchestrator.createTesters([file])
+    })
 
-    const orchestrators = [...browser.state.orchestrators.entries()]
+    browser.state.onError((_, error) => {
+      browser.state.cleanListeners()
+      defer.reject(error)
+    })
 
-    const promises: Promise<void>[] = []
+    const startPromises: Promise<void>[] = []
 
-    chunks.forEach((files, index) => {
-      if (orchestrators[index]) {
-        const [contextId, orchestrator] = orchestrators[index]
-        debug?.(
-          'Reusing orchestrator (context %s) for files: %s',
-          contextId,
-          [...files.map(f => relative(project.config.root, f))].join(', '),
-        )
-        const promise = waitForTests(method, contextId, project, files)
-        promises.push(promise)
-        orchestrator.createTesters(files)
-      }
-      else {
+    if (!orchestrators.length) {
+      files.splice(0, threadsCount).forEach((file) => {
         const contextId = crypto.randomUUID()
-        const waitPromise = waitForTests(method, contextId, project, files)
+        const waitPromise = waitForTests(method, contextId, project, [file])
         debug?.(
           'Opening a new context %s for files: %s',
           contextId,
-          [...files.map(f => relative(project.config.root, f))].join(', '),
+          file,
         )
         const url = new URL('/', origin)
         url.searchParams.set('contextId', contextId)
         const page = provider
-          .openPage(contextId, url.toString(), () => setBreakpoint(contextId, files[0]))
-        promises.push(page, waitPromise)
-      }
-    })
+          .openPage(contextId, url.toString(), () => setBreakpoint(contextId, file))
+        startPromises.push(page, waitPromise)
+      })
+    }
 
-    await Promise.all(promises)
+    await Promise.all([
+      defer,
+      ...startPromises,
+    ])
   }
 
   const runWorkspaceTests = async (method: 'run' | 'collect', specs: WorkspaceSpec[]) => {
@@ -116,19 +113,16 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
       groupedFiles.set(project, files)
     }
 
-    let isCancelled = false
+    const defer = createDefer<void>()
     ctx.onCancel(() => {
-      isCancelled = true
+      defer.reject(new Error('Tests cancelled'))
     })
 
-    // TODO: paralellize tests instead of running them sequentially (based on CPU?)
     for (const [project, files] of groupedFiles.entries()) {
-      if (isCancelled) {
-        break
-      }
-
-      await executeTests(method, project, files)
+      executeTests(defer, method, project, files)
     }
+
+    await defer
   }
 
   const numCpus

packages/browser/src/node/state.ts

@@ -9,6 +9,8 @@ export class BrowserServerState implements IBrowserServerState {
   public readonly cdps = new Map<string, BrowserServerCDPHandler>()
 
   private contexts = new Map<string, BrowserServerStateContext>()
+  private _onReady: ((contextId: string, orchestrator: WebSocketBrowserRPC) => void) | undefined
+  private _onError: ((contextId: string, error: unknown) => void) | undefined
 
   getContext(contextId: string) {
     return this.contexts.get(contextId)
@@ -22,12 +24,32 @@ export class BrowserServerState implements IBrowserServerState {
       resolve: () => {
         defer.resolve()
         this.contexts.delete(contextId)
+        const orchestrator = this.orchestrators.get(contextId)
+        if (orchestrator) {
+          this._onReady?.(contextId, orchestrator)
+        }
+      },
+      reject: (err) => {
+        this._onError?.(contextId, err)
+        defer.reject(err)
       },
-      reject: defer.reject,
     })
     return defer
   }
 
+  onReady(cb: (contextId: string, orchestrator: WebSocketBrowserRPC) => void) {
+    this._onReady = cb
+  }
+
+  onError(cb: (contextId: string, error: unknown) => void): void {
+    this._onError = cb
+  }
+
+  cleanListeners() {
+    this._onReady = undefined
+    this._onError = undefined
+  }
+
   async removeCDPHandler(sessionId: string) {
     this.cdps.delete(sessionId)
   }

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

@@ -198,6 +198,9 @@ export interface BrowserOrchestrator {
 
 export interface BrowserServerState {
   orchestrators: Map<string, BrowserOrchestrator>
+  onReady: (cb: (contextId: string, orchestrator: BrowserOrchestrator) => void) => void
+  onError: (cb: (contextId: string, error: unknown) => void) => void
+  cleanListeners: () => void
   getContext: (contextId: string) => BrowserServerStateContext | undefined
   createAsyncContext: (method: 'collect' | 'run', contextId: string, files: string[]) => Promise<void>
 }