Bug 1645180 - Don't under-report use counters for documents that hit the shared sheet cache. r=firefox-style-system-reviewers,heycam

This makes us refactor a bit more our setup but I think the end it's
worth it.

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

layout/style/GeckoBindings.cpp

@@ -1476,11 +1476,7 @@ void Gecko_StyleSheet_FinishAsyncParse(
                  counters = std::move(useCounters)]() mutable {
         MOZ_ASSERT(NS_IsMainThread());
         SheetLoadData* data = d->get();
-        if (Document* doc = data->mLoader->GetDocument()) {
-          if (const auto* docCounters = doc->GetStyleUseCounters()) {
-            Servo_UseCounters_Merge(docCounters, counters.get());
-          }
-        }
+        data->mUseCounters = std::move(counters);
         data->mSheet->FinishAsyncParse(contents.forget());
       }));
 }

layout/style/Loader.cpp

@@ -822,8 +822,25 @@ nsresult Loader::CheckContentPolicy(nsIPrincipal* aLoadingPrincipal,
   return NS_OK;
 }
 
-void Loader::DidHitCompleteSheetCache(SheetLoadData& aData) {
-  mLoadsPerformed.PutEntry(SheetLoadDataHashKey(aData));
+static void RecordUseCountersIfNeeded(Document* aDoc,
+                                      const StyleUseCounters* aCounters) {
+  if (!aDoc || !aCounters) {
+    return;
+  }
+  const StyleUseCounters* docCounters = aDoc->GetStyleUseCounters();
+  if (!docCounters) {
+    return;
+  }
+  Servo_UseCounters_Merge(docCounters, aCounters);
+  aDoc->MaybeWarnAboutZoom();
+}
+
+void Loader::DidHitCompleteSheetCache(const SheetLoadDataHashKey& aKey,
+                                      const StyleUseCounters* aCounters) {
+  MOZ_ASSERT(mDocument);
+  if (mLoadsPerformed.EnsureInserted(aKey)) {
+    RecordUseCountersIfNeeded(mDocument, aCounters);
+  }
 }
 
 /**
@@ -1414,6 +1431,8 @@ Loader::Completed Loader::ParseSheet(const nsACString& aBytes,
 }
 
 void Loader::NotifyObservers(SheetLoadData& aData, nsresult aStatus) {
+  RecordUseCountersIfNeeded(mDocument, aData.mUseCounters.get());
+
   // Constructable sheets do get here via StyleSheet::Replace, but they don't
   // count like a regular sheet load.
   //
@@ -1460,11 +1479,6 @@ void Loader::NotifyObservers(SheetLoadData& aData, nsresult aStatus) {
 void Loader::SheetComplete(SheetLoadData& aLoadData, nsresult aStatus) {
   LOG(("css::Loader::SheetComplete, status: 0x%" PRIx32,
        static_cast<uint32_t>(aStatus)));
-
-  if (mDocument) {
-    mDocument->MaybeWarnAboutZoom();
-  }
-
   SharedStyleSheetCache::LoadCompleted(mSheets.get(), aLoadData, aStatus);
 }
 
@@ -1686,7 +1700,6 @@ Result<Loader::LoadSheetResult, nsresult> Loader::LoadStyleLink(
       aInfo.mReferrerInfo, context);
   if (state == SheetState::Complete) {
     LOG(("  Sheet already complete: 0x%p", sheet.get()));
-    DidHitCompleteSheetCache(*data);
     if (aObserver || !mObservers.IsEmpty() || aInfo.mContent) {
       rv = PostLoadEvent(std::move(data));
       if (NS_FAILED(rv)) {
@@ -1842,7 +1855,6 @@ nsresult Loader::LoadChildSheet(StyleSheet& aParentSheet,
 
   if (state == SheetState::Complete) {
     LOG(("  Sheet already complete"));
-    DidHitCompleteSheetCache(*data);
     // We're completely done.  No need to notify, even, since the
     // @import rule addition/modification will trigger the right style
     // changes automatically.
@@ -1932,7 +1944,6 @@ Result<RefPtr<StyleSheet>, nsresult> Loader::InternalLoadNonDocumentSheet(
       mDocument);
   if (state == SheetState::Complete) {
     LOG(("  Sheet already complete"));
-    DidHitCompleteSheetCache(*data);
     if (aObserver || !mObservers.IsEmpty()) {
       rv = PostLoadEvent(std::move(data));
       if (NS_FAILED(rv)) {

layout/style/Loader.h

@@ -361,7 +361,8 @@ class Loader final {
   /**
    * Called when we hit the internal memory cache with a complete stylesheet.
    */
-  void DidHitCompleteSheetCache(SheetLoadData&);
+  void DidHitCompleteSheetCache(const SheetLoadDataHashKey&,
+                                const StyleUseCounters* aCounters);
 
   enum class UseSystemPrincipal { No, Yes };
 

layout/style/SharedStyleSheetCache.cpp

@@ -9,6 +9,7 @@
 #include "mozilla/MemoryReporting.h"
 #include "mozilla/StyleSheet.h"
 #include "mozilla/css/SheetLoadData.h"
+#include "mozilla/ServoBindings.h"
 #include "nsXULPrototypeCache.h"
 
 extern mozilla::LazyLogModule sCssLoaderLog;
@@ -104,6 +105,7 @@ SharedStyleSheetCache::CacheResult SharedStyleSheetCache::Lookup(
         // We need to check the parsing mode manually because the XUL cache only
         // keys off the URI. But we should fix that!
         if (sheet->ParsingMode() == aKey.ParsingMode()) {
+          aLoader.DidHitCompleteSheetCache(aKey, nullptr);
           return {CloneSheet(*sheet), SheetState::Complete};
         }
 
@@ -135,6 +137,8 @@ SharedStyleSheetCache::CacheResult SharedStyleSheetCache::Lookup(
       RefPtr<StyleSheet> clone = CloneSheet(cachedSheet);
       MOZ_ASSERT(!clone->HasForcedUniqueInner());
       MOZ_ASSERT(!clone->HasModifiedRules());
+
+      aLoader.DidHitCompleteSheetCache(aKey, completeSheet.mUseCounters.get());
       return {std::move(clone), SheetState::Complete};
     }
   }
@@ -204,7 +208,8 @@ size_t SharedStyleSheetCache::SizeOfIncludingThis(
 
   n += mCompleteSheets.ShallowSizeOfExcludingThis(aMallocSizeOf);
   for (auto iter = mCompleteSheets.ConstIter(); !iter.Done(); iter.Next()) {
-    n += iter.UserData().mSheet->SizeOfIncludingThis(aMallocSizeOf);
+    n += iter.Data().mSheet->SizeOfIncludingThis(aMallocSizeOf);
+    n += aMallocSizeOf(iter.Data().mUseCounters.get());
   }
 
   // Measurement of the following members may be added later if DMD finds it is
@@ -396,7 +401,15 @@ void SharedStyleSheetCache::InsertIntoCompleteCacheIfNeeded(
     }
 #endif
 
-    mCompleteSheets.Put(key, {aData.mExpirationTime, std::move(sheet)});
+    UniquePtr<StyleUseCounters> counters;
+    if (aData.mUseCounters) {
+      // TODO(emilio): Servo_UseCounters_Clone() or something?
+      counters = Servo_UseCounters_Create().Consume();
+      Servo_UseCounters_Merge(counters.get(), aData.mUseCounters.get());
+    }
+
+    mCompleteSheets.Put(
+        key, {aData.mExpirationTime, std::move(counters), std::move(sheet)});
   }
 }
 
