Bug 1702055 - Replace CallRestoreTabContentComplete with a Promise, r=nika

This should help avoid message ordering issues in the next patches. To ensure
that we're not firing events too early, we now do a "full" restore even if we
don't have form or scroll data, including for about:blank URIs.

This also moves all restore state on the CanonicalBrowsingContext into a
separate `mRestoreState` struct.

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

browser/components/sessionstore/SessionStore.jsm

@@ -519,10 +519,6 @@ var SessionStore = {
   finishTabRemotenessChange(aTab, aSwitchId) {
     SessionStoreInternal.finishTabRemotenessChange(aTab, aSwitchId);
   },
-
-  restoreTabContentComplete(aBrowser, aData) {
-    SessionStoreInternal._restoreTabContentComplete(aBrowser, aData);
-  },
 };
 
 // Freeze the SessionStore object. We don't want anyone to modify it.
@@ -5620,7 +5616,7 @@ var SessionStoreInternal = {
           .get(browser.permanentKey)
           .uninstall();
       }
-      SessionStoreUtils.setRestoreData(browser.browsingContext, null);
+      browser.browsingContext.clearRestoreState();
     }
 
     // Keep the tab's previous state for later in this method
@@ -5911,41 +5907,31 @@ var SessionStoreInternal = {
     this._restoreTabContentStarted(browser, restoreData);
 
     let { tabData } = restoreData;
-    let uri = null;
-    let loadFlags = null;
+    let uri = "about:blank";
+    let loadFlags = Ci.nsIWebNavigation.LOAD_FLAGS_BYPASS_HISTORY;
 
     if (tabData?.userTypedValue && tabData?.userTypedClear) {
       uri = tabData.userTypedValue;
       loadFlags = Ci.nsIWebNavigation.LOAD_FLAGS_ALLOW_THIRD_PARTY_FIXUP;
     } else if (tabData?.entries.length) {
-      uri = tabData.entries[tabData.index - 1].url;
-      let willRestoreContent = SessionStoreUtils.setRestoreData(
+      let promise = SessionStoreUtils.initializeRestore(
         browser.browsingContext,
         this.buildRestoreData(tabData.formdata, tabData.scroll)
       );
-      // We'll manually call RestoreTabContentComplete when the restore is done,
-      // so we only want to create the listener below if we're not restoring tab
-      // content.
-      if (willRestoreContent) {
+      promise.then(() => {
+        if (TAB_STATE_FOR_BROWSER.get(browser) === TAB_STATE_RESTORING) {
+          this._restoreTabContentComplete(browser, restoreData);
+        }
+      });
       return;
     }
-    } else {
-      uri = "about:blank";
-      loadFlags = Ci.nsIWebNavigation.LOAD_FLAGS_BYPASS_HISTORY;
-    }
 
-    if (uri) {
     this.addProgressListenerForRestore(browser, {
       onStopRequest: (request, listener) => {
         let requestURI = request.QueryInterface(Ci.nsIChannel)?.originalURI;
         // FIXME: We sometimes see spurious STATE_STOP events for about:blank
         // URIs, so we have to manually drop those here (unless we're actually
         // expecting an about:blank load).
-          //
-          // In the case where we're firing _restoreTabContentComplete due to
-          // a normal load (i.e. !willRestoreContent), we could perhaps just not
-          // wait for the load here, and instead fix tests that depend on this
-          // behavior.
         if (requestURI?.spec !== "about:blank" || uri === "about:blank") {
           listener.uninstall();
           this._restoreTabContentComplete(browser, restoreData);
@@ -5953,15 +5939,10 @@ var SessionStoreInternal = {
       },
     });
 
-      if (loadFlags) {
     browser.browsingContext.loadURI(uri, {
       loadFlags,
       triggeringPrincipal: Services.scriptSecurityManager.getSystemPrincipal(),
     });
-      } else {
-        browser.browsingContext.sessionHistory.reloadCurrentEntry();
-      }
-    }
   },
 
   _sendRestoreTabContent(browser, options) {

docshell/base/CanonicalBrowsingContext.cpp

@@ -182,14 +182,6 @@ void CanonicalBrowsingContext::ReplacedBy(
   }
   aNewContext->mWebProgress = std::move(mWebProgress);
 
-  bool hasRestoreData = !!mRestoreData && !mRestoreData->IsEmpty();
-  aNewContext->mRestoreData = std::move(mRestoreData);
-  aNewContext->mRequestedContentRestores = mRequestedContentRestores;
-  aNewContext->mCompletedContentRestores = mCompletedContentRestores;
-  SetRestoreData(nullptr);
-  mRequestedContentRestores = 0;
-  mCompletedContentRestores = 0;
-
   // Use the Transaction for the fields which need to be updated whether or not
   // the new context has been attached before.
   // SetWithoutSyncing can be used if context hasn't been attached.
@@ -197,13 +189,16 @@ void CanonicalBrowsingContext::ReplacedBy(
   txn.SetBrowserId(GetBrowserId());
   txn.SetHistoryID(GetHistoryID());
   txn.SetExplicitActive(GetExplicitActive());
-  txn.SetHasRestoreData(hasRestoreData);
+  txn.SetHasRestoreData(GetHasRestoreData());
   if (aNewContext->EverAttached()) {
     MOZ_ALWAYS_SUCCEEDS(txn.Commit(aNewContext));
   } else {
     txn.CommitWithoutSyncing(aNewContext);
   }
 
+  aNewContext->mRestoreState = mRestoreState.forget();
+  MOZ_ALWAYS_SUCCEEDS(SetHasRestoreData(false));
+
   // XXXBFCache name handling is still a bit broken in Fission in general,
   // at least in case name should be cleared.
   if (aRemotenessOptions.mTryUseBFCache) {
@@ -1913,49 +1908,103 @@ void CanonicalBrowsingContext::ResetScalingZoom() {
   }
 }
 
-void CanonicalBrowsingContext::SetRestoreData(SessionStoreRestoreData* aData) {
-  MOZ_DIAGNOSTIC_ASSERT(!mRestoreData || !aData,
-                        "must either be clearing or initializing");
-  MOZ_DIAGNOSTIC_ASSERT(
-      !aData || mCompletedContentRestores == mRequestedContentRestores,
-      "must not start restore in an unstable state");
-  mRestoreData = aData;
-  MOZ_ALWAYS_SUCCEEDS(SetHasRestoreData(mRestoreData));
+void CanonicalBrowsingContext::SetRestoreData(SessionStoreRestoreData* aData,
+                                              ErrorResult& aError) {
+  MOZ_DIAGNOSTIC_ASSERT(aData);
+
+  nsCOMPtr<nsIGlobalObject> global = do_QueryInterface(GetParentObject());
+  RefPtr<Promise> promise = Promise::Create(global, aError);
+  if (aError.Failed()) {
+    return;
+  }
+
+  if (NS_WARN_IF(NS_FAILED(SetHasRestoreData(true)))) {
+    aError.Throw(NS_ERROR_DOM_INVALID_STATE_ERR);
+    return;
+  }
+
+  mRestoreState = new RestoreState();
+  mRestoreState->mData = aData;
+  mRestoreState->mPromise = promise;
+}
+
+already_AddRefed<Promise> CanonicalBrowsingContext::GetRestorePromise() {
+  if (mRestoreState) {
+    return do_AddRef(mRestoreState->mPromise);
+  }
+  return nullptr;
+}
+
+void CanonicalBrowsingContext::ClearRestoreState() {
+  if (!mRestoreState) {
+    MOZ_DIAGNOSTIC_ASSERT(!GetHasRestoreData());
+    return;
+  }
+  if (mRestoreState->mPromise) {
+    mRestoreState->mPromise->MaybeRejectWithUndefined();
+    mRestoreState->mPromise = nullptr;
+  }
+  mRestoreState = nullptr;
+  MOZ_ALWAYS_SUCCEEDS(SetHasRestoreData(false));
+}
+
+void CanonicalBrowsingContext::ResolveAndClearRestoreState(
+    RestoreState* aState) {
+  MOZ_DIAGNOSTIC_ASSERT(aState->mPromise);
+  aState->mPromise->MaybeResolveWithUndefined();
+  aState->mPromise = nullptr;
+  // We may have started a new restore before the previous one had resolved.
+  // This ensures that we aren't clearing state for the new restore.
+  if (aState == mRestoreState) {
+    ClearRestoreState();
+  }
 }
 
 void CanonicalBrowsingContext::RequestRestoreTabContent(
     WindowGlobalParent* aWindow) {
   MOZ_DIAGNOSTIC_ASSERT(IsTop());
 
-  if (IsDiscarded() || !mRestoreData || mRestoreData->IsEmpty()) {
+  if (IsDiscarded() || !mRestoreState || !mRestoreState->mData) {
     return;
   }
 
   CanonicalBrowsingContext* context = aWindow->GetBrowsingContext();
   MOZ_DIAGNOSTIC_ASSERT(!context->IsDiscarded());
 
-  RefPtr<SessionStoreRestoreData> data = mRestoreData->FindChild(context);
+  RefPtr<SessionStoreRestoreData> data =
+      mRestoreState->mData->FindDataForChild(context);
 
-  // We'll only arrive here for a toplevel context after we've already sent
-  // down data for any out-of-process descendants, so it's fine to clear our
-  // data now.
   if (context->IsTop()) {
     MOZ_DIAGNOSTIC_ASSERT(context == this);
-    SetRestoreData(nullptr);
+
+    // We need to wait until the appropriate load event has fired before we
+    // can "complete" the restore process, so if we're holding an empty data
+    // object, just resolve the promise immediately.
+    if (mRestoreState->mData->IsEmpty()) {
+      MOZ_DIAGNOSTIC_ASSERT(!data || data->IsEmpty());
+      ResolveAndClearRestoreState(mRestoreState);
+      return;
+    }
+
+    // Since we're following load event order, we'll only arrive here for a
+    // toplevel context after we've already sent down data for all child frames,
+    // so it's safe to clear this reference now. The completion callback below
+    // relies on the mData field being null to determine if all requests have
+    // been sent out.
+    mRestoreState->ClearData();
+    MOZ_ALWAYS_SUCCEEDS(SetHasRestoreData(false));
   }
 
   if (data && !data->IsEmpty()) {
-    auto onTabRestoreComplete = [self = RefPtr{this}](auto) {
-      self->mCompletedContentRestores++;
-      if (!self->mRestoreData &&
-          self->mCompletedContentRestores == self->mRequestedContentRestores) {
-        if (Element* browser = self->GetEmbedderElement()) {
-          SessionStoreUtils::CallRestoreTabContentComplete(browser);
-        }
+    auto onTabRestoreComplete = [self = RefPtr{this},
+                                 state = RefPtr{mRestoreState}](auto) {
+      state->mResolves++;
+      if (!state->mData && state->mRequests == state->mResolves) {
+        self->ResolveAndClearRestoreState(state);
       }
     };
 
-    mRequestedContentRestores++;
+    mRestoreState->mRequests++;
 
     if (data->CanRestoreInto(aWindow->GetDocumentURI())) {
       if (!aWindow->IsInProcess()) {

docshell/base/CanonicalBrowsingContext.h

@@ -10,6 +10,7 @@
 #include "mozilla/dom/BrowsingContext.h"
 #include "mozilla/dom/MediaControlKeySource.h"
 #include "mozilla/dom/BrowsingContextWebProgress.h"
+#include "mozilla/dom/Promise.h"
 #include "mozilla/dom/SessionStoreRestoreData.h"
 #include "mozilla/dom/SessionStoreUtils.h"
 #include "mozilla/dom/ipc/IdType.h"
@@ -286,8 +287,10 @@ class CanonicalBrowsingContext final : public BrowsingContext {
     return mContainerFeaturePolicy;
   }
 
-  void SetRestoreData(SessionStoreRestoreData* aData);
+  void SetRestoreData(SessionStoreRestoreData* aData, ErrorResult& aError);
+  void ClearRestoreState();
   void RequestRestoreTabContent(WindowGlobalParent* aWindow);
+  already_AddRefed<Promise> GetRestorePromise();
 
   // Called when a BrowserParent for this BrowsingContext has been fully
   // destroyed (i.e. `ActorDestroy` was called).
@@ -349,6 +352,20 @@ class CanonicalBrowsingContext final : public BrowsingContext {
     RemotenessChangeOptions mOptions;
   };
 
+  struct RestoreState {
+    NS_INLINE_DECL_REFCOUNTING(RestoreState)
+
+    void ClearData() { mData = nullptr; }
+
+    RefPtr<SessionStoreRestoreData> mData;
+    RefPtr<Promise> mPromise;
+    uint32_t mRequests = 0;
+    uint32_t mResolves = 0;
+
+   private:
+    ~RestoreState() = default;
+  };
+
   friend class net::DocumentLoadListener;
   // Called when a DocumentLoadListener is created to start a load for
   // this browsing context. Returns false if a higher priority load is
@@ -376,6 +393,8 @@ class CanonicalBrowsingContext final : public BrowsingContext {
   // has become unloaded for one reason or another.
   void ShowSubframeCrashedUI(BrowserBridgeParent* aBridge);
 
+  void ResolveAndClearRestoreState(RestoreState* aState);
+
   // XXX(farre): Store a ContentParent pointer here rather than mProcessId?
   // Indicates which process owns the docshell.
   uint64_t mProcessId;
@@ -424,9 +443,7 @@ class CanonicalBrowsingContext final : public BrowsingContext {
 
   RefPtr<FeaturePolicy> mContainerFeaturePolicy;
 
-  RefPtr<SessionStoreRestoreData> mRestoreData;
-  uint32_t mRequestedContentRestores = 0;
-  uint32_t mCompletedContentRestores = 0;
+  RefPtr<RestoreState> mRestoreState;
 };
 
 }  // namespace dom

dom/chrome-webidl/BrowsingContext.webidl

@@ -287,6 +287,8 @@ interface CanonicalBrowsingContext : BrowsingContext {
   // The current URI loaded in this BrowsingContext according to nsDocShell.
   // This may not match the current window global's document URI in some cases.
   readonly attribute URI? currentURI;
+
+  void clearRestoreState();
 };
 
 [Exposed=Window, ChromeOnly]

dom/chrome-webidl/SessionStoreUtils.webidl

@@ -126,7 +126,8 @@ namespace SessionStoreUtils {
 
    nsISessionStoreRestoreData constructSessionStoreRestoreData();
 
-   boolean setRestoreData(CanonicalBrowsingContext browsingContext,
+   [Throws]
+   Promise<void> initializeRestore(CanonicalBrowsingContext browsingContext,
                                    nsISessionStoreRestoreData? data);
 };
 

toolkit/components/sessionstore/RestoreTabContentObserver.cpp

@@ -80,9 +80,7 @@ RestoreTabContentObserver::Observe(nsISupports* aSubject, const char* aTopic,
   }
 
   nsCOMPtr<nsIURI> uri = inner->GetDocumentURI();
-  // We'll never need to restore data into an about:blank document, so we can
-  // ignore those here.
-  if (!uri || NS_IsAboutBlank(uri)) {
+  if (!uri) {
     return NS_OK;
   }
 

toolkit/components/sessionstore/SessionStoreFunctions.idl

@@ -17,8 +17,6 @@ interface nsISessionStoreFunctions : nsISupports {
     in uint32_t aEpoch,
     in jsval aData, in boolean aCollectSHistory);
 
-  void RestoreTabContentComplete(in Element aBrowser);
-  
   void UpdateSessionStoreForWindow(
     in Element aBrowser, in BrowsingContext aBrowsingContext,
     in uint32_t aEpoch, in jsval aData);

toolkit/components/sessionstore/SessionStoreFunctions.jsm

@@ -10,10 +10,6 @@ XPCOMUtils.defineLazyModuleGetters(this, {
   SessionStore: "resource:///modules/sessionstore/SessionStore.jsm",
 });
 
-function RestoreTabContentComplete(aBrowser) {
-  SessionStore.restoreTabContentComplete(aBrowser, {});
-}
-
 function UpdateSessionStore(
   aBrowser,
   aBrowsingContext,
@@ -44,11 +40,7 @@ function UpdateSessionStoreForWindow(
   );
 }
 
-var EXPORTED_SYMBOLS = [
-  "RestoreTabContentComplete",
-  "UpdateSessionStore",
-  "UpdateSessionStoreForWindow",
-];
+var EXPORTED_SYMBOLS = ["UpdateSessionStore", "UpdateSessionStoreForWindow"];
 
 var SessionStoreFuncInternal = {
   updateStorage: function SSF_updateStorage(aOrigins, aKeys, aValues) {

toolkit/components/sessionstore/SessionStoreRestoreData.cpp

@@ -17,7 +17,7 @@ bool SessionStoreRestoreData::IsEmpty() {
           mEntries.IsEmpty() && mChildren.IsEmpty());
 }
 
-SessionStoreRestoreData* SessionStoreRestoreData::FindChild(
+SessionStoreRestoreData* SessionStoreRestoreData::FindDataForChild(
     BrowsingContext* aContext) {
   nsTArray<uint32_t> offsets;
   for (const BrowsingContext* current = aContext;

toolkit/components/sessionstore/SessionStoreRestoreData.h

@@ -26,7 +26,7 @@ class SessionStoreRestoreData final : public nsISessionStoreRestoreData {
  public:
   SessionStoreRestoreData() = default;
   bool IsEmpty();
-  SessionStoreRestoreData* FindChild(BrowsingContext* aContext);
+  SessionStoreRestoreData* FindDataForChild(BrowsingContext* aContext);
   bool CanRestoreInto(nsIURI* aDocumentURI);
   bool RestoreInto(RefPtr<BrowsingContext> aContext);
 

toolkit/components/sessionstore/SessionStoreUtils.cpp

@@ -1563,34 +1563,26 @@ SessionStoreUtils::ConstructSessionStoreRestoreData(
 
 /* static */
 MOZ_CAN_RUN_SCRIPT
-bool SessionStoreUtils::SetRestoreData(const GlobalObject& aGlobal,
-                                       CanonicalBrowsingContext& aContext,
-                                       nsISessionStoreRestoreData* aData) {
+already_AddRefed<Promise> SessionStoreUtils::InitializeRestore(
+    const GlobalObject& aGlobal, CanonicalBrowsingContext& aContext,
+    nsISessionStoreRestoreData* aData, ErrorResult& aError) {
   if (!mozilla::SessionHistoryInParent()) {
     MOZ_CRASH("why were we called?");
   }
 
   MOZ_DIAGNOSTIC_ASSERT(aContext.IsTop());
 
+  MOZ_DIAGNOSTIC_ASSERT(aData);
   nsCOMPtr<SessionStoreRestoreData> data = do_QueryInterface(aData);
-  aContext.SetRestoreData(data);
-
-  if (data && !data->IsEmpty()) {
-    if (nsISHistory* shistory = aContext.GetSessionHistory()) {
-      shistory->ReloadCurrentEntry();
-      return true;
-    }
+  aContext.SetRestoreData(data, aError);
+  if (aError.Failed()) {
+    return nullptr;
   }
 
-  return false;
-}
+  MOZ_DIAGNOSTIC_ASSERT(aContext.GetSessionHistory());
+  aContext.GetSessionHistory()->ReloadCurrentEntry();
 
-/* static */
-nsresult SessionStoreUtils::CallRestoreTabContentComplete(Element* aBrowser) {
-  nsCOMPtr<nsISessionStoreFunctions> funcs =
-      do_ImportModule("resource://gre/modules/SessionStoreFunctions.jsm");
-  NS_ENSURE_TRUE(funcs, NS_ERROR_FAILURE);
-  return funcs->RestoreTabContentComplete(aBrowser);
+  return aContext.GetRestorePromise();
 }
 
 /* static */

toolkit/components/sessionstore/SessionStoreUtils.h

@@ -113,11 +113,9 @@ class SessionStoreUtils {
   static already_AddRefed<nsISessionStoreRestoreData>
   ConstructSessionStoreRestoreData(const GlobalObject& aGlobal);
 
-  static bool SetRestoreData(const GlobalObject& aGlobal,
-                             CanonicalBrowsingContext& aContext,
-                             nsISessionStoreRestoreData* aData);
-
-  static nsresult CallRestoreTabContentComplete(Element* aBrowser);
+  static already_AddRefed<Promise> InitializeRestore(
+      const GlobalObject& aGlobal, CanonicalBrowsingContext& aContext,
+      nsISessionStoreRestoreData* aData, ErrorResult& aError);
 
   static nsresult ConstructFormDataValues(
       JSContext* aCx, const nsTArray<sessionstore::FormEntry>& aValues,