Skip to content

Commit

Permalink
refactor(runner): limit max concurrency without p-limit as a dependen…
Browse files Browse the repository at this point in the history
…cy (#5980)
  • Loading branch information
jviide authored Jun 30, 2024
1 parent be32317 commit d066ba9
Show file tree
Hide file tree
Showing 7 changed files with 70 additions and 18 deletions.
1 change: 0 additions & 1 deletion packages/browser/src/node/plugin.ts
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,6 @@ export default (browserServer: BrowserServer, base = '/'): Plugin[] => {
'vitest > pretty-format > ansi-regex',
'vitest > chai',
'vitest > chai > loupe',
'vitest > @vitest/runner > p-limit',
'vitest > @vitest/utils > diff-sequences',
'vitest > @vitest/utils > loupe',
'@vitest/browser > @testing-library/user-event',
Expand Down
1 change: 0 additions & 1 deletion packages/runner/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,6 @@
},
"dependencies": {
"@vitest/utils": "workspace:*",
"p-limit": "^5.0.0",
"pathe": "^1.1.2"
}
}
6 changes: 3 additions & 3 deletions packages/runner/src/run.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import limit from 'p-limit'
import type { Awaitable } from '@vitest/utils'
import { getSafeTimers, shuffle } from '@vitest/utils'
import { processError } from '@vitest/utils/error'
Expand All @@ -20,6 +19,7 @@ import type {
Test,
} from './types'
import { partitionSuiteChildren } from './utils/suite'
import { limitConcurrency } from './utils/limit-concurrency'
import { getFn, getHooks } from './map'
import { collectTests } from './collect'
import { setCurrentTest } from './test-state'
Expand Down Expand Up @@ -466,7 +466,7 @@ export async function runSuite(suite: Suite, runner: VitestRunner) {
}
}

let limitMaxConcurrency: ReturnType<typeof limit>
let limitMaxConcurrency: ReturnType<typeof limitConcurrency>

async function runSuiteChild(c: Task, runner: VitestRunner) {
if (c.type === 'test' || c.type === 'custom') {
Expand All @@ -478,7 +478,7 @@ async function runSuiteChild(c: Task, runner: VitestRunner) {
}

export async function runFiles(files: File[], runner: VitestRunner) {
limitMaxConcurrency ??= limit(runner.config.maxConcurrency)
limitMaxConcurrency ??= limitConcurrency(runner.config.maxConcurrency)

for (const file of files) {
if (!file.tasks.length && !runner.config.passWithNoTests) {
Expand Down
1 change: 1 addition & 0 deletions packages/runner/src/utils/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,3 +2,4 @@ export * from './collect'
export * from './suite'
export * from './tasks'
export * from './chain'
export * from './limit-concurrency'
57 changes: 57 additions & 0 deletions packages/runner/src/utils/limit-concurrency.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
// A compact (code-wise, probably not memory-wise) singly linked list node.
type QueueNode<T> = [value: T, next?: QueueNode<T>]

/**
* Return a function for running multiple async operations with limited concurrency.
*/
export function limitConcurrency(concurrency = Infinity): <Args extends unknown[], T>(func: (...args: Args) => PromiseLike<T> | T, ...args: Args) => Promise<T> {
// The number of currently active + pending tasks.
let count = 0

// The head and tail of the pending task queue, built using a singly linked list.
// Both head and tail are initially undefined, signifying an empty queue.
// They both become undefined again whenever there are no pending tasks.
let head: undefined | QueueNode<() => void>
let tail: undefined | QueueNode<() => void>

// A bookkeeping function executed whenever a task has been run to completion.
const finish = () => {
count--

// Check if there are further pending tasks in the queue.
if (head) {
// Allow the next pending task to run and pop it from the queue.
head[0]()
head = head[1]

// The head may now be undefined if there are no further pending tasks.
// In that case, set tail to undefined as well.
tail = head && tail
}
}

return (func, ...args) => {
// Create a promise chain that:
// 1. Waits for its turn in the task queue (if necessary).
// 2. Runs the task.
// 3. Allows the next pending task (if any) to run.
return new Promise<void>((resolve) => {
if (count++ < concurrency) {
// No need to queue if fewer than maxConcurrency tasks are running.
resolve()
}
else if (tail) {
// There are pending tasks, so append to the queue.
tail = tail[1] = [resolve]
}
else {
// No other pending tasks, initialize the queue with a new tail and head.
head = tail = [resolve]
}
}).then(() => {
// Running func here ensures that even a non-thenable result or an
// immediately thrown error gets wrapped into a Promise.
return func(...args)
}).finally(finish)
}
}
11 changes: 1 addition & 10 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

11 changes: 8 additions & 3 deletions test/cli/test/console.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,8 @@ test('can run custom pools with Vitest', async () => {
`)
const stderrArray = vitest.stderr.split('\n')
// remove stack trace
const stderr = stderrArray.slice(0, -9).join('\n')
const stackStderr = stderrArray.slice(-9).join('\n')
const stderr = stderrArray.slice(0, -14).join('\n')
const stackStderr = stderrArray.slice(-14).join('\n')
expect(stderr).toMatchInlineSnapshot(`
"stderr | trace.test.ts > logging to stdout
warn with trace
Expand All @@ -64,14 +64,19 @@ test('can run custom pools with Vitest', async () => {
expect(stackStderr).not.toMatch('❯ ')
if (process.platform !== 'win32') {
const root = resolve(process.cwd(), '../..')
expect(stackStderr.replace(new RegExp(root, 'g'), '<root>').replace(/\d+:\d+/g, 'ln:cl')).toMatchInlineSnapshot(`
expect(stackStderr.replace(new RegExp(root, 'g'), '<root>').replace(/\d+:\d+/g, 'ln:cl').replace(/\.\w+\.js:/g, '.<chunk>.js:')).toMatchInlineSnapshot(`
"stderr | trace.test.ts > logging to stdout
Trace: trace with trace
at <root>/test/cli/fixtures/console/trace.test.ts:ln:cl
at file://<root>/packages/runner/dist/index.js:ln:cl
at file://<root>/packages/runner/dist/index.js:ln:cl
at runTest (file://<root>/packages/runner/dist/index.js:ln:cl)
at processTicksAndRejections (node:internal/process/task_queues:ln:cl)
at runSuite (file://<root>/packages/runner/dist/index.js:ln:cl)
at runFiles (file://<root>/packages/runner/dist/index.js:ln:cl)
at startTests (file://<root>/packages/runner/dist/index.js:ln:cl)
at file://<root>/packages/vitest/dist/chunks/runtime-runBaseTests.<chunk>.js:ln:cl
at withEnv (file://<root>/packages/vitest/dist/chunks/runtime-runBaseTests.<chunk>.js:ln:cl)
"
`)
Expand Down

0 comments on commit d066ba9

Please sign in to comment.