chore: rework wait task to accept arbitrary task on dom world

.gitignore

@@ -1,6 +1,7 @@
 /node_modules/
 /test/output-chromium
 /test/output-firefox
+/test/output-webkit
 /test/test-user-data-dir*
 /.local-chromium/
 /.local-browser/

docs/api.md

@@ -1828,7 +1828,7 @@ await element.setInputFiles('/tmp/myfile.pdf');
 #### page.waitForFunction(pageFunction[, options[, ...args]])
 - `pageFunction` <[function]|[string]> Function to be evaluated in browser context
 - `options` <[Object]> Optional waiting parameters
-  - `polling` <[string]|[number]> An interval at which the `pageFunction` is executed, defaults to `raf`. If `polling` is a number, then it is treated as an interval in milliseconds at which the function would be executed. If `polling` is a string, then it can be one of the following values:
+  - `polling` <[number]|"raf"|"mutation"> An interval at which the `pageFunction` is executed, defaults to `raf`. If `polling` is a number, then it is treated as an interval in milliseconds at which the function would be executed. If `polling` is a string, then it can be one of the following values:
     - `raf` - to constantly execute `pageFunction` in `requestAnimationFrame` callback. This is the tightest polling mode which is suitable to observe styling changes.
     - `mutation` - to execute `pageFunction` on every DOM mutation.
   - `timeout` <[number]> maximum time to wait for in milliseconds. Defaults to `30000` (30 seconds). Pass `0` to disable timeout. The default value can be changed by using the [page.setDefaultTimeout(timeout)](#pagesetdefaulttimeouttimeout) method.
@@ -2864,7 +2864,7 @@ await page.waitFor(selector => !!document.querySelector(selector), {}, selector)
 #### frame.waitForFunction(pageFunction[, options[, ...args]])
 - `pageFunction` <[function]|[string]> Function to be evaluated in browser context
 - `options` <[Object]> Optional waiting parameters
-  - `polling` <[string]|[number]> An interval at which the `pageFunction` is executed, defaults to `raf`. If `polling` is a number, then it is treated as an interval in milliseconds at which the function would be executed. If `polling` is a string, then it can be one of the following values:
+  - `polling` <[number]|"raf"|"mutation"> An interval at which the `pageFunction` is executed, defaults to `raf`. If `polling` is a number, then it is treated as an interval in milliseconds at which the function would be executed. If `polling` is a string, then it can be one of the following values:
     - `raf` - to constantly execute `pageFunction` in `requestAnimationFrame` callback. This is the tightest polling mode which is suitable to observe styling changes.
     - `mutation` - to execute `pageFunction` on every DOM mutation.
   - `timeout` <[number]> maximum time to wait for in milliseconds. Defaults to `30000` (30 seconds). Pass `0` to disable timeout. The default value can be changed by using the [page.setDefaultTimeout(timeout)](#pagesetdefaulttimeouttimeout) method.

src/chromium/Page.ts

@@ -579,10 +579,7 @@ export class Page extends EventEmitter {
     return this.mainFrame().waitForXPath(xpath, options);
   }
 
-  waitForFunction(pageFunction: Function, options: {
-      polling?: string | number;
-      timeout?: number; } = {},
-  ...args: any[]): Promise<js.JSHandle> {
+  waitForFunction(pageFunction: Function | string, options: dom.WaitForFunctionOptions, ...args: any[]): Promise<js.JSHandle> {
     return this.mainFrame().waitForFunction(pageFunction, options, ...args);
   }
 }

src/dom.ts

@@ -10,7 +10,6 @@ import * as cssSelectorEngineSource from './generated/cssSelectorEngineSource';
 import * as xpathSelectorEngineSource from './generated/xpathSelectorEngineSource';
 import { assert, helper } from './helper';
 import Injected from './injected/injected';
-import { WaitTaskParams } from './waitTask';
 
 export interface DOMWorldDelegate {
   keyboard: input.Keyboard;
@@ -307,36 +306,57 @@ function selectorToString(selector: Selector): string {
   return `:scope >> ${selector.selector}`;
 }
 
+export type Task = (domWorld: DOMWorld) => Promise<js.JSHandle>;
+
+export type Polling = 'raf' | 'mutation' | number;
+export type WaitForFunctionOptions = { polling?: Polling, timeout?: number };
+
+export function waitForFunctionTask(pageFunction: Function | string, options: WaitForFunctionOptions, ...args: any[]) {
+  const { polling = 'raf' } = options;
+  if (helper.isString(polling))
+    assert(polling === 'raf' || polling === 'mutation', 'Unknown polling option: ' + polling);
+  else if (helper.isNumber(polling))
+    assert(polling > 0, 'Cannot poll with non-positive interval: ' + polling);
+  else
+    throw new Error('Unknown polling options: ' + polling);
+  const predicateBody = helper.isString(pageFunction) ? 'return (' + pageFunction + ')' : 'return (' + pageFunction + ')(...args)';
+
+  return async (domWorld: DOMWorld) => domWorld.context.evaluateHandle((injected: Injected, predicateBody: string, polling: Polling, timeout: number, ...args) => {
+    const predicate = new Function('...args', predicateBody);
+    if (polling === 'raf')
+      return injected.pollRaf(predicate, timeout, ...args);
+    if (polling === 'mutation')
+      return injected.pollMutation(predicate, timeout, ...args);
+    return injected.pollInterval(polling, predicate, timeout, ...args);
+  }, await domWorld.injected(), predicateBody, polling, options.timeout, ...args);
+}
+
 export type WaitForSelectorOptions = { visible?: boolean, hidden?: boolean, timeout?: number };
 
-export function waitForSelectorTask(selector: string, options: WaitForSelectorOptions): WaitTaskParams {
-  const { visible: waitForVisible = false, hidden: waitForHidden = false, timeout } = options;
-  const polling = waitForVisible || waitForHidden ? 'raf' : 'mutation';
-  const title = `selector "${selector}"${waitForHidden ? ' to be hidden' : ''}`;
-  const params: WaitTaskParams = {
-    predicateBody: predicate,
-    title,
-    polling,
-    timeout,
-    args: [normalizeSelector(selector), waitForVisible, waitForHidden],
-    passInjected: true
-  };
-  return params;
-
-  function predicate(injected: Injected, selector: string, waitForVisible: boolean, waitForHidden: boolean): (Node | boolean) | null {
-    const element = injected.querySelector(selector, document);
-    if (!element)
-      return waitForHidden;
-    if (!waitForVisible && !waitForHidden)
-      return element;
-    const style = window.getComputedStyle(element);
-    const isVisible = style && style.visibility !== 'hidden' && hasVisibleBoundingBox();
-    const success = (waitForVisible === isVisible || waitForHidden === !isVisible);
-    return success ? element : null;
-
-    function hasVisibleBoundingBox(): boolean {
-      const rect = element.getBoundingClientRect();
-      return !!(rect.top || rect.bottom || rect.width || rect.height);
+export function waitForSelectorTask(selector: string, options: WaitForSelectorOptions): Task {
+  const { visible: waitForVisible = false, hidden: waitForHidden = false } = options;
+  selector = normalizeSelector(selector);
+
+  return async (domWorld: DOMWorld) => domWorld.context.evaluateHandle((injected: Injected, selector: string, waitForVisible: boolean, waitForHidden: boolean, timeout: number) => {
+    if (waitForVisible || waitForHidden)
+      return injected.pollRaf(predicate, timeout);
+    return injected.pollMutation(predicate, timeout);
+
+    function predicate(): Element | boolean {
+      const element = injected.querySelector(selector, document);
+      if (!element)
+        return waitForHidden;
+      if (!waitForVisible && !waitForHidden)
+        return element;
+      const style = window.getComputedStyle(element);
+      const isVisible = style && style.visibility !== 'hidden' && hasVisibleBoundingBox();
+      const success = (waitForVisible === isVisible || waitForHidden === !isVisible);
+      return success ? element : false;
+
+      function hasVisibleBoundingBox(): boolean {
+        const rect = element.getBoundingClientRect();
+        return !!(rect.top || rect.bottom || rect.width || rect.height);
+      }
     }
-  }
+  }, await domWorld.injected(), selector, waitForVisible, waitForHidden, options.timeout);
 }

src/firefox/Page.ts

@@ -493,7 +493,7 @@ export class Page extends EventEmitter {
     return this._frameManager.mainFrame().waitFor(selectorOrFunctionOrTimeout, options, ...args);
   }
 
-  waitForFunction(pageFunction: Function | string, options: { polling?: string | number; timeout?: number; } | undefined = {}, ...args): Promise<js.JSHandle> {
+  waitForFunction(pageFunction: Function | string, options: dom.WaitForFunctionOptions, ...args): Promise<js.JSHandle> {
     return this._frameManager.mainFrame().waitForFunction(pageFunction, options, ...args);
   }
 

src/frames.ts

@@ -22,8 +22,8 @@ import * as dom from './dom';
 import * as network from './network';
 import { helper, assert } from './helper';
 import { ClickOptions, MultiClickOptions, PointerActionOptions, SelectOption } from './input';
-import { WaitTaskParams, WaitTask } from './waitTask';
 import { TimeoutSettings } from './TimeoutSettings';
+import { TimeoutError } from './Errors';
 
 const readFileAsync = helper.promisify(fs.readFile);
 
@@ -32,7 +32,7 @@ type World = {
   contextPromise: Promise<js.ExecutionContext>;
   contextResolveCallback: (c: js.ExecutionContext) => void;
   context: js.ExecutionContext | null;
-  waitTasks: Set<WaitTask>;
+  rerunnableTasks: Set<RerunnableTask>;
 };
 
 export type NavigateOptions = {
@@ -65,8 +65,8 @@ export class Frame {
     this._timeoutSettings = timeoutSettings;
     this._parentFrame = parentFrame;
 
-    this._worlds.set('main', { contextPromise: new Promise(() => {}), contextResolveCallback: () => {}, context: null, waitTasks: new Set() });
-    this._worlds.set('utility', { contextPromise: new Promise(() => {}), contextResolveCallback: () => {}, context: null, waitTasks: new Set() });
+    this._worlds.set('main', { contextPromise: new Promise(() => {}), contextResolveCallback: () => {}, context: null, rerunnableTasks: new Set() });
+    this._worlds.set('utility', { contextPromise: new Promise(() => {}), contextResolveCallback: () => {}, context: null, rerunnableTasks: new Set() });
     this._setContext('main', null);
     this._setContext('utility', null);
 
@@ -386,8 +386,9 @@ export class Frame {
   }
 
   async waitForSelector(selector: string, options: dom.WaitForSelectorOptions = {}): Promise<dom.ElementHandle | null> {
-    const params = dom.waitForSelectorTask(selector, { timeout: this._timeoutSettings.timeout(), ...options });
-    const handle = await this._scheduleWaitTask(params, 'utility');
+    const task = dom.waitForSelectorTask(selector, { timeout: this._timeoutSettings.timeout(), ...options });
+    const title = `selector "${selector}"${options.hidden ? ' to be hidden' : ''}`;
+    const handle = await this._scheduleRerunnableTask(task, 'utility', options.timeout, title);
     if (!handle.asElement()) {
       await handle.dispose();
       return null;
@@ -404,22 +405,10 @@ export class Frame {
     return this.waitForSelector('xpath=' + xpath, options);
   }
 
-  waitForFunction(
-    pageFunction: Function | string,
-    options: { polling?: string | number; timeout?: number; } = {},
-    ...args): Promise<js.JSHandle> {
-    const {
-      polling = 'raf',
-      timeout = this._timeoutSettings.timeout(),
-    } = options;
-    const params: WaitTaskParams = {
-      predicateBody: pageFunction,
-      title: 'function',
-      polling,
-      timeout,
-      args
-    };
-    return this._scheduleWaitTask(params, 'main');
+  waitForFunction(pageFunction: Function | string, options: dom.WaitForFunctionOptions = {}, ...args: any[]): Promise<js.JSHandle> {
+    options = { timeout: this._timeoutSettings.timeout(), ...options };
+    const task = dom.waitForFunctionTask(pageFunction, options, ...args);
+    return this._scheduleRerunnableTask(task, 'main', options.timeout);
   }
 
   async title(): Promise<string> {
@@ -435,30 +424,31 @@ export class Frame {
   _detach() {
     this._detached = true;
     for (const world of this._worlds.values()) {
-      for (const waitTask of world.waitTasks)
-        waitTask.terminate(new Error('waitForFunction failed: frame got detached.'));
+      for (const rerunnableTask of world.rerunnableTasks)
+        rerunnableTask.terminate(new Error('waitForFunction failed: frame got detached.'));
     }
     if (this._parentFrame)
       this._parentFrame._childFrames.delete(this);
     this._parentFrame = null;
   }
 
-  private _scheduleWaitTask(params: WaitTaskParams, worldType: WorldType): Promise<js.JSHandle> {
+  private _scheduleRerunnableTask(task: dom.Task, worldType: WorldType, timeout?: number, title?: string): Promise<js.JSHandle> {
     const world = this._worlds.get(worldType);
-    const task = new WaitTask(params, () => world.waitTasks.delete(task));
-    world.waitTasks.add(task);
+    const rerunnableTask = new RerunnableTask(world, task, timeout, title);
+    world.rerunnableTasks.add(rerunnableTask);
     if (world.context)
-      task.rerun(world.context);
-    return task.promise;
+      rerunnableTask.rerun(world.context._domWorld);
+    return rerunnableTask.promise;
   }
 
   private _setContext(worldType: WorldType, context: js.ExecutionContext | null) {
     const world = this._worlds.get(worldType);
     world.context = context;
     if (context) {
+      assert(context._domWorld, 'Frame context must have a dom world');
       world.contextResolveCallback.call(null, context);
-      for (const waitTask of world.waitTasks)
-        waitTask.rerun(context);
+      for (const rerunnableTask of world.rerunnableTasks)
+        rerunnableTask.rerun(context._domWorld);
     } else {
       world.contextPromise = new Promise(fulfill => {
         world.contextResolveCallback = fulfill;
@@ -482,3 +472,83 @@ export class Frame {
     }
   }
 }
+
+class RerunnableTask {
+  readonly promise: Promise<js.JSHandle>;
+  private _world: World;
+  private _task: dom.Task;
+  private _runCount: number;
+  private _resolve: (result: js.JSHandle) => void;
+  private _reject: (reason: Error) => void;
+  private _timeoutTimer: NodeJS.Timer;
+  private _terminated: boolean;
+
+  constructor(world: World, task: dom.Task, timeout?: number, title?: string) {
+    this._world = world;
+    this._task = task;
+    this._runCount = 0;
+    this.promise = new Promise<js.JSHandle>((resolve, reject) => {
+      this._resolve = resolve;
+      this._reject = reject;
+    });
+    // Since page navigation requires us to re-install the pageScript, we should track
+    // timeout on our end.
+    if (timeout) {
+      const timeoutError = new TimeoutError(`waiting for ${title || 'function'} failed: timeout ${timeout}ms exceeded`);
+      this._timeoutTimer = setTimeout(() => this.terminate(timeoutError), timeout);
+    }
+  }
+
+  terminate(error: Error) {
+    this._terminated = true;
+    this._reject(error);
+    this._doCleanup();
+  }
+
+  async rerun(domWorld: dom.DOMWorld) {
+    const runCount = ++this._runCount;
+    let success: js.JSHandle | null = null;
+    let error = null;
+    try {
+      success = await this._task(domWorld);
+    } catch (e) {
+      error = e;
+    }
+
+    if (this._terminated || runCount !== this._runCount) {
+      if (success)
+        await success.dispose();
+      return;
+    }
+
+    // Ignore timeouts in pageScript - we track timeouts ourselves.
+    // If execution context has been already destroyed, `context.evaluate` will
+    // throw an error - ignore this predicate run altogether.
+    if (!error && await domWorld.context.evaluate(s => !s, success).catch(e => true)) {
+      await success.dispose();
+      return;
+    }
+
+    // When the page is navigated, the promise is rejected.
+    // We will try again in the new execution context.
+    if (error && error.message.includes('Execution context was destroyed'))
+      return;
+
+    // We could have tried to evaluate in a context which was already
+    // destroyed.
+    if (error && error.message.includes('Cannot find context with specified id'))
+      return;
+
+    if (error)
+      this._reject(error);
+    else
+      this._resolve(success);
+
+    this._doCleanup();
+  }
+
+  _doCleanup() {
+    clearTimeout(this._timeoutTimer);
+    this._world.rerunnableTasks.delete(this);
+  }
+}