Skip to content

Commit

Permalink
feat(watch): test run cancelling
Browse files Browse the repository at this point in the history
  • Loading branch information
AriPerkkio committed Apr 10, 2023
1 parent 0ce8364 commit 2f85b5c
Show file tree
Hide file tree
Showing 12 changed files with 113 additions and 24 deletions.
17 changes: 17 additions & 0 deletions packages/runner/src/run.ts
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,14 @@ async function callCleanupHooks(cleanups: HookCleanupCallback[]) {
}

export async function runTest(test: Test, runner: VitestRunner) {
// @ts-expect-error untyped global
if (globalThis.__vitest_worker__.isCancelling) {
test.mode = 'skip'
test.result = { ...test.result, state: 'skip' }
updateTask(test, runner)
return
}

await runner.onBeforeRunTest?.(test)

if (test.mode !== 'run')
Expand Down Expand Up @@ -227,6 +235,15 @@ function markTasksAsSkipped(suite: Suite, runner: VitestRunner) {
}

export async function runSuite(suite: Suite, runner: VitestRunner) {
// @ts-expect-error untyped global
if (globalThis.__vitest_worker__.isCancelling) {
suite.result = { state: 'skip' }
suite.mode = 'skip'
markTasksAsSkipped(suite, runner)
updateTask(suite, runner)
return
}

await runner.onBeforeRunSuite?.(suite)

if (suite.result?.state === 'fail') {
Expand Down
19 changes: 18 additions & 1 deletion packages/vitest/src/node/core.ts
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ export class Vitest {
filenamePattern?: string
runningPromise?: Promise<void>
closingPromise?: Promise<void>
isCancelling = false

isFirstRun = true
restartsCount = 0
Expand All @@ -64,6 +65,7 @@ export class Vitest {

private _onRestartListeners: OnServerRestartHandler[] = []
private _onSetServer: OnServerRestartHandler[] = []
private _onCancelListeners: (() => Promise<void> | void)[] = []

async setServer(options: UserConfig, server: ViteDevServer, cliOptions: UserConfig) {
this.unregisterWatcher?.()
Expand Down Expand Up @@ -394,13 +396,13 @@ export class Vitest {

async runFiles(paths: WorkspaceSpec[]) {
const filepaths = paths.map(([, file]) => file)

this.state.collectPaths(filepaths)

await this.report('onPathsCollected', filepaths)

// previous run
await this.runningPromise
this.resetCancel()

// schedule the new run
this.runningPromise = (async () => {
Expand Down Expand Up @@ -431,11 +433,22 @@ export class Vitest {
const specs = Array.from(new Set(paths.map(([, p]) => p)))
await this.report('onFinished', this.state.getFiles(specs), this.state.getUnhandledErrors())
this.runningPromise = undefined
this.resetCancel()
})

return await this.runningPromise
}

async cancelCurrentRun() {
this.isCancelling = true
await Promise.all(this._onCancelListeners.splice(0).map(listener => listener()))
}

resetCancel() {
this._onCancelListeners = []
this.isCancelling = false
}

async rerunFiles(files: string[] = this.state.getFilepaths(), trigger?: string) {
if (this.filenamePattern) {
const filteredFiles = await this.globTestFiles([this.filenamePattern])
Expand Down Expand Up @@ -765,4 +778,8 @@ export class Vitest {
onAfterSetServer(fn: OnServerRestartHandler) {
this._onSetServer.push(fn)
}

onCancel(fn: () => void) {
this._onCancelListeners.push(fn)
}
}
3 changes: 3 additions & 0 deletions packages/vitest/src/node/pools/browser.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,9 @@ export function createBrowserPool(ctx: Vitest): ProcessPool {
const isolate = project.config.isolate
if (isolate) {
for (const path of paths) {
if (ctx.isCancelling)
break

const url = new URL('/', origin)
url.searchParams.append('path', path)
url.searchParams.set('id', path)
Expand Down
7 changes: 5 additions & 2 deletions packages/vitest/src/node/pools/child.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ import { fork } from 'node:child_process'
import { fileURLToPath, pathToFileURL } from 'node:url'
import { createBirpc } from 'birpc'
import { resolve } from 'pathe'
import type { ContextTestEnvironment, ResolvedConfig, RuntimeRPC, Vitest } from '../../types'
import type { ContextTestEnvironment, ResolvedConfig, RunnerRPC, RuntimeRPC, Vitest } from '../../types'
import type { ChildContext } from '../../types/child'
import type { PoolProcessOptions, ProcessPool, WorkspaceSpec } from '../pool'
import { distDir } from '../../paths'
Expand All @@ -16,9 +16,10 @@ import { createMethodsRPC } from './rpc'
const childPath = fileURLToPath(pathToFileURL(resolve(distDir, './child.js')).href)

function setupChildProcessChannel(project: WorkspaceProject, fork: ChildProcess): void {
createBirpc<{}, RuntimeRPC>(
const rpc = createBirpc<RunnerRPC, RuntimeRPC>(
createMethodsRPC(project),
{
eventNames: ['onCancel'],
serialize: v8.serialize,
deserialize: v => v8.deserialize(Buffer.from(v)),
post(v) {
Expand All @@ -29,6 +30,8 @@ function setupChildProcessChannel(project: WorkspaceProject, fork: ChildProcess)
},
},
)

project.ctx.onCancel(() => rpc.onCancel())
}

function stringifyRegex(input: RegExp | string): string {
Expand Down
39 changes: 31 additions & 8 deletions packages/vitest/src/node/pools/threads.ts
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
import { MessageChannel } from 'node:worker_threads'
import { cpus } from 'node:os'
import { pathToFileURL } from 'node:url'
import EventEmitter from 'node:events'
import { createBirpc } from 'birpc'
import { resolve } from 'pathe'
import type { Options as TinypoolOptions } from 'tinypool'
import Tinypool from 'tinypool'
import { distDir } from '../../paths'
import type { ContextTestEnvironment, ResolvedConfig, RuntimeRPC, Vitest, WorkerContext } from '../../types'
import type { ContextTestEnvironment, ResolvedConfig, RunnerRPC, RuntimeRPC, Vitest, WorkerContext } from '../../types'
import type { PoolProcessOptions, ProcessPool, RunWithFiles } from '../pool'
import { envsOrder, groupFilesByEnv } from '../../utils/test-helpers'
import { groupBy } from '../../utils/base'
Expand All @@ -20,9 +21,10 @@ function createWorkerChannel(project: WorkspaceProject) {
const port = channel.port2
const workerPort = channel.port1

createBirpc<{}, RuntimeRPC>(
const rpc = createBirpc<RunnerRPC, RuntimeRPC>(
createMethodsRPC(project),
{
eventNames: ['onCancel'],
post(v) {
port.postMessage(v)
},
Expand All @@ -32,6 +34,8 @@ function createWorkerChannel(project: WorkspaceProject) {
},
)

project.ctx.onCancel(() => rpc.onCancel())

return { workerPort, port }
}

Expand Down Expand Up @@ -74,7 +78,7 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt
const runWithFiles = (name: string): RunWithFiles => {
let id = 0

async function runFiles(project: WorkspaceProject, config: ResolvedConfig, files: string[], environment: ContextTestEnvironment, invalidates: string[] = []) {
async function runFiles(project: WorkspaceProject, config: ResolvedConfig, files: string[], environment: ContextTestEnvironment, invalidates: string[] = [], signal: EventEmitter) {
ctx.state.clearFiles(project, files)
const { workerPort, port } = createWorkerChannel(project)
const workerId = ++id
Expand All @@ -87,14 +91,28 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt
workerId,
}
try {
await pool.run(data, { transferList: [workerPort], name })
await pool.run(data, { transferList: [workerPort], name, signal })
}
catch (error) {
// Worker got stuck and won't terminate - this may cause process to hang
if (error instanceof Error && /Failed to terminate worker/.test(error.message))
if (error instanceof Error && /Failed to terminate worker/.test(error.message)) {
ctx.state.addProcessTimeoutCause(`Failed to terminate worker while running ${files.join(', ')}.`)
else
}
else if (ctx.isCancelling && error instanceof Error && /The task has been aborted/.test(error.message)) {
// Intentionally cancelled
// TODO: This could be a "ctx.state.cancelFiles(files: string[])" instead
ctx.state.collectFiles(files.map(filepath => ({
filepath,
id: filepath,
mode: 'skip',
name: filepath,
tasks: [],
type: 'suite',
})))
}
else {
throw error
}
}
finally {
port.close()
Expand Down Expand Up @@ -123,6 +141,11 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt
workspaceMap.set(file, workspaceFiles)
}

// TODO: This signal could also be in "ctx.cancelSignal" etc. Though browser and child_process do not need it.
const signal = new EventEmitter()
signal.setMaxListeners(ctx.state.getFiles().length)
ctx.onCancel(() => signal.emit('abort'))

// it's possible that project defines a file that is also defined by another project
const { shard } = ctx.config

Expand All @@ -138,7 +161,7 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt
const filesByEnv = await groupFilesByEnv(multipleThreads)
const promises = Object.values(filesByEnv).flat()
const results = await Promise.allSettled(promises
.map(({ file, environment, project }) => runFiles(project, getConfig(project), [file], environment, invalidates)))
.map(({ file, environment, project }) => runFiles(project, getConfig(project), [file], environment, invalidates, signal)))

const errors = results.filter((r): r is PromiseRejectedResult => r.status === 'rejected').map(r => r.reason)
if (errors.length > 0)
Expand All @@ -162,7 +185,7 @@ export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOpt

const promises = Object.values(filesByOptions).map(async (files) => {
const filenames = files.map(f => f.file)
await runFiles(files[0].project, getConfig(files[0].project), filenames, files[0].environment, invalidates)
await runFiles(files[0].project, getConfig(files[0].project), filenames, files[0].environment, invalidates, signal)
})

await Promise.all(promises)
Expand Down
7 changes: 6 additions & 1 deletion packages/vitest/src/node/reporters/base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ const HELP_QUITE = `${c.dim('press ')}${c.bold('q')}${c.dim(' to quit')}`

const WAIT_FOR_CHANGE_PASS = `\n${c.bold(c.inverse(c.green(' PASS ')))}${c.green(' Waiting for file changes...')}`
const WAIT_FOR_CHANGE_FAIL = `\n${c.bold(c.inverse(c.red(' FAIL ')))}${c.red(' Tests failed. Watching for file changes...')}`
const WAIT_FOR_CHANGE_CANCELLED = `\n${c.bold(c.inverse(c.red(' CANCELLED ')))}${c.red(' Test run cancelled. Watching for file changes...')}`

const LAST_RUN_LOG_TIMEOUT = 1_500

Expand Down Expand Up @@ -100,9 +101,13 @@ export abstract class BaseReporter implements Reporter {
async onWatcherStart(files = this.ctx.state.getFiles(), errors = this.ctx.state.getUnhandledErrors()) {
this.resetLastRunLog()

const cancelled = this.ctx.isCancelling
const failed = errors.length > 0 || hasFailed(files)
const failedSnap = hasFailedSnapshot(files)
if (failed)

if (cancelled)
this.ctx.logger.log(WAIT_FOR_CHANGE_CANCELLED)
else if (failed)
this.ctx.logger.log(WAIT_FOR_CHANGE_FAIL)
else
this.ctx.logger.log(WAIT_FOR_CHANGE_PASS)
Expand Down
14 changes: 8 additions & 6 deletions packages/vitest/src/node/stdin.ts
Original file line number Diff line number Diff line change
Expand Up @@ -37,12 +37,14 @@ export function registerConsoleShortcuts(ctx: Vitest) {
return
}

// is running, ignore keypress
if (ctx.runningPromise)
return

const name = key?.name

if (ctx.runningPromise) {
if (['space', 'c'].includes(name) || keys.map(key => key[0]).includes(name))
await ctx.cancelCurrentRun()
return
}

// quit
if (name === 'q')
return ctx.exit(true)
Expand Down Expand Up @@ -83,8 +85,8 @@ export function registerConsoleShortcuts(ctx: Vitest) {
message: 'Input test name pattern (RegExp)',
initial: ctx.configOverride.testNamePattern?.source || '',
}])
await ctx.changeNamePattern(filter.trim(), undefined, 'change pattern')
on()
await ctx.changeNamePattern(filter.trim(), undefined, 'change pattern')
}

async function inputFilePattern() {
Expand All @@ -96,8 +98,8 @@ export function registerConsoleShortcuts(ctx: Vitest) {
initial: latestFilename,
}])
latestFilename = filter.trim()
await ctx.changeFilenamePattern(filter.trim())
on()
await ctx.changeFilenamePattern(filter.trim())
}

let rl: readline.Interface | undefined
Expand Down
11 changes: 8 additions & 3 deletions packages/vitest/src/runtime/child.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@ import v8 from 'node:v8'
import { createBirpc } from 'birpc'
import { parseRegexp } from '@vitest/utils'
import type { ResolvedConfig } from '../types'
import type { RuntimeRPC } from '../types/rpc'
import type { RunnerRPC, RuntimeRPC } from '../types/rpc'
import type { ChildContext } from '../types/child'
import { getWorkerState } from '../utils'
import { mockMap, moduleCache, startViteNode } from './execute'
import { rpcDone } from './rpc'
import { setupInspect } from './inspector'
Expand All @@ -26,8 +27,12 @@ function init(ctx: ChildContext) {
environment: 0,
prepare: performance.now(),
},
rpc: createBirpc<RuntimeRPC>(
{},
rpc: createBirpc<RuntimeRPC, RunnerRPC>(
{
onCancel() {
getWorkerState().isCancelling = true
},
},
{
eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onWorkerExit'],
serialize: v8.serialize,
Expand Down
4 changes: 4 additions & 0 deletions packages/vitest/src/runtime/entry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,10 @@ export async function run(files: string[], config: ResolvedConfig, environment:
workerState.durations.environment = performance.now() - workerState.durations.environment

for (const file of files) {
if (workerState.isCancelling) {
// TODO: Mark as skipped and break
}

// it doesn't matter if running with --threads
// if running with --no-threads, we usually want to reset everything before running a test
// but we have --isolate option to disable this
Expand Down
10 changes: 7 additions & 3 deletions packages/vitest/src/runtime/worker.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { createBirpc } from 'birpc'
import { workerId as poolId } from 'tinypool'
import type { RuntimeRPC, WorkerContext } from '../types'
import type { RunnerRPC, RuntimeRPC, WorkerContext } from '../types'
import { getWorkerState } from '../utils/global'
import { mockMap, moduleCache, startViteNode } from './execute'
import { setupInspect } from './inspector'
Expand Down Expand Up @@ -28,8 +28,12 @@ function init(ctx: WorkerContext) {
environment: 0,
prepare: performance.now(),
},
rpc: createBirpc<RuntimeRPC>(
{},
rpc: createBirpc<RuntimeRPC, RunnerRPC>(
{
onCancel() {
getWorkerState().isCancelling = true
},
},
{
eventNames: ['onUserConsoleLog', 'onFinished', 'onCollected', 'onWorkerExit'],
post(v) { port.postMessage(v) },
Expand Down
5 changes: 5 additions & 0 deletions packages/vitest/src/types/rpc.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,11 @@ export interface RuntimeRPC {
resolveSnapshotPath: (testPath: string) => string
}

export interface RunnerRPC {
// TODO: This could be "(reason: 'INPUT' | 'BAIL') => void" instead
onCancel: () => void
}

export interface ContextTestEnvironment {
name: VitestEnvironment
options: EnvironmentOptions | null
Expand Down
1 change: 1 addition & 0 deletions packages/vitest/src/types/worker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ export interface WorkerGlobalState {
current?: Test
filepath?: string
environmentTeardownRun?: boolean
isCancelling?: boolean
moduleCache: ModuleCacheMap
mockMap: MockMap
durations: {
Expand Down

0 comments on commit 2f85b5c

Please sign in to comment.