From d16306beed731082b649a99161a5db92a7a6b3c0 Mon Sep 17 00:00:00 2001 From: Vladimir Sheremet Date: Tue, 5 Nov 2024 12:00:50 +0100 Subject: [PATCH] perf: orchestrate browser pool on demand --- packages/browser/src/node/pool.ts | 92 +++++++++++------------ packages/browser/src/node/state.ts | 24 +++++- packages/vitest/src/node/types/browser.ts | 3 + 3 files changed, 69 insertions(+), 50 deletions(-) diff --git a/packages/browser/src/node/pool.ts b/packages/browser/src/node/pool.ts index c0bc9f0f8828..8f2b1f103015 100644 --- a/packages/browser/src/node/pool.ts +++ b/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() - const executeTests = async (method: 'run' | 'collect', project: WorkspaceProject, files: string[]) => { + const executeTests = async ( + defer: DeferPromise, + 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[] = [] + const startPromises: Promise[] = [] - 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() 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 diff --git a/packages/browser/src/node/state.ts b/packages/browser/src/node/state.ts index 5e07bd51c33d..cdef7f160b7f 100644 --- a/packages/browser/src/node/state.ts +++ b/packages/browser/src/node/state.ts @@ -9,6 +9,8 @@ export class BrowserServerState implements IBrowserServerState { public readonly cdps = new Map() private contexts = new Map() + 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) } diff --git a/packages/vitest/src/node/types/browser.ts b/packages/vitest/src/node/types/browser.ts index 3d607c6dae47..58b0380ec83a 100644 --- a/packages/vitest/src/node/types/browser.ts +++ b/packages/vitest/src/node/types/browser.ts @@ -198,6 +198,9 @@ export interface BrowserOrchestrator { export interface BrowserServerState { orchestrators: Map + 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 }