Skip to content

Commit 3545fe7

Browse files
sheremet-vasheremet-va
authored
fix: add handshake timeout to iframe communication (#10656)
1 parent 90c4ed4 commit 3545fe7

4 files changed

Lines changed: 243 additions & 70 deletions

File tree

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,8 @@ export type IframeChannelEvent
7070
= | IframeChannelIncomingEvent
7171
| IframeChannelOutgoingEvent
7272

73+
export type IframeReceivedEvent = IframeChannelEvent & { messageId: number }
74+
7375
export const channel: BroadcastChannel = new BroadcastChannel(
7476
`vitest:${getBrowserState().sessionId}`,
7577
)

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

Lines changed: 75 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import type { Context as OTELContext } from '@opentelemetry/api'
2-
import type { GlobalChannelIncomingEvent, IframeChannelEvent, IframeChannelOutgoingEvent, IframeViewportDoneEvent, IframeViewportFailEvent } from '@vitest/browser/client'
2+
import type { GlobalChannelIncomingEvent, IframeChannelEvent, IframeChannelOutgoingEvent, IframeReceivedEvent, IframeViewportDoneEvent, IframeViewportFailEvent } from '@vitest/browser/client'
33
import type { BrowserTesterOptions, SerializedConfig } from 'vitest'
44
import type { FileSpecification } from 'vitest/internal/browser'
55
import { channel, client, globalChannel } from '@vitest/browser/client'
@@ -17,7 +17,8 @@ export class IframeOrchestrator {
1717
private recreateNonIsolatedIframe = false
1818
private iframes = new Map<string, HTMLIFrameElement>()
1919
private readyIframes = new Set<string>()
20-
private readyWaiters = new Map<string, () => void>()
20+
private readyWaiters = new Map<string, { resolve: () => void; reject: (error: Error) => void }>()
21+
private messageId = 0
2122

2223
public eventTarget: EventTarget = new EventTarget()
2324

@@ -290,7 +291,7 @@ export class IframeOrchestrator {
290291
const waiter = this.readyWaiters.get(iframeId)
291292
if (waiter) {
292293
this.readyWaiters.delete(iframeId)
293-
waiter()
294+
waiter.resolve()
294295
}
295296
}
296297

@@ -299,16 +300,41 @@ export class IframeOrchestrator {
299300
return Promise.resolve()
300301
}
301302

302-
return new Promise((resolve) => {
303-
this.readyWaiters.set(iframeId, resolve)
303+
return new Promise<void>((resolve, reject) => {
304+
const timeout = getIframeTimeout()
305+
// the tester reports readiness as soon as its module evaluates; if it
306+
// never does (e.g. it threw during bootstrap), don't wait forever
307+
const timer = setTimeout(() => {
308+
this.readyWaiters.delete(iframeId)
309+
reject(new Error(
310+
`The iframe "${iframeId}" did not become ready within ${timeout}ms. `
311+
+ `The tester likely failed to initialize, check the browser console for errors.`,
312+
))
313+
}, timeout)
314+
315+
this.readyWaiters.set(iframeId, {
316+
resolve: () => {
317+
clearTimeout(timer)
318+
resolve()
319+
},
320+
reject: (error) => {
321+
clearTimeout(timer)
322+
reject(error)
323+
},
324+
})
304325
})
305326
}
306327

307328
private removeIframe(iframeId: string) {
308329
const iframe = this.iframes.get(iframeId)
309330
this.iframes.delete(iframeId)
310331
this.readyIframes.delete(iframeId)
311-
this.readyWaiters.delete(iframeId)
332+
const waiter = this.readyWaiters.get(iframeId)
333+
if (waiter) {
334+
this.readyWaiters.delete(iframeId)
335+
// surface an error instead of silently abandoning whoever awaits readiness
336+
waiter.reject(new Error(`The iframe "${iframeId}" was removed before it became ready.`))
337+
}
312338
iframe?.remove()
313339
}
314340

@@ -432,10 +458,11 @@ export class IframeOrchestrator {
432458
break
433459
}
434460
default: {
435-
// ignore responses
461+
// ignore acknowledgements and responses to events we sent
462+
const event = e.data.event
436463
if (
437-
typeof e.data.event === 'string'
438-
&& (e.data.event as string).startsWith('response:')
464+
typeof event === 'string'
465+
&& (event.startsWith('response:') || event.startsWith('ack:'))
439466
) {
440467
break
441468
}
@@ -465,25 +492,51 @@ export class IframeOrchestrator {
465492
}
466493
events.add(event.event)
467494

468-
channel.postMessage(event)
495+
const messageId = this.messageId++
496+
channel.postMessage({ ...event, messageId } satisfies IframeReceivedEvent)
497+
469498
return new Promise<void>((resolve, reject) => {
499+
let ackTimer: ReturnType<typeof setTimeout>
500+
470501
const cleanupEvents = () => {
502+
clearTimeout(ackTimer)
471503
channel.removeEventListener('message', onReceived)
472504
this.eventTarget.removeEventListener('iframeerror', onError)
505+
events!.delete(event.event)
473506
}
474507

508+
// The tester acknowledges the message as soon as it receives it, then
509+
// sends the actual response once the work is done. We only time out
510+
// waiting for the acknowledgement: it proves the tester is alive, after
511+
// which the work (e.g. running a whole test file) may take any amount of
512+
// time, so there is intentionally no deadline on the response itself.
513+
const timeout = getIframeTimeout()
514+
ackTimer = setTimeout(() => {
515+
cleanupEvents()
516+
reject(new Error(
517+
`The iframe "${event.iframeId}" did not acknowledge the "${event.event}" message within ${timeout}ms. `
518+
+ `The tester might have crashed, been removed, or be blocked by a long synchronous task.`,
519+
))
520+
}, timeout)
521+
475522
function onReceived(e: MessageEvent) {
476-
if (e.data.iframeId === event.iframeId && e.data.event === `response:${event.event}`) {
477-
resolve()
523+
if (e.data.iframeId !== event.iframeId || e.data.messageId !== messageId) {
524+
return
525+
}
526+
if (e.data.event === `ack:${event.event}`) {
527+
// alive and processing: wait for the response without a deadline
528+
clearTimeout(ackTimer)
529+
return
530+
}
531+
if (e.data.event === `response:${event.event}`) {
478532
cleanupEvents()
479-
events!.delete(event.event)
533+
resolve()
480534
}
481535
}
482536

483537
function onError(e: Event) {
484-
reject((e as CustomEvent).detail)
485538
cleanupEvents()
486-
events!.delete(event.event)
539+
reject((e as CustomEvent).detail)
487540
}
488541

489542
this.eventTarget.addEventListener('iframeerror', onError)
@@ -551,3 +604,10 @@ function debug(...args: unknown[]) {
551604
client.rpc.debug(...args.map(String))
552605
}
553606
}
607+
608+
// Liveness timeout for tester iframes (readiness and message acknowledgement),
609+
// not a timeout for the test work itself. Overridable via the `VITEST_BROWSER_IFRAME_TIMEOUT`
610+
// env in case a tester legitimately needs longer to boot or acknowledge.
611+
function getIframeTimeout(): number {
612+
return Number(getConfig().env.VITEST_BROWSER_IFRAME_TIMEOUT) || 60_000
613+
}

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

Lines changed: 76 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { BrowserRPC, IframeChannelEvent } from '@vitest/browser/client'
1+
import type { BrowserRPC, IframeReceivedEvent } from '@vitest/browser/client'
22
import type { FileSpecification } from 'vitest/internal/browser'
33
import { channel, client, onCancel, registerPageMarkHandler } from '@vitest/browser/client'
44
import { parse } from 'flatted'
@@ -35,12 +35,10 @@ let rootTesterSpan: ReturnType<Traces['startContextSpan']> | undefined
3535
getBrowserState().traces = traces
3636

3737
channel.addEventListener('message', async (e) => {
38-
await client.waitForConnection()
39-
4038
const data = e.data
41-
debug?.('event from orchestrator', JSON.stringify(e.data))
4239

4340
if (!isEvent(data)) {
41+
await client.waitForConnection()
4442
const error = new Error(`Unknown message: ${JSON.stringify(e.data)}`)
4543
unhandledError(error, 'Unknown Iframe Message')
4644
return
@@ -51,60 +49,83 @@ channel.addEventListener('message', async (e) => {
5149
return
5250
}
5351

54-
switch (data.event) {
55-
case 'execute': {
56-
const { method, files, context, concurrencyId, workerId } = data
57-
const state = getWorkerState()
58-
const parsedContext = parse(context)
59-
60-
state.ctx.concurrencyId = concurrencyId
61-
state.ctx.workerId = workerId
62-
state.ctx.providedContext = parsedContext
63-
state.providedContext = parsedContext
64-
state.metaEnv.VITEST_POOL_ID = String(concurrencyId)
65-
state.metaEnv.VITEST_WORKER_ID = String(workerId)
66-
67-
if (method === 'collect') {
68-
await executeTests('collect', files).catch(err => unhandledError(err, 'Collect Error'))
52+
// tell the orchestrator we received the event before doing any work (which
53+
// may be long-running or gated on the connection), so it can tell a busy
54+
// tester apart from a crashed one. See `sendEventToIframe` in orchestrator.ts.
55+
channel.postMessage({
56+
event: `ack:${data.event}`,
57+
iframeId: data.iframeId,
58+
messageId: data.messageId,
59+
})
60+
61+
await client.waitForConnection()
62+
debug?.('event from orchestrator', JSON.stringify(e.data))
63+
64+
try {
65+
switch (data.event) {
66+
case 'execute': {
67+
const { method, files, context, concurrencyId, workerId } = data
68+
const state = getWorkerState()
69+
const parsedContext = parse(context)
70+
71+
state.ctx.concurrencyId = concurrencyId
72+
state.ctx.workerId = workerId
73+
state.ctx.providedContext = parsedContext
74+
state.providedContext = parsedContext
75+
state.metaEnv.VITEST_POOL_ID = String(concurrencyId)
76+
state.metaEnv.VITEST_WORKER_ID = String(workerId)
77+
78+
if (method === 'collect') {
79+
await executeTests('collect', files).catch(err => unhandledError(err, 'Collect Error'))
80+
}
81+
else {
82+
await executeTests('run', files).catch(err => unhandledError(err, 'Run Error'))
83+
}
84+
break
6985
}
70-
else {
71-
await executeTests('run', files).catch(err => unhandledError(err, 'Run Error'))
86+
case 'cleanup': {
87+
await cleanup().catch(err => unhandledError(err, 'Cleanup Error'))
88+
rootTesterSpan?.span.end()
89+
await traces.finish()
90+
break
91+
}
92+
case 'prepare': {
93+
await traces.waitInit()
94+
const tracesContext = traces.getContextFromCarrier(data.otelCarrier)
95+
traces.recordInitSpan(tracesContext)
96+
rootTesterSpan = traces.startContextSpan(
97+
`vitest.browser.tester.run`,
98+
tracesContext,
99+
)
100+
traces.bind(rootTesterSpan.context)
101+
await prepare(data).catch(err => unhandledError(err, 'Prepare Error'))
102+
break
103+
}
104+
case 'viewport:done':
105+
case 'viewport:fail':
106+
case 'viewport': {
107+
break
108+
}
109+
default: {
110+
const error = new Error(`Unknown event: ${(data as any).event}`)
111+
unhandledError(error, 'Unknown Event')
72112
}
73-
break
74-
}
75-
case 'cleanup': {
76-
await cleanup().catch(err => unhandledError(err, 'Cleanup Error'))
77-
rootTesterSpan?.span.end()
78-
await traces.finish()
79-
break
80-
}
81-
case 'prepare': {
82-
await traces.waitInit()
83-
const tracesContext = traces.getContextFromCarrier(data.otelCarrier)
84-
traces.recordInitSpan(tracesContext)
85-
rootTesterSpan = traces.startContextSpan(
86-
`vitest.browser.tester.run`,
87-
tracesContext,
88-
)
89-
traces.bind(rootTesterSpan.context)
90-
await prepare(data).catch(err => unhandledError(err, 'Prepare Error'))
91-
break
92-
}
93-
case 'viewport:done':
94-
case 'viewport:fail':
95-
case 'viewport': {
96-
break
97-
}
98-
default: {
99-
const error = new Error(`Unknown event: ${(data as any).event}`)
100-
unhandledError(error, 'Unknown Event')
101113
}
102114
}
103-
104-
channel.postMessage({
105-
event: `response:${data.event}`,
106-
iframeId: getBrowserState().iframeId!,
107-
})
115+
catch (error: any) {
116+
// errors not handled by the cases above (e.g. tracing setup/teardown) must
117+
// not stop us from responding, otherwise the orchestrator would wait forever
118+
await unhandledError(error, 'Tester Error')
119+
}
120+
finally {
121+
// always let the orchestrator know the event was handled so its
122+
// `sendEventToIframe` promise resolves, even if the work above threw
123+
channel.postMessage({
124+
event: `response:${data.event}`,
125+
iframeId: data.iframeId,
126+
messageId: data.messageId,
127+
})
128+
}
108129
})
109130

110131
const url = new URL(location.href)
@@ -307,6 +328,6 @@ function unhandledError(e: Error, type: string) {
307328
stack: e.stack,
308329
}, type).catch(() => {})
309330
}
310-
function isEvent(data: unknown): data is IframeChannelEvent {
331+
function isEvent(data: unknown): data is IframeReceivedEvent {
311332
return typeof data === 'object' && !!data && 'event' in data
312333
}

0 commit comments

Comments
 (0)