chore(wk, ff): simplify target management (#1279)

src/firefox/ffBrowser.ts

@@ -24,19 +24,19 @@ import { Page, PageBinding, PageEvent } from '../page';
 import * as platform from '../platform';
 import { ConnectionTransport, SlowMoTransport } from '../transport';
 import * as types from '../types';
-import { ConnectionEvents, FFConnection, FFSession, FFSessionEvents } from './ffConnection';
+import { ConnectionEvents, FFConnection } from './ffConnection';
 import { headersArray } from './ffNetworkManager';
 import { FFPage } from './ffPage';
 import { Protocol } from './protocol';
 
-const kAttachedToTarget = Symbol('kAttachedToTarget');
-
 export class FFBrowser extends platform.EventEmitter implements Browser {
   _connection: FFConnection;
-  _targets: Map<string, Target>;
+  readonly _ffPages: Map<string, FFPage>;
   readonly _defaultContext: FFBrowserContext;
   readonly _contexts: Map<string, FFBrowserContext>;
   private _eventListeners: RegisteredListener[];
+  readonly _firstPagePromise: Promise<void>;
+  private _firstPageCallback = () => {};
 
   static async connect(transport: ConnectionTransport, attachToDefaultContext: boolean, slowMo?: number): Promise<FFBrowser> {
     const connection = new FFConnection(SlowMoTransport.wrap(transport, slowMo));
@@ -48,7 +48,7 @@ export class FFBrowser extends platform.EventEmitter implements Browser {
   constructor(connection: FFConnection) {
     super();
     this._connection = connection;
-    this._targets = new Map();
+    this._ffPages = new Map();
 
     this._defaultContext = new FFBrowserContext(this, null, validateBrowserContextOptions({}));
     this._contexts = new Map();
@@ -61,6 +61,7 @@ export class FFBrowser extends platform.EventEmitter implements Browser {
       helper.addEventListener(this._connection, 'Browser.attachedToTarget', this._onAttachedToTarget.bind(this)),
       helper.addEventListener(this._connection, 'Browser.detachedFromTarget', this._onDetachedFromTarget.bind(this)),
     ];
+    this._firstPagePromise = new Promise(f => this._firstPageCallback = f);
   }
 
   isConnected(): boolean {
@@ -112,53 +113,26 @@ export class FFBrowser extends platform.EventEmitter implements Browser {
     return createPageInNewContext(this, options);
   }
 
-  async _waitForTarget(predicate: (target: Target) => boolean, options: { timeout?: number; } = {}): Promise<Target> {
-    const {
-      timeout = 30000
-    } = options;
-    const existingTarget = this._allTargets().find(predicate);
-    if (existingTarget)
-      return existingTarget;
-    let resolve: (t: Target) => void;
-    const targetPromise = new Promise<Target>(x => resolve = x);
-    this.on(kAttachedToTarget, check);
-    try {
-      if (!timeout)
-        return await targetPromise;
-      return await helper.waitWithTimeout(targetPromise, 'target', timeout);
-    } finally {
-      this.removeListener(kAttachedToTarget, check);
-    }
-
-    function check(target: Target) {
-      if (predicate(target))
-        resolve(target);
-    }
-  }
-
-  _allTargets() {
-    return Array.from(this._targets.values());
-  }
-
   _onDetachedFromTarget(payload: Protocol.Browser.detachedFromTargetPayload) {
-    const {targetId} = payload;
+    const ffPage = this._ffPages.get(payload.targetId)!;
-    const target = this._targets.get(targetId)!;
+    this._ffPages.delete(payload.targetId);
-    this._targets.delete(targetId);
+    ffPage.didClose();
-    target._didClose();
   }
 
   async _onAttachedToTarget(payload: Protocol.Browser.attachedToTargetPayload) {
     const {targetId, browserContextId, openerId, type} = payload.targetInfo;
+    assert(type === 'page');
     const context = browserContextId ? this._contexts.get(browserContextId)! : this._defaultContext;
-    const target = new Target(this, context, type, '', openerId);
+    const session = this._connection.createSession(payload.sessionId, type);
-    this._targets.set(targetId, target);
+    const opener = openerId ? this._ffPages.get(openerId)! : null;
-    target._initPagePromise(this._connection.createSession(payload.sessionId, type));
+    const ffPage = new FFPage(session, context, opener);
+    this._ffPages.set(targetId, ffPage);
 
-    const pageEvent = new PageEvent(target.pageOrError());
+    const pageEvent = new PageEvent(ffPage.pageOrError());
-    target.context().emit(Events.BrowserContext.Page, pageEvent);
+    context.emit(Events.BrowserContext.Page, pageEvent);
-    this.emit(kAttachedToTarget, target);
+
+    ffPage.pageOrError().then(() => this._firstPageCallback());
 
-    const opener = target.opener();
     if (!opener)
       return;
     const openerPage = await opener.pageOrError();
@@ -179,83 +153,6 @@ export class FFBrowser extends platform.EventEmitter implements Browser {
   }
 }
 
-class Target {
-  _pagePromise?: Promise<Page | Error>;
-  _ffPage: FFPage | null = null;
-  private readonly _browser: FFBrowser;
-  private readonly _context: FFBrowserContext;
-  private readonly _type: 'page' | 'browser';
-  _url: string;
-  private readonly _openerId: string | undefined;
-  private _session?: FFSession;
-
-  constructor(browser: FFBrowser, context: FFBrowserContext, type: 'page' | 'browser', url: string, openerId: string | undefined) {
-    this._browser = browser;
-    this._context = context;
-    this._type = type;
-    this._url = url;
-    this._openerId = openerId;
-  }
-
-  _didClose() {
-    if (this._session)
-      this._session.dispose();
-    if (this._ffPage)
-      this._ffPage.didClose();
-  }
-
-  opener(): Target | null {
-    return this._openerId ? this._browser._targets.get(this._openerId)! : null;
-  }
-
-  type(): 'page' | 'browser' {
-    return this._type;
-  }
-
-  url() {
-    return this._url;
-  }
-
-  context(): FFBrowserContext {
-    return this._context;
-  }
-
-  async pageOrError(): Promise<Page | Error> {
-    if (this._type !== 'page')
-      throw new Error(`Cannot create page for "${this._type}" target`);
-    return this._pagePromise!;
-  }
-
-  _initPagePromise(session: FFSession) {
-    this._session = session;
-    this._pagePromise = new Promise(async f => {
-      this._ffPage = new FFPage(session, this._context, async () => {
-        const openerTarget = this.opener();
-        if (!openerTarget)
-          return null;
-        const result = await openerTarget.pageOrError();
-        if (result instanceof Page && !result.isClosed())
-          return result;
-        return null;
-      });
-      const page = this._ffPage._page;
-      session.once(FFSessionEvents.Disconnected, () => page._didDisconnect());
-      let pageOrError: Page | Error;
-      try {
-        await this._ffPage._initialize();
-        pageOrError = page;
-      } catch (e) {
-        pageOrError = e;
-      }
-      f(pageOrError);
-    });
-  }
-
-  browser() {
-    return this._browser;
-  }
-}
-
 export class FFBrowserContext extends BrowserContextBase {
   readonly _browser: FFBrowser;
   readonly _browserContextId: string | null;
@@ -281,13 +178,12 @@ export class FFBrowserContext extends BrowserContextBase {
       await this.setHTTPCredentials(this._options.httpCredentials);
   }
 
-  _existingPages(): Page[] {
+  _ffPages(): FFPage[] {
-    const pages: Page[] = [];
+    return Array.from(this._browser._ffPages.values()).filter(ffPage => ffPage._browserContext === this);
-    for (const target of this._browser._allTargets()) {
-      if (target.context() === this && target._ffPage)
-        pages.push(target._ffPage._page);
   }
-    return pages;
+
+  _existingPages(): Page[] {
+    return this._ffPages().map(ffPage => ffPage._initializedPage()).filter(pageOrNull => !!pageOrNull) as Page[];
   }
 
   setDefaultNavigationTimeout(timeout: number) {
@@ -299,9 +195,8 @@ export class FFBrowserContext extends BrowserContextBase {
   }
 
   async pages(): Promise<Page[]> {
-    const targets = this._browser._allTargets().filter(target => target.context() === this && target.type() === 'page');
+    const pagesOrErrors = await Promise.all(this._ffPages().map(ffPage => ffPage.pageOrError()));
-    const pages = await Promise.all(targets.map(target => target.pageOrError()));
+    return pagesOrErrors.filter(pageOrError => pageOrError instanceof Page && !pageOrError.isClosed()) as Page[];
-    return pages.filter(page => page instanceof Page && !page.isClosed()) as Page[];
   }
 
   async newPage(): Promise<Page> {
@@ -309,14 +204,14 @@ export class FFBrowserContext extends BrowserContextBase {
     const { targetId } = await this._browser._connection.send('Browser.newPage', {
       browserContextId: this._browserContextId || undefined
     });
-    const target = this._browser._targets.get(targetId)!;
+    const ffPage = this._browser._ffPages.get(targetId)!;
-    const result = await target.pageOrError();
+    const pageOrError = await ffPage.pageOrError();
-    if (result instanceof Page) {
+    if (pageOrError instanceof Page) {
-      if (result.isClosed())
+      if (pageOrError.isClosed())
         throw new Error('Page has been closed.');
-      return result;
+      return pageOrError;
     }
-    throw result;
+    throw pageOrError;
   }
 
   async cookies(urls?: string | string[]): Promise<network.NetworkCookie[]> {

src/firefox/ffPage.ts

@@ -26,7 +26,7 @@ import { kScreenshotDuringNavigationError } from '../screenshotter';
 import * as types from '../types';
 import { getAccessibilityTree } from './ffAccessibility';
 import { FFBrowserContext } from './ffBrowser';
-import { FFSession } from './ffConnection';
+import { FFSession, FFSessionEvents } from './ffConnection';
 import { FFExecutionContext } from './ffExecutionContext';
 import { RawKeyboardImpl, RawMouseImpl } from './ffInput';
 import { FFNetworkManager, headersArray } from './ffNetworkManager';
@@ -40,17 +40,22 @@ export class FFPage implements PageDelegate {
   readonly _session: FFSession;
   readonly _page: Page;
   readonly _networkManager: FFNetworkManager;
-  private readonly _openerResolver: () => Promise<Page | null>;
+  readonly _browserContext: FFBrowserContext;
+  private _pagePromise: Promise<Page | Error>;
+  private _pageCallback: (pageOrError: Page | Error) => void = () => {};
+  private _initialized = false;
+  private readonly _opener: FFPage | null;
   private readonly _contextIdToContext: Map<string, dom.FrameExecutionContext>;
   private _eventListeners: RegisteredListener[];
   private _workers = new Map<string, { frameId: string, session: FFSession }>();
 
-  constructor(session: FFSession, browserContext: FFBrowserContext, openerResolver: () => Promise<Page | null>) {
+  constructor(session: FFSession, browserContext: FFBrowserContext, opener: FFPage | null) {
     this._session = session;
-    this._openerResolver = openerResolver;
+    this._opener = opener;
     this.rawKeyboard = new RawKeyboardImpl(session);
     this.rawMouse = new RawMouseImpl(session);
     this._contextIdToContext = new Map();
+    this._browserContext = browserContext;
     this._page = new Page(this, browserContext);
     this._networkManager = new FFNetworkManager(session, this._page);
     this._page.on(Events.Page.FrameDetached, frame => this._removeContextsForFrame(frame));
@@ -75,9 +80,13 @@ export class FFPage implements PageDelegate {
       helper.addEventListener(this._session, 'Page.dispatchMessageFromWorker', this._onDispatchMessageFromWorker.bind(this)),
       helper.addEventListener(this._session, 'Page.crashed', this._onCrashed.bind(this)),
     ];
+    this._pagePromise = new Promise(f => this._pageCallback = f);
+    session.once(FFSessionEvents.Disconnected, () => this._page._didDisconnect());
+    this._initialize();
   }
 
   async _initialize() {
+    try {
       await Promise.all([
         this._session.send('Page.addScriptToEvaluateOnNewDocument', {
           script: '',
@@ -85,6 +94,19 @@ export class FFPage implements PageDelegate {
         }),
         new Promise(f => this._session.once('Page.ready', f)),
       ]);
+      this._pageCallback(this._page);
+    } catch (e) {
+      this._pageCallback(e);
+    }
+    this._initialized = true;
+  }
+
+  _initializedPage(): Page | null {
+    return this._initialized ? this._page : null;
+  }
+
+  async pageOrError(): Promise<Page | Error> {
+    return this._pagePromise;
   }
 
   _onExecutionContextCreated(payload: Protocol.Runtime.executionContextCreatedPayload) {
@@ -249,6 +271,7 @@ export class FFPage implements PageDelegate {
   }
 
   didClose() {
+    this._session.dispose();
     helper.removeEventListeners(this._eventListeners);
     this._networkManager.dispose();
     this._page._didClose();
@@ -289,7 +312,12 @@ export class FFPage implements PageDelegate {
   }
 
   async opener(): Promise<Page | null> {
-    return await this._openerResolver();
+    if (!this._opener)
+      return null;
+    const result = await this._opener.pageOrError();
+    if (result instanceof Page && !result.isClosed())
+      return result;
+    return null;
   }
 
   async reload(): Promise<void> {

src/server/firefox.ts

@@ -83,7 +83,7 @@ export class Firefox implements BrowserType {
     const browser = await platform.connectToWebsocket(browserServer.wsEndpoint()!, transport => {
       return FFBrowser.connect(transport, true);
     });
-    await helper.waitWithTimeout(browser._waitForTarget(t => t.type() === 'page'), 'first page', timeout);
+    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();

src/webkit/wkBrowser.ts

@@ -139,7 +139,7 @@ export class WKBrowser extends platform.EventEmitter implements Browser {
     const wkPage = this._wkPages.get(pageProxyId);
     if (!wkPage)
       return;
-    wkPage.didClose(false);
+    wkPage.didClose();
     wkPage.dispose();
     this._wkPages.delete(pageProxyId);
   }
@@ -202,16 +202,12 @@ export class WKBrowserContext extends BrowserContextBase {
       await this.setHTTPCredentials(this._options.httpCredentials);
   }
 
-  _existingPages(): Page[] {
+  _wkPages(): WKPage[] {
-    const pages: Page[] = [];
+    return Array.from(this._browser._wkPages.values()).filter(wkPage => wkPage._browserContext === this);
-    for (const wkPage of this._browser._wkPages.values()) {
-      if (wkPage._browserContext !== this)
-        continue;
-      const page = wkPage._initializedPage();
-      if (page)
-        pages.push(page);
   }
-    return pages;
+
+  _existingPages(): Page[] {
+    return this._wkPages().map(wkPage => wkPage._initializedPage()).filter(pageOrNull => !!pageOrNull) as Page[];
   }
 
   async pages(): Promise<Page[]> {

src/webkit/wkPage.ts

@@ -78,8 +78,8 @@ export class WKPage implements PageDelegate {
     this._pagePromise = new Promise(f => this._pagePromiseCallback = f);
   }
 
-  _initializedPage(): Page | undefined {
+  _initializedPage(): Page | null {
-    return this._initialized ? this._page : undefined;
+    return this._initialized ? this._page : null;
   }
 
   private async _initializePageProxySession() {
@@ -179,21 +179,19 @@ export class WKPage implements PageDelegate {
       this._provisionalPage = null;
     } else if (this._session.sessionId === targetId) {
       this._session.dispose();
-      if (crashed)
-        this.didClose(crashed);
-    }
-  }
-
-  didClose(crashed: boolean) {
       helper.removeEventListeners(this._sessionListeners);
       if (crashed)
         this._page._didCrash();
-    else
+    }
+  }
+
+  didClose() {
     this._page._didClose();
   }
 
   dispose() {
     this._pageProxySession.dispose();
+    helper.removeEventListeners(this._sessionListeners);
     helper.removeEventListeners(this._eventListeners);
     if (this._session)
       this._session.dispose();