Bug 1751954 - [remote] Allow to return from waitForInitialNavigationCompleted when load started. r=webdriver-reviewers,jdescottes

Differential Revision: https://phabricator.services.mozilla.com/D137103

remote/marionette/driver.js

@@ -464,7 +464,7 @@ GeckoDriver.prototype.newSession = async function(cmd) {
       const browsingContext = this.curBrowser.contentBrowser.browsingContext;
       this.currentSession.contentBrowsingContext = browsingContext;
 
-      await waitForInitialNavigationCompleted(browsingContext);
+      await waitForInitialNavigationCompleted(browsingContext.webProgress);
 
       this.curBrowser.contentBrowser.focus();
     }
@@ -2023,7 +2023,9 @@ GeckoDriver.prototype.newWindow = async function(cmd) {
 
   // Actors need the new window to be loaded to safely execute queries.
   // Wait until the initial page load has been finished.
-  await waitForInitialNavigationCompleted(contentBrowser.browsingContext);
+  await waitForInitialNavigationCompleted(
+    contentBrowser.browsingContext.webProgress
+  );
 
   const id = TabManager.getIdForBrowser(contentBrowser);
 

remote/shared/Navigate.jsm

@@ -23,14 +23,26 @@ XPCOMUtils.defineLazyGetter(this, "logger", () =>
 const webProgressListeners = new Set();
 
 /**
- * Wait until the initial load of the given browsing context is done.
+ * Wait until the initial load of the given WebProgress is done.
  *
- * @param {BrowsingContext} browsingContext
- *     The browsing context to check.
+ * @param {WebProgress} webProgress
+ *     The WebProgress instance to observe.
+ * @param {Object=} options
+ * @param {Boolean=} options.resolveWhenStarted
+ *     Flag to indicate that the Promise has to be resolved when the
+ *     page load has been started. Otherwise wait until the page has
+ *     finished loading. Defaults to `false`.
+ *
+ * @returns {Promise}
+ *     Promise which resolves when the page load is in the expected state.
  */
-function waitForInitialNavigationCompleted(browsingContext) {
+function waitForInitialNavigationCompleted(webProgress, options = {}) {
+  const { resolveWhenStarted = false } = options;
+
+  const browsingContext = webProgress.browsingContext;
+
   // Start the listener right away to avoid race conditions.
-  const listener = new ProgressListener(browsingContext.webProgress);
+  const listener = new ProgressListener(webProgress, { resolveWhenStarted });
   const navigated = listener.start();
 
   // Right after a browsing context has been attached it could happen that
@@ -48,33 +60,46 @@ function waitForInitialNavigationCompleted(browsingContext) {
       truncate`[${browsingContext.id}] Document already finished loading: ${browsingContext.currentURI?.spec}`
     );
 
+    // Will resolve the navigated promise.
     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.
+ * WebProgressListener to observe for page loads.
  */
 class ProgressListener {
-  #resolve;
+  #resolveWhenStarted;
   #unloadTimeout;
+
+  #resolve;
+  #seenStartFlag;
   #unloadTimer;
   #webProgress;
 
+  /**
+   * Create a new WebProgressListener instance.
+   *
+   * @param {WebProgress} webProgress
+   *     The web progress to attach the listener to.
+   * @param {Object=} options
+   * @param {Boolean=} options.resolveWhenStarted
+   *     Flag to indicate that the Promise has to be resolved when the
+   *     page load has been started. Otherwise wait until the page has
+   *     finished loading. Defaults to `false`.
+   * @param {Number=} options.unloadTimeout
+   *     Time to allow before the page gets unloaded. Defaults to 200ms.
+   */
   constructor(webProgress, options = {}) {
-    const { unloadTimeout = 200 } = options;
+    const { resolveWhenStarted = false, unloadTimeout = 200 } = options;
+
+    this.#resolveWhenStarted = resolveWhenStarted;
+    this.#unloadTimeout = unloadTimeout;
 
     this.#resolve = null;
-    this.#unloadTimeout = unloadTimeout;
+    this.#seenStartFlag = false;
     this.#unloadTimer = null;
     this.#webProgress = webProgress;
   }
@@ -91,27 +116,44 @@ class ProgressListener {
     return this.#webProgress.isLoadingDocument;
   }
 
-  onStateChange(progress, request, flag, status) {
-    const isStart = flag & Ci.nsIWebProgressListener.STATE_START;
-    const isStop = flag & Ci.nsIWebProgressListener.STATE_STOP;
+  #checkLoadingState(options = {}) {
+    const { isStart = false, isStop = false } = options;
+
+    if (isStart && !this.#seenStartFlag) {
+      this.#seenStartFlag = true;
 
-    if (isStart) {
-      this.#unloadTimer?.cancel();
       logger.trace(
-        truncate`[${this.browsingContext.id}] Web progress state=start: ${this.currentURI?.spec}`
+        truncate`[${this.browsingContext.id}] ${this.constructor.name} state=start: ${this.currentURI.spec}`
       );
+
+      if (this.#unloadTimer) {
+        this.#unloadTimer.cancel();
+        this.#unloadTimer = null;
+      }
+
+      if (this.#resolveWhenStarted) {
+        // Return immediately when the load should not be awaited.
+        this.stop();
+        return;
+      }
     }
 
-    if (isStop) {
+    if (isStop && this.#seenStartFlag) {
       logger.trace(
-        truncate`this.browsingContext.id}] Web progress state=stop: ${this.currentURI?.spec}`
+        truncate`[${this.browsingContext.id}] ${this.constructor.name} state=stop: ${this.currentURI.spec}`
       );
 
-      this.#resolve();
       this.stop();
     }
   }
 
+  onStateChange(progress, request, flag, status) {
+    this.#checkLoadingState({
+      isStart: flag & Ci.nsIWebProgressListener.STATE_START,
+      isStop: flag & Ci.nsIWebProgressListener.STATE_STOP,
+    });
+  }
+
   /**
    * Start observing web progress changes.
    *
@@ -123,17 +165,26 @@ class ProgressListener {
       throw new Error(`Progress listener already started`);
     }
 
+    if (this.#webProgress.isLoadingDocument && this.#resolveWhenStarted) {
+      // Resolve immediately when the page is already loading and there
+      // is no requirement to wait for it to finish.
+      return Promise.resolve();
+    }
+
     const promise = new Promise(resolve => (this.#resolve = resolve));
 
+    // Enable all state notifications to get informed about an upcoming load
+    // as early as possible.
     this.#webProgress.addProgressListener(
       this,
-      Ci.nsIWebProgress.NOTIFY_STATE_WINDOW |
-        Ci.nsIWebProgress.NOTIFY_STATE_DOCUMENT
+      Ci.nsIWebProgress.NOTIFY_STATE_ALL
     );
 
     webProgressListeners.add(this);
 
-    if (!this.#webProgress.isLoadingDocument) {
+    if (this.#webProgress.isLoadingDocument) {
+      this.#checkLoadingState({ isStart: true });
+    } else {
       // If the document is not loading yet wait some time for the navigation
       // to be started.
       this.#unloadTimer = Cc["@mozilla.org/timer;1"].createInstance(
@@ -144,7 +195,6 @@ class ProgressListener {
           logger.trace(
             truncate`[${this.browsingContext.id}] No navigation detected: ${this.currentURI?.spec}`
           );
-          this.#resolve();
           this.stop();
         },
         this.#unloadTimeout,
@@ -166,17 +216,17 @@ class ProgressListener {
     this.#unloadTimer?.cancel();
     this.#unloadTimer = null;
 
+    this.#resolve();
     this.#resolve = null;
 
     this.#webProgress.removeProgressListener(
       this,
-      Ci.nsIWebProgress.NOTIFY_STATE_WINDOW |
-        Ci.nsIWebProgress.NOTIFY_STATE_DOCUMENT
+      Ci.nsIWebProgress.NOTIFY_STATE_ALL
     );
     webProgressListeners.delete(this);
   }
 
-  get toString() {
+  toString() {
     return `[object ${this.constructor.name}]`;
   }
 

remote/shared/test/xpcshell/test_Navigate.js

@@ -2,6 +2,8 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this file,
  * You can obtain one at http://mozilla.org/MPL/2.0/. */
 
+const { Services } = ChromeUtils.import("resource://gre/modules/Services.jsm");
+
 const { waitForInitialNavigationCompleted } = ChromeUtils.import(
   "chrome://remote/content/shared/Navigate.jsm"
 );
@@ -67,6 +69,7 @@ class MockWebProgress {
 
 class MockTopContext {
   constructor(webProgress = null) {
+    this.currentURI = Services.io.newURI("http://foo.bar");
     this.currentWindowGlobal = { isInitialDocument: true };
     this.id = 7;
     this.top = this;
@@ -77,7 +80,7 @@ class MockTopContext {
 add_test(
   async function test_waitForInitialNavigation_initialDocumentFinishedLoading() {
     const browsingContext = new MockTopContext();
-    await waitForInitialNavigationCompleted(browsingContext);
+    await waitForInitialNavigationCompleted(browsingContext.webProgress);
 
     ok(
       !browsingContext.webProgress.isLoadingDocument,
@@ -97,7 +100,9 @@ add_test(
     const browsingContext = new MockTopContext();
     browsingContext.webProgress.sendStartState({ isInitial: true });
 
-    const completed = waitForInitialNavigationCompleted(browsingContext);
+    const completed = waitForInitialNavigationCompleted(
+      browsingContext.webProgress
+    );
     browsingContext.webProgress.sendStopState();
     await completed;
 
@@ -119,7 +124,9 @@ add_test(
     const browsingContext = new MockTopContext();
     delete browsingContext.currentWindowGlobal;
 
-    const completed = waitForInitialNavigationCompleted(browsingContext);
+    const completed = waitForInitialNavigationCompleted(
+      browsingContext.webProgress
+    );
     browsingContext.webProgress.sendStartState({ isInitial: true });
     browsingContext.webProgress.sendStopState();
     await completed;
@@ -143,7 +150,7 @@ add_test(
     browsingContext.webProgress.sendStartState({ isInitial: false });
     browsingContext.webProgress.sendStopState();
 
-    await waitForInitialNavigationCompleted(browsingContext);
+    await waitForInitialNavigationCompleted(browsingContext.webProgress);
 
     ok(
       !browsingContext.webProgress.isLoadingDocument,
@@ -163,7 +170,9 @@ add_test(
     const browsingContext = new MockTopContext();
     browsingContext.webProgress.sendStartState({ isInitial: false });
 
-    const completed = waitForInitialNavigationCompleted(browsingContext);
+    const completed = waitForInitialNavigationCompleted(
+      browsingContext.webProgress
+    );
     browsingContext.webProgress.sendStopState();
     await completed;
 
@@ -180,6 +189,30 @@ add_test(
   }
 );
 
+add_test(async function test_waitForInitialNavigation_resolveWhenStarted() {
+  const browsingContext = new MockTopContext();
+  browsingContext.webProgress.sendStartState();
+
+  const completed = waitForInitialNavigationCompleted(
+    browsingContext.webProgress,
+    {
+      resolveWhenStarted: true,
+    }
+  );
+  await completed;
+
+  ok(
+    browsingContext.webProgress.isLoadingDocument,
+    "Document is still loading"
+  );
+  ok(
+    !browsingContext.currentWindowGlobal.isInitialDocument,
+    "Is not initial document"
+  );
+
+  run_next_test();
+});
+
 add_test(
   async function test_waitForInitialNavigation_crossOriginAlreadyLoading() {
     const browsingContext = new MockTopContext();
@@ -187,7 +220,9 @@ add_test(
 
     browsingContext.webProgress.sendStartState({ coop: true });
 
-    const completed = waitForInitialNavigationCompleted(browsingContext);
+    const completed = waitForInitialNavigationCompleted(
+      browsingContext.webProgress
+    );
     browsingContext.webProgress.sendStopState();
     await completed;