fix: prevent rpc timeout on slow thread blocking synchronous methods (#8297)

Co-authored-by: Vladimir Sheremet <[email protected]>

Author: AriPerkkio

packages/browser/src/client/client.ts

@@ -109,7 +109,6 @@ function createClient() {
     {
       post: msg => ctx.ws.send(msg),
       on: fn => (onMessage = fn),
-      timeout: -1, // createTesters can take a while
       serialize: e =>
         stringify(e, (_, v) => {
           if (v instanceof Error) {
@@ -122,9 +121,7 @@ function createClient() {
           return v
         }),
       deserialize: parse,
-      onTimeoutError(functionName) {
-        throw new Error(`[vitest-browser]: Timeout calling "${functionName}"`)
-      },
+      timeout: -1, // createTesters can take a while
     },
   )
 

packages/browser/src/node/rpc.ts

@@ -337,11 +337,8 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject, defaultMocke
         on: fn => ws.on('message', fn),
         eventNames: ['onCancel', 'cdpEvent'],
         serialize: (data: any) => stringify(data, stringifyReplace),
-        timeout: -1, // createTesters can take a long time
         deserialize: parse,
-        onTimeoutError(functionName) {
-          throw new Error(`[vitest-api]: Timeout calling "${functionName}"`)
-        },
+        timeout: -1, // createTesters can take a long time
       },
     )
 

packages/vitest/src/api/setup.ts

@@ -1,5 +1,4 @@
 import type { File, TaskEventPack, TaskResultPack, TestAnnotation } from '@vitest/runner'
-
 import type { IncomingMessage } from 'node:http'
 import type { ViteDevServer } from 'vite'
 import type { WebSocket } from 'ws'
@@ -142,16 +141,15 @@ export function setup(ctx: Vitest, _server?: ViteDevServer): void {
         ],
         serialize: (data: any) => stringify(data, stringifyReplace),
         deserialize: parse,
-        onTimeoutError(functionName) {
-          throw new Error(`[vitest-api]: Timeout calling "${functionName}"`)
-        },
+        timeout: -1,
       },
     )
 
     clients.set(ws, rpc)
 
     ws.on('close', () => {
       clients.delete(ws)
+      rpc.$close(new Error('[vitest-api]: Pending methods while closing rpc'))
     })
   }
 

packages/vitest/src/node/pools/forks.ts

@@ -19,13 +19,7 @@ import { createMethodsRPC } from './rpc'
 
 function createChildProcessChannel(project: TestProject, collect = false) {
   const emitter = new EventEmitter()
-
   const events = { message: 'message', response: 'response' }
-  const channel: TinypoolChannel = {
-    onMessage: callback => emitter.on(events.message, callback),
-    postMessage: message => emitter.emit(events.response, message),
-    onClose: () => emitter.removeAllListeners(),
-  }
 
   const rpc = createBirpc<RunnerRPC, RuntimeRPC>(createMethodsRPC(project, { cacheFs: true, collect }), {
     eventNames: ['onCancel'],
@@ -51,13 +45,20 @@ function createChildProcessChannel(project: TestProject, collect = false) {
     on(fn) {
       emitter.on(events.response, fn)
     },
-    onTimeoutError(functionName) {
-      throw new Error(`[vitest-pool]: Timeout calling "${functionName}"`)
-    },
+    timeout: -1,
   })
 
   project.vitest.onCancel(reason => rpc.onCancel(reason))
 
+  const channel: TinypoolChannel = {
+    onMessage: callback => emitter.on(events.message, callback),
+    postMessage: message => emitter.emit(events.response, message),
+    onClose: () => {
+      emitter.removeAllListeners()
+      rpc.$close(new Error('[vitest-pool]: Pending methods while closing rpc'))
+    },
+  }
+
   return channel
 }
 

packages/vitest/src/node/pools/threads.ts

@@ -29,14 +29,18 @@ function createWorkerChannel(project: TestProject, collect: boolean) {
     on(fn) {
       port.on('message', fn)
     },
-    onTimeoutError(functionName) {
-      throw new Error(`[vitest-pool]: Timeout calling "${functionName}"`)
-    },
+    timeout: -1,
   })
 
   project.vitest.onCancel(reason => rpc.onCancel(reason))
 
