From eb89720704c7d5c2d79713d33c7d53a3e5ce3392 Mon Sep 17 00:00:00 2001
From: jrandolf <101637635+jrandolf@users.noreply.github.com>
Date: Mon, 11 Sep 2023 14:07:52 +0200
Subject: [PATCH] chore: reimplement document caching (#10878)

---
 packages/puppeteer-core/src/api/Frame.ts      | 42 ++++++++++++++-----
 .../src/common/IsolatedWorld.ts               |  3 ++
 .../puppeteer-core/src/common/bidi/Realm.ts   |  1 +
 3 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/packages/puppeteer-core/src/api/Frame.ts b/packages/puppeteer-core/src/api/Frame.ts
index b8f185b9..5c45ff82 100644
--- a/packages/puppeteer-core/src/api/Frame.ts
+++ b/packages/puppeteer-core/src/api/Frame.ts
@@ -322,14 +322,31 @@ export abstract class Frame extends EventEmitter {
    */
   abstract isolatedRealm(): Realm;
 
+  #_document: Promise<ElementHandle<Document>> | undefined;
+
   /**
    * @internal
    */
-  async document(): Promise<ElementHandle<Document>> {
-    // TODO(#10813): Implement document caching.
-    return await this.evaluateHandle(() => {
-      return document;
-    });
+  #document(): Promise<ElementHandle<Document>> {
+    if (!this.#_document) {
+      this.#_document = this.isolatedRealm()
+        .evaluateHandle(() => {
+          return document;
+        })
+        .then(handle => {
+          return this.mainRealm().transferHandle(handle);
+        });
+    }
+    return this.#_document;
+  }
+
+  /**
+   * Used to clear the document handle that has been destroyed.
+   *
+   * @internal
+   */
+  clearDocumentHandle(): void {
+    this.#_document = undefined;
   }
 
   /**
@@ -441,7 +458,8 @@ export abstract class Frame extends EventEmitter {
   async $<Selector extends string>(
     selector: Selector
   ): Promise<ElementHandle<NodeFor<Selector>> | null> {
-    using document = await this.document();
+    // eslint-disable-next-line rulesdir/use-using -- This is cached.
+    const document = await this.#document();
     return await document.$(selector);
   }
 
@@ -456,7 +474,8 @@ export abstract class Frame extends EventEmitter {
   async $$<Selector extends string>(
     selector: Selector
   ): Promise<Array<ElementHandle<NodeFor<Selector>>>> {
-    using document = await this.document();
+    // eslint-disable-next-line rulesdir/use-using -- This is cached.
+    const document = await this.#document();
     return await document.$$(selector);
   }
 
@@ -494,7 +513,8 @@ export abstract class Frame extends EventEmitter {
     ...args: Params
   ): Promise<Awaited<ReturnType<Func>>> {
     pageFunction = withSourcePuppeteerURLIfNone(this.$eval.name, pageFunction);
-    using document = await this.document();
+    // eslint-disable-next-line rulesdir/use-using -- This is cached.
+    const document = await this.#document();
     return await document.$eval(selector, pageFunction, ...args);
   }
 
@@ -532,7 +552,8 @@ export abstract class Frame extends EventEmitter {
     ...args: Params
   ): Promise<Awaited<ReturnType<Func>>> {
     pageFunction = withSourcePuppeteerURLIfNone(this.$$eval.name, pageFunction);
-    using document = await this.document();
+    // eslint-disable-next-line rulesdir/use-using -- This is cached.
+    const document = await this.#document();
     return await document.$$eval(selector, pageFunction, ...args);
   }
 
@@ -548,7 +569,8 @@ export abstract class Frame extends EventEmitter {
    */
   @throwIfDetached
   async $x(expression: string): Promise<Array<ElementHandle<Node>>> {
-    using document = await this.document();
+    // eslint-disable-next-line rulesdir/use-using -- This is cached.
+    const document = await this.#document();
     return await document.$x(expression);
   }
 
diff --git a/packages/puppeteer-core/src/common/IsolatedWorld.ts b/packages/puppeteer-core/src/common/IsolatedWorld.ts
index b5950fe1..90c34439 100644
--- a/packages/puppeteer-core/src/common/IsolatedWorld.ts
+++ b/packages/puppeteer-core/src/common/IsolatedWorld.ts
@@ -126,6 +126,9 @@ export class IsolatedWorld extends Realm {
 
   clearContext(): void {
     this.#context = Deferred.create();
+    if (this.#frameOrWorker instanceof CDPFrame) {
+      this.#frameOrWorker.clearDocumentHandle();
+    }
   }
 
   setContext(context: ExecutionContext): void {
diff --git a/packages/puppeteer-core/src/common/bidi/Realm.ts b/packages/puppeteer-core/src/common/bidi/Realm.ts
index dfb04f9b..1ede1946 100644
--- a/packages/puppeteer-core/src/common/bidi/Realm.ts
+++ b/packages/puppeteer-core/src/common/bidi/Realm.ts
@@ -49,6 +49,7 @@ export class Realm extends EventEmitter {
       // Note: The Realm is destroyed, so in theory the handle should be as
       // well.
       this.internalPuppeteerUtil = undefined;
+      this.#sandbox.environment.clearDocumentHandle();
     }
   };