Skip to content

Commit bea8746

Browse files
AriPerkkiosheremet-va
AriPerkkio
and
sheremet-va
authored
fix: prevent rpc timeout on slow thread blocking synchronous methods (#8297)
Co-authored-by: Vladimir Sheremet <sleuths.slews0s@icloud.com>
1 parent c16abe7 commit bea8746

File tree

12 files changed

+65
-73
lines changed

12 files changed

+65
-73
lines changed

‎packages/browser/src/client/client.ts

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -109,7 +109,6 @@ function createClient() {
109109
{
110110
post: msg => ctx.ws.send(msg),
111111
on: fn => (onMessage = fn),
112-
timeout: -1, // createTesters can take a while
113112
serialize: e =>
114113
stringify(e, (_, v) => {
115114
if (v instanceof Error) {
@@ -122,9 +121,7 @@ function createClient() {
122121
return v
123122
}),
124123
deserialize: parse,
125-
onTimeoutError(functionName) {
126-
throw new Error(`[vitest-browser]: Timeout calling "${functionName}"`)
127-
},
124+
timeout: -1, // createTesters can take a while
128125
},
129126
)
130127

‎packages/browser/src/node/rpc.ts

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -337,11 +337,8 @@ export function setupBrowserRpc(globalServer: ParentBrowserProject, defaultMocke
337337
on: fn => ws.on('message', fn),
338338
eventNames: ['onCancel', 'cdpEvent'],
339339
serialize: (data: any) => stringify(data, stringifyReplace),
340-
timeout: -1, // createTesters can take a long time
341340
deserialize: parse,
342-
onTimeoutError(functionName) {
343-
throw new Error(`[vitest-api]: Timeout calling "${functionName}"`)
344-
},
341+
timeout: -1, // createTesters can take a long time
345342
},
346343
)
347344

‎packages/vitest/src/api/setup.ts

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import type { File, TaskEventPack, TaskResultPack, TestAnnotation } from '@vitest/runner'
2-
32
import type { IncomingMessage } from 'node:http'
43
import type { ViteDevServer } from 'vite'
54
import type { WebSocket } from 'ws'
@@ -142,16 +141,15 @@ export function setup(ctx: Vitest, _server?: ViteDevServer): void {
142141
],
143142
serialize: (data: any) => stringify(data, stringifyReplace),
144143
deserialize: parse,
145-
onTimeoutError(functionName) {
146-
throw new Error(`[vitest-api]: Timeout calling "${functionName}"`)
147-
},
144+
timeout: -1,
148145
},
149146
)
150147

151148
clients.set(ws, rpc)
152149

153150
ws.on('close', () => {
154151
clients.delete(ws)
152+
rpc.$close(new Error('[vitest-api]: Pending methods while closing rpc'))
155153
})
156154
}
157155

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

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,7 @@ import { createMethodsRPC } from './rpc'
1919

2020
function createChildProcessChannel(project: TestProject, collect = false) {
2121
const emitter = new EventEmitter()
22-
2322
const events = { message: 'message', response: 'response' }
24-
const channel: TinypoolChannel = {
25-
onMessage: callback => emitter.on(events.message, callback),
26-
postMessage: message => emitter.emit(events.response, message),
27-
onClose: () => emitter.removeAllListeners(),
28-
}
2923

3024
const rpc = createBirpc<RunnerRPC, RuntimeRPC>(createMethodsRPC(project, { cacheFs: true, collect }), {
3125
eventNames: ['onCancel'],
@@ -51,13 +45,20 @@ function createChildProcessChannel(project: TestProject, collect = false) {
5145
on(fn) {
5246
emitter.on(events.response, fn)
5347
},
54-
onTimeoutError(functionName) {
55-
throw new Error(`[vitest-pool]: Timeout calling "${functionName}"`)
56-
},
48+
timeout: -1,
5749
})
5850

5951
project.vitest.onCancel(reason => rpc.onCancel(reason))
6052

53+
const channel: TinypoolChannel = {
54+
onMessage: callback => emitter.on(events.message, callback),
55+
postMessage: message => emitter.emit(events.response, message),
56+
onClose: () => {
57+
emitter.removeAllListeners()
58+
rpc.$close(new Error('[vitest-pool]: Pending methods while closing rpc'))
59+
},
60+
}
61+
6162
return channel
6263
}
6364

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

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -29,14 +29,18 @@ function createWorkerChannel(project: TestProject, collect: boolean) {
2929
on(fn) {
3030
port.on('message', fn)
3131
},
32-
onTimeoutError(functionName) {
33-
throw new Error(`[vitest-pool]: Timeout calling "${functionName}"`)
34-
},
32+
timeout: -1,
3533
})
3634