@@ -405,7 +418,7 @@ void SharedStyleSheetCache::StartDeferredLoadsForLoader(
   LoadDataArray arr;
   for (auto iter = mPendingDatas.Iter(); !iter.Done(); iter.Next()) {
     bool startIt = false;
-    SheetLoadData* data = iter.UserData();
+    SheetLoadData* data = iter.Data();
     do {
       if (data->mLoader == &aLoader) {
         // Note that we don't want to affect what the selected style set is, so
@@ -434,7 +447,7 @@ void SharedStyleSheetCache::CancelDeferredLoadsForLoader(css::Loader& aLoader) {
   for (auto iter = mPendingDatas.Iter(); !iter.Done(); iter.Next()) {
     RefPtr<SheetLoadData>& first = iter.Data();
     SheetLoadData* prev = nullptr;
-    SheetLoadData* current = iter.UserData();
+    SheetLoadData* current = iter.Data();
     do {
       if (current->mLoader != &aLoader) {
         prev = current;
@@ -473,7 +486,7 @@ void SharedStyleSheetCache::CancelLoadsForLoader(css::Loader& aLoader) {
   // We can't stop in-progress loads because some other loader may care about
   // them.
   for (auto iter = mLoadingDatas.Iter(); !iter.Done(); iter.Next()) {
-    SheetLoadData* data = iter.UserData();
+    SheetLoadData* data = iter.Data();
     do {
       if (data->mLoader == &aLoader) {
         data->mIsCancelled = true;

layout/style/SharedStyleSheetCache.h

@@ -114,6 +114,7 @@ class SharedStyleSheetCache final : public nsIMemoryReporter {
 
   struct CompleteSheet {
     uint32_t mExpirationTime = 0;
+    UniquePtr<StyleUseCounters> mUseCounters;
     RefPtr<StyleSheet> mSheet;
 
     bool Expired() const;

layout/style/SheetLoadData.h

@@ -11,6 +11,7 @@
 #include "mozilla/css/SheetParsingMode.h"
 #include "mozilla/Encoding.h"
 #include "mozilla/NotNull.h"
+#include "mozilla/UniquePtr.h"
 #include "nsIThreadInternal.h"
 #include "nsProxyRelease.h"
 
@@ -22,6 +23,7 @@ class nsINode;
 class nsIPrincipal;
 class nsIURI;
 class nsIReferrerInfo;
+struct StyleUseCounters;
 
 namespace mozilla {
 namespace css {
@@ -209,6 +211,10 @@ class SheetLoadData final : public nsIRunnable, public nsIThreadObserver {
   // The encoding guessed from attributes and the document character set.
   const NotNull<const Encoding*> mGuessedEncoding;
 
+  // If we've parsed the stylesheet, the use counters for the properties parsed
+  // in this styleshetet.
+  UniquePtr<StyleUseCounters> mUseCounters;
+
   // The quirks mode of the loader at the time the load was triggered.
   const nsCompatibility mCompatMode;
 

layout/style/StyleSheet.cpp

@@ -1180,31 +1180,33 @@ RefPtr<StyleSheetParsePromise> StyleSheet::ParseSheet(
   RefPtr<StyleSheetParsePromise> p = mParsePromise.Ensure(__func__);
   SetURLExtraData();
 
-  const StyleUseCounters* useCounters =
-      aLoader.GetDocument() ? aLoader.GetDocument()->GetStyleUseCounters()
-                            : nullptr;
   // @import rules are disallowed due to this decision:
   // https://github.com/WICG/construct-stylesheets/issues/119#issuecomment-588352418
   // We may allow @import rules again in the future.
   auto allowImportRules = SelfOrAncestorIsConstructed()
                               ? StyleAllowImportRules::No
                               : StyleAllowImportRules::Yes;
+  const bool shouldRecordCounters =
+      aLoader.GetDocument() && aLoader.GetDocument()->GetStyleUseCounters();
   if (!AllowParallelParse(aLoader, GetSheetURI())) {
+    UniquePtr<StyleUseCounters> counters =
+        shouldRecordCounters ? Servo_UseCounters_Create().Consume() : nullptr;
+
     RefPtr<RawServoStyleSheetContents> contents =
         Servo_StyleSheet_FromUTF8Bytes(
             &aLoader, this, &aLoadData, &aBytes, mParsingMode, Inner().mURLData,
             aLoadData.mLineNumber, aLoadData.mCompatMode,
-            /* reusable_sheets = */ nullptr, useCounters, allowImportRules,
+            /* reusable_sheets = */ nullptr, counters.get(), allowImportRules,
             StyleSanitizationKind::None,
             /* sanitized_output = */ nullptr)
             .Consume();
+    aLoadData.mUseCounters = std::move(counters);
     FinishAsyncParse(contents.forget());
   } else {
     auto holder = MakeRefPtr<css::SheetLoadDataHolder>(__func__, &aLoadData);
     Servo_StyleSheet_FromUTF8BytesAsync(
         holder, Inner().mURLData, &aBytes, mParsingMode, aLoadData.mLineNumber,
-        aLoadData.mCompatMode,
-        /* should_record_use_counters = */ !!useCounters, allowImportRules);
+        aLoadData.mCompatMode, shouldRecordCounters, allowImportRules);
   }
 
   return p;

layout/style/StyleSheet.h

@@ -119,8 +119,7 @@ class StyleSheet final : public nsICSSLoaderObserver, public nsWrapperCache {
 
   // Common code that needs to be called after servo finishes parsing. This is
   // shared between the parallel and sequential paths.
-  void FinishAsyncParse(
-      already_AddRefed<RawServoStyleSheetContents> aSheetContents);
+  void FinishAsyncParse(already_AddRefed<RawServoStyleSheetContents>);
 
   // Similar to `ParseSheet`, but guarantees that
   // parsing will be performed synchronously.