From e3207f3af3bb9e6faa5415c6b22b5dc2fb49599b Mon Sep 17 00:00:00 2001
From: Tim Huang <tihuang@mozilla.com>
Date: Thu, 12 Sep 2024 16:21:19 +0000
Subject: [PATCH] Bug 1894380 - Queue the AsyncApplyUpdate request if there is
 one executing. r=dimi

Differential Revision: https://phabricator.services.mozilla.com/D220249
---
 .../components/url-classifier/Classifier.cpp  | 38 +++++++++++++++++++
 .../components/url-classifier/Classifier.h    | 10 +++++
 2 files changed, 48 insertions(+)

diff --git a/toolkit/components/url-classifier/Classifier.cpp b/toolkit/components/url-classifier/Classifier.cpp
index 52fc84d958b3..3b0f532bce04 100644
--- a/toolkit/components/url-classifier/Classifier.cpp
+++ b/toolkit/components/url-classifier/Classifier.cpp
@@ -124,6 +124,7 @@ nsresult Classifier::GetPrivateStoreDirectory(
 
 Classifier::Classifier()
     : mIsTableRequestResultOutdated(true),
+      mAsyncUpdateInProgress(false),
       mUpdateInterrupted(true),
       mIsClosed(false) {
   // Make a lazy thread for any IO
@@ -697,6 +698,8 @@ void Classifier::FlushAndDisableAsyncUpdate() {
 
   mUpdateThread->Shutdown();
   mUpdateThread = nullptr;
+  mPendingUpdates.Clear();
+  mAsyncUpdateInProgress = false;
 }
 
 nsresult Classifier::AsyncApplyUpdates(const TableUpdateArray& aUpdates,
@@ -708,6 +711,22 @@ nsresult Classifier::AsyncApplyUpdates(const TableUpdateArray& aUpdates,
     return NS_ERROR_FAILURE;
   }
 
+  if (mAsyncUpdateInProgress) {
+    mPendingUpdates.AppendElement(NS_NewRunnableFunction(
+        "safebrowsing::Classifier::AsyncApplyUpdates",
+        [self = RefPtr{this}, aUpdates = aUpdates.Clone(),
+         aCallback]() mutable {
+          nsresult rv = self->AsyncApplyUpdates(aUpdates, aCallback);
+
+          // Calling the callback if we got an failure here to notify update
+          // observers.
+          if (NS_FAILED(rv)) {
+            aCallback(rv);
+          }
+        }));
+    return NS_OK;
+  }
+
   //         Caller thread      |       Update thread
   // --------------------------------------------------------
   //                            |    ApplyUpdatesBackground
@@ -718,6 +737,7 @@ nsresult Classifier::AsyncApplyUpdates(const TableUpdateArray& aUpdates,
   MOZ_ASSERT(mNewLookupCaches.IsEmpty(),
              "There should be no leftovers from a previous update.");
 
+  mAsyncUpdateInProgress = true;
   mUpdateInterrupted = false;
   nsresult rv =
       mRootStoreDirectory->Clone(getter_AddRefs(mRootStoreDirectoryForUpdate));
@@ -773,6 +793,8 @@ nsresult Classifier::AsyncApplyUpdates(const TableUpdateArray& aUpdates,
 
               LOG(("Step 3. Updates applied! Fire callback."));
               aCallback(rv);
+
+              classifier->AsyncUpdateFinished();
             });
 
         callerThread->Dispatch(fgRunnable, NS_DISPATCH_NORMAL);
@@ -781,6 +803,22 @@ nsresult Classifier::AsyncApplyUpdates(const TableUpdateArray& aUpdates,
   return mUpdateThread->Dispatch(bgRunnable, NS_DISPATCH_NORMAL);
 }
 
+void Classifier::AsyncUpdateFinished() {
+  MOZ_ASSERT(!OnUpdateThread(),
+             "AsyncUpdateFinished() MUST NOT be called on update thread");
+  MOZ_ASSERT(!NS_IsMainThread(),
+             "AsyncUpdateFinished() must be called on the worker thread");
+
+  mAsyncUpdateInProgress = false;
+
+  // If there are pending updates, run the first one.
+  if (!mPendingUpdates.IsEmpty()) {
+    auto& runnable = mPendingUpdates.ElementAt(0);
+    runnable->Run();
+    mPendingUpdates.RemoveElementAt(0);
+  }
+}
+
 nsresult Classifier::ApplyUpdatesBackground(
     TableUpdateArray& aUpdates, nsTArray<nsCString>& aFailedTableNames) {
   // |mUpdateInterrupted| is guaranteed to have been unset.
diff --git a/toolkit/components/url-classifier/Classifier.h b/toolkit/components/url-classifier/Classifier.h
index a9af2736d2fb..8995a2ff604b 100644
--- a/toolkit/components/url-classifier/Classifier.h
+++ b/toolkit/components/url-classifier/Classifier.h
@@ -205,6 +205,10 @@ class Classifier {
   nsresult ApplyUpdatesForeground(nsresult aBackgroundRv,
                                   const nsTArray<nsCString>& aFailedTableNames);
 
+  // Notify the worker thread that the async update is finished to kick off the
+  // pending updates.
+  void AsyncUpdateFinished();
+
   // Used by worker thread and update thread to abort current operation.
   bool ShouldAbort() const;
 
@@ -237,6 +241,12 @@ class Classifier {
   // The copy of mLookupCaches for update only.
   LookupCacheArray mNewLookupCaches;
 
+  // Whether an async update is in progress.
+  bool mAsyncUpdateInProgress;
+
+  // Pending updates to be executed after the current async update is finished.
+  nsTArray<nsCOMPtr<nsIRunnable>> mPendingUpdates;
+
   // True when Reset() is called.
   bool mUpdateInterrupted;