chore: refactor frame tree management (#8952)

This PR removes the strong references between
parent and child frames and stores the frame in a
lazy FrameTree with weak references.

Drive-by: fix --no-coverage flag in test runner

src/common/Frame.ts

@@ -1,3 +1,19 @@
+/**
+ * Copyright 2017 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 import {Protocol} from 'devtools-protocol';
 import {assert} from '../util/assert.js';
 import {isErrorLike} from '../util/ErrorLike.js';
@@ -150,7 +166,6 @@ export interface FrameAddStyleTagOptions {
  * @public
  */
 export class Frame {
-  #parentFrame: Frame | null;
   #url = '';
   #detached = false;
   #client!: CDPSession;
@@ -186,30 +201,25 @@ export class Frame {
   /**
    * @internal
    */
-  _childFrames: Set<Frame>;
+  _parentId?: string;
 
   /**
    * @internal
    */
   constructor(
     frameManager: FrameManager,
-    parentFrame: Frame | null,
     frameId: string,
+    parentFrameId: string | undefined,
     client: CDPSession
   ) {
     this._frameManager = frameManager;
-    this.#parentFrame = parentFrame ?? null;
     this.#url = '';
     this._id = frameId;
+    this._parentId = parentFrameId;
     this.#detached = false;
 
     this._loaderId = '';
 
-    this._childFrames = new Set();
-    if (this.#parentFrame) {
-      this.#parentFrame._childFrames.add(this);
-    }
-
     this.updateClient(client);
   }
 
@@ -720,14 +730,14 @@ export class Frame {
    * @returns The parent frame, if any. Detached and main frames return `null`.
    */
   parentFrame(): Frame | null {
-    return this.#parentFrame;
+    return this._frameManager._frameTree.parentFrame(this._id) || null;
   }
 
   /**
    * @returns An array of child frames.
    */
   childFrames(): Frame[] {
-    return Array.from(this._childFrames);
+    return this._frameManager._frameTree.childFrames(this._id);
   }
 
   /**
@@ -1090,9 +1100,5 @@ export class Frame {
     this.#detached = true;
     this.worlds[MAIN_WORLD]._detach();
     this.worlds[PUPPETEER_WORLD]._detach();
-    if (this.#parentFrame) {
-      this.#parentFrame._childFrames.delete(this);
-    }
-    this.#parentFrame = null;
   }
 }

src/common/FrameManager.ts

@@ -16,13 +16,12 @@
 
 import {Protocol} from 'devtools-protocol';
 import {assert} from '../util/assert.js';
-import {createDebuggableDeferredPromise} from '../util/DebuggableDeferredPromise.js';
-import {DeferredPromise} from '../util/DeferredPromise.js';
 import {isErrorLike} from '../util/ErrorLike.js';
 import {CDPSession, isTargetClosedError} from './Connection.js';
 import {EventEmitter} from './EventEmitter.js';
 import {EVALUATION_SCRIPT_URL, ExecutionContext} from './ExecutionContext.js';
 import {Frame} from './Frame.js';
+import {FrameTree} from './FrameTree.js';
 import {IsolatedWorld, MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorld.js';
 import {NetworkManager} from './NetworkManager.js';
 import {Page} from './Page.js';
@@ -60,20 +59,13 @@ export class FrameManager extends EventEmitter {
   #page: Page;
   #networkManager: NetworkManager;
   #timeoutSettings: TimeoutSettings;
-  #frames = new Map<string, Frame>();
   #contextIdToContext = new Map<string, ExecutionContext>();
   #isolatedWorlds = new Set<string>();
-  #mainFrame?: Frame;
   #client: CDPSession;
   /**
-   * Keeps track of OOPIF targets/frames (target ID == frame ID for OOPIFs)
+   * @internal
-   * that are being initialized.
    */
-  #framesPendingTargetInit = new Map<string, DeferredPromise<void>>();
+  _frameTree = new FrameTree();
-  /**
-   * Keeps track of frames that are in the process of being attached in #onFrameAttached.
-   */
-  #framesPendingAttachment = new Map<string, DeferredPromise<void>>();
 
   get timeoutSettings(): TimeoutSettings {
     return this.#timeoutSettings;
@@ -140,19 +132,8 @@ export class FrameManager extends EventEmitter {
     });
   }
 
-  async initialize(
+  async initialize(client: CDPSession = this.#client): Promise<void> {
-    targetId: string,
-    client: CDPSession = this.#client
-  ): Promise<void> {
     try {
-      if (!this.#framesPendingTargetInit.has(targetId)) {
-        this.#framesPendingTargetInit.set(
-          targetId,
-          createDebuggableDeferredPromise(
-            `Waiting for target frame ${targetId} failed`
-          )
-        );
-      }
       const result = await Promise.all([
         client.send('Page.enable'),
         client.send('Page.getFrameTree'),
@@ -177,9 +158,6 @@ export class FrameManager extends EventEmitter {
       }
 
       throw error;
-    } finally {
-      this.#framesPendingTargetInit.get(targetId)?.resolve();
-      this.#framesPendingTargetInit.delete(targetId);
     }
   }
 
@@ -198,16 +176,17 @@ export class FrameManager extends EventEmitter {
   }
 
   mainFrame(): Frame {
-    assert(this.#mainFrame, 'Requesting main frame too early!');
+    const mainFrame = this._frameTree.getMainFrame();
-    return this.#mainFrame;
+    assert(mainFrame, 'Requesting main frame too early!');
+    return mainFrame;
   }
 
   frames(): Frame[] {
-    return Array.from(this.#frames.values());
+    return Array.from(this._frameTree.frames());
   }
 
   frame(frameId: string): Frame | null {
-    return this.#frames.get(frameId) || null;
+    return this._frameTree.getById(frameId) || null;
   }
 
   onAttachedToTarget(target: Target): void {
@@ -215,16 +194,16 @@ export class FrameManager extends EventEmitter {
       return;
     }
 
-    const frame = this.#frames.get(target._getTargetInfo().targetId);
+    const frame = this.frame(target._getTargetInfo().targetId);
     if (frame) {
       frame.updateClient(target._session()!);
     }
     this.setupEventListeners(target._session()!);
-    this.initialize(target._getTargetInfo().targetId, target._session());
+    this.initialize(target._session());
   }
 
   onDetachedFromTarget(target: Target): void {
-    const frame = this.#frames.get(target._targetId);
+    const frame = this.frame(target._targetId);
     if (frame && frame.isOOPFrame()) {
       // When an OOP iframe is removed from the page, it
       // will only get a Target.detachedFromTarget event.
@@ -233,7 +212,7 @@ export class FrameManager extends EventEmitter {
   }
 
   #onLifecycleEvent(event: Protocol.Page.LifecycleEventEvent): void {
-    const frame = this.#frames.get(event.frameId);
+    const frame = this.frame(event.frameId);
     if (!frame) {
       return;
     }
@@ -242,7 +221,7 @@ export class FrameManager extends EventEmitter {
   }
 
   #onFrameStartedLoading(frameId: string): void {
-    const frame = this.#frames.get(frameId);
+    const frame = this.frame(frameId);
     if (!frame) {
       return;
     }
@@ -250,7 +229,7 @@ export class FrameManager extends EventEmitter {
   }
 
   #onFrameStoppedLoading(frameId: string): void {
-    const frame = this.#frames.get(frameId);
+    const frame = this.frame(frameId);
     if (!frame) {
       return;
     }
@@ -284,8 +263,8 @@ export class FrameManager extends EventEmitter {
     frameId: string,
     parentFrameId: string
   ): void {
-    if (this.#frames.has(frameId)) {
+    let frame = this.frame(frameId);
-      const frame = this.#frames.get(frameId)!;
+    if (frame) {
       if (session && frame.isOOPFrame()) {
         // If an OOP iframes becomes a normal iframe again
         // it is first attached to the parent page before
@@ -294,86 +273,41 @@ export class FrameManager extends EventEmitter {
       }
       return;
     }
-    const parentFrame = this.#frames.get(parentFrameId);
 
-    const complete = (parentFrame: Frame) => {
+    frame = new Frame(this, frameId, parentFrameId, session);
-      assert(parentFrame, `Parent frame ${parentFrameId} not found`);
+    this._frameTree.addFrame(frame);
-      const frame = new Frame(this, parentFrame, frameId, session);
+    this.emit(FrameManagerEmittedEvents.FrameAttached, frame);
-      this.#frames.set(frame._id, frame);
-      this.emit(FrameManagerEmittedEvents.FrameAttached, frame);
-    };
-
-    if (parentFrame) {
-      return complete(parentFrame);
-    }
-
-    const frame = this.#framesPendingTargetInit.get(parentFrameId);
-    if (frame) {
-      if (!this.#framesPendingAttachment.has(frameId)) {
-        this.#framesPendingAttachment.set(
-          frameId,
-          createDebuggableDeferredPromise(
-            `Waiting for frame ${frameId} to attach failed`
-          )
-        );
-      }
-      frame.then(() => {
-        complete(this.#frames.get(parentFrameId)!);
-        this.#framesPendingAttachment.get(frameId)?.resolve();
-        this.#framesPendingAttachment.delete(frameId);
-      });
-      return;
-    }
-
-    throw new Error(`Parent frame ${parentFrameId} not found`);
   }
 
-  #onFrameNavigated(framePayload: Protocol.Page.Frame): void {
+  async #onFrameNavigated(framePayload: Protocol.Page.Frame): Promise<void> {
     const frameId = framePayload.id;
     const isMainFrame = !framePayload.parentId;
-    const frame = isMainFrame ? this.#mainFrame : this.#frames.get(frameId);
 
-    const complete = (frame?: Frame) => {
+    let frame = this._frameTree.getById(frameId);
-      assert(
-        isMainFrame || frame,
-        `Missing frame isMainFrame=${isMainFrame}, frameId=${frameId}`
-      );
 
-      // Detach all child frames first.
+    // Detach all child frames first.
-      if (frame) {
+    if (frame) {
-        for (const child of frame.childFrames()) {
+      for (const child of frame.childFrames()) {
-          this.#removeFramesRecursively(child);
+        this.#removeFramesRecursively(child);
-        }
       }
-
-      // Update or create main frame.
-      if (isMainFrame) {
-        if (frame) {
-          // Update frame id to retain frame identity on cross-process navigation.
-          this.#frames.delete(frame._id);
-          frame._id = frameId;
-        } else {
-          // Initial main frame navigation.
-          frame = new Frame(this, null, frameId, this.#client);
-        }
-        this.#frames.set(frameId, frame);
-        this.#mainFrame = frame;
-      }
-
-      // Update frame payload.
-      assert(frame);
-      frame._navigated(framePayload);
-
-      this.emit(FrameManagerEmittedEvents.FrameNavigated, frame);
-    };
-    const pendingFrame = this.#framesPendingAttachment.get(frameId);
-    if (pendingFrame) {
-      pendingFrame.then(() => {
-        complete(isMainFrame ? this.#mainFrame : this.#frames.get(frameId));
-      });
-    } else {
-      complete(frame);
     }
+
+    // Update or create main frame.
+    if (isMainFrame) {
+      if (frame) {
+        // Update frame id to retain frame identity on cross-process navigation.
+        this._frameTree.removeFrame(frame);
+        frame._id = frameId;
+      } else {
+        // Initial main frame navigation.
+        frame = new Frame(this, frameId, undefined, this.#client);
+      }
+      this._frameTree.addFrame(frame);
+    }
+
+    frame = await this._frameTree.waitForFrame(frameId);
+    frame._navigated(framePayload);
+    this.emit(FrameManagerEmittedEvents.FrameNavigated, frame);
   }
 
   async #createIsolatedWorld(session: CDPSession, name: string): Promise<void> {
@@ -410,7 +344,7 @@ export class FrameManager extends EventEmitter {
   }
 
   #onFrameNavigatedWithinDocument(frameId: string, url: string): void {
-    const frame = this.#frames.get(frameId);
+    const frame = this.frame(frameId);
     if (!frame) {
       return;
     }
@@ -423,7 +357,7 @@ export class FrameManager extends EventEmitter {
     frameId: string,
     reason: Protocol.Page.FrameDetachedEventReason
   ): void {
-    const frame = this.#frames.get(frameId);
+    const frame = this.frame(frameId);
     if (reason === 'remove') {
       // Only remove the frame if the reason for the detached event is
       // an actual removement of the frame.
@@ -442,8 +376,7 @@ export class FrameManager extends EventEmitter {
   ): void {
     const auxData = contextPayload.auxData as {frameId?: string} | undefined;
     const frameId = auxData && auxData.frameId;
-    const frame =
+    const frame = typeof frameId === 'string' ? this.frame(frameId) : undefined;
-      typeof frameId === 'string' ? this.#frames.get(frameId) : undefined;
     let world: IsolatedWorld | undefined;
     if (frame) {
       // Only care about execution contexts created for the current session.
@@ -509,7 +442,7 @@ export class FrameManager extends EventEmitter {
       this.#removeFramesRecursively(child);
     }
     frame._detach();
-    this.#frames.delete(frame._id);
+    this._frameTree.removeFrame(frame);
     this.emit(FrameManagerEmittedEvents.FrameDetached, frame);
   }
 }

src/common/FrameTree.ts

@@ -0,0 +1,111 @@
+/**
+ * Copyright 2022 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import {
+  createDeferredPromise,
+  DeferredPromise,
+} from '../util/DeferredPromise.js';
+import type {Frame} from './Frame.js';
+
+/**
+ * Keeps track of the page frame tree and it's is managed by
+ * {@link FrameManager}. FrameTree uses frame IDs to reference frame and it
+ * means that referenced frames might not be in the tree anymore. Thus, the tree
+ * structure is eventually consistent.
+ * @internal
+ */
+export class FrameTree {
+  #frames = new Map<string, Frame>();
+  // frameID -> parentFrameID
+  #parentIds = new Map<string, string>();
+  // frameID -> childFrameIDs
+  #childIds = new Map<string, Set<string>>();
+  #mainFrame?: Frame;
+  #waitRequests = new Map<string, Set<DeferredPromise<Frame>>>();
+
+  getMainFrame(): Frame | undefined {
+    return this.#mainFrame;
+  }
+
+  getById(frameId: string): Frame | undefined {
+    return this.#frames.get(frameId);
+  }
+
+  /**
+   * Returns a promise that is resolved once the frame with
+   * the given ID is added to the tree.
+   */
+  waitForFrame(frameId: string): Promise<Frame> {
+    const frame = this.getById(frameId);
+    if (frame) {
+      return Promise.resolve(frame);
+    }
+    const deferred = createDeferredPromise<Frame>();
+    const callbacks =
+      this.#waitRequests.get(frameId) || new Set<DeferredPromise<Frame>>();
+    callbacks.add(deferred);
+    return deferred;
+  }
+
+  frames(): Frame[] {
+    return Array.from(this.#frames.values());
+  }
+
+  addFrame(frame: Frame): void {
+    this.#frames.set(frame._id, frame);
+    if (frame._parentId) {
+      this.#parentIds.set(frame._id, frame._parentId);
+      if (!this.#childIds.has(frame._parentId)) {
+        this.#childIds.set(frame._parentId, new Set());
+      }
+      this.#childIds.get(frame._parentId)!.add(frame._id);
+    } else {
+      this.#mainFrame = frame;
+    }
+    this.#waitRequests.get(frame._id)?.forEach(request => {
+      return request.resolve(frame);
+    });
+  }
+
+  removeFrame(frame: Frame): void {
+    this.#frames.delete(frame._id);
+    this.#parentIds.delete(frame._id);
+    if (frame._parentId) {
+      this.#childIds.get(frame._parentId)?.delete(frame._id);
+    } else {
+      this.#mainFrame = undefined;
+    }
+  }
+
+  childFrames(frameId: string): Frame[] {
+    const childIds = this.#childIds.get(frameId);
+    if (!childIds) {
+      return [];
+    }
+    return Array.from(childIds)
+      .map(id => {
+        return this.getById(id);
+      })
+      .filter((frame): frame is Frame => {
+        return frame !== undefined;
+      });
+  }
+
+  parentFrame(frameId: string): Frame | undefined {
+    const parentId = this.#parentIds.get(frameId);
+    return parentId ? this.getById(parentId) : undefined;
+  }
+}

src/common/Page.ts

@@ -660,7 +660,7 @@ export class Page extends EventEmitter {
   async #initialize(): Promise<void> {
     try {
       await Promise.all([
-        this.#frameManager.initialize(this.#target._targetId),
+        this.#frameManager.initialize(),
         this.#client.send('Performance.enable'),
         this.#client.send('Log.enable'),
       ]);

src/types.ts

@@ -24,6 +24,7 @@ export * from './common/FileChooser.js';
 export * from './common/FirefoxTargetManager.js';
 export * from './common/Frame.js';
 export * from './common/FrameManager.js';
+export * from './common/FrameTree.js';
 export * from './common/HTTPRequest.js';
 export * from './common/HTTPResponse.js';
 export * from './common/Input.js';

utils/mochaRunner/src/main.ts

@@ -27,7 +27,7 @@ import {
 import path from 'path';
 import fs from 'fs';
 import os from 'os';
-import {spawn} from 'node:child_process';
+import {spawn, SpawnOptions} from 'node:child_process';
 import {
   extendProcessEnv,
   filterByPlatform,
@@ -126,8 +126,14 @@ async function main() {
         '-O',
         'output=' + tmpFilename,
       ];
+      const spawnArgs: SpawnOptions = {
+        shell: true,
+        cwd: process.cwd(),
+        stdio: 'inherit',
+        env,
+      };
       const handle = noCoverage
-        ? spawn('npx', ['mocha', ...args])
+        ? spawn('npx', ['mocha', ...args], spawnArgs)
         : spawn(
             'npx',
             [
@@ -138,12 +144,7 @@ async function main() {
               'npx mocha',
               ...args,
             ],
-            {
+            spawnArgs
-              shell: true,
-              cwd: process.cwd(),
-              stdio: 'inherit',
-              env,
-            }
           );
       await new Promise<void>((resolve, reject) => {
         handle.on('error', err => {