fix: wait for the process to close when closing the browser (#1629)

JoelEinbinder · web-flow · commit 3d6d9db44a03 · 2020-04-02T16:57:12.000-07:00
diff --git a/src/browser.ts b/src/browser.ts
@@ -24,6 +24,7 @@ export interface Browser extends EventEmitter {
   newPage(options?: BrowserContextOptions): Promise<Page>;
   isConnected(): boolean;
   close(): Promise<void>;
+  _disconnect(): Promise<void>;
   _setDebugFunction(debugFunction: (message: string) => void): void;
 }
 
diff --git a/src/chromium/crBrowser.ts b/src/chromium/crBrowser.ts
@@ -30,6 +30,7 @@ import { Events } from './events';
 import { Protocol } from './protocol';
 import { CRExecutionContext } from './crExecutionContext';
 import { EventEmitter } from 'events';
+import type { BrowserServer } from '../server/browserServer';
 
 export class CRBrowser extends EventEmitter implements Browser {
   readonly _connection: CRConnection;
@@ -46,6 +47,7 @@ export class CRBrowser extends EventEmitter implements Browser {
   private _tracingRecording = false;
   private _tracingPath: string | null = '';
   private _tracingClient: CRSession | undefined;
+  _ownedServer: BrowserServer | null = null;
 
   static async connect(transport: ConnectionTransport, isPersistent: boolean, slowMo?: number): Promise<CRBrowser> {
     const connection = new CRConnection(SlowMoTransport.wrap(transport, slowMo));
@@ -183,13 +185,20 @@ export class CRBrowser extends EventEmitter implements Browser {
     await this._session.send('Target.closeTarget', { targetId: crPage._targetId });
   }
 
-  async close() {
+  async _disconnect() {
     const disconnected = new Promise(f => this._connection.once(ConnectionEvents.Disconnected, f));
     await Promise.all(this.contexts().map(context => context.close()));
     this._connection.close();
     await disconnected;
   }
 
+  async close() {
+    if (this._ownedServer)
+      await this._ownedServer.close();
+    else
+      await this._disconnect();
+  }
+
   async newBrowserCDPSession(): Promise<CRSession> {
     return await this._connection.createBrowserSession();
   }
diff --git a/src/firefox/ffBrowser.ts b/src/firefox/ffBrowser.ts
@@ -28,6 +28,7 @@ import { headersArray } from './ffNetworkManager';
 import { FFPage } from './ffPage';
 import { Protocol } from './protocol';
 import { EventEmitter } from 'events';
+import type { BrowserServer } from '../server/browserServer';
 
 export class FFBrowser extends EventEmitter implements Browser {
   _connection: FFConnection;
@@ -37,6 +38,7 @@ export class FFBrowser extends EventEmitter implements Browser {
   private _eventListeners: RegisteredListener[];
   readonly _firstPagePromise: Promise<void>;
   private _firstPageCallback = () => {};
+  _ownedServer: BrowserServer | null = null;
 
   static async connect(transport: ConnectionTransport, attachToDefaultContext: boolean, slowMo?: number): Promise<FFBrowser> {
     const connection = new FFConnection(SlowMoTransport.wrap(transport, slowMo));
@@ -140,14 +142,21 @@ export class FFBrowser extends EventEmitter implements Browser {
     });
   }
 
-  async close() {
+  async _disconnect() {
     await Promise.all(this.contexts().map(context => context.close()));
     helper.removeEventListeners(this._eventListeners);
     const disconnected = new Promise(f => this.once(Events.Browser.Disconnected, f));
     this._connection.close();
     await disconnected;
   }
 
+  async close() {
+    if (this._ownedServer)
+      await this._ownedServer.close();
+    else
+      await this._disconnect();
+  }
+
   _setDebugFunction(debugFunction: debug.IDebugger) {
     this._connection._debugProtocol = debugFunction;
   }
diff --git a/src/server/chromium.ts b/src/server/chromium.ts
@@ -22,7 +22,7 @@ import * as util from 'util';
 import { debugError, helper, assert } from '../helper';
 import { CRBrowser } from '../chromium/crBrowser';
 import * as ws from 'ws';
-import { launchProcess } from '../server/processLauncher';
+import { launchProcess } from './processLauncher';
 import { kBrowserCloseMessageId } from '../chromium/crConnection';
 import { PipeTransport } from './pipeTransport';
 import { LaunchOptions, BrowserArgOptions, BrowserType, ConnectOptions, LaunchServerOptions } from './browserType';
@@ -49,7 +49,7 @@ export class Chromium implements BrowserType<CRBrowser> {
     assert(!(options as any).userDataDir, 'userDataDir option is not supported in `browserType.launch`. Use `browserType.launchPersistentContext` instead');
     const { browserServer, transport } = await this._launchServer(options, 'local');
     const browser = await CRBrowser.connect(transport!, false, options.slowMo);
-    (browser as any)['__server__'] = browserServer;
+    browser._ownedServer = browserServer;
     return browser;
   }
 
@@ -62,8 +62,9 @@ export class Chromium implements BrowserType<CRBrowser> {
       timeout = 30000,
       slowMo = 0
     } = options;
-    const { transport } = await this._launchServer(options, 'persistent', userDataDir);
+    const { transport, browserServer } = await this._launchServer(options, 'persistent', userDataDir);
     const browser = await CRBrowser.connect(transport!, true, slowMo);
+    browser._ownedServer = browserServer;
     await helper.waitWithTimeout(browser._firstPagePromise, 'first page', timeout);
     return browser._defaultContext;
   }
@@ -110,8 +111,7 @@ export class Chromium implements BrowserType<CRBrowser> {
       pipe: true,
       tempDir: temporaryUserDataDir || undefined,
       attemptToGracefullyClose: async () => {
-        if (!browserServer)
-          return Promise.reject();
+        assert(browserServer);
         // 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.
@@ -127,7 +127,7 @@ export class Chromium implements BrowserType<CRBrowser> {
 
     let transport: PipeTransport | undefined = undefined;
     let browserServer: BrowserServer | undefined = undefined;
-    transport = new PipeTransport(launchedProcess.stdio[3] as NodeJS.WritableStream, launchedProcess.stdio[4] as NodeJS.ReadableStream, () => browserServer!.close());
+    transport = new PipeTransport(launchedProcess.stdio[3] as NodeJS.WritableStream, launchedProcess.stdio[4] as NodeJS.ReadableStream);
     browserServer = new BrowserServer(launchedProcess, gracefullyClose, launchType === 'server' ? wrapTransportWithWebSocket(transport, port) : null);
     return { browserServer, transport };
   }
diff --git a/src/server/firefox.ts b/src/server/firefox.ts
@@ -53,9 +53,7 @@ export class Firefox implements BrowserType<FFBrowser> {
     const browser = await WebSocketTransport.connect(browserServer.wsEndpoint()!, transport => {
       return FFBrowser.connect(transport, false, options.slowMo);
     });
-    // Hack: for typical launch scenario, ensure that close waits for actual process termination.
-    browser.close = () => browserServer.close();
-    (browser as any)['__server__'] = browserServer;
+    browser._ownedServer = browserServer;
     return browser;
   }
 
@@ -72,10 +70,9 @@ export class Firefox implements BrowserType<FFBrowser> {
     const browser = await WebSocketTransport.connect(browserServer.wsEndpoint()!, transport => {
       return FFBrowser.connect(transport, true, slowMo);
     });
+    browser._ownedServer = browserServer;
     await helper.waitWithTimeout(browser._firstPagePromise, '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;
   }
 
@@ -128,11 +125,8 @@ export class Firefox implements BrowserType<FFBrowser> {
       pipe: false,
       tempDir: temporaryProfileDir || undefined,
       attemptToGracefullyClose: async () => {
-        if (!browserServer)
-          return Promise.reject();
+        assert(browserServer);
         // 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 WebSocketTransport.connect(browserWSEndpoint!, async transport => transport);
         const message = { method: 'Browser.close', params: {}, id: kBrowserCloseMessageId };
         await transport.send(message);
diff --git a/src/server/pipeTransport.ts b/src/server/pipeTransport.ts
@@ -23,14 +23,12 @@ export class PipeTransport implements ConnectionTransport {
   private _pendingMessage = '';
   private _eventListeners: RegisteredListener[];
   private _waitForNextTask = helper.makeWaitForNextTask();
-  private readonly _closeCallback: () => void;
 
   onmessage?: (message: ProtocolResponse) => void;
   onclose?: () => void;
 
-  constructor(pipeWrite: NodeJS.WritableStream, pipeRead: NodeJS.ReadableStream, closeCallback: () => void) {
+  constructor(pipeWrite: NodeJS.WritableStream, pipeRead: NodeJS.ReadableStream) {
     this._pipeWrite = pipeWrite;
-    this._closeCallback = closeCallback;
     this._eventListeners = [
       helper.addEventListener(pipeRead, 'data', buffer => this._dispatch(buffer)),
       helper.addEventListener(pipeRead, 'close', () => {
@@ -51,7 +49,7 @@ export class PipeTransport implements ConnectionTransport {
   }
 
   close() {
-    this._closeCallback();
+    throw new Error('unimplemented');
   }
 
   _dispatch(buffer: Buffer) {
diff --git a/src/server/processLauncher.ts b/src/server/processLauncher.ts
@@ -135,7 +135,7 @@ export async function launchProcess(options: LaunchProcessOptions): Promise<Laun
     }
     gracefullyClosing = true;
     debugBrowser(`<gracefully close start>`);
-    options.attemptToGracefullyClose().catch(() => killProcess());
+    await options.attemptToGracefullyClose().catch(() => killProcess());
     await waitForProcessToClose;
     debugBrowser(`<gracefully close end>`);
   }
diff --git a/src/server/webkit.ts b/src/server/webkit.ts
@@ -48,8 +48,8 @@ export class WebKit implements BrowserType<WKBrowser> {
   async launch(options: LaunchOptions = {}): Promise<WKBrowser> {
     assert(!(options as any).userDataDir, 'userDataDir option is not supported in `browserType.launch`. Use `browserType.launchPersistentContext` instead');
     const { browserServer, transport } = await this._launchServer(options, 'local');
-    const browser = await WKBrowser.connect(transport!, options.slowMo);
-    (browser as any)['__server__'] = browserServer;
+    const browser = await WKBrowser.connect(transport!, options.slowMo, false, () => browserServer.close());
+    browser._ownedServer = browserServer;
     return browser;
   }
 
@@ -62,8 +62,8 @@ export class WebKit implements BrowserType<WKBrowser> {
       timeout = 30000,
       slowMo = 0,
     } = options;
-    const { transport } = await this._launchServer(options, 'persistent', userDataDir);
-    const browser = await WKBrowser.connect(transport!, slowMo, true);
+    const { transport, browserServer } = await this._launchServer(options, 'persistent', userDataDir);
+    const browser = await WKBrowser.connect(transport!, slowMo, true, () => browserServer.close());
     await helper.waitWithTimeout(browser._waitForFirstPageTarget(), 'first page', timeout);
     return browser._defaultContext;
   }
@@ -111,12 +111,11 @@ export class WebKit implements BrowserType<WKBrowser> {
       pipe: true,
       tempDir: temporaryUserDataDir || undefined,
       attemptToGracefullyClose: async () => {
-        if (!transport)
-          return Promise.reject();
+        assert(transport);
         // 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.
-        transport.send({method: 'Playwright.close', params: {}, id: kBrowserCloseMessageId});
+        await transport.send({method: 'Playwright.close', params: {}, id: kBrowserCloseMessageId});
       },
       onkill: (exitCode, signal) => {
         if (browserServer)
@@ -127,7 +126,7 @@ export class WebKit implements BrowserType<WKBrowser> {
     // For local launch scenario close will terminate the browser process.
     let transport: ConnectionTransport | undefined = undefined;
     let browserServer: BrowserServer | undefined = undefined;
-    transport = new PipeTransport(launchedProcess.stdio[3] as NodeJS.WritableStream, launchedProcess.stdio[4] as NodeJS.ReadableStream, () => browserServer!.close());
+    transport = new PipeTransport(launchedProcess.stdio[3] as NodeJS.WritableStream, launchedProcess.stdio[4] as NodeJS.ReadableStream);
     browserServer = new BrowserServer(launchedProcess, gracefullyClose, launchType === 'server' ? wrapTransportWithWebSocket(transport, port || 0) : null);
     return { browserServer, transport };
   }
diff --git a/src/webkit/wkBrowser.ts b/src/webkit/wkBrowser.ts
@@ -27,6 +27,7 @@ import { Protocol } from './protocol';
 import { kPageProxyMessageReceived, PageProxyMessageReceivedPayload, WKConnection, WKSession } from './wkConnection';
 import { WKPage } from './wkPage';
 import { EventEmitter } from 'events';
+import type { BrowserServer } from '../server/browserServer';
 
 const DEFAULT_USER_AGENT = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_2) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/13.0.4 Safari/605.1.15';
 
@@ -39,19 +40,22 @@ export class WKBrowser extends EventEmitter implements Browser {
   readonly _wkPages = new Map<string, WKPage>();
   private readonly _eventListeners: RegisteredListener[];
   private _popupOpeners: string[] = [];
+  private _closeOverride?: () => Promise<void>;
 
   private _firstPageCallback: () => void = () => {};
   private readonly _firstPagePromise: Promise<void>;
+  _ownedServer: BrowserServer | null = null;
 
-  static async connect(transport: ConnectionTransport, slowMo: number = 0, attachToDefaultContext: boolean = false): Promise<WKBrowser> {
-    const browser = new WKBrowser(SlowMoTransport.wrap(transport, slowMo), attachToDefaultContext);
+  static async connect(transport: ConnectionTransport, slowMo: number = 0, attachToDefaultContext: boolean = false, closeOverride?: () => Promise<void>): Promise<WKBrowser> {
+    const browser = new WKBrowser(SlowMoTransport.wrap(transport, slowMo), attachToDefaultContext, closeOverride);
     return browser;
   }
 
-  constructor(transport: ConnectionTransport, attachToDefaultContext: boolean) {
+  constructor(transport: ConnectionTransport, attachToDefaultContext: boolean, closeOverride?: () => Promise<void>) {
     super();
     this._connection = new WKConnection(transport, this._onDisconnect.bind(this));
     this._attachToDefaultContext = attachToDefaultContext;
+    this._closeOverride = closeOverride;
     this._browserSession = this._connection.browserSession;
 
     this._defaultContext = new WKBrowserContext(this, undefined, validateBrowserContextOptions({}));
@@ -178,14 +182,21 @@ export class WKBrowser extends EventEmitter implements Browser {
     return !this._connection.isClosed();
   }
 
-  async close() {
+  async _disconnect() {
     helper.removeEventListeners(this._eventListeners);
     const disconnected = new Promise(f => this.once(Events.Browser.Disconnected, f));
     await Promise.all(this.contexts().map(context => context.close()));
     this._connection.close();
     await disconnected;
   }
 
+  async close() {
+    if (this._closeOverride)
+      await this._closeOverride();
+    else
+      await this._disconnect();
+  }
+
   _setDebugFunction(debugFunction: debug.IDebugger) {
     this._connection._debugProtocol = debugFunction;
   }
diff --git a/test/playwright.spec.js b/test/playwright.spec.js
@@ -160,7 +160,7 @@ module.exports.addPlaywrightTests = ({testRunner, platform, products, playwright
       describe('', function() {
         beforeAll(async state => {
           state.browser = await browserType.launch(defaultBrowserOptions);
-          state.browserServer = state.browser.__server__;
+          state.browserServer = state.browser._ownedServer;
           state._stdout = readline.createInterface({ input: state.browserServer.process().stdout });
           state._stderr = readline.createInterface({ input: state.browserServer.process().stderr });
         });
diff --git a/utils/doclint/check_public_api/JSBuilder.js b/utils/doclint/check_public_api/JSBuilder.js
@@ -112,7 +112,7 @@ function checkSources(sources) {
       if (ts.isImportDeclaration(node) && ts.isStringLiteral(node.moduleSpecifier)) {
         const module = node.moduleSpecifier.text;
         const isServerDependency = path.resolve(path.dirname(fileName), module).includes('src/server');
-        if (isServerDependency) {
+        if (isServerDependency && !node.importClause.isTypeOnly) {
           const lac = ts.getLineAndCharacterOfPosition(node.getSourceFile(), node.moduleSpecifier.pos);
           errors.push(`Disallowed import "${module}" at ${node.getSourceFile().fileName}:${lac.line + 1}`);
         }

Original file line number	Diff line number	Diff line change
`@@ -24,6 +24,7 @@ export interface Browser extends EventEmitter {`
`24`	`24`	`newPage(options?: BrowserContextOptions): Promise<Page>;`
`25`	`25`	`isConnected(): boolean;`
`26`	`26`	`close(): Promise<void>;`
	`27`	`+ _disconnect(): Promise<void>;`
`27`	`28`	`_setDebugFunction(debugFunction: (message: string) => void): void;`
`28`	`29`	`}`
`29`	`30`
Original file line number	Diff line number	Diff line change
`@@ -135,7 +135,7 @@ export async function launchProcess(options: LaunchProcessOptions): Promise<Laun`
`135`	`135`	`}`
`136`	`136`	`gracefullyClosing = true;`
`137`	`137`	debugBrowser(`<gracefully close start>`);
`138`		`- options.attemptToGracefullyClose().catch(() => killProcess());`
	`138`	`+ await options.attemptToGracefullyClose().catch(() => killProcess());`
`139`	`139`	`await waitForProcessToClose;`
`140`	`140`	debugBrowser(`<gracefully close end>`);
`141`	`141`	`}`