From 2f7d42c99c59ce8acdf7bd5bae6bb7434b24068e Mon Sep 17 00:00:00 2001
From: Diogo Torres <diogo.rey@hotmail.com>
Date: Wed, 20 Sep 2023 07:58:19 +0100
Subject: [PATCH] chore: update threadsCount logic to use availableParallelism
 (#4126)

---
 packages/vitest/src/node/pools/threads.ts    | 11 ++++++++---
 packages/vitest/src/node/pools/vm-threads.ts | 11 ++++++++---
 packages/vitest/src/utils/memory-limit.ts    | 11 ++++++++---
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/packages/vitest/src/node/pools/threads.ts b/packages/vitest/src/node/pools/threads.ts
index b405f2052..53e0b993a 100644
--- 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
index 89aecc1e3..3a97c237f 100644
--- 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
index 581b19f8c..623b19ad0 100644
--- 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) {
-- 
2.51.2