chore: refactor deferred promise (#8759)

src/common/DOMWorld.ts

@@ -28,7 +28,9 @@ import {getQueryHandlerAndSelector} from './QueryHandler.js';
 import {TimeoutSettings} from './TimeoutSettings.js';
 import {EvaluateFunc, HandleFor, NodeFor} from './types.js';
 import {
+  createDeferredPromise,
   debugError,
+  DeferredPromise,
   importFS,
   isNumber,
   isString,
@@ -75,8 +77,7 @@ export class DOMWorld {
   #frame: Frame;
   #timeoutSettings: TimeoutSettings;
   #documentPromise: Promise<ElementHandle<Document>> | null = null;
-  #contextPromise: Promise<ExecutionContext> | null = null;
-  #contextResolveCallback: ((x: ExecutionContext) => void) | null = null;
+  #contextPromise: DeferredPromise<ExecutionContext> = createDeferredPromise();
   #detached = false;
 
   // Set of bindings that have been registered in the current context.
@@ -110,7 +111,6 @@ export class DOMWorld {
     this.#frameManager = frameManager;
     this.#frame = frame;
     this.#timeoutSettings = timeoutSettings;
-    this._setContext(null);
     this.#client.on('Runtime.bindingCalled', this.#onBindingCalled);
   }
 
@@ -118,28 +118,25 @@ export class DOMWorld {
     return this.#frame;
   }
 
-  async _setContext(context: ExecutionContext | null): Promise<void> {
-    if (context) {
-      assert(
-        this.#contextResolveCallback,
-        `ExecutionContext ${context._contextId} has already been set.`
-      );
-      this.#ctxBindings.clear();
-      this.#contextResolveCallback?.call(null, context);
-      this.#contextResolveCallback = null;
-      for (const waitTask of this._waitTasks) {
-        waitTask.rerun();
-      }
-    } else {
-      this.#documentPromise = null;
-      this.#contextPromise = new Promise(fulfill => {
-        this.#contextResolveCallback = fulfill;
-      });
+  clearContext(): void {
+    this.#documentPromise = null;
+    this.#contextPromise = createDeferredPromise();
+  }
+
+  setContext(context: ExecutionContext): void {
+    assert(
+      this.#contextPromise,
+      `ExecutionContext ${context._contextId} has already been set.`
+    );
+    this.#ctxBindings.clear();
+    this.#contextPromise.resolve(context);
+    for (const waitTask of this._waitTasks) {
+      waitTask.rerun();
     }
   }
 
-  _hasContext(): boolean {
-    return !this.#contextResolveCallback;
+  hasContext(): boolean {
+    return this.#contextPromise.resolved();
   }
 
   _detach(): void {
@@ -619,7 +616,7 @@ export class DOMWorld {
     event: Protocol.Runtime.BindingCalledEvent
   ): Promise<void> => {
     let payload: {type: string; name: string; seq: number; args: unknown[]};
-    if (!this._hasContext()) {
+    if (!this.hasContext()) {
       return;
     }
     const context = await this.executionContext();

src/common/FrameManager.ts

@@ -397,7 +397,8 @@ export class FrameManager extends EventEmitter {
       return complete(parentFrame);
     }
 
-    if (this.#framesPendingTargetInit.has(parentFrameId)) {
+    const frame = this.#framesPendingTargetInit.get(parentFrameId);
+    if (frame) {
       if (!this.#framesPendingAttachment.has(frameId)) {
         this.#framesPendingAttachment.set(
           frameId,
@@ -406,7 +407,7 @@ export class FrameManager extends EventEmitter {
           )
         );
       }
-      this.#framesPendingTargetInit.get(parentFrameId)!.promise.then(() => {
+      frame.then(() => {
         complete(this.#frames.get(parentFrameId)!);
         this.#framesPendingAttachment.get(frameId)?.resolve();
         this.#framesPendingAttachment.delete(frameId);
@@ -455,8 +456,9 @@ export class FrameManager extends EventEmitter {
 
       this.emit(FrameManagerEmittedEvents.FrameNavigated, frame);
     };
-    if (this.#framesPendingAttachment.has(frameId)) {
-      this.#framesPendingAttachment.get(frameId)!.promise.then(() => {
+    const pendingFrame = this.#framesPendingAttachment.get(frameId);
+    if (pendingFrame) {
+      pendingFrame.then(() => {
         complete(isMainFrame ? this.#mainFrame : this.#frames.get(frameId));
       });
     } else {
@@ -539,7 +541,7 @@ export class FrameManager extends EventEmitter {
         world = frame._mainWorld;
       } else if (
         contextPayload.name === UTILITY_WORLD_NAME &&
-        !frame._secondaryWorld._hasContext()
+        !frame._secondaryWorld.hasContext()
       ) {
         // In case of multiple sessions to the same target, there's a race between
         // connections so we might end up creating multiple isolated worlds.
@@ -553,7 +555,7 @@ export class FrameManager extends EventEmitter {
       world
     );
     if (world) {
-      world._setContext(context);
+      world.setContext(context);
     }
     const key = `${session.id()}:${contextPayload.id}`;
     this.#contextIdToContext.set(key, context);
@@ -570,7 +572,7 @@ export class FrameManager extends EventEmitter {
     }
     this.#contextIdToContext.delete(key);
     if (context._world) {
-      context._world._setContext(null);
+      context._world.clearContext();
     }
   }
 
@@ -582,7 +584,7 @@ export class FrameManager extends EventEmitter {
         continue;
       }
       if (context._world) {
-        context._world._setContext(null);
+        context._world.clearContext();
       }
       this.#contextIdToContext.delete(key);
     }

src/common/NetworkManager.ts

@@ -143,7 +143,7 @@ export class NetworkManager extends EventEmitter {
    */
   initialize(): Promise<void> {
     if (this.#deferredInitPromise) {
-      return this.#deferredInitPromise.promise;
+      return this.#deferredInitPromise;
     }
     this.#deferredInitPromise = createDeferredPromiseWithTimer<void>(
       'NetworkManager initialization timed out',
@@ -157,14 +157,15 @@ export class NetworkManager extends EventEmitter {
         : null,
       this.#client.send('Network.enable'),
     ]);
+    const deferredInitPromise = this.#deferredInitPromise;
     init
       .then(() => {
-        this.#deferredInitPromise?.resolve();
+        deferredInitPromise.resolve();
       })
       .catch(err => {
-        return this.#deferredInitPromise?.reject(err);
+        deferredInitPromise.reject(err);
       });
-    return this.#deferredInitPromise.promise;
+    return this.#deferredInitPromise;
   }
 
   async authenticate(credentials?: Credentials): Promise<void> {

src/common/util.ts

@@ -527,11 +527,11 @@ export function isErrnoException(obj: unknown): obj is NodeJS.ErrnoException {
 /**
  * @internal
  */
-export type DeferredPromise<T> = {
-  promise: Promise<T>;
+export interface DeferredPromise<T> extends Promise<T> {
+  resolved: () => boolean;
   resolve: (_: T) => void;
   reject: (_: Error) => void;
-};
+}
 
 /**
  * Creates an returns a promise along with the resolve/reject functions.
@@ -545,6 +545,7 @@ export function createDeferredPromiseWithTimer<T>(
   timeoutMessage: string,
   timeout = 5000
 ): DeferredPromise<T> {
+  let isResolved = false;
   let resolver = (_: T): void => {};
   let rejector = (_: Error) => {};
   const taskPromise = new Promise<T>((resolve, reject) => {
@@ -554,15 +555,45 @@ export function createDeferredPromiseWithTimer<T>(
   const timeoutId = setTimeout(() => {
     rejector(new Error(timeoutMessage));
   }, timeout);
-  return {
-    promise: taskPromise,
+  return Object.assign(taskPromise, {
+    resolved: () => {
+      return isResolved;
+    },
     resolve: (value: T) => {
       clearTimeout(timeoutId);
+      isResolved = true;
       resolver(value);
     },
     reject: (err: Error) => {
       clearTimeout(timeoutId);
       rejector(err);
     },
-  };
+  });
+}
+
+/**
+ * Creates an returns a promise along with the resolve/reject functions.
+ *
+ * @internal
+ */
+export function createDeferredPromise<T>(): DeferredPromise<T> {
+  let isResolved = false;
+  let resolver = (_: T): void => {};
+  let rejector = (_: Error) => {};
+  const taskPromise = new Promise<T>((resolve, reject) => {
+    resolver = resolve;
+    rejector = reject;
+  });
+  return Object.assign(taskPromise, {
+    resolved: () => {
+      return isResolved;
+    },
+    resolve: (value: T) => {
+      isResolved = true;
+      resolver(value);
+    },
+    reject: (err: Error) => {
+      rejector(err);
+    },
+  });
 }