From 2673c3bb2d2f944858629fb1230700e28c5ebad2 Mon Sep 17 00:00:00 2001
From: Vladimir <sleuths.slews0s@icloud.com>
Date: Tue, 3 Sep 2024 17:57:19 +0200
Subject: [PATCH] fix(vitest): dispose vmForks listeners to avoid memory leak
 (#6448)

---
 packages/vitest/src/runtime/worker.ts        |  3 +++
 packages/vitest/src/runtime/workers/utils.ts | 18 ++++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/packages/vitest/src/runtime/worker.ts b/packages/vitest/src/runtime/worker.ts
index c35f7a6733b7..3df04d3dab48 100644
--- a/packages/vitest/src/runtime/worker.ts
+++ b/packages/vitest/src/runtime/worker.ts
@@ -7,6 +7,7 @@ import type { ContextRPC, WorkerGlobalState } from '../types/worker'
 import { setupInspect } from './inspector'
 import { createRuntimeRpc, rpcDone } from './rpc'
 import type { VitestWorker } from './workers/types'
+import { disposeInternalListeners } from './workers/utils'
 
 if (isChildProcess()) {
   setProcessTitle(`vitest ${poolId}`)
@@ -14,6 +15,8 @@ if (isChildProcess()) {
 
 // this is what every pool executes when running tests
 async function execute(mehtod: 'run' | 'collect', ctx: ContextRPC) {
+  disposeInternalListeners()
+
   const prepareStart = performance.now()
 
   const inspectorCleanup = setupInspect(ctx)
diff --git a/packages/vitest/src/runtime/workers/utils.ts b/packages/vitest/src/runtime/workers/utils.ts
index 82b518c2b746..eca86051f040 100644
--- a/packages/vitest/src/runtime/workers/utils.ts
+++ b/packages/vitest/src/runtime/workers/utils.ts
@@ -9,6 +9,8 @@ const REGEXP_WRAP_PREFIX = '$$vitest:'
 // Store global APIs in case process is overwritten by tests
 const processSend = process.send?.bind(process)
 const processOn = process.on?.bind(process)
+const processOff = process.off?.bind(process)
+const dispose: (() => void)[] = []
 
 export function createThreadsRpcOptions({
   port,
@@ -23,6 +25,16 @@ export function createThreadsRpcOptions({
   }
 }
 
+export function disposeInternalListeners() {
+  for (const fn of dispose) {
+    try {
+      fn()
+    }
+    catch {}
+  }
+  dispose.length = 0
+}
+
 export function createForksRpcOptions(
   nodeV8: typeof import('v8'),
 ): WorkerRpcOptions {
@@ -33,14 +45,16 @@ export function createForksRpcOptions(
       processSend!(v)
     },
     on(fn) {
-      processOn('message', (message: any, ...extras: any) => {
+      const handler = (message: any, ...extras: any) => {
         // Do not react on Tinypool's internal messaging
         if ((message as TinypoolWorkerMessage)?.__tinypool_worker_message__) {
           return
         }
 
         return fn(message, ...extras)
-      })
+      }
+      processOn('message', handler)
+      dispose.push(() => processOff('message', handler))
     },
   }
 }