Bug 1754448 - Fix snapshot usage calculation in LSSnapshot::Clear; r=dom-storage-reviewers,jari,webidl,smaug

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

dom/localstorage/LSDatabase.cpp

@@ -323,6 +323,15 @@ bool LSDatabase::HasSnapshot() const {
   return !!mSnapshot;
 }
 
+int64_t LSDatabase::GetSnapshotUsage() const {
+  AssertIsOnOwningThread();
+  MOZ_ASSERT(mActor);
+  MOZ_ASSERT(!mAllowedToClose);
+  MOZ_ASSERT(mSnapshot);
+
+  return mSnapshot->GetUsage();
+}
+
 nsresult LSDatabase::EnsureSnapshot(LSObject* aObject, const nsAString& aKey,
                                     bool aExplicit) {
   MOZ_ASSERT(aObject);

dom/localstorage/LSDatabase.h

@@ -88,6 +88,8 @@ class LSDatabase final {
 
   bool HasSnapshot() const;
 
+  int64_t GetSnapshotUsage() const;
+
  private:
   ~LSDatabase();
 

dom/localstorage/LSObject.cpp

@@ -879,6 +879,28 @@ bool LSObject::GetHasSnapshot(nsIPrincipal& aSubjectPrincipal,
   return mDatabase->HasSnapshot();
 }
 
+int64_t LSObject::GetSnapshotUsage(nsIPrincipal& aSubjectPrincipal,
+                                   ErrorResult& aError) {
+  AssertIsOnOwningThread();
+
+  if (!CanUseStorage(aSubjectPrincipal)) {
+    aError.Throw(NS_ERROR_DOM_SECURITY_ERR);
+    return 0;
+  }
+
+  if (!mDatabase || mDatabase->IsAllowedToClose()) {
+    aError.Throw(NS_ERROR_NOT_AVAILABLE);
+    return 0;
+  }
+
+  if (!mDatabase->HasSnapshot()) {
+    aError.Throw(NS_ERROR_NOT_AVAILABLE);
+    return 0;
+  }
+
+  return mDatabase->GetSnapshotUsage();
+}
+
 NS_IMPL_ADDREF_INHERITED(LSObject, Storage)
 NS_IMPL_RELEASE_INHERITED(LSObject, Storage)
 

dom/localstorage/LSObject.h

@@ -185,6 +185,9 @@ class LSObject final : public Storage {
   bool GetHasSnapshot(nsIPrincipal& aSubjectPrincipal,
                       ErrorResult& aError) override;
 
+  int64_t GetSnapshotUsage(nsIPrincipal& aSubjectPrincipal,
+                           ErrorResult& aError) override;
+
   //////////////////////////////////////////////////////////////////////////////
 
   NS_DECL_ISUPPORTS_INHERITED

dom/localstorage/LSSnapshot.cpp

@@ -510,7 +510,16 @@ nsresult LSSnapshot::Clear(LSNotifyInfo& aNotifyInfo) {
   } else {
     changed = true;
 
-    DebugOnly<nsresult> rv = UpdateUsage(-mExactUsage);
+    int64_t delta = 0;
+    for (const auto& entry : mValues) {
+      const nsAString& key = entry.GetKey();
+      const nsString& value = entry.GetData();
+
+      delta += -static_cast<int64_t>(key.Length()) -
+               static_cast<int64_t>(value.Length());
+    }
+
+    DebugOnly<nsresult> rv = UpdateUsage(delta);
     MOZ_ASSERT(NS_SUCCEEDED(rv));
 
     mValues.Clear();
@@ -584,6 +593,15 @@ nsresult LSSnapshot::End() {
   return NS_OK;
 }
 
+int64_t LSSnapshot::GetUsage() const {
+  AssertIsOnOwningThread();
+  MOZ_ASSERT(mActor);
+  MOZ_ASSERT(mInitialized);
+  MOZ_ASSERT(!mSentFinish);
+
+  return mExactUsage;
+}
+
 void LSSnapshot::ScheduleStableStateCallback() {
   AssertIsOnOwningThread();
   MOZ_ASSERT(mIdleTimer);

dom/localstorage/LSSnapshot.h

@@ -157,6 +157,8 @@ class LSSnapshot final : public nsIRunnable {
 
   nsresult End();
 
+  int64_t GetUsage() const;
+
  private:
   ~LSSnapshot();
 

dom/storage/Storage.cpp

@@ -62,6 +62,12 @@ bool Storage::StoragePrefIsEnabled() {
   return mozilla::Preferences::GetBool(kStorageEnabled);
 }
 
+int64_t Storage::GetSnapshotUsage(nsIPrincipal& aSubjectPrincipal,
+                                  ErrorResult& aRv) {
+  aRv.Throw(NS_ERROR_NOT_IMPLEMENTED);
+  return 0;
+}
+
 bool Storage::CanUseStorage(nsIPrincipal& aSubjectPrincipal) {
   if (!StoragePrefIsEnabled()) {
     return false;

dom/storage/Storage.h

@@ -126,6 +126,9 @@ class Storage : public nsISupports, public nsWrapperCache {
     return false;
   }
 
+  virtual int64_t GetSnapshotUsage(nsIPrincipal& aSubjectPrincipal,
+                                   ErrorResult& aRv);
+
   //////////////////////////////////////////////////////////////////////////////
 
   // Dispatch storage notification events on all impacted pages in the current

dom/tests/browser/browser_localStorage_snapshotting_e10s.js

@@ -77,6 +77,28 @@ async function endExplicitSnapshot(knownTab) {
   });
 }
 
+async function verifyHasSnapshot(knownTab, expectedHasSnapshot) {
+  let hasSnapshot = await SpecialPowers.spawn(
+    knownTab.tab.linkedBrowser,
+    [],
+    function() {
+      return content.wrappedJSObject.getHasSnapshot();
+    }
+  );
+  is(hasSnapshot, expectedHasSnapshot, "Correct has snapshot");
+}
+
+async function verifySnapshotUsage(knownTab, expectedSnapshotUsage) {
+  let snapshotUsage = await SpecialPowers.spawn(
+    knownTab.tab.linkedBrowser,
+    [],
+    function() {
+      return content.wrappedJSObject.getSnapshotUsage();
+    }
+  );
+  is(snapshotUsage, expectedSnapshotUsage, "Correct snapshot usage");
+}
+
 // We spin up a ton of child processes.
 requestLongerTimeout(4);
 
@@ -542,3 +564,87 @@ add_task(async function() {
 
   clearOriginStorageEnsuringNoPreload(HELPER_PAGE_ORIGIN);
 });
+
+/**
+ * Verify that snapshot usage is correctly updated after each operation.
+ */
+add_task(async function() {
+  await SpecialPowers.pushPrefEnv({
+    set: [
+      // Force multiple web and webIsolated content processes so that the
+      // multi-e10s logic works correctly.
+      ["dom.ipc.processCount", 4],
+      ["dom.ipc.processCount.webIsolated", 2],
+      // Disable snapshot peak usage pre-incrementation to make the testing
+      // easier.
+      ["dom.storage.snapshot_peak_usage.initial_preincrement", 0],
+      ["dom.storage.snapshot_peak_usage.reduced_initial_preincrement", 0],
+      ["dom.storage.snapshot_peak_usage.gradual_preincrement", 0],
+      ["dom.storage.snapshot_peak_usage.reuced_gradual_preincrement", 0],
+      // Enable LocalStorage's testing API so we can explicitly create
+      // snapshots when needed.
+      ["dom.storage.testing", true],
+    ],
+  });
+
+  // Ensure that there is no localstorage data by forcing the origin to be
+  // cleared prior to the start of our test.
+  await clearOriginStorageEnsuringNoPreload(HELPER_PAGE_ORIGIN);
+
+  // Open tabs.  Don't configure any of them yet.
+  const knownTabs = new KnownTabs();
+  const writerTab1 = await openTestTab(
+    HELPER_PAGE_URL,
+    "writer1",
+    knownTabs,
+    true
+  );
+  const writerTab2 = await openTestTab(
+    HELPER_PAGE_URL,
+    "writer2",
+    knownTabs,
+    true
+  );
+
+  // Apply the initial mutation using an explicit snapshot. The explicit
+  // snapshot here ensures that the parent process have received the changes.
+  await beginExplicitSnapshot(writerTab1);
+  await verifySnapshotUsage(writerTab1, 0);
+  await applyMutations(writerTab1, [["key", "something"]]);
+  await verifySnapshotUsage(writerTab1, 12);
+  await endExplicitSnapshot(writerTab1);
+  await verifyHasSnapshot(writerTab1, false);
+
+  // Begin an explicit snapshot in writerTab1 and apply the first mutatation
+  // in it.
+  await beginExplicitSnapshot(writerTab1);
+  await verifySnapshotUsage(writerTab1, 12);
+  await applyMutations(writerTab1, [["key", "somethingBigger"]]);
+  await verifySnapshotUsage(writerTab1, 18);
+
+  // Begin an explicit snapshot in writerTab2 and apply the second mutatation
+  // in it.
+  await beginExplicitSnapshot(writerTab2);
+  await verifySnapshotUsage(writerTab2, 18);
+  await applyMutations(writerTab2, [[null, null]]);
+  await verifySnapshotUsage(writerTab2, 6);
+
+  // End explicit snapshots in both tabs.
+  await endExplicitSnapshot(writerTab1);
+  await verifyHasSnapshot(writerTab1, false);
+  await endExplicitSnapshot(writerTab2);
+  await verifyHasSnapshot(writerTab2, false);
+
+  // Verify the final state, it should match the state after the second
+  // mutation has been applied and "commited". An explicit snapshot is used.
+  await beginExplicitSnapshot(writerTab1);
+  await verifySnapshotUsage(writerTab1, 0);
+  await verifyState(writerTab1, {});
+  await endExplicitSnapshot(writerTab1);
+  await verifyHasSnapshot(writerTab1, false);
+
+  // Clean up.
+  await cleanupTabs(knownTabs);
+
+  clearOriginStorageEnsuringNoPreload(HELPER_PAGE_ORIGIN);
+});

dom/tests/browser/page_localstorage_snapshotting_e10s.html

@@ -49,6 +49,15 @@ function beginExplicitSnapshot() {
 function endExplicitSnapshot() {
   localStorage.endExplicitSnapshot();
 }
+
+function getHasSnapshot() {
+  return localStorage.hasSnapshot;
+}
+
+function getSnapshotUsage() {
+  return localStorage.snapshotUsage;
+}
+
 </script>
 </head>
 <body><h2 id="pageNameH"></h2></body>

dom/webidl/Storage.webidl

@@ -79,4 +79,14 @@ partial interface Storage {
    */
   [Throws, NeedsSubjectPrincipal, Pref="dom.storage.testing"]
   readonly attribute boolean hasSnapshot;
+
+  /**
+   * Returns snapshot usage.
+   *
+   * @throws NS_ERROR_NOT_AVAILABLE if the underlying database hasn't been
+   *         opened or the database is being closed or it doesn't have a
+   *         snapshot.
+   */
+  [Throws, NeedsSubjectPrincipal, Pref="dom.storage.testing"]
+  readonly attribute long long snapshotUsage;
 };