test: add web socket leak test coverage

src/server/browserServer.ts

@@ -17,25 +17,53 @@
 import { ChildProcess, execSync } from 'child_process';
 import * as platform from '../platform';
 
+export class WebSocketWrapper {
+  readonly wsEndpoint: string;
+  private _bindings: (Map<any, any> | Set<any>)[];
+  constructor(wsEndpoint: string, bindings: (Map<any, any>|Set<any>)[]) {
+    this.wsEndpoint = wsEndpoint;
+    this._bindings = bindings;
+  }
+
+  async checkLeaks() {
+    let counter = 0;
+    return new Promise((fulfill, reject) => {
+      const check = () => {
+        const filtered = this._bindings.filter(entry => entry.size);
+        if (!filtered.length) {
+          fulfill();
+          return;
+        }
+
+        if (++counter >= 50) {
+          reject(new Error('Web socket leak ' + filtered.map(entry => [...entry.keys()].join(':')).join('|')));
+          return;
+        }
+        setTimeout(check, 100);
+      };
+      check();
+    });
+  }
+}
+
 export class BrowserServer extends platform.EventEmitter {
   private _process: ChildProcess;
   private _gracefullyClose: () => Promise<void>;
-  private _browserWSEndpoint: string = '';
+  private _webSocketWrapper: WebSocketWrapper | null;
 
-  constructor(process: ChildProcess, gracefullyClose: () => Promise<void>, wsEndpoint: string|null) {
+  constructor(process: ChildProcess, gracefullyClose: () => Promise<void>, webSocketWrapper: WebSocketWrapper | null) {
     super();
     this._process = process;
     this._gracefullyClose = gracefullyClose;
-    if (wsEndpoint)
-      this._browserWSEndpoint = wsEndpoint;
+    this._webSocketWrapper = webSocketWrapper;
   }
 
   process(): ChildProcess {
     return this._process;
   }
 
   wsEndpoint(): string {
-    return this._browserWSEndpoint;
+    return this._webSocketWrapper ? this._webSocketWrapper.wsEndpoint : '';
   }
 
   kill() {
@@ -54,4 +82,9 @@ export class BrowserServer extends platform.EventEmitter {
   async close(): Promise<void> {
     await this._gracefullyClose();
   }
+
+  async _checkLeaks(): Promise<void> {
+    if (this._webSocketWrapper)
+      await this._webSocketWrapper.checkLeaks();
+  }
 }

src/server/chromium.ts

@@ -27,7 +27,7 @@ import { kBrowserCloseMessageId } from '../chromium/crConnection';
 import { PipeTransport } from './pipeTransport';
 import { LaunchOptions, BrowserArgOptions, BrowserType } from './browserType';
 import { ConnectOptions, LaunchType } from '../browser';
-import { BrowserServer } from './browserServer';
+import { BrowserServer, WebSocketWrapper } from './browserServer';
 import { Events } from '../events';
 import { ConnectionTransport, ProtocolRequest } from '../transport';
 import { BrowserContext } from '../browserContext';
@@ -172,7 +172,7 @@ export class Chromium implements BrowserType<CRBrowser> {
   }
 }
 
-function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number): string {
+function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number): WebSocketWrapper {
   const server = new ws.Server({ port });
   const guid = platform.guid();
 
@@ -285,9 +285,8 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number
   });
 
   const address = server.address();
-  if (typeof address === 'string')
-    return address + '/' + guid;
-  return 'ws://127.0.0.1:' + address.port + '/' + guid;
+  const wsEndpoint = typeof address === 'string' ? `${address}/${guid}` : `ws://127.0.0.1:${address.port}/${guid}`;
+  return new WebSocketWrapper(wsEndpoint, [awaitingBrowserTarget, sessionToSocket, socketToBrowserSession, browserSessions]);
 }
 
 

src/server/firefox.ts

