diff --git a/apps/playwright-browser-tunnel/src/PlaywrightBrowserTunnel.ts b/apps/playwright-browser-tunnel/src/PlaywrightBrowserTunnel.ts index 1fa0627b531..aff0a1b8d3c 100644 --- a/apps/playwright-browser-tunnel/src/PlaywrightBrowserTunnel.ts +++ b/apps/playwright-browser-tunnel/src/PlaywrightBrowserTunnel.ts @@ -96,11 +96,12 @@ export class PlaywrightTunnel { private readonly _listenPort: number | undefined; private readonly _playwrightInstallPath: string; private _status: TunnelStatus = 'stopped'; - private _initWsPromise?: Promise; + private _initWsPromise?: Promise; private _keepRunning: boolean = false; private _ws?: WebSocket; private _mode: TunnelMode; private _pendingConnectionAttempt?: Promise; + private _cancelPendingConnection?: () => void; private _pollInterval?: NodeJS.Timeout; public constructor(options: IPlaywrightTunnelOptions) { @@ -144,9 +145,14 @@ export class PlaywrightTunnel { public async waitForCloseAsync(): Promise { const terminal: ITerminal = this._terminal; - const initWsPromise: Promise | undefined = this._initWsPromise; + const initWsPromise: Promise | undefined = this._initWsPromise; if (initWsPromise) { - const ws: WebSocket = await initWsPromise; + const ws: WebSocket | undefined = await initWsPromise; + if (!ws) { + terminal.writeDebugLine('WebSocket connection was cancelled before it was established.'); + this._initWsPromise = undefined; + return; + } await once(ws, 'close'); terminal.writeDebugLine('WebSocket connection closed. resolving init promise.'); this._initWsPromise = undefined; @@ -173,6 +179,15 @@ export class PlaywrightTunnel { clearInterval(this._pollInterval); this._pollInterval = undefined; } + if (!this._ws) { + this._cancelPendingConnection?.(); + this._cancelPendingConnection = undefined; + this._pendingConnectionAttempt = undefined; + this._initWsPromise = undefined; + this.status = 'stopped'; + return; + } + await this._initWsPromise?.finally(() => { this._ws?.close(WebSocketCloseCode.NORMAL_CLOSURE, 'Tunnel stopped'); }); @@ -269,9 +284,28 @@ export class PlaywrightTunnel { // TODO: Only supporting one test at a time. // Need to support multiple simultaneous connections for parallel tests. - private async _pollConnectionAsync(): Promise { + private async _pollConnectionAsync(): Promise { this._terminal.writeLine(`Waiting for WebSocket connection`); - return await new Promise((resolve, reject) => { + return await new Promise((resolve) => { + let settled: boolean = false; + const cleanup = (): void => { + if (this._pollInterval) { + clearInterval(this._pollInterval); + this._pollInterval = undefined; + } + this._pendingConnectionAttempt = undefined; + this._cancelPendingConnection = undefined; + }; + + this._cancelPendingConnection = (): void => { + if (settled) { + return; + } + settled = true; + cleanup(); + resolve(undefined); + }; + this._pollInterval = setInterval(() => { if (this._pendingConnectionAttempt) { return; // Skip if a connection attempt is already in progress @@ -280,10 +314,14 @@ export class PlaywrightTunnel { this._pendingConnectionAttempt = connectionPromise; connectionPromise .then((ws: WebSocket) => { - clearInterval(this._pollInterval); - this._pollInterval = undefined; + if (settled || !this._keepRunning) { + ws.close(WebSocketCloseCode.NORMAL_CLOSURE, 'Tunnel stopped'); + return; + } + settled = true; + cleanup(); + this._ws = ws; ws.removeAllListeners(); - this._pendingConnectionAttempt = undefined; resolve(ws); }) .catch(() => { @@ -519,17 +557,24 @@ export class PlaywrightTunnel { * and setting up the browser server. * Returns when the handshake is complete and the browser server is running. */ - private async _initPlaywrightBrowserTunnelAsync(): Promise { + private async _initPlaywrightBrowserTunnelAsync(): Promise { let handshake: IHandshake | undefined = undefined; let client: WebSocket | undefined = undefined; let browserServer: BrowserServer | undefined = undefined; this.status = 'waiting-for-connection'; - const ws: WebSocket = + const ws: WebSocket | undefined = this._mode === 'poll-connection' ? await this._pollConnectionAsync() : await this._waitForIncomingConnectionAsync(); + if (!ws) { + this._terminal.writeLine('Playwright tunnel start cancelled before a WebSocket connected.'); + this._initWsPromise = undefined; + this.status = 'stopped'; + return undefined; + } + ws.on('open', () => { this._terminal.writeLine(`WebSocket connection established`); handshake = undefined; @@ -543,6 +588,7 @@ export class PlaywrightTunnel { const reasonStr: string = reason.toString() || 'no reason provided'; const codeDescription: string = getWebSocketCloseReason(code); this._initWsPromise = undefined; + this._ws = undefined; this.status = 'stopped'; this._terminal.writeLine( `WebSocket connection closed - code: ${code} (${codeDescription}), reason: ${reasonStr}` @@ -601,6 +647,7 @@ export class PlaywrightTunnel { } this.status = 'browser-server-running'; + this._ws = ws; // Send ack so that the counterpart also knows to start forwarding messages. // NOTE: The 1-second delay is an intentional workaround. In the current diff --git a/apps/playwright-browser-tunnel/tests/pollStop.test.cjs b/apps/playwright-browser-tunnel/tests/pollStop.test.cjs new file mode 100644 index 00000000000..0d31f5e3783 --- /dev/null +++ b/apps/playwright-browser-tunnel/tests/pollStop.test.cjs @@ -0,0 +1,95 @@ +const assert = require('node:assert/strict'); +const { mkdtempSync } = require('node:fs'); +const { tmpdir } = require('node:os'); +const { join } = require('node:path'); +const test = require('node:test'); +const { WebSocketServer } = require('ws'); + +const { PlaywrightTunnel } = require('../lib-commonjs/index.js'); + +function createTerminal() { + return { + writeLine() {}, + writeWarningLine() {}, + writeErrorLine() {}, + writeDebugLine() {} + }; +} + +async function createUnusedWsEndpointAsync() { + const server = new WebSocketServer({ port: 0 }); + await new Promise((resolve) => server.once('listening', resolve)); + const { port } = server.address(); + await new Promise((resolve) => server.close(resolve)); + return `ws://127.0.0.1:${port}`; +} + +async function settledWithinAsync(promise, milliseconds) { + let settled = false; + promise.then( + () => { + settled = true; + }, + () => { + settled = true; + } + ); + await new Promise((resolve) => setTimeout(resolve, milliseconds)); + return settled; +} + +test('stopAsync settles while waiting for a connection', async () => { + const tunnel = new PlaywrightTunnel({ + mode: 'poll-connection', + wsEndpoint: await createUnusedWsEndpointAsync(), + terminal: createTerminal(), + playwrightInstallPath: mkdtempSync(join(tmpdir(), 'rushstack-playwright-browser-tunnel-')), + onStatusChange() {} + }); + + const startPromise = tunnel.startAsync(); + assert.equal(tunnel.status, 'waiting-for-connection'); + + const stopPromise = tunnel.stopAsync(); + assert.equal(await settledWithinAsync(stopPromise, 250), true); + assert.equal(tunnel.status, 'stopped'); + assert.equal(await settledWithinAsync(startPromise, 250), true); + + const restartPromise = tunnel.startAsync(); + assert.equal(tunnel.status, 'waiting-for-connection'); + + const restartStopPromise = tunnel.stopAsync(); + assert.equal(await settledWithinAsync(restartStopPromise, 250), true); + assert.equal(await settledWithinAsync(restartPromise, 250), true); +}); + +test('stopAsync closes an active websocket connection', async () => { + const server = new WebSocketServer({ port: 0 }); + await new Promise((resolve) => server.once('listening', resolve)); + + const connectionPromise = new Promise((resolve) => server.once('connection', resolve)); + const closePromise = new Promise((resolve) => { + server.once('connection', (ws) => ws.once('close', resolve)); + }); + + const { port } = server.address(); + const tunnel = new PlaywrightTunnel({ + mode: 'poll-connection', + wsEndpoint: `ws://127.0.0.1:${port}`, + terminal: createTerminal(), + playwrightInstallPath: mkdtempSync(join(tmpdir(), 'rushstack-playwright-browser-tunnel-')), + onStatusChange() {} + }); + + const startPromise = tunnel.startAsync(); + assert.equal(tunnel.status, 'waiting-for-connection'); + await connectionPromise; + + const stopPromise = tunnel.stopAsync(); + assert.equal(await settledWithinAsync(stopPromise, 250), true); + assert.equal(tunnel.status, 'stopped'); + assert.equal(await settledWithinAsync(startPromise, 250), true); + assert.equal(await settledWithinAsync(closePromise, 250), true); + + await new Promise((resolve) => server.close(resolve)); +});