refactor: implement reverse argument binding (#9651)

packages/puppeteer-core/src/common/Binding.ts

@@ -0,0 +1,119 @@
+import {JSHandle} from '../api/JSHandle.js';
+import {isErrorLike} from '../util/ErrorLike.js';
+import {ExecutionContext} from './ExecutionContext.js';
+import {debugError} from './util.js';
+
+/**
+ * @internal
+ */
+export class Binding {
+  #name: string;
+  #fn: (...args: unknown[]) => unknown;
+  constructor(name: string, fn: (...args: unknown[]) => unknown) {
+    this.#name = name;
+    this.#fn = fn;
+  }
+
+  /**
+   *
+   * @param context - Context to run the binding in; the context should have
+   * the binding added to it beforehand.
+   * @param id - ID of the call. This should come from the CDP
+   * `onBindingCalled` response.
+   * @param args - Plain arguments from CDP.
+   */
+  async run(
+    context: ExecutionContext,
+    id: number,
+    args: unknown[],
+    isTrivial: boolean
+  ): Promise<void> {
+    const garbage = [];
+    try {
+      if (!isTrivial) {
+        // Getting non-trivial arguments.
+        const handles = await context.evaluateHandle(
+          (name, seq) => {
+            // @ts-expect-error Code is evaluated in a different context.
+            return globalThis[name].args.get(seq);
+          },
+          this.#name,
+          id
+        );
+        try {
+          const properties = await handles.getProperties();
+          for (const [index, handle] of properties) {
+            // This is not straight-forward since some arguments can stringify, but
+            // aren't plain objects so add subtypes when the use-case arises.
+            if (index in args) {
+              switch (handle.remoteObject().subtype) {
+                case 'node':
+                  args[+index] = handle;
+                  break;
+                default:
+                  garbage.push(handle.dispose());
+              }
+            } else {
+              garbage.push(handle.dispose());
+            }
+          }
+        } finally {
+          await handles.dispose();
+        }
+      }
+
+      await context.evaluate(
+        (name, seq, result) => {
+          // @ts-expect-error Code is evaluated in a different context.
+          const callbacks = globalThis[name].callbacks;
+          callbacks.get(seq).resolve(result);
+          callbacks.delete(seq);
+        },
+        this.#name,
+        id,
+        await this.#fn(...args)
+      );
+
+      for (const arg of args) {
+        if (arg instanceof JSHandle) {
+          garbage.push(arg.dispose());
+        }
+      }
+    } catch (error) {
+      if (isErrorLike(error)) {
+        await context
+          .evaluate(
+            (name, seq, message, stack) => {
+              const error = new Error(message);
+              error.stack = stack;
+              // @ts-expect-error Code is evaluated in a different context.
+              const callbacks = globalThis[name].callbacks;
+              callbacks.get(seq).reject(error);
+              callbacks.delete(seq);
+            },
+            this.#name,
+            id,
+            error.message,
+            error.stack
+          )
+          .catch(debugError);
+      } else {
+        await context
+          .evaluate(
+            (name, seq, error) => {
+              // @ts-expect-error Code is evaluated in a different context.
+              const callbacks = globalThis[name].callbacks;
+              callbacks.get(seq).reject(error);
+              callbacks.delete(seq);
+            },
+            this.#name,
+            id,
+            error
+          )
+          .catch(debugError);
+      }
+    } finally {
+      await Promise.all(garbage);
+    }
+  }
+}

packages/puppeteer-core/src/common/IsolatedWorld.ts

@@ -28,16 +28,18 @@ import {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js';
 import {LifecycleWatcher, PuppeteerLifeCycleEvent} from './LifecycleWatcher.js';
 import {TimeoutSettings} from './TimeoutSettings.js';
 import {
+  BindingPayload,
   EvaluateFunc,
   EvaluateFuncWith,
   HandleFor,
   InnerLazyParams,
   NodeFor,
 } from './types.js';
-import {createJSHandle, debugError, pageBindingInitString} from './util.js';
+import {addPageBinding, createJSHandle, debugError} from './util.js';
 import {TaskManager, WaitTask} from './WaitTask.js';
 
 import type {ElementHandle} from '../api/ElementHandle.js';
+import {Binding} from './Binding.js';
 import {LazyArg} from './LazyArg.js';
 
 /**
@@ -95,24 +97,20 @@ export class IsolatedWorld {
   #detached = false;
 
   // Set of bindings that have been registered in the current context.
-  #ctxBindings = new Set<string>();
+  #contextBindings = new Set<string>();
 
   // Contains mapping from functions that should be bound to Puppeteer functions.
-  #boundFunctions = new Map<string, Function>();
+  #bindings = new Map<string, Binding>();
   #taskManager = new TaskManager();
 
   get taskManager(): TaskManager {
     return this.#taskManager;
   }
 
-  get _boundFunctions(): Map<string, Function> {
-    return this.#boundFunctions;
+  get _bindings(): Map<string, Binding> {
+    return this.#bindings;
   }
 
-  static #bindingIdentifier = (name: string, contextId: number) => {
-    return `${name}_${contextId}`;
-  };
-
   constructor(frame: Frame) {
     // Keep own reference to client because it might differ from the FrameManager's
     // client for OOP iframes.
@@ -142,7 +140,7 @@ export class IsolatedWorld {
   }
 
   setContext(context: ExecutionContext): void {
-    this.#ctxBindings.clear();
+    this.#contextBindings.clear();
     this.#context.resolve(context);
     this.#taskManager.rerunAll();
   }
@@ -354,118 +352,72 @@ export class IsolatedWorld {
 
   // If multiple waitFor are set up asynchronously, we need to wait for the
   // first one to set up the binding in the page before running the others.
-  #settingUpBinding: Promise<void> | null = null;
-
+  #mutex = new Mutex();
   async _addBindingToContext(
     context: ExecutionContext,
     name: string
   ): Promise<void> {
-    // Previous operation added the binding so we are done.
-    if (
-      this.#ctxBindings.has(
-        IsolatedWorld.#bindingIdentifier(name, context._contextId)
-      )
-    ) {
+    if (this.#contextBindings.has(name)) {
       return;
     }
-    // Wait for other operation to finish
-    if (this.#settingUpBinding) {
-      await this.#settingUpBinding;
-      return this._addBindingToContext(context, name);
-    }
 
-    const bind = async (name: string) => {
-      const expression = pageBindingInitString('internal', name);
-      try {
-        // TODO: In theory, it would be enough to call this just once
-        await context._client.send('Runtime.addBinding', {
-          name,
-          executionContextName: context._contextName,
-        });
-        await context.evaluate(expression);
-      } catch (error) {
-        // We could have tried to evaluate in a context which was already
-        // destroyed. This happens, for example, if the page is navigated while
-        // we are trying to add the binding
-        if (error instanceof Error) {
-          // Destroyed context.
-          if (error.message.includes('Execution context was destroyed')) {
-            return;
-          }
-          // Missing context.
-          if (error.message.includes('Cannot find context with specified id')) {
-            return;
-          }
-        }
+    await this.#mutex.acquire();
+    try {
+      await context._client.send('Runtime.addBinding', {
+        name,
+        executionContextName: context._contextName,
+      });
 
-        debugError(error);
-        return;
+      await context.evaluate(addPageBinding, 'internal', name);
+
+      this.#contextBindings.add(name);
+    } catch (error) {
+      // We could have tried to evaluate in a context which was already
+      // destroyed. This happens, for example, if the page is navigated while
+      // we are trying to add the binding
+      if (error instanceof Error) {
+        // Destroyed context.
+        if (error.message.includes('Execution context was destroyed')) {
+          return;
+        }
+        // Missing context.
+        if (error.message.includes('Cannot find context with specified id')) {
+          return;
+        }
       }
-      this.#ctxBindings.add(
-        IsolatedWorld.#bindingIdentifier(name, context._contextId)
-      );
-    };
 
-    this.#settingUpBinding = bind(name);
-    await this.#settingUpBinding;
-    this.#settingUpBinding = null;
+      debugError(error);
+    } finally {
+      this.#mutex.release();
+    }
   }
 
   #onBindingCalled = async (
     event: Protocol.Runtime.BindingCalledEvent
   ): Promise<void> => {
-    let payload: {type: string; name: string; seq: number; args: unknown[]};
-    if (!this.hasContext()) {
-      return;
-    }
-    const context = await this.executionContext();
+    let payload: BindingPayload;
     try {
       payload = JSON.parse(event.payload);
     } catch {
       // The binding was either called by something in the page or it was
       // called before our wrapper was initialized.
       return;
     }
-    const {type, name, seq, args} = payload;
-    if (
-      type !== 'internal' ||
-      !this.#ctxBindings.has(
-        IsolatedWorld.#bindingIdentifier(name, context._contextId)
-      )
-    ) {
+    const {type, name, seq, args, isTrivial} = payload;
+    if (type !== 'internal') {
       return;
     }
-    if (context._contextId !== event.executionContextId) {
+    if (!this.#contextBindings.has(name)) {
       return;
     }
-    try {
-      const fn = this._boundFunctions.get(name);
-      if (!fn) {
-        throw new Error(`Bound function $name is not found`);
-      }
-      const result = await fn(...args);
-      await context.evaluate(
-        (name: string, seq: number, result: unknown) => {
-          // @ts-expect-error Code is evaluated in a different context.
-          const callbacks = self[name].callbacks;
-          callbacks.get(seq).resolve(result);
-          callbacks.delete(seq);
-        },
-        name,
-        seq,
-        result
-      );
-    } catch (error) {
-      // The WaitTask may already have been resolved by timing out, or the
-      // execution context may have been destroyed.
-      // In both caes, the promises above are rejected with a protocol error.
-      // We can safely ignores these, as the WaitTask is re-installed in
-      // the next execution context if needed.
-      if ((error as Error).message.includes('Protocol error')) {
-        return;
-      }
-      debugError(error);
+
+    const context = await this.#context;
+    if (event.executionContextId !== context._contextId) {
+      return;
     }
+
+    const binding = this._bindings.get(name);
+    await binding?.run(context, seq, args, isTrivial);
   };
 
   async _waitForSelectorInPage(
@@ -598,3 +550,31 @@ export class IsolatedWorld {
     return result;
   }
 }
+
+class Mutex {
+  #locked = false;
+  #acquirers: Array<() => void> = [];
+
+  // This is FIFO.
+  acquire(): Promise<void> {
+    if (!this.#locked) {
+      this.#locked = true;
+      return Promise.resolve();
+    }
+    let resolve!: () => void;
+    const promise = new Promise<void>(res => {
+      resolve = res;
+    });
+    this.#acquirers.push(resolve);
+    return promise;
+  }
+
+  release(): void {
+    const resolve = this.#acquirers.shift();
+    if (!resolve) {
+      this.#locked = false;
+      return;
+    }
+    resolve();
+  }
+}