Bug 1747359 - [marionette] Improve checks for tab's initial page load…

…. r=webdriver-reviewers,jdescottes Differential Revision: https://phabricator.services.mozilla.com/D135594
dothq · Jan 17, 2022 · 65e979f · 65e979f
1 parent d1b5c1b
commit 65e979f
Show file tree

Hide file tree

Showing 2 changed files with 305 additions and 99 deletions.
diff --git a/remote/shared/Navigate.jsm b/remote/shared/Navigate.jsm
@@ -28,51 +28,153 @@ const webProgressListeners = new Set();
  *     The browsing context to check.
  */
 function waitForInitialNavigationCompleted(browsingContext) {
-  // Keep a reference to the webProgress instance in case the browsing context
-  // gets replaced by a cross-group navigation.
-  const webProgress = browsingContext.webProgress;
-
-  let listener;
-  return new Promise(resolve => {
-    listener = new ProgressListener(resolve);
-    // Keep a reference to the weak listener so it's not getting gc'ed.
-    webProgressListeners.add(listener);
-
-    // Monitor the webprogress listener before checking isLoadingDocument to
-    // avoid race conditions.
-    webProgress.addProgressListener(
-      listener,
-      Ci.nsIWebProgress.NOTIFY_STATE_WINDOW |
-        Ci.nsIWebProgress.NOTIFY_STATE_DOCUMENT
-    );
+  // Start the listener right away to avoid race conditions.
+  const listener = new ProgressListener(browsingContext.webProgress);
+  const navigated = listener.start();
 
-    // Right after its creation a browsing context doesn't have a window global.
-    const isInitial = !browsingContext.currentWindowGlobal;
+  // Right after a browsing context has been attached it could happen that
+  // no window global has been set yet. Consider this as nothing has been
+  // loaded yet.
+  let isInitial = true;
+  if (browsingContext.currentWindowGlobal) {
+    isInitial = browsingContext.currentWindowGlobal.isInitialDocument;
+  }
 
-    if (!webProgress.isLoadingDocument && !isInitial) {
-      logger.trace("Initial navigation already completed");
-      resolve();
-    }
-  }).finally(() => {
-    webProgress.removeProgressListener(
-      listener,
-      Ci.nsIWebProgress.NOTIFY_STATE_WINDOW |
-        Ci.nsIWebProgress.NOTIFY_STATE_DOCUMENT
+  // If the current document is not the initial "about:blank" and is also
+  // no longer loading, assume the navigation is done and return.
+  if (!isInitial && !listener.isLoadingDocument) {
+    logger.trace(
+      `[${browsingContext.id}] Document already finished loading: ${browsingContext.currentURI?.spec}`
     );
-    webProgressListeners.delete(listener);
-  });
+
+    listener.stop();
+    return Promise.resolve();
+  }
+
+  return navigated;
 }
 
+/**
+ * WebProgressListener to observe page loads.
+ *
+ * @param {WebProgress} webProgress
+ *     The web progress to attach the listener to.
+ * @param {Object=} options
+ * @param {Number=} options.unloadTimeout
+ *     Time to allow before the page gets unloaded. Defaults to 200ms.
+ */
 class ProgressListener {
-  constructor(resolve) {
-    this.resolve = resolve;
+  #resolve;
+  #unloadTimeout;
+  #unloadTimer;
+  #webProgress;
+
+  constructor(webProgress, options = {}) {
+    const { unloadTimeout = 200 } = options;
+
+    this.#resolve = null;
+    this.#unloadTimeout = unloadTimeout;
+    this.#unloadTimer = null;
+    this.#webProgress = webProgress;
+  }
+
+  get browsingContext() {
+    return this.#webProgress.browsingContext;
+  }
+
+  get currentURI() {
+    return this.#webProgress.browsingContext.currentURI;
+  }
+
+  get isLoadingDocument() {
+    return this.#webProgress.isLoadingDocument;
   }
 
   onStateChange(progress, request, flag, status) {
+    const isStart = flag & Ci.nsIWebProgressListener.STATE_START;
     const isStop = flag & Ci.nsIWebProgressListener.STATE_STOP;
+
+    if (isStart) {
+      this.#unloadTimer?.cancel();
+      logger.trace(
+        `[${this.browsingContext.id}] Web progress state=start: ${this.currentURI?.spec}`
+      );
+    }
+
     if (isStop) {
-      this.resolve();
+      logger.trace(
+        `[${this.browsingContext.id}] Web progress state=stop: ${this.currentURI?.spec}`
+      );
+
+      this.#resolve();
+      this.stop();
+    }
+  }
+
+  /**
+   * Start observing web progress changes.
+   *
+   * @returns {Promise}
+   *     A promise that will resolve when the navigation has been finished.
+   */
+  start() {
+    if (this.#resolve) {
+      throw new Error(`Progress listener already started`);
+    }
+
+    const promise = new Promise(resolve => (this.#resolve = resolve));
+
+    this.#webProgress.addProgressListener(
+      this,
+      Ci.nsIWebProgress.NOTIFY_STATE_WINDOW |
+        Ci.nsIWebProgress.NOTIFY_STATE_DOCUMENT
+    );
+
+    webProgressListeners.add(this);
+
+    if (!this.#webProgress.isLoadingDocument) {
+      // If the document is not loading yet wait some time for the navigation
+      // to be started.
+      this.#unloadTimer = Cc["@mozilla.org/timer;1"].createInstance(
+        Ci.nsITimer
+      );
+      this.#unloadTimer.initWithCallback(
+        () => {
+          logger.trace("No initial navigation detected");
+          this.#resolve();
+          this.stop();
+        },
+        this.#unloadTimeout,
+        Ci.nsITimer.TYPE_ONE_SHOT
+      );
     }
+
+    return promise;
+  }
+
+  /**
+   * Stop observing web progress changes.
+   */
+  stop() {
+    if (!this.#resolve) {
+      throw new Error(`Progress listener not yet started`);
+    }
+
+    this.#unloadTimer?.cancel();
+    this.#unloadTimer = null;
+
+    this.#resolve = null;
+
+    this.#webProgress.removeProgressListener(
+      this,
+      Ci.nsIWebProgress.NOTIFY_STATE_WINDOW |
+        Ci.nsIWebProgress.NOTIFY_STATE_DOCUMENT
+    );
+    webProgressListeners.delete(this);
+  }
+
+  get toString() {
+    return `[object ${this.constructor.name}]`;
   }
 
   get QueryInterface() {