Refactor Frame.waitForSelector method

Refactor Frame.waitForSelector to make room for Frame.waitForFunction
implementation.
This patch:
- removes AwaitedElement class which proved to be confusing, and
  introduces a more straight-forward WaitTask.
- refactors the mutation observer to return true in case of successful
  waiting or false in case of timeout.

References #91

lib/FrameManager.js

@@ -172,93 +172,6 @@ class FrameManager extends EventEmitter {
       throw new Error('Evaluation failed: ' + helper.getExceptionMessage(exceptionDetails));
     return await helper.serializeRemoteObject(this._client, remoteObject);
   }
-
-  /**
-   * @param {!Frame} frame
-   * @param {string} selector
-   * @param {boolean} waitForVisible
-   * @param {number} timeout
-   * @return {!Promise<undefined>}
-   */
-  async _waitForSelector(frame, selector, waitForVisible, timeout) {
-    let contextId = undefined;
-    if (!frame.isMainFrame()) {
-      contextId = this._frameIdToExecutionContextId.get(frame._id);
-      console.assert(contextId, 'Frame does not have default context to evaluate in!');
-    }
-    let { exceptionDetails } = await this._client.send('Runtime.evaluate', {
-      expression: helper.evaluationString(inPageWatchdog, selector, waitForVisible, timeout),
-      contextId,
-      awaitPromise: true,
-      returnByValue: false,
-    });
-    if (exceptionDetails)
-      throw new Error('Evaluation failed: ' + helper.getExceptionMessage(exceptionDetails));
-
-    /**
-     * @param {string} selector
-     * @param {boolean} waitForVisible
-     * @param {number} timeout
-     * @return {!Promise}
-     */
-    async function inPageWatchdog(selector, visible, timeout) {
-      const resultPromise = visible ? waitForVisible(selector) : waitInDOM(selector);
-      const timeoutPromise = new Promise((resolve, reject) => {
-        setTimeout(reject.bind(null, new Error(`waitForSelector failed: timeout ${timeout}ms exceeded.`)), timeout);
-      });
-      await Promise.race([resultPromise, timeoutPromise]);
-
-      /**
-       * @param {string} selector
-       * @return {!Promise<!Element>}
-       */
-      function waitInDOM(selector) {
-        let node = document.querySelector(selector);
-        if (node)
-          return Promise.resolve(node);
-
-        let fulfill;
-        const result = new Promise(x => fulfill = x);
-        const observer = new MutationObserver(mutations => {
-          const node = document.querySelector(selector);
-          if (node) {
-            observer.disconnect();
-            fulfill(node);
-          }
-        });
-        observer.observe(document, {
-          childList: true,
-          subtree: true
-        });
-        return result;
-      }
-
-      /**
-       * @param {string} selector
-       * @return {!Promise<!Element>}
-       */
-      async function waitForVisible(selector) {
-        let fulfill;
-        const result = new Promise(x => fulfill = x);
-        onRaf();
-        return result;
-
-        async function onRaf() {
-          const node = await waitInDOM(selector);
-          if (!node) {
-            fulfill(null);
-            return;
-          }
-          const style = window.getComputedStyle(node);
-          if (style.display === 'none' || style.visibility === 'hidden') {
-            requestAnimationFrame(onRaf);
-            return;
-          }
-          fulfill(node);
-        }
-      }
-    }
-  }
 }
 
 /** @enum {string} */
