chore: update threadsCount logic to use availableParallelism (#4126)

vitest-dev · Sep 20, 2023 · 2f7d42c · 2f7d42c
1 parent c21c0ef
commit 2f7d42c
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 9 deletions.
diff --git a/packages/vitest/src/node/pools/threads.ts b/packages/vitest/src/node/pools/threads.ts
@@ -1,5 +1,5 @@
 import { MessageChannel } from 'node:worker_threads'
-import { cpus } from 'node:os'
+import * as nodeos from 'node:os'
 import { pathToFileURL } from 'node:url'
 import { createBirpc } from 'birpc'
 import { resolve } from 'pathe'
@@ -39,9 +39,14 @@ function createWorkerChannel(project: WorkspaceProject) {
 }
 
 export function createThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOptions): ProcessPool {
+  const numCpus
+    = typeof nodeos.availableParallelism === 'function'
+      ? nodeos.availableParallelism()
+      : nodeos.cpus().length
+
   const threadsCount = ctx.config.watch
-    ? Math.max(Math.floor(cpus().length / 2), 1)
-    : Math.max(cpus().length - 1, 1)
+    ? Math.max(Math.floor(numCpus / 2), 1)
+    : Math.max(numCpus - 1, 1)
 
   const maxThreads = ctx.config.maxThreads ?? threadsCount
   const minThreads = ctx.config.minThreads ?? threadsCount

diff --git a/packages/vitest/src/node/pools/vm-threads.ts b/packages/vitest/src/node/pools/vm-threads.ts
@@ -1,5 +1,5 @@
 import { MessageChannel } from 'node:worker_threads'
-import { cpus } from 'node:os'
+import * as nodeos from 'node:os'
 import { pathToFileURL } from 'node:url'
 import { createBirpc } from 'birpc'
 import { resolve } from 'pathe'
@@ -40,9 +40,14 @@ function createWorkerChannel(project: WorkspaceProject) {
 }
 
 export function createVmThreadsPool(ctx: Vitest, { execArgv, env }: PoolProcessOptions): ProcessPool {
+  const numCpus
+    = typeof nodeos.availableParallelism === 'function'
+      ? nodeos.availableParallelism()
+      : nodeos.cpus().length
+
   const threadsCount = ctx.config.watch
-    ? Math.max(Math.floor(cpus().length / 2), 1)
-    : Math.max(cpus().length - 1, 1)
+    ? Math.max(Math.floor(numCpus / 2), 1)
+    : Math.max(numCpus - 1, 1)
 
   const maxThreads = ctx.config.maxThreads ?? threadsCount
   const minThreads = ctx.config.minThreads ?? threadsCount

diff --git a/packages/vitest/src/utils/memory-limit.ts b/packages/vitest/src/utils/memory-limit.ts
@@ -5,13 +5,18 @@
  * LICENSE file in the root directory of facebook/jest GitHub project tree.
  */
 
-import { cpus } from 'node:os'
+import * as nodeos from 'node:os'
 import type { ResolvedConfig } from '../types'
 
 function getDefaultThreadsCount(config: ResolvedConfig) {
+  const numCpus
+    = typeof nodeos.availableParallelism === 'function'
+      ? nodeos.availableParallelism()
+      : nodeos.cpus().length
+
   return config.watch
-    ? Math.max(Math.floor(cpus().length / 2), 1)
-    : Math.max(cpus().length - 1, 1)
+    ? Math.max(Math.floor(numCpus / 2), 1)
+    : Math.max(numCpus - 1, 1)
 }
 
 export function getWorkerMemoryLimit(config: ResolvedConfig) {