chore: further unify launching and connection (#2320)

Author: pavelfeldman

docs/api.md

@@ -3976,6 +3976,7 @@ const { chromium } = require('playwright');  // Or 'firefox' or 'webkit'.
   - `wsEndpoint` <[string]> A browser websocket endpoint to connect to. **required**
   - `slowMo` <[number]> Slows down Playwright operations by the specified amount of milliseconds. Useful so that you can see what is going on. Defaults to 0.
   - `logger` <[Logger]> Logger sink for Playwright logging.
+  - `timeout` <[number]> Maximum time in milliseconds to wait for the connection to be established. Defaults to `30000` (30 seconds). Pass `0` to disable timeout.
 - returns: <[Promise]<[Browser]>>
 
 This methods attaches Playwright to an existing browser instance.

src/helper.ts

@@ -21,6 +21,7 @@ import * as fs from 'fs';
 import * as util from 'util';
 import { TimeoutError } from './errors';
 import * as types from './types';
+import { ChildProcess, execSync } from 'child_process';
 
 // NOTE: update this to point to playwright/lib when moving this file.
 const PLAYWRIGHT_LIB_PATH = __dirname;
@@ -332,6 +333,19 @@ class Helper {
   static optionsWithUpdatedTimeout<T extends types.TimeoutOptions>(options: T | undefined, deadline: number): T {
     return { ...(options || {}) as T, timeout: this.timeUntilDeadline(deadline) };
   }
+
+  static killProcess(proc: ChildProcess) {
+    if (proc.pid && !proc.killed) {
+      try {
+        if (process.platform === 'win32')
+          execSync(`taskkill /pid ${proc.pid} /T /F`);
+        else
+          process.kill(-proc.pid, 'SIGKILL');
+      } catch (e) {
+        // the process might have already stopped
+      }
+    }
+  }
 }
 
 export function assert(value: any, message?: string): asserts value {

src/logger.ts

@@ -62,12 +62,16 @@ export class RootLogger implements InnerLogger {
     this._logger.add(`launch`, new RecordingLogger('browser'));
   }
 
-  stopLaunchRecording(): string {
-    const logger = this._logger.remove(`launch`) as RecordingLogger;
+  launchRecording(): string {
+    const logger = this._logger.get(`launch`) as RecordingLogger;
     if (logger)
       return logger.recording();
     return '';
   }
+
+  stopLaunchRecording() {
+    this._logger.remove(`launch`);
+  }
 }
 
 const colorMap = new Map<string, number>([
@@ -87,10 +91,12 @@ class MultiplexingLogger implements Logger {
     this._loggers.set(id, logger);
   }
 
-  remove(id: string): Logger | undefined {
-    const logger = this._loggers.get(id);
+  get(id: string): Logger | undefined {
+    return this._loggers.get(id);
+  }
+
+  remove(id: string) {
     this._loggers.delete(id);
-    return logger;
   }
 
   isEnabled(name: string, severity: LoggerSeverity): boolean {

src/server/browserServer.ts

@@ -14,12 +14,12 @@
  * limitations under the License.
  */
 
-import { ChildProcess, execSync } from 'child_process';
+import { ChildProcess } from 'child_process';
 import { EventEmitter } from 'events';
 import { helper } from '../helper';
 import { RootLogger } from '../logger';
-import { TimeoutSettings } from '../timeoutSettings';
-import { LaunchOptions } from './browserType';
+import { LaunchOptions, processBrowserArgOptions } from './browserType';
+import { ConnectionTransport } from '../transport';
 
 export class WebSocketWrapper {
   readonly wsEndpoint: string;
@@ -51,82 +51,52 @@ export class WebSocketWrapper {
 }
 
 export class BrowserServer extends EventEmitter {
-  private _process: ChildProcess | undefined;
-  private _gracefullyClose: (() => Promise<void>) | undefined;
+  private _process: ChildProcess;
+  private _gracefullyClose: (() => Promise<void>);
   private _webSocketWrapper: WebSocketWrapper | null = null;
   readonly _launchOptions: LaunchOptions;
   readonly _logger: RootLogger;
-  readonly _launchDeadline: number;
+  readonly _downloadsPath: string;
+  readonly _transport: ConnectionTransport;
+  readonly _headful: boolean;
 
-  constructor(options: LaunchOptions) {
+  constructor(options: LaunchOptions, process: ChildProcess, gracefullyClose: () => Promise<void>, transport: ConnectionTransport, downloadsPath: string, webSocketWrapper: WebSocketWrapper | null) {
     super();
     this._launchOptions = options;
+    this._headful = !processBrowserArgOptions(options).headless;
     this._logger = new RootLogger(options.logger);
-    this._launchDeadline = TimeoutSettings.computeDeadline(typeof options.timeout === 'number' ? options.timeout : 30000);
-  }
 
-  _initialize(process: ChildProcess, gracefullyClose: () => Promise<void>, webSocketWrapper: WebSocketWrapper | null) {
     this._process = process;
     this._gracefullyClose = gracefullyClose;
+    this._transport = transport;
+    this._downloadsPath = downloadsPath;
     this._webSocketWrapper = webSocketWrapper;
   }
 
-  _isInitialized(): boolean {
-    return !!this._process;
-  }
-
   process(): ChildProcess {
-    return this._process!;
+    return this._process;
   }
 
   wsEndpoint(): string {
     return this._webSocketWrapper ? this._webSocketWrapper.wsEndpoint : '';
   }
 
   kill() {
-    if (this._process!.pid && !this._process!.killed) {
-      try {
-        if (process.platform === 'win32')
-          execSync(`taskkill /pid ${this._process!.pid} /T /F`);
-        else
-          process.kill(-this._process!.pid, 'SIGKILL');
-      } catch (e) {
-        // the process might have already stopped
-      }
-    }
+    helper.killProcess(this._process);
   }
 
   async close(): Promise<void> {
-    await this._gracefullyClose!();
+    await this._gracefullyClose();
   }
 
   async _checkLeaks(): Promise<void> {
     if (this._webSocketWrapper)
       await this._webSocketWrapper.checkLeaks();
   }
 
-  async _initializeOrClose<T>(init: () => Promise<T>): Promise<T> {
-    try {
-      let promise: Promise<T>;
-      if ((this._launchOptions as any).__testHookBeforeCreateBrowser)
-        promise = (this._launchOptions as any).__testHookBeforeCreateBrowser().then(init);
-      else
-        promise = init();
-      const result = await helper.waitWithDeadline(promise, 'the browser to launch', this._launchDeadline, 'pw:browser*');
-      this._logger.stopLaunchRecording();
-      return result;
-    } catch (e) {
-      e.message += '\n=============== Process output during launch: ===============\n' +
-          this._logger.stopLaunchRecording() +
-          '\n=============================================================';
-      await this._closeOrKill();
-      throw e;
-    }
-  }
-
-  private async _closeOrKill(): Promise<void> {
+  async _closeOrKill(deadline: number): Promise<void> {
     try {
-      await helper.waitWithDeadline(this.close(), '', this._launchDeadline, ''); // The error message is ignored.
+      await helper.waitWithDeadline(this.close(), '', deadline, ''); // The error message is ignored.
     } catch (ignored) {
       this.kill();
     }

src/server/browserType.ts

@@ -20,7 +20,8 @@ import * as browserPaths from '../install/browserPaths';
 import { Logger, RootLogger } from '../logger';
 import { ConnectionTransport, WebSocketTransport } from '../transport';
 import { BrowserBase, BrowserOptions, Browser } from '../browser';
-import { assert } from '../helper';
+import { assert, helper } from '../helper';
+import { TimeoutSettings } from '../timeoutSettings';
 
 export type BrowserArgOptions = {
   headless?: boolean,
@@ -48,6 +49,7 @@ export type ConnectOptions = {
   wsEndpoint: string,
   slowMo?: number,
   logger?: Logger,
+  timeout?: number,
 };
 export type LaunchType = 'local' | 'server' | 'persistent';
 export type LaunchOptions = LaunchOptionsBase & { slowMo?: number };
@@ -86,42 +88,88 @@ export abstract class BrowserTypeBase implements BrowserType {
 
   async launch(options: LaunchOptions = {}): Promise<Browser> {
     assert(!(options as any).userDataDir, 'userDataDir option is not supported in `browserType.launch`. Use `browserType.launchPersistentContext` instead');
-    const browserServer = new BrowserServer(options);
-    const { transport, downloadsPath } = await this._launchServer(options, 'local', browserServer);
-    return await browserServer._initializeOrClose(async () => {
-      return this._connectToServer(browserServer, false, transport!, downloadsPath);
-    });
+    return this._innerLaunch('local', options);
   }
 
   async launchPersistentContext(userDataDir: string, options: LaunchOptions = {}): Promise<BrowserContext> {
-    const browserServer = new BrowserServer(options);
-    const { transport, downloadsPath } = await this._launchServer(options, 'persistent', browserServer, userDataDir);
+    const browser = await this._innerLaunch('persistent', options, userDataDir);
+    return browser._defaultContext!;
+  }
+
+  async _innerLaunch(launchType: LaunchType, options: LaunchOptions, userDataDir?: string): Promise<BrowserBase> {
+    const deadline = TimeoutSettings.computeDeadline(options.timeout, 30000);
+    const logger = new RootLogger(options.logger);
+    logger.startLaunchRecording();
+
+    let browserServer: BrowserServer | undefined;
+    try {
+      browserServer = await this._launchServer(options, launchType, logger, deadline, userDataDir);
+      const promise = this._innerLaunchPromise(browserServer, launchType, options);
+      const browser = await helper.waitWithDeadline(promise, 'the browser to launch', deadline, 'pw:browser*');
+      return browser;
+    } catch (e) {
+      e.message += '\n=============== Process output during launch: ===============\n' +
+          logger.launchRecording() +
+          '\n=============================================================';
+      if (browserServer)
+        await browserServer._closeOrKill(deadline);
+      throw e;
+    } finally {
+      logger.stopLaunchRecording();
+    }
+  }
 
-    return await browserServer._initializeOrClose(async () => {
-      const browser = await this._connectToServer(browserServer, true, transport!, downloadsPath);
+  async _innerLaunchPromise(browserServer: BrowserServer, launchType: LaunchType, options: LaunchOptions): Promise<BrowserBase> {
+    if ((options as any).__testHookBeforeCreateBrowser)
+      await (options as any).__testHookBeforeCreateBrowser();
+
+    const browser = await this._connectToServer(browserServer, launchType === 'persistent');
+    if (launchType === 'persistent' && (!options.ignoreDefaultArgs || Array.isArray(options.ignoreDefaultArgs))) {
       const context = browser._defaultContext!;
-      if (!options.ignoreDefaultArgs || Array.isArray(options.ignoreDefaultArgs))
-        await context._loadDefaultContext();
-      return context;
-    });
+      await context._loadDefaultContext();
+    }
+    return browser;
   }
 
   async launchServer(options: LaunchServerOptions = {}): Promise<BrowserServer> {
-    const browserServer = new BrowserServer(options);
-    await this._launchServer(options, 'server', browserServer);
-    return browserServer;
+    const logger = new RootLogger(options.logger);
+    return this._launchServer(options, 'server', logger, TimeoutSettings.computeDeadline(options.timeout, 30000));
   }
 
   async connect(options: ConnectOptions): Promise<Browser> {
+    const deadline = TimeoutSettings.computeDeadline(options.timeout, 30000);
     const logger = new RootLogger(options.logger);
-    return await WebSocketTransport.connect(options.wsEndpoint, async transport => {
-      if ((options as any).__testHookBeforeCreateBrowser)
-        await (options as any).__testHookBeforeCreateBrowser();
-      return this._connectToTransport(transport, { slowMo: options.slowMo, logger, downloadsPath: '' });
-    }, logger);
+    logger.startLaunchRecording();
+
+    let transport: ConnectionTransport | undefined;
+    try {
+      transport = await WebSocketTransport.connect(options.wsEndpoint, logger, deadline);
+      const promise = this._innerConnectPromise(transport, options, logger);
+      const browser = await helper.waitWithDeadline(promise, 'connect to browser', deadline, 'pw:browser*');
+      logger.stopLaunchRecording();
+      return browser;
+    } catch (e) {
+      e.message += '\n=============== Process output during connect: ===============\n' +
+          logger.launchRecording() +
+          '\n=============================================================';
+      try {
+        if (transport)
+          transport.close();
+      } catch (e) {
+      }
+      throw e;
+    } finally {
+      logger.stopLaunchRecording();
+    }
+  }
+
+  async _innerConnectPromise(transport: ConnectionTransport, options: ConnectOptions, logger: RootLogger): Promise<Browser> {
+    if ((options as any).__testHookBeforeCreateBrowser)
+      await (options as any).__testHookBeforeCreateBrowser();
+    return this._connectToTransport(transport, { slowMo: options.slowMo, logger, downloadsPath: '' });
   }
 
-  abstract _launchServer(options: LaunchServerOptions, launchType: LaunchType, browserServer: BrowserServer, userDataDir?: string): Promise<{ transport?: ConnectionTransport, downloadsPath: string }>;
-  abstract _connectToServer(browserServer: BrowserServer, persistent: boolean, transport: ConnectionTransport, downloadsPath: string): Promise<BrowserBase>;
+  abstract _launchServer(options: LaunchServerOptions, launchType: LaunchType, logger: RootLogger, deadline: number, userDataDir?: string): Promise<BrowserServer>;
+  abstract _connectToServer(browserServer: BrowserServer, persistent: boolean): Promise<BrowserBase>;
   abstract _connectToTransport(transport: ConnectionTransport, options: BrowserOptions): Promise<BrowserBase>;
 }