@@ -27,7 +27,7 @@ import { FFBrowser } from '../firefox/ffBrowser';
 import { kBrowserCloseMessageId } from '../firefox/ffConnection';
 import { helper } from '../helper';
 import * as platform from '../platform';
-import { BrowserServer } from './browserServer';
+import { BrowserServer, WebSocketWrapper } from './browserServer';
 import { BrowserArgOptions, BrowserType, LaunchOptions } from './browserType';
 import { launchProcess, waitForLine } from './processLauncher';
 import { ConnectionTransport, SequenceNumberMixer } from '../transport';
@@ -129,7 +129,7 @@ export class Firefox implements BrowserType<FFBrowser> {
         // We try to gracefully close to prevent crash reporting and core dumps.
         // Note that it's fine to reuse the pipe transport, since
         // our connection ignores kBrowserCloseMessageId.
-        const transport = await platform.connectToWebsocket(browserWSEndpoint, async transport => transport);
+        const transport = await platform.connectToWebsocket(browserWSEndpoint!, async transport => transport);
         const message = { method: 'Browser.close', params: {}, id: kBrowserCloseMessageId };
         await transport.send(message);
       },
@@ -144,8 +144,10 @@ export class Firefox implements BrowserType<FFBrowser> {
     const innerEndpoint = match[1];
 
     let browserServer: BrowserServer | undefined = undefined;
-    const browserWSEndpoint = launchType === 'server' ? (await platform.connectToWebsocket(innerEndpoint, t => wrapTransportWithWebSocket(t, port || 0))) : innerEndpoint;
-    browserServer = new BrowserServer(launchedProcess, gracefullyClose, browserWSEndpoint);
+    let browserWSEndpoint: string | undefined = undefined;
+    const webSocketWrapper = launchType === 'server' ? (await platform.connectToWebsocket(innerEndpoint, t => wrapTransportWithWebSocket(t, port || 0))) : new WebSocketWrapper(innerEndpoint, []);
+    browserWSEndpoint = webSocketWrapper.wsEndpoint;
+    browserServer = new BrowserServer(launchedProcess, gracefullyClose, webSocketWrapper);
     return browserServer;
   }
 
@@ -189,7 +191,7 @@ export class Firefox implements BrowserType<FFBrowser> {
   }
 }
 
-function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number): string {
+function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number): WebSocketWrapper {
   const server = new ws.Server({ port });
   const guid = platform.guid();
   const idMixer = new SequenceNumberMixer<{id: number, socket: ws}>();
@@ -311,8 +313,7 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number
   });
 
   const address = server.address();
-  if (typeof address === 'string')
-    return address + '/' + guid;
-  return 'ws://127.0.0.1:' + address.port + '/' + guid;
+  const wsEndpoint = typeof address === 'string' ? `${address}/${guid}` : `ws://127.0.0.1:${address.port}/${guid}`;
+  return new WebSocketWrapper(wsEndpoint,
+      [pendingBrowserContextCreations, pendingBrowserContextDeletions, browserContextIds, sessionToSocket, sockets]);
 }
-

src/server/webkit.ts

@@ -28,7 +28,7 @@ import { LaunchOptions, BrowserArgOptions, BrowserType } from './browserType';
 import { ConnectionTransport, SequenceNumberMixer } from '../transport';
 import * as ws from 'ws';
 import { ConnectOptions, LaunchType } from '../browser';
-import { BrowserServer } from './browserServer';
+import { BrowserServer, WebSocketWrapper } from './browserServer';
 import { Events } from '../events';
 import { BrowserContext } from '../browserContext';
 
@@ -163,7 +163,7 @@ const mkdtempAsync = platform.promisify(fs.mkdtemp);
 
 const WEBKIT_PROFILE_PATH = path.join(os.tmpdir(), 'playwright_dev_profile-');
 