@@ -283,8 +196,8 @@ class Frame {
     this._parentFrame = parentFrame;
     this._url = '';
     this._id = frameId;
-    /** @type {!Set<!AwaitedElement>} */
-    this._awaitedElements = new Set();
+    /** @type {!Set<!WaitTask>} */
+    this._waitTasks = new Set();
 
     this._adoptPayload(payload);
 
@@ -365,15 +278,14 @@ class Frame {
    */
   waitForSelector(selector, options = {}) {
     const timeout = options.timeout || 30000;
-    const awaitedElement = new AwaitedElement(() => this._frameManager._waitForSelector(this, selector, !!options.visible, timeout));
-    // Since navigation will re-install page watchdogs, we should timeout on our
-    // end as well.
-    setTimeout(() => awaitedElement.terminate(new Error(`waitForSelector failed: timeout ${timeout}ms exceeded`)), timeout);
+    const waitForVisible = !!options.visible;
+    const pageScript = helper.evaluationString(waitForSelectorPageFunction, selector, waitForVisible, timeout);
+    const waitTask = new WaitTask(this._frameManager, this, pageScript, timeout);
 
-    this._awaitedElements.add(awaitedElement);
-    let cleanup = () => this._awaitedElements.delete(awaitedElement);
-    awaitedElement.promise.then(cleanup, cleanup);
-    return awaitedElement.promise;
+    this._waitTasks.add(waitTask);
+    let cleanup = () => this._waitTasks.delete(waitTask);
+    waitTask.promise.then(cleanup, cleanup);
+    return waitTask.promise;
   }
 
   /**
@@ -451,13 +363,13 @@ class Frame {
     };
     this._name = framePayload.name;
     this._url = framePayload.url;
-    for (let awaitedElement of this._awaitedElements)
-      awaitedElement.startWaiting();
+    for (let waitTask of this._waitTasks)
+      waitTask.run();
   }
 
   _detach() {
-    for (let awaitedElement of this._awaitedElements)
-      awaitedElement.terminate(new Error('waitForSelector failed: frame got detached.'));
+    for (let waitTask of this._waitTasks)
+      waitTask.terminate(new Error('waitForSelector failed: frame got detached.'));
     this._detached = true;
     if (this._parentFrame)
       this._parentFrame._childFrames.delete(this);
@@ -466,18 +378,26 @@ class Frame {
 }
 helper.tracePublicAPI(Frame);
 
-class AwaitedElement {
+class WaitTask {
   /**
-   * @param {function():!Promise} waitInPageCallback
+   * @param {!FrameManager} frameManager
+   * @param {!Frame} frame
+   * @param {string} pageScript
+   * @param {number} timeout
    */
-  constructor(waitInPageCallback) {
+  constructor(frameManager, frame, pageScript, timeout) {
+    this._frameManager = frameManager;
+    this._frame = frame;
+    this._pageScript = pageScript;
+    this._runningTask = null;
     this.promise = new Promise((resolve, reject) => {
       this._resolve = resolve;
       this._reject = reject;
     });
-    this._waitInPageCallback = waitInPageCallback;
-    this._waitPromise = null;
-    this.startWaiting();
+    // Since page navigation requires us to re-install the pageScript, we should track
+    // timeout on our end.
+    this._timeoutTimer = setTimeout(() => this.terminate(new Error(`waiting failed: timeout ${timeout}ms exceeded`)), timeout);
+    this.run();
   }
 
   /**
@@ -485,23 +405,94 @@ class AwaitedElement {
    */
   terminate(error) {
     this._reject(error);
-    this._waitTaskPromise = null;
+    this._cleanup();
   }
 
-  startWaiting() {
-    let waitPromise = this._waitInPageCallback.call(null).then(finish.bind(this), finish.bind(this));
-    this._waitPromise = waitPromise;
+  run() {
+    let runningTask = this._frameManager._evaluateOnFrame(this._frame, this._pageScript).then(finish.bind(this), finish.bind(this, false));
+    this._runningTask = runningTask;
 
     /**
-     * @param {?Error} error
+     * @param {boolean} success
+     * @param {?Error=} error
      */
-    function finish(error) {
-      if (this._waitPromise !== waitPromise)
+    function finish(success, error) {
+      if (runningTask !== this._runningTask)
+        return;
+      // Ignore timeouts in pageScript - we track timeouts ourselves.
+      if (!success && !error)
         return;
       if (error)
         this._reject(error);
       else
         this._resolve();
+      this._cleanup();
+    }
+  }
+
+  _cleanup() {
+    clearTimeout(this._timeoutTimer);
+    this._runningTask = null;
+  }
+}
+
+/**
+ * @param {string} selector
+ * @param {boolean} waitForVisible
+ * @param {number} timeout
+ * @return {!Promise<boolean>}
+ */
+function waitForSelectorPageFunction(selector, visible, timeout) {
+  const resultPromise = visible ? waitForVisible(selector) : waitInDOM(selector);
+  const timeoutPromise = new Promise(fulfill => setTimeout(fulfill, timeout));
+  return Promise.race([
+    resultPromise.then(() => true),
+    timeoutPromise.then(() => false)
+  ]);
+
+  /**
+   * @param {string} selector
+   * @return {!Promise<!Element>}
+   */
+  function waitInDOM(selector) {
+    let node = document.querySelector(selector);
+    if (node)
+      return Promise.resolve(node);
+
+    let fulfill;
+    const result = new Promise(x => fulfill = x);
+    const observer = new MutationObserver(mutations => {
+      const node = document.querySelector(selector);
+      if (node) {
+        observer.disconnect();
+        fulfill(node);
+      }
+    });
+    observer.observe(document, {
+      childList: true,
+      subtree: true
+    });
+    return result;
+  }
+
+  /**
+   * @param {string} selector
+   * @return {!Promise<!Element>}
+   */
+  async function waitForVisible(selector) {
+    let fulfill;
+    const result = new Promise(x => fulfill = x);
+    onRaf();
+    return result;
+
+    async function onRaf() {
+      const node = await waitInDOM(selector);
+      const style = window.getComputedStyle(node);
+      if (style.display === 'none' || style.visibility === 'hidden') {
+        requestAnimationFrame(onRaf);
+        return;
+      }
+      fulfill(node);
     }
   }
 }

test/test.js

@@ -284,7 +284,7 @@ describe('Puppeteer', function() {
       let error = null;
       await page.waitForSelector('div', {timeout: 10}).catch(e => error = e);
       expect(error).toBeTruthy();
-      expect(error.message).toContain('waitForSelector failed: timeout');
+      expect(error.message).toContain('waiting failed: timeout');
     }));
   });
 

utils/doclint/lint.js

@@ -3,7 +3,6 @@ const mdBuilder = require('./MDBuilder');
 const Documentation = require('./Documentation');
 
 const EXCLUDE_CLASSES = new Set([
-  'AwaitedElement',
   'Connection',
   'EmulationManager',
   'FrameManager',
@@ -12,6 +11,7 @@ const EXCLUDE_CLASSES = new Set([
   'NetworkManager',
   'ProxyStream',
   'TaskQueue',
+  'WaitTask',
 ]);
 
 const EXCLUDE_METHODS = new Set([