From 348a0088575a33f10926380687043f3283d2ef49 Mon Sep 17 00:00:00 2001
From: Vladimir <sheremet.va@icloud.com>
Date: Tue, 14 Jun 2022 20:57:56 +0300
Subject: [PATCH] feat: add `maxConcurrency` option (#1483)

---
 docs/config/index.md                | 9 +++++++++
 packages/vitest/package.json        | 1 +
 packages/vitest/src/defaults.ts     | 1 +
 packages/vitest/src/runtime/run.ts  | 8 +++++---
 packages/vitest/src/types/config.ts | 6 ++++++
 pnpm-lock.yaml                      | 2 ++
 6 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/docs/config/index.md b/docs/config/index.md
index a96789e5d019..b0764ea364b5 100644
--- a/docs/config/index.md
+++ b/docs/config/index.md
@@ -533,3 +533,12 @@ Vitest will not fail, if no tests will be found.
 - **Default**: `false`
 
 Show heap usage after each test. Useful for debugging memory leaks.
+
+### maxConcurrency
+
+- **Type**: `number`
+- **Default**: `5`
+
+A number of tests that are allowed to run at the same time marked with `test.concurrent`.
+
+Test above this limit will be queued to run when available slot appears.
diff --git a/packages/vitest/package.json b/packages/vitest/package.json
index 02baea1ed5a2..89128b746ce0 100644
--- a/packages/vitest/package.json
+++ b/packages/vitest/package.json
@@ -119,6 +119,7 @@
     "micromatch": "^4.0.5",
     "mlly": "^0.5.2",
     "natural-compare": "^1.4.0",
+    "p-limit": "^4.0.0",
     "pathe": "^0.2.0",
     "picocolors": "^1.0.0",
     "pkg-types": "^0.3.2",
diff --git a/packages/vitest/src/defaults.ts b/packages/vitest/src/defaults.ts
index 9521ee7d60f6..7f7a9e6694e9 100644
--- a/packages/vitest/src/defaults.ts
+++ b/packages/vitest/src/defaults.ts
@@ -72,6 +72,7 @@ const config = {
   open: true,
   coverage: coverageConfigDefaults,
   fakeTimers: fakeTimersDefaults,
+  maxConcurrency: 5,
 }
 
 export const configDefaults: Required<Pick<UserConfig, keyof typeof config>> = Object.freeze(config)
diff --git a/packages/vitest/src/runtime/run.ts b/packages/vitest/src/runtime/run.ts
index 875f49837d3c..a12529c4fd81 100644
--- a/packages/vitest/src/runtime/run.ts
+++ b/packages/vitest/src/runtime/run.ts
@@ -1,3 +1,4 @@
+import limit from 'p-limit'
 import type { File, HookCleanupCallback, HookListener, ResolvedConfig, Suite, SuiteHooks, Task, TaskResult, TaskState, Test } from '../types'
 import { vi } from '../integrations/vi'
 import { getSnapshotClient } from '../integrations/snapshot/chai'
@@ -195,6 +196,8 @@ export async function runSuite(suite: Suite) {
 
   updateTask(suite)
 
+  const workerState = getWorkerState()
+
   if (suite.mode === 'skip') {
     suite.result.state = 'skip'
   }
@@ -207,7 +210,8 @@ export async function runSuite(suite: Suite) {
 
       for (const tasksGroup of partitionSuiteChildren(suite)) {
         if (tasksGroup[0].concurrent === true) {
-          await Promise.all(tasksGroup.map(c => runSuiteChild(c)))
+          const mutex = limit(workerState.config.maxConcurrency)
+          await Promise.all(tasksGroup.map(c => mutex(() => runSuiteChild(c))))
         }
         else {
           for (const c of tasksGroup)
@@ -225,8 +229,6 @@ export async function runSuite(suite: Suite) {
   }
   suite.result.duration = now() - start
 
-  const workerState = getWorkerState()
-
   if (workerState.config.logHeapUsage)
     suite.result.heap = process.memoryUsage().heapUsed
 
diff --git a/packages/vitest/src/types/config.ts b/packages/vitest/src/types/config.ts
index 1c133f281283..26499cd46c14 100644
--- a/packages/vitest/src/types/config.ts
+++ b/packages/vitest/src/types/config.ts
@@ -344,6 +344,12 @@ export interface InlineConfig {
    * Return `false` to ignore the log.
    */
   onConsoleLog?: (log: string, type: 'stdout' | 'stderr') => false | void
+
+  /**
+   * A number of tests that are allowed to run at the same time marked with `test.concurrent`.
+   * @default 5
+   */
+  maxConcurrency?: number
 }
 
 export interface UserConfig extends InlineConfig {
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 8a4b361e0f4b..7d4d0ff34fe6 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -676,6 +676,7 @@ importers:
       micromatch: ^4.0.5
       mlly: ^0.5.2
       natural-compare: ^1.4.0
+      p-limit: ^4.0.0
       pathe: ^0.2.0
       picocolors: ^1.0.0
       pkg-types: ^0.3.2
@@ -727,6 +728,7 @@ importers:
       micromatch: 4.0.5
       mlly: 0.5.2
       natural-compare: 1.4.0
+      p-limit: 4.0.0
       pathe: 0.2.0
       picocolors: 1.0.0
       pkg-types: 0.3.2