-function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number): string {
+function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number): WebSocketWrapper {
   const server = new ws.Server({ port });
   const guid = platform.guid();
   const idMixer = new SequenceNumberMixer<{id: number, socket: ws}>();
@@ -296,7 +296,8 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number
   });
 
   const address = server.address();
-  if (typeof address === 'string')
-    return address + '/' + guid;
-  return 'ws://127.0.0.1:' + address.port + '/' + guid;
+  const wsEndpoint = typeof address === 'string' ? `${address}/${guid}` : `ws://127.0.0.1:${address.port}/${guid}`;
+
+  return new WebSocketWrapper(wsEndpoint,
+      [pendingBrowserContextCreations, pendingBrowserContextDeletions, browserContextIds, pageProxyIds, sockets]);
 }

test/launcher.spec.js

@@ -142,6 +142,7 @@ module.exports.describe = function({testRunner, expect, defaultBrowserOptions, p
       expect(remote.isConnected()).toBe(true);
       await remote.close();
       expect(remote.isConnected()).toBe(false);
+      await browserServer._checkLeaks();
       await browserServer.close();
     });
     it('should throw when used after isConnected returns false', async({server}) => {
@@ -169,6 +170,7 @@ module.exports.describe = function({testRunner, expect, defaultBrowserOptions, p
       await remote.close();
       const error = await navigationPromise;
       expect(error.message).toContain('Navigation failed because browser has disconnected!');
+      await browserServer._checkLeaks();
       await browserServer.close();
     });
     it('should reject waitForSelector when browser closes', async({server}) => {
@@ -184,6 +186,7 @@ module.exports.describe = function({testRunner, expect, defaultBrowserOptions, p
       await remote.close();
       const error = await watchdog;
       expect(error.message).toContain('Protocol error');
+      await browserServer._checkLeaks();
       await browserServer.close();
     });
     it('should throw if used after disconnect', async({server}) => {
@@ -193,6 +196,7 @@ module.exports.describe = function({testRunner, expect, defaultBrowserOptions, p
       await remote.close();
       const error = await page.evaluate('1 + 1').catch(e => e);
       expect(error.message).toContain('has been closed');
+      await browserServer._checkLeaks();
       await browserServer.close();
     });
     it('should emit close events on pages and contexts', async({server}) => {
@@ -246,6 +250,8 @@ module.exports.describe = function({testRunner, expect, defaultBrowserOptions, p
       const page = await browserContext.newPage();
       expect(await page.evaluate('11 * 11')).toBe(121);
       await page.close();
+      await browser.close();
+      await browserServer._checkLeaks();
       await browserServer.close();
     });
     it('should fire "disconnected" when closing with webSocket', async() => {
@@ -274,6 +280,7 @@ module.exports.describe = function({testRunner, expect, defaultBrowserOptions, p
         await page.goto(server.EMPTY_PAGE);
         await browser.close();
       }
+      await browserServer._checkLeaks();
       await browserServer.close();
     });
   });

test/multiclient.spec.js

@@ -39,6 +39,11 @@ module.exports.describe = function({testRunner, expect, defaultBrowserOptions, b
       expect(browser2.contexts().length).toBe(1);
 
       expect(browser1.contexts().length).toBe(1);
+
+      await browser1.close();
+      await browser2.close();
+
+      await browserServer._checkLeaks();
       await browserServer.close();
     });
   });
@@ -90,6 +95,8 @@ module.exports.describe = function({testRunner, expect, defaultBrowserOptions, b
 
       const page2 = await browser2.newPage();
       expect(await page2.evaluate(() => 7 * 6)).toBe(42, 'original browser should still work');
+      await browser2.close();
+      await browserServer._checkLeaks();
       await browserServer.close();
     });
     it('should not be able to close remote browser', async() => {
@@ -104,6 +111,8 @@ module.exports.describe = function({testRunner, expect, defaultBrowserOptions, b
         await remote.newContext();
         await remote.close();
       }
+      await browserServer._checkLeaks();
+      await browserServer.close();
     });
   });
 };