From 0eae3d1bb6a4722b527c56fc9ea91db45d97a6e8 Mon Sep 17 00:00:00 2001
From: Alex Rudenko <OrKoN@users.noreply.github.com>
Date: Tue, 7 May 2024 09:50:56 +0200
Subject: [PATCH] refactor: move execution context events (#12402)

---
 .../src/cdp/ExecutionContext.ts               | 28 ++++++++----
 .../puppeteer-core/src/cdp/FrameManager.ts    | 44 +++----------------
 2 files changed, 27 insertions(+), 45 deletions(-)

diff --git a/packages/puppeteer-core/src/cdp/ExecutionContext.ts b/packages/puppeteer-core/src/cdp/ExecutionContext.ts
index 3b1c53ce8e8..1b3e3d1bd3b 100644
--- a/packages/puppeteer-core/src/cdp/ExecutionContext.ts
+++ b/packages/puppeteer-core/src/cdp/ExecutionContext.ts
@@ -6,7 +6,7 @@
 
 import type {Protocol} from 'devtools-protocol';
 
-import type {CDPSession, CDPSessionEvents} from '../api/CDPSession.js';
+import type {CDPSession} from '../api/CDPSession.js';
 import type {ElementHandle} from '../api/ElementHandle.js';
 import type {JSHandle} from '../api/JSHandle.js';
 import {EventEmitter} from '../common/EventEmitter.js';
@@ -62,31 +62,42 @@ const ariaQuerySelectorAllBinding = new Binding(
 /**
  * @internal
  */
-export class ExecutionContext implements Disposable {
+export class ExecutionContext
+  extends EventEmitter<{
+    /** Emitted when this execution context is disposed. */
+    disposed: undefined;
+  }>
+  implements Disposable
+{
   _client: CDPSession;
   _world: IsolatedWorld;
   _contextId: number;
   _contextName?: string;
 
   readonly #disposables = new DisposableStack();
-  readonly #clientEmitter: EventEmitter<CDPSessionEvents>;
 
   constructor(
     client: CDPSession,
     contextPayload: Protocol.Runtime.ExecutionContextDescription,
     world: IsolatedWorld
   ) {
+    super();
     this._client = client;
     this._world = world;
     this._contextId = contextPayload.id;
     if (contextPayload.name) {
       this._contextName = contextPayload.name;
     }
-    this.#clientEmitter = this.#disposables.use(new EventEmitter(this._client));
-    this.#clientEmitter.on(
-      'Runtime.bindingCalled',
-      this.#onBindingCalled.bind(this)
-    );
+    const clientEmitter = this.#disposables.use(new EventEmitter(this._client));
+    clientEmitter.on('Runtime.bindingCalled', this.#onBindingCalled.bind(this));
+    clientEmitter.on('Runtime.executionContextDestroyed', async event => {
+      if (event.executionContextId === this._contextId) {
+        this[disposeSymbol]();
+      }
+    });
+    clientEmitter.on('Runtime.executionContextsCleared', async () => {
+      this[disposeSymbol]();
+    });
   }
 
   // Contains mapping from functions that should be bound to Puppeteer functions.
@@ -456,6 +467,7 @@ export class ExecutionContext implements Disposable {
 
   [disposeSymbol](): void {
     this.#disposables.dispose();
+    this.emit('disposed', undefined);
   }
 }
 
diff --git a/packages/puppeteer-core/src/cdp/FrameManager.ts b/packages/puppeteer-core/src/cdp/FrameManager.ts
index 13822fc6aa3..f7c7d0a99f8 100644
--- 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('disposed', () => {
+      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 {