refactor: move execution context events

puppeteer · May 6, 2024 · 1b9fa5a · 1b9fa5a
1 parent 3eab6b5
commit 1b9fa5a
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 38 deletions.
diff --git a/packages/puppeteer-core/src/cdp/ExecutionContext.ts b/packages/puppeteer-core/src/cdp/ExecutionContext.ts
@@ -9,6 +9,7 @@ import type {Protocol} from 'devtools-protocol';
 import type {CDPSession, CDPSessionEvents} from '../api/CDPSession.js';
 import type {ElementHandle} from '../api/ElementHandle.js';
 import type {JSHandle} from '../api/JSHandle.js';
+import type {EventType} from '../common/EventEmitter.js';
 import {EventEmitter} from '../common/EventEmitter.js';
 import {LazyArg} from '../common/LazyArg.js';
 import {scriptInjector} from '../common/ScriptInjector.js';
@@ -59,10 +60,17 @@ const ariaQuerySelectorAllBinding = new Binding(
   }) as (...args: unknown[]) => unknown
 );
 
+interface ExecutionContextEvents extends Record<EventType, unknown> {
+  dispose: undefined;
+}
+
 /**
  * @internal
  */
-export class ExecutionContext implements Disposable {
+export class ExecutionContext
+  extends EventEmitter<ExecutionContextEvents>
+  implements Disposable
+{
   _client: CDPSession;
   _world: IsolatedWorld;
   _contextId: number;
@@ -76,6 +84,7 @@ export class ExecutionContext implements Disposable {
     contextPayload: Protocol.Runtime.ExecutionContextDescription,
     world: IsolatedWorld
   ) {
+    super();
     this._client = client;
     this._world = world;
     this._contextId = contextPayload.id;
@@ -87,6 +96,14 @@ export class ExecutionContext implements Disposable {
       'Runtime.bindingCalled',
       this.#onBindingCalled.bind(this)
     );
+    this.#clientEmitter.on('Runtime.executionContextDestroyed', async event => {
+      if (event.executionContextId === this._contextId) {
+        this[disposeSymbol]();
+      }
+    });
+    this.#clientEmitter.on('Runtime.executionContextsCleared', async () => {
+      this[disposeSymbol]();
+    });
   }
 
   // Contains mapping from functions that should be bound to Puppeteer functions.
@@ -456,6 +473,7 @@ export class ExecutionContext implements Disposable {
 
   [disposeSymbol](): void {
     this.#disposables.dispose();
+    this.emit('dispose', undefined);
   }
 }
 

diff --git a/packages/puppeteer-core/src/cdp/FrameManager.ts b/packages/puppeteer-core/src/cdp/FrameManager.ts
@@ -122,8 +122,6 @@ export class FrameManager extends EventEmitter<FrameManagerEvents> {
    * its frame tree and ID.
    */
   async swapFrameTree(client: CDPSession): Promise<void> {
-    this.#onExecutionContextsCleared(this.#client);
-
     this.#client = client;
     assert(
       this.#client instanceof CdpCDPSession,
@@ -188,14 +186,6 @@ export class FrameManager extends EventEmitter<FrameManagerEvents> {
       await this.#frameTreeHandled?.valueOrThrow();
       this.#onExecutionContextCreated(event.context, session);
     });
-    session.on('Runtime.executionContextDestroyed', async event => {
-      await this.#frameTreeHandled?.valueOrThrow();
-      this.#onExecutionContextDestroyed(event.executionContextId, session);
-    });
-    session.on('Runtime.executionContextsCleared', async () => {
-      await this.#frameTreeHandled?.valueOrThrow();
-      this.#onExecutionContextsCleared(session);
-    });
     session.on('Page.lifecycleEvent', async event => {
       await this.#frameTreeHandled?.valueOrThrow();
       this.#onLifecycleEvent(event);
@@ -505,35 +495,15 @@ export class FrameManager extends EventEmitter<FrameManagerEvents> {
     }
     const key = `${session.id()}:${contextPayload.id}`;
     this.#contextIdToContext.set(key, context);
-  }
-
-  #onExecutionContextDestroyed(
-    executionContextId: number,
-    session: CDPSession
-  ): void {
-    const key = `${session.id()}:${executionContextId}`;
-    const context = this.#contextIdToContext.get(key);
-    if (!context) {
-      return;
-    }
-    this.#contextIdToContext.delete(key);
-    if (context._world) {
-      context._world.clearContext();
-    }
-  }
-
-  #onExecutionContextsCleared(session: CDPSession): void {
-    for (const [key, context] of this.#contextIdToContext.entries()) {
-      // Make sure to only clear execution contexts that belong
-      // to the current session.
-      if (context._client !== session) {
-        continue;
-      }
-      if (context._world) {
-        context._world.clearContext();
+    context.once('dispose', () => {
+      const key = `${session.id()}:${contextPayload.id}`;
+      const context = this.#contextIdToContext.get(key);
+      if (!context) {
+        return;
       }
       this.#contextIdToContext.delete(key);
-    }
+      context._world.clearContext();
+    });
   }
 
   #removeFramesRecursively(frame: CdpFrame): void {