fix(vitest): adjust getWorkerMemoryLimit priority for vmForks (#7960)

pengooseDev · web-flow · commit 5a91eca20de0 · 2025-05-12T20:52:58.000+02:00
diff --git a/packages/vitest/src/node/pools/vmForks.ts b/packages/vitest/src/node/pools/vmForks.ts
@@ -213,7 +213,7 @@ export function createVmForksPool(
 
 function getMemoryLimit(config: ResolvedConfig) {
   const memory = nodeos.totalmem()
-  const limit = getWorkerMemoryLimit(config)
+  const limit = getWorkerMemoryLimit(config, 'vmForks')
 
   if (typeof memory === 'number') {
     return stringToBytes(limit, config.watch ? memory / 2 : memory)
diff --git a/packages/vitest/src/node/pools/vmThreads.ts b/packages/vitest/src/node/pools/vmThreads.ts
@@ -206,7 +206,7 @@ export function createVmThreadsPool(
 
 function getMemoryLimit(config: ResolvedConfig) {
   const memory = nodeos.totalmem()
-  const limit = getWorkerMemoryLimit(config)
+  const limit = getWorkerMemoryLimit(config, 'vmThreads')
 
   if (typeof memory === 'number') {
     return stringToBytes(limit, config.watch ? memory / 2 : memory)
diff --git a/packages/vitest/src/utils/memory-limit.ts b/packages/vitest/src/utils/memory-limit.ts
@@ -19,18 +19,25 @@ function getDefaultThreadsCount(config: ResolvedConfig) {
     : Math.max(numCpus - 1, 1)
 }
 
-export function getWorkerMemoryLimit(config: ResolvedConfig): string | number {
-  const memoryLimit = config.poolOptions?.vmThreads?.memoryLimit
+export function getWorkerMemoryLimit(config: ResolvedConfig, pool: 'vmThreads' | 'vmForks'): string | number {
+  if (pool === 'vmForks') {
+    const opts = config.poolOptions?.vmForks ?? {}
+    if (opts.memoryLimit) {
+      return opts.memoryLimit
+    }
+    const workers = opts.maxForks ?? getDefaultThreadsCount(config)
 
-  if (memoryLimit) {
-    return memoryLimit
+    return 1 / workers
   }
+  else {
+    const opts = config.poolOptions?.vmThreads ?? {}
+    if (opts.memoryLimit) {
+      return opts.memoryLimit
+    }
+    const workers = opts.maxThreads ?? getDefaultThreadsCount(config)
 
-  return (
-    1
-    / (config.poolOptions?.vmThreads?.maxThreads
-      ?? getDefaultThreadsCount(config))
-  )
+    return 1 / workers
+  }
 }
 
 /**
diff --git a/test/core/test/memory-limit.test.ts b/test/core/test/memory-limit.test.ts
@@ -0,0 +1,46 @@
+import type { PoolOptions, ResolvedConfig } from 'vitest/node'
+import { describe, expect, it } from 'vitest'
+import { getWorkerMemoryLimit } from 'vitest/src/utils/memory-limit.js'
+
+function makeConfig(poolOptions: PoolOptions): ResolvedConfig {
+  return {
+    poolOptions: {
+      vmForks: {
+        maxForks: poolOptions.maxForks,
+        memoryLimit: poolOptions.memoryLimit,
+      },
+      vmThreads: {
+        maxThreads: poolOptions.maxThreads,
+        memoryLimit: poolOptions.memoryLimit,
+      },
+    },
+  } as ResolvedConfig
+}
+
+describe('getWorkerMemoryLimit', () => {
+  it('should prioritize vmThreads.memoryLimit when pool is vmThreads', () => {
+    const config = {
+      poolOptions: {
+        vmForks: { memoryLimit: undefined },
+        vmThreads: { memoryLimit: '256MB' },
+      },
+    } as ResolvedConfig
+
+    expect(getWorkerMemoryLimit(config, 'vmThreads')).toBe('256MB')
+  })
+
+  it('should prioritize vmForks.memoryLimit when pool is vmForks', () => {
+    const config = makeConfig({ memoryLimit: '512MB' })
+    expect(getWorkerMemoryLimit(config, 'vmForks')).toBe('512MB')
+  })
+
+  it('should calculate 1/maxThreads when vmThreads.memoryLimit is unset', () => {
+    const config = makeConfig({ maxThreads: 4 })
+    expect(getWorkerMemoryLimit(config, 'vmThreads')).toBe(1 / 4)
+  })
+
+  it('should calculate 1/maxForks when vmForks.memoryLimit is unset', () => {
+    const config = makeConfig({ maxForks: 4 })
+    expect(getWorkerMemoryLimit(config, 'vmForks')).toBe(1 / 4)
+  })
+})