-  return { workerPort, port }
+  const onClose = () => {
+    port.close()
+    workerPort.close()
+    rpc.$close(new Error('[vitest-pool]: Pending methods while closing rpc'))
+  }
+
+  return { workerPort, port, onClose }
 }
 
 export function createThreadsPool(
@@ -104,11 +108,7 @@ export function createThreadsPool(
       const paths = files.map(f => f.filepath)
       vitest.state.clearFiles(project, paths)
 
-      const { workerPort, port } = createWorkerChannel(project, name === 'collect')
-      const onClose = () => {
-        port.close()
-        workerPort.close()
-      }
+      const { workerPort, onClose } = createWorkerChannel(project, name === 'collect')
 
       const workerId = ++id
       const data: WorkerContext = {

packages/vitest/src/node/pools/vmForks.ts

@@ -22,13 +22,8 @@ const suppressWarningsPath = resolve(rootDir, './suppress-warnings.cjs')
 
 function createChildProcessChannel(project: TestProject, collect: boolean) {
   const emitter = new EventEmitter()
-  const cleanup = () => emitter.removeAllListeners()
 
   const events = { message: 'message', response: 'response' }
-  const channel: TinypoolChannel = {
-    onMessage: callback => emitter.on(events.message, callback),
-    postMessage: message => emitter.emit(events.response, message),
-  }
 
   const rpc = createBirpc<RunnerRPC, RuntimeRPC>(
     createMethodsRPC(project, { cacheFs: true, collect }),
@@ -56,15 +51,22 @@ function createChildProcessChannel(project: TestProject, collect: boolean) {
       on(fn) {
         emitter.on(events.response, fn)
       },
-      onTimeoutError(functionName) {
-        throw new Error(`[vitest-pool]: Timeout calling "${functionName}"`)
-      },
+      timeout: -1,
     },
   )
 
   project.vitest.onCancel(reason => rpc.onCancel(reason))
 
-  return { channel, cleanup }
+  const channel = {
+    onMessage: callback => emitter.on(events.message, callback),
+    postMessage: message => emitter.emit(events.response, message),
+    onClose: () => {
+      emitter.removeAllListeners()
+      rpc.$close(new Error('[vitest-pool]: Pending methods while closing rpc'))
+    },
+  } satisfies TinypoolChannel
+
+  return { channel }
 }
 
 export function createVmForksPool(
@@ -131,7 +133,7 @@ export function createVmForksPool(
       const paths = files.map(f => f.filepath)
       vitest.state.clearFiles(project, paths)
 
-      const { channel, cleanup } = createChildProcessChannel(project, name === 'collect')
+      const { channel } = createChildProcessChannel(project, name === 'collect')
       const workerId = ++id
       const data: ContextRPC = {
         pool: 'forks',
@@ -170,7 +172,7 @@ export function createVmForksPool(
         }
       }
       finally {
-        cleanup()
+        channel.onClose()
       }
     }
 

packages/vitest/src/node/pools/vmThreads.ts

@@ -32,14 +32,18 @@ function createWorkerChannel(project: TestProject, collect: boolean) {
     on(fn) {
       port.on('message', fn)
     },
-    onTimeoutError(functionName) {
-      throw new Error(`[vitest-pool]: Timeout calling "${functionName}"`)
-    },
+    timeout: -1,
   })
 
   project.vitest.onCancel(reason => rpc.onCancel(reason))
 
-  return { workerPort, port }
+  function onClose() {
+    workerPort.close()
+    port.close()
+    rpc.$close(new Error('[vitest-pool]: Pending methods while closing rpc'))
+  }
+
+  return { workerPort, onClose }
 }
 
 export function createVmThreadsPool(
@@ -108,7 +112,7 @@ export function createVmThreadsPool(
       const paths = files.map(f => f.filepath)
       vitest.state.clearFiles(project, paths)
 
-      const { workerPort, port } = createWorkerChannel(project, name === 'collect')
+      const { workerPort, onClose } = createWorkerChannel(project, name === 'collect')
       const workerId = ++id
       const data: WorkerContext = {
         pool: 'vmThreads',
@@ -150,8 +154,7 @@ export function createVmThreadsPool(
         }
       }
       finally {
-        port.close()
-        workerPort.close()
+        onClose()
       }
     }
 

packages/vitest/src/runtime/rpc.ts

@@ -83,24 +83,7 @@ export function createRuntimeRpc(
           'onCollected',
           'onCancel',
         ],
-        onTimeoutError(functionName, args) {
-          let message = `[vitest-worker]: Timeout calling "${functionName}"`
-
-          if (
-            functionName === 'fetch'
-            || functionName === 'transform'
-            || functionName === 'resolveId'
-          ) {
-            message += ` with "${JSON.stringify(args)}"`
-          }
-
-          // JSON.stringify cannot serialize Error instances
-          if (functionName === 'onUnhandledError') {
-            message += ` with "${args[0]?.message || args[0]}"`
-          }
-
-          throw new Error(message)
-        },
+        timeout: -1,
         ...options,
       },
     ),
@@ -115,6 +98,11 @@ export function createRuntimeRpc(
 export function createSafeRpc(rpc: WorkerRPC): WorkerRPC {
   return new Proxy(rpc, {
     get(target, p, handler) {
+      // keep $rejectPendingCalls as sync function
+      if (p === '$rejectPendingCalls') {
+        return rpc.$rejectPendingCalls
+      }
+
       const sendCall = get(target, p, handler)
       const safeSendCall = (...args: any[]) =>
         withSafeTimers(async () => {

packages/vitest/src/runtime/worker.ts

@@ -36,7 +36,7 @@ async function execute(method: 'run' | 'collect', ctx: ContextRPC) {
 
   const prepareStart = performance.now()
 
-  const inspectorCleanup = setupInspect(ctx)
+  const cleanups: (() => void | Promise<void>)[] = [setupInspect(ctx)]
 
   process.env.VITEST_WORKER_ID = String(ctx.workerId)
   process.env.VITEST_POOL_ID = String(poolId)
@@ -73,6 +73,13 @@ async function execute(method: 'run' | 'collect', ctx: ContextRPC) {
     // RPC is used to communicate between worker (be it a thread worker or child process or a custom implementation) and the main thread
     const { rpc, onCancel } = createRuntimeRpc(worker.getRpcOptions(ctx))
 
+    // do not close the RPC channel so that we can get the error messages sent to the main thread
+    cleanups.push(async () => {
+      await Promise.all(rpc.$rejectPendingCalls(({ method, reject }) => {
+        reject(new Error(`[vitest-worker]: Closing rpc while "${method}" was pending`))
+      }))
+    })
+
     const beforeEnvironmentTime = performance.now()
     const environment = await loadEnvironment(ctx, rpc)
     if (ctx.environment.transformMode) {
@@ -110,8 +117,9 @@ async function execute(method: 'run' | 'collect', ctx: ContextRPC) {
     await worker[methodName](state)
   }
   finally {
+    await Promise.all(cleanups.map(fn => fn()))
+
     await rpcDone().catch(() => {})
-    inspectorCleanup()
   }
 }
 

packages/ws-client/src/index.ts

@@ -99,9 +99,7 @@ export function createClient(url: string, options: VitestClientOptions = {}): Vi
         return v
       }),
     deserialize: parse,
-    onTimeoutError(functionName) {
-      throw new Error(`[vitest-ws-client]: Timeout calling "${functionName}"`)
-    },
+    timeout: -1,
   }
 
   ctx.rpc = createBirpc<WebSocketHandlers, WebSocketEvents>(