3735
project.vitest.onCancel(reason => rpc.onCancel(reason))
3836

39-
return { workerPort, port }
37+
const onClose = () => {
38+
port.close()
39+
workerPort.close()
40+
rpc.$close(new Error('[vitest-pool]: Pending methods while closing rpc'))
41+
}
42+
43+
return { workerPort, port, onClose }
4044
}
4145

4246
export function createThreadsPool(
@@ -104,11 +108,7 @@ export function createThreadsPool(
104108
const paths = files.map(f => f.filepath)
105109
vitest.state.clearFiles(project, paths)
106110

107-
const { workerPort, port } = createWorkerChannel(project, name === 'collect')
108-
const onClose = () => {
109-
port.close()
110-
workerPort.close()
111-
}
111+
const { workerPort, onClose } = createWorkerChannel(project, name === 'collect')
112112

113113
const workerId = ++id
114114
const data: WorkerContext = {

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

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -22,13 +22,8 @@ const suppressWarningsPath = resolve(rootDir, './suppress-warnings.cjs')
2222

2323
function createChildProcessChannel(project: TestProject, collect: boolean) {
2424
const emitter = new EventEmitter()
25-
const cleanup = () => emitter.removeAllListeners()
2625

2726
const events = { message: 'message', response: 'response' }
28-
const channel: TinypoolChannel = {
29-
onMessage: callback => emitter.on(events.message, callback),
30-
postMessage: message => emitter.emit(events.response, message),
31-
}
3227

3328
const rpc = createBirpc<RunnerRPC, RuntimeRPC>(
3429
createMethodsRPC(project, { cacheFs: true, collect }),
@@ -56,15 +51,22 @@ function createChildProcessChannel(project: TestProject, collect: boolean) {
5651
on(fn) {
5752
emitter.on(events.response, fn)
5853
},
59-
onTimeoutError(functionName) {
60-
throw new Error(`[vitest-pool]: Timeout calling "${functionName}"`)
61-
},
54+
timeout: -1,
6255
},
6356
)
6457

6558
project.vitest.onCancel(reason => rpc.onCancel(reason))
6659

67-
return { channel, cleanup }
60+
const channel = {
61+
onMessage: callback => emitter.on(events.message, callback),
62+
postMessage: message => emitter.emit(events.response, message),
63+
onClose: () => {
64+
emitter.removeAllListeners()
65+
rpc.$close(new Error('[vitest-pool]: Pending methods while closing rpc'))
66+
},
67+
} satisfies TinypoolChannel
68+
69+
return { channel }
6870
}
6971

7072
export function createVmForksPool(
@@ -131,7 +133,7 @@ export function createVmForksPool(
131133
const paths = files.map(f => f.filepath)
132134
vitest.state.clearFiles(project, paths)
133135

134-
const { channel, cleanup } = createChildProcessChannel(project, name === 'collect')
136+
const { channel } = createChildProcessChannel(project, name === 'collect')
135137
const workerId = ++id
136138
const data: ContextRPC = {
137139
pool: 'forks',
@@ -170,7 +172,7 @@ export function createVmForksPool(
170172
}
171173
}
172174
finally {
173-
cleanup()
175+
channel.onClose()
174176
}
175177
}
176178

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

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -32,14 +32,18 @@ function createWorkerChannel(project: TestProject, collect: boolean) {
3232
on(fn) {
3333
port.on('message', fn)
3434
},
35-
onTimeoutError(functionName) {
36-
throw new Error(`[vitest-pool]: Timeout calling "${functionName}"`)
37-
},
35+
timeout: -1,
3836
})
3937

4038
project.vitest.onCancel(reason => rpc.onCancel(reason))
4139

42-
return { workerPort, port }
40+
function onClose() {
41+
workerPort.close()
42+
port.close()
43+
rpc.$close(new Error('[vitest-pool]: Pending methods while closing rpc'))
44+
}
45+
46+
return { workerPort, onClose }
4347
}
4448

4549
export function createVmThreadsPool(
@@ -108,7 +112,7 @@ export function createVmThreadsPool(
108112
const paths = files.map(f => f.filepath)
109113
vitest.state.clearFiles(project, paths)
110114

111-
const { workerPort, port } = createWorkerChannel(project, name === 'collect')
115+
const { workerPort, onClose } = createWorkerChannel(project, name === 'collect')
112116
const workerId = ++id
113117
const data: WorkerContext = {
114118
pool: 'vmThreads',
@@ -150,8 +154,7 @@ export function createVmThreadsPool(
150154
}
151155
}
152156
finally {
153-
port.close()
154-
workerPort.close()
157+
onClose()
155158
}
156159
}
157160

‎packages/vitest/src/runtime/rpc.ts

Lines changed: 6 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -83,24 +83,7 @@ export function createRuntimeRpc(
8383
'onCollected',
8484
'onCancel',
8585
],
86-
onTimeoutError(functionName, args) {
87-
let message = `[vitest-worker]: Timeout calling "${functionName}"`
88-
89-
if (
90-
functionName === 'fetch'
91-
|| functionName === 'transform'
92-
|| functionName === 'resolveId'
93-
) {
94-
message += ` with "${JSON.stringify(args)}"`
95-
}
96-
97-
// JSON.stringify cannot serialize Error instances
98-
if (functionName === 'onUnhandledError') {
99-
message += ` with "${args[0]?.message || args[0]}"`
100-
}
101-
102-
throw new Error(message)
103-
},
86+
timeout: -1,
10487
...options,
10588
},
10689
),
@@ -115,6 +98,11 @@ export function createRuntimeRpc(
11598
export function createSafeRpc(rpc: WorkerRPC): WorkerRPC {
11699
return new Proxy(rpc, {
117100
get(target, p, handler) {
101+
// keep $rejectPendingCalls as sync function
102+
if (p === '$rejectPendingCalls') {
103+
return rpc.$rejectPendingCalls
104+
}
105+
118106
const sendCall = get(target, p, handler)
119107
const safeSendCall = (...args: any[]) =>
120108
withSafeTimers(async () => {

‎packages/vitest/src/runtime/worker.ts

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ async function execute(method: 'run' | 'collect', ctx: ContextRPC) {
3636

3737
const prepareStart = performance.now()
3838

39-
const inspectorCleanup = setupInspect(ctx)
39+
const cleanups: (() => void | Promise<void>)[] = [setupInspect(ctx)]
4040

4141
process.env.VITEST_WORKER_ID = String(ctx.workerId)
4242
process.env.VITEST_POOL_ID = String(poolId)
@@ -73,6 +73,13 @@ async function execute(method: 'run' | 'collect', ctx: ContextRPC) {
7373
// RPC is used to communicate between worker (be it a thread worker or child process or a custom implementation) and the main thread
7474
const { rpc, onCancel } = createRuntimeRpc(worker.getRpcOptions(ctx))
7575

76+
// do not close the RPC channel so that we can get the error messages sent to the main thread
77+
cleanups.push(async () => {
78+
await Promise.all(rpc.$rejectPendingCalls(({ method, reject }) => {
79+
reject(new Error(`[vitest-worker]: Closing rpc while "${method}" was pending`))
80+
}))
81+
})
82+
7683
const beforeEnvironmentTime = performance.now()
7784
const environment = await loadEnvironment(ctx, rpc)
7885
if (ctx.environment.transformMode) {
@@ -110,8 +117,9 @@ async function execute(method: 'run' | 'collect', ctx: ContextRPC) {
110117
await worker[methodName](state)
111118
}
112119
finally {
120+
await Promise.all(cleanups.map(fn => fn()))
121+
113122
await rpcDone().catch(() => {})
114-
inspectorCleanup()
115123
}
116124
}
117125

‎packages/ws-client/src/index.ts

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -99,9 +99,7 @@ export function createClient(url: string, options: VitestClientOptions = {}): Vi
9999
return v
100100
}),
101101
deserialize: parse,
102-
onTimeoutError(functionName) {
103-
throw new Error(`[vitest-ws-client]: Timeout calling "${functionName}"`)
104-
},
102+
timeout: -1,
105103
}
106104

107105
ctx.rpc = createBirpc<WebSocketHandlers, WebSocketEvents>(

0 commit comments

Comments
 (0)