Bug 1618289 - Have PreloadService hanging off dom::Document, r=smaug

Depends on D67479

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

dom/base/Document.cpp

@@ -10486,6 +10486,9 @@ void Document::Destroy() {
   // Manually break cycles via promise's global object pointer.
   mReadyForIdle = nullptr;
   mOrientationPendingPromise = nullptr;
+
+  // To break cycles.
+  mPreloadService.ClearAllPreloads();
 }
 
 void Document::RemovedFromDocShell() {

dom/base/Document.h

@@ -57,6 +57,7 @@
 #include "mozilla/HashTable.h"
 #include "mozilla/LinkedList.h"
 #include "mozilla/NotNull.h"
+#include "mozilla/PreloadService.h"
 #include "mozilla/SegmentedVector.h"
 #include "mozilla/TimeStamp.h"
 #include "mozilla/UniquePtr.h"
@@ -1522,6 +1523,8 @@ class Document : public nsINode,
   // parser inserted form control element.
   int32_t GetNextControlNumber() { return mNextControlNumber++; }
 
+  PreloadService& Preloads() { return mPreloadService; }
+
  protected:
   friend class nsUnblockOnloadEvent;
 
@@ -5088,6 +5091,9 @@ class Document : public nsINode,
   int32_t mNextFormNumber;
   int32_t mNextControlNumber;
 
+  // Scope preloads per document.  This is used by speculative loading as well.
+  PreloadService mPreloadService;
+
  public:
   // Needs to be public because the bindings code pokes at it.
   js::ExpandoAndGeneration mExpandoAndGeneration;

uriloader/preload/PreloadService.cpp

@@ -0,0 +1,37 @@
+/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * 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/. */
+
+#include "PreloadService.h"
+
+namespace mozilla {
+
+bool PreloadService::RegisterPreload(PreloadHashKey* aKey,
+                                     PreloaderBase* aPreload) {
+  if (PreloadExists(aKey)) {
+    return false;
+  }
+
+  mPreloads.Put(aKey, RefPtr{aPreload});
+  return true;
+}
+
+void PreloadService::DeregisterPreload(PreloadHashKey* aKey) {
+  mPreloads.Remove(aKey);
+}
+
+void PreloadService::ClearAllPreloads() { mPreloads.Clear(); }
+
+bool PreloadService::PreloadExists(PreloadHashKey* aKey) {
+  bool found;
+  mPreloads.GetWeak(aKey, &found);
+  return found;
+}
+
+already_AddRefed<PreloaderBase> PreloadService::LookupPreload(
+    PreloadHashKey* aKey) const {
+  return mPreloads.Get(aKey);
+}
+
+}  // namespace mozilla

uriloader/preload/PreloadService.h

@@ -0,0 +1,51 @@
+/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * 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/. */
+
+#ifndef PreloadService_h__
+#define PreloadService_h__
+
+#include "nsRefPtrHashtable.h"
+#include "PreloaderBase.h"
+
+namespace mozilla {
+
+/**
+ * Intended to scope preloads and speculative loads under one roof.  This class
+ * is intended to be a member of dom::Document. Provides registration of
+ * speculative loads via a `key` which is defined to consist of the URL,
+ * resource type, and resource-specific attributes that are further
+ * distinguishing loads with otherwise same type and url.
+ */
+class PreloadService {
+ public:
+  // Called by resource loaders to register a running resource load.  This is
+  // called for a speculative load when it's started the first time, being it
+  // either regular speculative load or a preload.
+  //
+  // Returns false and does nothing if a preload is already registered under
+  // this key, true otherwise.
+  bool RegisterPreload(PreloadHashKey* aKey, PreloaderBase* aPreload);
+
+  // Called when the load is about to be cancelled.  Exact behavior is to be
+  // determined yet.
+  void DeregisterPreload(PreloadHashKey* aKey);
+
+  // Called when the scope is to go away.
+  void ClearAllPreloads();
+
+  // True when there is a preload registered under the key.
+  bool PreloadExists(PreloadHashKey* aKey);
+
+  // Returns an existing preload under the key or null, when there is none
+  // registered under the key.
+  already_AddRefed<PreloaderBase> LookupPreload(PreloadHashKey* aKey) const;
+
+ private:
+  nsRefPtrHashtable<PreloadHashKey, PreloaderBase> mPreloads;
+};
+
+}  // namespace mozilla
+
+#endif

uriloader/preload/PreloaderBase.cpp

@@ -83,7 +83,14 @@ void PreloaderBase::AddLoadBackgroundFlag(nsIChannel* aChannel) {
 
 void PreloaderBase::NotifyOpen(PreloadHashKey* aKey, nsIChannel* aChannel,
                                dom::Document* aDocument, bool aIsPreload) {
-  // * Register this preload in document's preload service.
+  if (aDocument && !aDocument->Preloads().RegisterPreload(aKey, this)) {
+    // This means there is already a preload registered under this key in this
+    // document.  We only allow replacement when this is a regular load.
+    // Otherwise, this should never happen and is a suspected misuse of the API.
+    MOZ_ASSERT(!aIsPreload);
+    aDocument->Preloads().DeregisterPreload(aKey);
+    aDocument->Preloads().RegisterPreload(aKey, this);
+  }
 
   mChannel = aChannel;
   mKey = *aKey;
@@ -130,7 +137,7 @@ void PreloaderBase::NotifyUsage() {
 
 void PreloaderBase::NotifyRestart(dom::Document* aDocument,
                                   PreloaderBase* aNewPreloader) {
-  // * Deregister this preload from Document's preloads
+  aDocument->Preloads().DeregisterPreload(&mKey);
   mKey = PreloadHashKey();
 
   if (aNewPreloader) {
@@ -206,7 +213,14 @@ void PreloaderBase::RemoveLinkPreloadNode(nsINode* aNode) {
   mNodes.RemoveElement(node);
 
   if (mNodes.Length() == 0 && !mIsUsed) {
-    // * Deregister this preload from document's preloads
+    // Keep a reference, because the following call may release us.  The caller
+    // may use a WeakPtr to access this.
+    RefPtr<PreloaderBase> self(this);
+
+    dom::Document* doc = aNode->OwnerDoc();
+    if (doc) {
+      doc->Preloads().DeregisterPreload(&mKey);
+    }
 
     if (mChannel) {
       mChannel->Cancel(NS_BINDING_ABORTED);

uriloader/preload/moz.build

@@ -10,11 +10,13 @@ with Files("**"):
 EXPORTS.mozilla += [
     'PreloaderBase.h',
     'PreloadHashKey.h',
+    'PreloadService.h',
 ]
 
 UNIFIED_SOURCES += [
     'PreloaderBase.cpp',
     'PreloadHashKey.cpp',
+    'PreloadService.cpp',
 ]
 
 FINAL_LIBRARY = 'xul'