fix(context): reliably fire BrowserContext.Close event when browser i…

…s closing (#1277)
microsoft · Mar 9, 2020 · 9bd3711 · 9bd3711
1 parent 27eb25a
commit 9bd3711
Show file tree

Hide file tree

Showing 9 changed files with 53 additions and 39 deletions.
diff --git a/src/chromium/crBrowser.ts b/src/chromium/crBrowser.ts
@@ -396,7 +396,12 @@ export class CRBrowserContext extends BrowserContextBase {
   async close() {
     if (this._closed)
       return;
-    assert(this._browserContextId, 'Non-incognito profiles cannot be closed!');
+    if (!this._browserContextId) {
+      // Default context is only created in 'persistent' mode and closing it should close
+      // the browser.
+      await this._browser.close();
+      return;
+    }
     await this._browser._session.send('Target.disposeBrowserContext', { browserContextId: this._browserContextId });
     this._browser._contexts.delete(this._browserContextId);
     this._didCloseInternal();

diff --git a/src/firefox/ffBrowser.ts b/src/firefox/ffBrowser.ts
@@ -293,7 +293,12 @@ export class FFBrowserContext extends BrowserContextBase {
   async close() {
     if (this._closed)
       return;
-    assert(this._browserContextId, 'Non-incognito profiles cannot be closed!');
+    if (!this._browserContextId) {
+      // Default context is only created in 'persistent' mode and closing it should close
+      // the browser.
+      await this._browser.close();
+      return;
+    }
     await this._browser._connection.send('Browser.removeBrowserContext', { browserContextId: this._browserContextId });
     this._browser._contexts.delete(this._browserContextId);
     this._didCloseInternal();

diff --git a/src/server/chromium.ts b/src/server/chromium.ts
@@ -56,8 +56,6 @@ export class Chromium implements BrowserType {
       throw new Error('userDataDir option is not supported in `browserType.launch`. Use `browserType.launchPersistent` instead');
     const { browserServer, transport } = await this._launchServer(options, 'local');
     const browser = await CRBrowser.connect(transport!, false, options && options.slowMo);
-    // Hack: for typical launch scenario, ensure that close waits for actual process termination.
-    browser.close = () => browserServer.close();
     (browser as any)['__server__'] = browserServer;
     return browser;
   }
@@ -68,7 +66,7 @@ export class Chromium implements BrowserType {
 
   async launchPersistent(userDataDir: string, options?: LaunchOptions): Promise<BrowserContext> {
     const { timeout = 30000 } = options || {};
-    const { browserServer, transport } = await this._launchServer(options, 'persistent', userDataDir);
+    const { transport } = await this._launchServer(options, 'persistent', userDataDir);
     const browser = await CRBrowser.connect(transport!, true);
     const browserContext = browser._defaultContext;
 
@@ -78,9 +76,6 @@ export class Chromium implements BrowserType {
     const firstTarget = targets().length ? Promise.resolve() : new Promise(f => browserContext.once('page', f));
     const firstPage = firstTarget.then(() => targets()[0].pageOrError());
     await helper.waitWithTimeout(firstPage, 'first page', timeout);
-
-    // Hack: for typical launch scenario, ensure that close waits for actual process termination.
-    browserContext.close = () => browserServer.close();
     return browserContext;
   }
 
@@ -145,18 +140,20 @@ export class Chromium implements BrowserType {
       },
     });
 
-    let transport: ConnectionTransport | undefined;
-    let browserWSEndpoint: string | null;
+    let transport: PipeTransport | undefined;
+    let browserWSEndpoint: string | undefined;
     if (launchType === 'server') {
       const timeoutError = new TimeoutError(`Timed out after ${timeout} ms while trying to connect to Chromium! The only Chromium revision guaranteed to work is r${this._revision}`);
       const match = await waitForLine(launchedProcess, launchedProcess.stderr, /^DevTools listening on (ws:\/\/.*)$/, timeout, timeoutError);
       browserWSEndpoint = match[1];
+      browserServer = new BrowserServer(launchedProcess, gracefullyClose, browserWSEndpoint);
+      return { browserServer };
     } else {
-      transport = new PipeTransport(launchedProcess.stdio[3] as NodeJS.WritableStream, launchedProcess.stdio[4] as NodeJS.ReadableStream);
-      browserWSEndpoint = null;
+      // For local launch scenario close will terminate the browser process.
+      transport = new PipeTransport(launchedProcess.stdio[3] as NodeJS.WritableStream, launchedProcess.stdio[4] as NodeJS.ReadableStream, () => browserServer!.close());
+      browserServer = new BrowserServer(launchedProcess, gracefullyClose, null);
+      return { browserServer, transport };
     }
-    browserServer = new BrowserServer(launchedProcess, gracefullyClose, browserWSEndpoint);
-    return { browserServer, transport };
   }
 
   async connect(options: ConnectOptions): Promise<CRBrowser> {

diff --git a/src/server/pipeTransport.ts b/src/server/pipeTransport.ts
@@ -24,14 +24,18 @@ export class PipeTransport implements ConnectionTransport {
   private _pendingMessage = '';
   private _eventListeners: RegisteredListener[];
   private _waitForNextTask = makeWaitForNextTask();
+  private readonly _closeCallback: () => void;
+
   onmessage?: (message: string) => void;
   onclose?: () => void;
 
-  constructor(pipeWrite: NodeJS.WritableStream, pipeRead: NodeJS.ReadableStream) {
+  constructor(pipeWrite: NodeJS.WritableStream, pipeRead: NodeJS.ReadableStream, closeCallback: () => void) {
     this._pipeWrite = pipeWrite;
+    this._closeCallback = closeCallback;
     this._eventListeners = [
       helper.addEventListener(pipeRead, 'data', buffer => this._dispatch(buffer)),
       helper.addEventListener(pipeRead, 'close', () => {
+        helper.removeEventListeners(this._eventListeners);
         if (this.onclose)
           this.onclose.call(null);
       }),
@@ -47,6 +51,10 @@ export class PipeTransport implements ConnectionTransport {
     this._pipeWrite!.write('\0');
   }
 
+  close() {
+    this._closeCallback();
+  }
+
   _dispatch(buffer: Buffer) {
     let end = buffer.indexOf('\0');
     if (end === -1) {
@@ -72,11 +80,4 @@ export class PipeTransport implements ConnectionTransport {
     }
     this._pendingMessage = buffer.toString(undefined, start);
   }
-
-  close() {
-    this._pipeWrite = null;
-    helper.removeEventListeners(this._eventListeners);
-    if (this.onclose)
-      this.onclose.call(null);
-  }
 }
diff --git a/src/server/webkit.ts b/src/server/webkit.ts
@@ -68,8 +68,6 @@ export class WebKit implements BrowserType {
       throw new Error('userDataDir option is not supported in `browserType.launch`. Use `browserType.launchPersistent` instead');
     const { browserServer, transport } = await this._launchServer(options, 'local');
     const browser = await WKBrowser.connect(transport!, options && options.slowMo);
-    // Hack: for typical launch scenario, ensure that close waits for actual process termination.
-    browser.close = () => browserServer.close();
     (browser as any)['__server__'] = browserServer;
     return browser;
   }
@@ -80,13 +78,10 @@ export class WebKit implements BrowserType {
 
   async launchPersistent(userDataDir: string, options?: LaunchOptions): Promise<BrowserContext> {
     const { timeout = 30000 } = options || {};
-    const { browserServer, transport } = await this._launchServer(options, 'persistent', userDataDir);
+    const { transport } = await this._launchServer(options, 'persistent', userDataDir);
     const browser = await WKBrowser.connect(transport!, undefined, true);
     await helper.waitWithTimeout(browser._waitForFirstPageTarget(), 'first page', timeout);
-    // Hack: for typical launch scenario, ensure that close waits for actual process termination.
-    const browserContext = browser._defaultContext;
-    browserContext.close = () => browserServer.close();
-    return browserContext;
+    return browser._defaultContext;
   }
 
   private async _launchServer(options: LaunchOptions = {}, launchType: LaunchType, userDataDir?: string, port?: number): Promise<{ browserServer: BrowserServer, transport?: ConnectionTransport }> {
@@ -150,8 +145,11 @@ export class WebKit implements BrowserType {
       },
     });
 
-    transport = new PipeTransport(launchedProcess.stdio[3] as NodeJS.WritableStream, launchedProcess.stdio[4] as NodeJS.ReadableStream);
+    // For local launch scenario close will terminate the browser process.
+    transport = new PipeTransport(launchedProcess.stdio[3] as NodeJS.WritableStream, launchedProcess.stdio[4] as NodeJS.ReadableStream, () => browserServer!.close());
     browserServer = new BrowserServer(launchedProcess, gracefullyClose, launchType === 'server' ? await wrapTransportWithWebSocket(transport, port || 0) : null);
+    if (launchType === 'server')
+      return { browserServer };
     return { browserServer, transport };
   }
 
@@ -378,7 +376,7 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number
         pendingBrowserContextDeletions.set(seqNum, params.browserContextId);
     });
 
-    socket.on('close', () => {
+    socket.on('close', (socket as any).__closeListener = () => {
       for (const [pageProxyId, s] of pageProxyIds) {
         if (s === socket)
           pageProxyIds.delete(pageProxyId);
@@ -398,8 +396,10 @@ function wrapTransportWithWebSocket(transport: ConnectionTransport, port: number
   });
 
   transport.onclose = () => {
-    for (const socket of sockets)
+    for (const socket of sockets) {
+      socket.removeListener('close', (socket as any).__closeListener);
       socket.close(undefined, 'Browser disconnected');
+    }
     server.close();
     transport.onmessage = undefined;
     transport.onclose = undefined;

diff --git a/src/webkit/wkBrowser.ts b/src/webkit/wkBrowser.ts
@@ -66,11 +66,11 @@ export class WKBrowser extends platform.EventEmitter implements Browser {
   }
 
   _onDisconnect() {
-    for (const context of this._contexts.values())
-      context._browserClosed();
     for (const wkPage of this._wkPages.values())
       wkPage.dispose();
     this._wkPages.clear();
+    for (const context of this._contexts.values())
+      context._browserClosed();
     this.emit(Events.Browser.Disconnected);
   }
 
@@ -318,7 +318,12 @@ export class WKBrowserContext extends BrowserContextBase {
   async close() {
     if (this._closed)
       return;
-    assert(this._browserContextId, 'Non-incognito profiles cannot be closed!');
+    if (!this._browserContextId) {
+      // Default context is only created in 'persistent' mode and closing it should close
+      // the browser.
+      await this._browser.close();
+      return;
+    }
     await this._browser._browserSession.send('Browser.deleteContext', { browserContextId: this._browserContextId });
     this._browser._contexts.delete(this._browserContextId);
     this._didCloseInternal();

diff --git a/src/webkit/wkConnection.ts b/src/webkit/wkConnection.ts
@@ -30,10 +30,10 @@ export const kPageProxyMessageReceived = 'kPageProxyMessageReceived';
 export type PageProxyMessageReceivedPayload = { pageProxyId: string, message: any };
 
 export class WKConnection {
-  private _lastId = 0;
   private readonly _transport: ConnectionTransport;
+  private readonly _onDisconnect: () => void;
+  private _lastId = 0;
   private _closed = false;
-  private _onDisconnect: () => void;
   _debugFunction: (message: string) => void = platform.debug('pw:protocol');
 
   readonly browserSession: WKSession;

diff --git a/test/evaluation.spec.js b/test/evaluation.spec.js
@@ -230,7 +230,8 @@ module.exports.describe = function({testRunner, expect, FFOX, CHROMIUM, WEBKIT})
       await page.evaluate(e => e.textContent, element).catch(e => error = e);
       expect(error.message).toContain('JSHandle is disposed');
     });
-    it('should simulate a user gesture', async({page, server}) => {
+    it.fail(FFOX)('should simulate a user gesture', async({page, server}) => {
+      // Flaky on Limux: https://github.com/microsoft/playwright/issues/1305
       const result = await page.evaluate(() => {
         document.body.appendChild(document.createTextNode('test'));
         document.execCommand('selectAll');

diff --git a/test/launcher.spec.js b/test/launcher.spec.js
@@ -226,7 +226,7 @@ module.exports.describe = function({testRunner, expect, defaultBrowserOptions, p
         expect(message).not.toContain('Timeout');
       }
     });
-    it.fail(WEBKIT)('should fire close event for all contexts', async() => {
+    it('should fire close event for all contexts', async(state, test) => {
       const browser = await playwright.launch(defaultBrowserOptions);
       const context = await browser.newContext();
       let closed = false;