Bug 1339050 - Asynchronously apply safebrowsing DB update. draft
authorHenry Chang <hchang@mozilla.com>
Thu, 06 Apr 2017 07:07:56 +0800
changeset 556495 06873c5553a35ebd12df1eef3e78dcab2f0fdc97
parent 556469 867df9483d5af4c8c12e19fab9b0de18bee30db7
child 622911 9af2771a133b2a48122ad415fec795ea165e4d1e
push id52572
push userhchang@mozilla.com
push dateThu, 06 Apr 2017 00:11:20 +0000
bugs1339050
milestone55.0a1
Bug 1339050 - Asynchronously apply safebrowsing DB update. A new function Classifier::AsyncApplyUpdates() is implemented for async update. Besides, all public Classifier interfaces become "worker thread only" and we remove DBServiceWorker::ApplyUpdatesBackground/Foreground. In DBServiceWorker::FinishUpdate, instead of calling Classifier::ApplyUpdates, we call Classifier::AsyncApplyUpdates and install a callback for notifying the update observer when update is finished. The callback will occur on the caller thread (i.e. worker thread.) As for the shutdown issue, when the main thread is notified to shut down, we at first *synchronously* dispatch an event to the worker thread to shut down the update thread. After getting synchronized with all other threads, we send last two events "CancelUpdate" and "CloseDb" to notify dangling update (i.e. BeginUpdate is called but FinishUpdate isn't) and do cleanup work. MozReview-Commit-ID: DXZvA2eFKlc
browser/components/safebrowsing/content/test/browser_bug400731.js
browser/components/safebrowsing/content/test/browser_mixedcontent_aboutblocked.js
browser/components/safebrowsing/content/test/head.js
toolkit/components/url-classifier/Classifier.cpp
toolkit/components/url-classifier/Classifier.h
toolkit/components/url-classifier/HashStore.cpp
toolkit/components/url-classifier/LookupCache.cpp
toolkit/components/url-classifier/LookupCache.h
toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
toolkit/components/url-classifier/nsUrlClassifierDBService.h
toolkit/components/url-classifier/tests/gtest/Common.cpp
toolkit/components/url-classifier/tests/gtest/Common.h
toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp
toolkit/components/url-classifier/tests/unit/test_partial.js
--- a/browser/components/safebrowsing/content/test/browser_bug400731.js
+++ b/browser/components/safebrowsing/content/test/browser_bug400731.js
@@ -22,18 +22,20 @@ function onDOMContentLoaded(callback) {
     addEventListener("DOMContentLoaded", listener);
   }
   mm.loadFrameScript("data:,(" + contentScript.toString() + ")();", true);
 }
 
 function test() {
   waitForExplicitFinish();
 
-  gBrowser.selectedTab = gBrowser.addTab("http://www.itisatrap.org/firefox/its-an-attack.html");
-  onDOMContentLoaded(testMalware);
+  waitForDBInit(() => {
+    gBrowser.selectedTab = gBrowser.addTab("http://www.itisatrap.org/firefox/its-an-attack.html");
+    onDOMContentLoaded(testMalware);
+  });
 }
 
 function testMalware(data) {
   ok(data.buttonPresent, "Ignore warning button should be present for malware");
 
   Services.prefs.setBoolPref("browser.safebrowsing.allowOverride", false);
 
   // Now launch the unwanted software test
--- a/browser/components/safebrowsing/content/test/browser_mixedcontent_aboutblocked.js
+++ b/browser/components/safebrowsing/content/test/browser_mixedcontent_aboutblocked.js
@@ -1,57 +1,15 @@
 /* Any copyright is dedicated to the Public Domain.
  * http://creativecommons.org/publicdomain/zero/1.0/ */
 
 const { classes: Cc, interfaces: Ci, results: Cr } = Components;
 
-// This url must sync with the table, url in SafeBrowsing.jsm addMozEntries
-const PHISH_TABLE = "test-phish-simple";
-const PHISH_URL = "https://www.itisatrap.org/firefox/its-a-trap.html";
-
 const SECURE_CONTAINER_URL = "https://example.com/browser/browser/components/safebrowsing/content/test/empty_file.html";
 
-// This function is mostly ported from classifierCommon.js
-// under toolkit/components/url-classifier/tests/mochitest.
-function waitForDBInit(callback) {
-  // Since there are two cases that may trigger the callback,
-  // we have to carefully avoid multiple callbacks and observer
-  // leaking.
-  let didCallback = false;
-  function callbackOnce() {
-    Services.obs.removeObserver(obsFunc, "mozentries-update-finished");
-    if (!didCallback) {
-      callback();
-    }
-    didCallback = true;
-  }
-
-  // The first part: listen to internal event.
-  function obsFunc() {
-    ok(true, "Received internal event!");
-    callbackOnce();
-  }
-  Services.obs.addObserver(obsFunc, "mozentries-update-finished", false);
-
-  // The second part: we might have missed the event. Just do
-  // an internal database lookup to confirm if the url has been
-  // added.
-  let principal = Services.scriptSecurityManager
-    .createCodebasePrincipal(Services.io.newURI(PHISH_URL), {});
-
-  let dbService = Cc["@mozilla.org/url-classifier/dbservice;1"]
-    .getService(Ci.nsIUrlClassifierDBService);
-  dbService.lookup(principal, PHISH_TABLE, value => {
-    if (value === PHISH_TABLE) {
-      ok(true, "DB lookup success!");
-      callbackOnce();
-    }
-  });
-}
-
 add_task(function* testNormalBrowsing() {
   yield BrowserTestUtils.withNewTab(SECURE_CONTAINER_URL, function* (browser) {
     // Before we load the phish url, we have to make sure the hard-coded
     // black list has been added to the database.
     yield new Promise(resolve => waitForDBInit(resolve));
 
     yield ContentTask.spawn(browser, PHISH_URL, function* (aPhishUrl) {
       return new Promise(resolve => {
--- a/browser/components/safebrowsing/content/test/head.js
+++ b/browser/components/safebrowsing/content/test/head.js
@@ -1,15 +1,19 @@
 Components.utils.import("resource://gre/modules/XPCOMUtils.jsm");
 
 XPCOMUtils.defineLazyModuleGetter(this, "Promise",
   "resource://gre/modules/Promise.jsm");
 XPCOMUtils.defineLazyModuleGetter(this, "Task",
   "resource://gre/modules/Task.jsm");
 
+// This url must sync with the table, url in SafeBrowsing.jsm addMozEntries
+const PHISH_TABLE = "test-phish-simple";
+const PHISH_URL = "https://www.itisatrap.org/firefox/its-a-trap.html";
+
 /**
  * Waits for a load (or custom) event to finish in a given tab. If provided
  * load an uri into the tab.
  *
  * @param tab
  *        The tab to load into.
  * @param [optional] url
  *        The url to load, or the current url.
@@ -43,12 +47,50 @@ function promiseTabLoadEvent(tab, url, e
   }
 
   if (url)
     BrowserTestUtils.loadURI(tab.linkedBrowser, url);
 
   return loaded;
 }
 
+// This function is mostly ported from classifierCommon.js
+// under toolkit/components/url-classifier/tests/mochitest.
+function waitForDBInit(callback) {
+  // Since there are two cases that may trigger the callback,
+  // we have to carefully avoid multiple callbacks and observer
+  // leaking.
+  let didCallback = false;
+  function callbackOnce() {
+    if (!didCallback) {
+      Services.obs.removeObserver(obsFunc, "mozentries-update-finished");
+      callback();
+    }
+    didCallback = true;
+  }
+
+  // The first part: listen to internal event.
+  function obsFunc() {
+    ok(true, "Received internal event!");
+    callbackOnce();
+  }
+  Services.obs.addObserver(obsFunc, "mozentries-update-finished", false);
+
+  // The second part: we might have missed the event. Just do
+  // an internal database lookup to confirm if the url has been
+  // added.
+  let principal = Services.scriptSecurityManager
+    .createCodebasePrincipal(Services.io.newURI(PHISH_URL), {});
+
+  let dbService = Cc["@mozilla.org/url-classifier/dbservice;1"]
+    .getService(Ci.nsIUrlClassifierDBService);
+  dbService.lookup(principal, PHISH_TABLE, value => {
+    if (value === PHISH_TABLE) {
+      ok(true, "DB lookup success!");
+      callbackOnce();
+    }
+  });
+}
+
 Services.prefs.setCharPref("urlclassifier.malwareTable", "test-malware-simple,test-unwanted-simple");
 Services.prefs.setCharPref("urlclassifier.phishTable", "test-phish-simple");
 Services.prefs.setCharPref("urlclassifier.blockedTable", "test-block-simple");
 SafeBrowsing.init();
--- a/toolkit/components/url-classifier/Classifier.cpp
+++ b/toolkit/components/url-classifier/Classifier.cpp
@@ -143,17 +143,20 @@ Classifier::GetPrivateStoreDirectory(nsI
 
   providerDirectory.forget(aPrivateStoreDirectory);
 
   return NS_OK;
 }
 
 Classifier::Classifier()
   : mIsTableRequestResultOutdated(true)
+  , mUpdateInterrupted(true)
 {
+  NS_NewNamedThread(NS_LITERAL_CSTRING("Classifier Update"),
+                    getter_AddRefs(mUpdateThread));
 }
 
 Classifier::~Classifier()
 {
   Close();
 }
 
 nsresult
@@ -268,27 +271,44 @@ void
 Classifier::Close()
 {
   DropStores();
 }
 
 void
 Classifier::Reset()
 {
-  DropStores();
+  MOZ_ASSERT(NS_GetCurrentThread() != mUpdateThread,
+             "Reset() MUST NOT be called on update thread");
+
+  LOG(("Reset() is called so we interrupt the update."));
+  mUpdateInterrupted = true;
+
+  auto resetFunc = [=] {
+    DropStores();
+
+    mRootStoreDirectory->Remove(true);
+    mBackupDirectory->Remove(true);
+    mUpdatingDirectory->Remove(true);
+    mToDeleteDirectory->Remove(true);
 
-  mRootStoreDirectory->Remove(true);
-  mBackupDirectory->Remove(true);
-  mUpdatingDirectory->Remove(true);
-  mToDeleteDirectory->Remove(true);
+    CreateStoreDirectory();
+
+    mTableFreshness.Clear();
+    RegenActiveTables();
+  };
 
-  CreateStoreDirectory();
+  if (!mUpdateThread) {
+    LOG(("Async update has been disabled. Just Reset() on worker thread."));
+    resetFunc();
+    return;
+  }
 
-  mTableFreshness.Clear();
-  RegenActiveTables();
+  nsCOMPtr<nsIRunnable> r = NS_NewRunnableFunction(resetFunc);
+  SyncRunnable::DispatchToThread(mUpdateThread, r);
 }
 
 void
 Classifier::ResetTables(ClearType aType, const nsTArray<nsCString>& aTables)
 {
   for (uint32_t i = 0; i < aTables.Length(); i++) {
     LOG(("Resetting table: %s", aTables[i].get()));
     // Spoil this table by marking it as no known freshness
@@ -606,16 +626,46 @@ Classifier::RemoveUpdateIntermediaries()
     // If the directory is locked from removal for some reason,
     // we will fail here and it doesn't matter until the next
     // update. (the next udpate will fail due to the removable
     // "safebrowsing-udpating" directory.)
     LOG(("Failed to remove updating directory."));
   }
 }
 
+void
+Classifier::MergeNewLookupCaches()
+{
+  MOZ_ASSERT(NS_GetCurrentThread() != mUpdateThread,
+             "MergeNewLookupCaches cannot be called on update thread "
+             "since it mutates mLookupCaches which is only safe on "
+             "worker thread.");
+
+  for (auto& newCache: mNewLookupCaches) {
+    // For each element in mNewLookCaches, it will be swapped with
+    //   - An old cache in mLookupCache with the same table name or
+    //   - nullptr (mLookupCache will be expaned) otherwise.
+    size_t swapIndex = 0;
+    for (; swapIndex < mLookupCaches.Length(); swapIndex++) {
+      if (mLookupCaches[swapIndex]->TableName() == newCache->TableName()) {
+        break;
+      }
+    }
+    if (swapIndex == mLookupCaches.Length()) {
+      mLookupCaches.AppendElement(nullptr);
+    }
+
+    Swap(mLookupCaches[swapIndex], newCache);
+    mLookupCaches[swapIndex]->UpdateRootDirHandle(mRootStoreDirectory);
+  }
+
+  // At this point, mNewLookupCaches's length remains the same but
+  // will contain either old cache (override) or nullptr (append).
+}
+
 nsresult
 Classifier::SwapInNewTablesAndCleanup()
 {
   nsresult rv;
 
   // Step 1. Swap in on-disk tables. The idea of using "safebrowsing-backup"
   // as the intermediary directory is we can get databases recovered if
   // crash occurred in any step of the swap. (We will recover from
@@ -625,21 +675,20 @@ Classifier::SwapInNewTablesAndCleanup()
                             mCacheDirectory,     // common parent dir
                             mBackupDirectory);   // intermediary dir for swap
   if (NS_FAILED(rv)) {
     LOG(("Failed to swap in on-disk tables."));
     RemoveUpdateIntermediaries();
     return rv;
   }
 
-  // Step 2. Swap in in-memory tables and update root directroy handles.
-  mLookupCaches.SwapElements(mNewLookupCaches);
-  for (auto c: mLookupCaches) {
-    c->UpdateRootDirHandle(mRootStoreDirectory);
-  }
+  // Step 2. Merge mNewLookupCaches into mLookupCaches. The outdated
+  // LookupCaches will be stored in mNewLookupCaches and be cleaned
+  // up later.
+  MergeNewLookupCaches();
 
   // Step 3. Re-generate active tables based on on-disk tables.
   rv = RegenActiveTables();
   if (NS_FAILED(rv)) {
     LOG(("Failed to re-generate active tables!"));
   }
 
   // Step 4. Clean up intermediaries for update.
@@ -648,32 +697,89 @@ Classifier::SwapInNewTablesAndCleanup()
   // Step 5. Invalidate cached tableRequest request.
   mIsTableRequestResultOutdated = true;
 
   LOG(("Done swap in updated tables."));
 
   return rv;
 }
 
+void Classifier::FlushAndDisableAsyncUpdate()
+{
+  LOG(("Classifier::FlushAndDisableAsyncUpdate [%p, %p]", this, mUpdateThread.get()));
+
+  if (!mUpdateThread) {
+    LOG(("Async update has been disabled."));
+    return;
+  }
+
+  mUpdateThread->Shutdown();
+  mUpdateThread = nullptr;
+}
+
 nsresult
-Classifier::ApplyUpdates(nsTArray<TableUpdate*>* aUpdates)
+Classifier::AsyncApplyUpdates(nsTArray<TableUpdate*>* aUpdates,
+                              AsyncUpdateCallback aCallback)
 {
-  nsCString failedTableName;
+  LOG(("Classifier::AsyncApplyUpdates"));
+
+  if (!mUpdateThread) {
+    LOG(("Async update has already been disabled."));
+    return NS_ERROR_FAILURE;
+  }
+
+  //         Caller thread      |       Update thread
+  // --------------------------------------------------------
+  //                            |    ApplyUpdatesBackground
+  //    (processing other task) |    (bg-update done. ping back to caller thread)
+  //    (processing other task) |    idle...
+  //    ApplyUpdatesForeground  |
+  //          callback          |
+
+  mUpdateInterrupted = false;
+  nsresult rv = mRootStoreDirectory->Clone(getter_AddRefs(mRootStoreDirectoryForUpdate));
+  if (NS_FAILED(rv)) {
+    LOG(("Failed to clone mRootStoreDirectory for update."));
+    return rv;
+  }
 
-  nsresult bgRv = ApplyUpdatesBackground(aUpdates, failedTableName);
-  return ApplyUpdatesForeground(bgRv, failedTableName);
+  nsCOMPtr<nsIThread> callerThread = NS_GetCurrentThread();
+  MOZ_ASSERT(callerThread != mUpdateThread);
+
+  nsCOMPtr<nsIRunnable> bgRunnable = NS_NewRunnableFunction([=] {
+    MOZ_ASSERT(NS_GetCurrentThread() == mUpdateThread, "MUST be on update thread");
+
+    LOG(("Step 1. ApplyUpdatesBackground on update thread."));
+    nsCString failedTableName;
+    nsresult bgRv = ApplyUpdatesBackground(aUpdates, failedTableName);
+
+    nsCOMPtr<nsIRunnable> fgRunnable = NS_NewRunnableFunction([=] {
+      MOZ_ASSERT(NS_GetCurrentThread() == callerThread, "MUST be on caller thread");
+
+      LOG(("Step 2. ApplyUpdatesForeground on caller thread"));
+      nsresult rv = ApplyUpdatesForeground(bgRv, failedTableName);;
+
+      LOG(("Step 3. Updates applied! Fire callback."));
+
+      aCallback(rv);
+    });
+    callerThread->Dispatch(fgRunnable, NS_DISPATCH_NORMAL);
+  });
+
+  return mUpdateThread->Dispatch(bgRunnable, NS_DISPATCH_NORMAL);
 }
 
 nsresult
 Classifier::ApplyUpdatesBackground(nsTArray<TableUpdate*>* aUpdates,
                                    nsACString& aFailedTableName)
 {
-  // Will run on the update thread after Bug 1339760.
-  // No lock is required except for CopyInUseDirForUpdate() since
-  // the table lookup may lead to database removal.
+  // |mUpdateInterrupted| is guaranteed to have been unset.
+  // If |mUpdateInterrupted| is set at any point, Reset() must have
+  // been called then we need to interrupt the update process.
+  // We only add checkpoints for non-trivial tasks.
 
   if (!aUpdates || aUpdates->Length() == 0) {
     return NS_OK;
   }
 
   nsCOMPtr<nsIUrlClassifierUtils> urlUtil =
     do_GetService(NS_URLCLASSIFIERUTILS_CONTRACTID);
 
@@ -690,54 +796,52 @@ Classifier::ApplyUpdatesBackground(nsTAr
   }
 
   nsresult rv;
 
   {
     ScopedUpdatesClearer scopedUpdatesClearer(aUpdates);
 
     {
-      // TODO: Bug 1339050. This section should be mutual exclusive to
-      // LookupCache::Open() and HashStore::Open() since those operations
-      // might fail and cause the database to be removed. It'll be fine
-      // until we move "apply update" code off the worker thread.
+      // Check point 1: Copying file takes time so we check here.
+      if (mUpdateInterrupted) {
+        LOG(("Update is interrupted. Don't copy files."));
+        return NS_OK;
+      }
+
       rv = CopyInUseDirForUpdate(); // i.e. mUpdatingDirectory will be setup.
       if (NS_FAILED(rv)) {
         LOG(("Failed to copy in-use directory for update."));
         return rv;
       }
-      rv = CopyInUseLookupCacheForUpdate(); // i.e. mNewLookupCaches will be setup.
-      if (NS_FAILED(rv)) {
-        LOG(("Failed to create lookup caches from copied files."));
-        return rv;
-      }
     }
 
     LOG(("Applying %" PRIuSIZE " table updates.", aUpdates->Length()));
 
     for (uint32_t i = 0; i < aUpdates->Length(); i++) {
       // Previous UpdateHashStore() may have consumed this update..
       if ((*aUpdates)[i]) {
         // Run all updates for one table
         nsCString updateTable(aUpdates->ElementAt(i)->TableName());
 
+        // Check point 2: Processing downloaded data takes time.
+        if (mUpdateInterrupted) {
+          LOG(("Update is interrupted. Stop building new tables."));
+          return NS_OK;
+        }
+
         // Will update the mirrored in-memory and on-disk databases.
         if (TableUpdate::Cast<TableUpdateV2>((*aUpdates)[i])) {
           rv = UpdateHashStore(aUpdates, updateTable);
         } else {
           rv = UpdateTableV4(aUpdates, updateTable);
         }
 
         if (NS_FAILED(rv)) {
           aFailedTableName = updateTable;
-          if (rv != NS_ERROR_OUT_OF_MEMORY) {
-#ifdef MOZ_SAFEBROWSING_DUMP_FAILED_UPDATES
-            DumpFailedUpdate();
-#endif
-          }
           RemoveUpdateIntermediaries();
           return rv;
         }
       }
     }
 
   } // End of scopedUpdatesClearer scope.
 
@@ -749,16 +853,21 @@ Classifier::ApplyUpdatesBackground(nsTAr
 
   return rv;
 }
 
 nsresult
 Classifier::ApplyUpdatesForeground(nsresult aBackgroundRv,
                                    const nsACString& aFailedTableName)
 {
+  if (mUpdateInterrupted) {
+    LOG(("Update is interrupted! Just remove update intermediaries."));
+    RemoveUpdateIntermediaries();
+    return NS_OK;
+  }
   if (NS_SUCCEEDED(aBackgroundRv)) {
     return SwapInNewTablesAndCleanup();
   }
   if (NS_ERROR_OUT_OF_MEMORY != aBackgroundRv) {
     ResetTables(Clear_All, nsTArray<nsCString> { nsCString(aFailedTableName) });
   }
   return aBackgroundRv;
 }
@@ -913,20 +1022,19 @@ Classifier::GetFailedUpdateDirectroy()
   }
 
   return failedUpdatekDirectory.forget();
 }
 
 nsresult
 Classifier::DumpRawTableUpdates(const nsACString& aRawUpdates)
 {
-  // DumpFailedUpdate() MUST be called first to create the
-  // "failed update" directory
+  LOG(("Dumping raw table updates..."));
 
-  LOG(("Dumping raw table updates..."));
+  DumpFailedUpdate();
 
   nsCOMPtr<nsIFile> failedUpdatekDirectory = GetFailedUpdateDirectroy();
 
   // Create tableupdate.bin and dump raw table update data.
   nsCOMPtr<nsIFile> rawTableUpdatesFile;
   nsCOMPtr<nsIOutputStream> outputStream;
   if (NS_FAILED(failedUpdatekDirectory->Clone(getter_AddRefs(rawTableUpdatesFile))) ||
       NS_FAILED(rawTableUpdatesFile->AppendNative(nsCString("tableupdates.bin"))) ||
@@ -985,46 +1093,27 @@ Classifier::CopyInUseDirForUpdate()
   // for updating.
 
   nsCString updatingDirName;
   nsresult rv = mUpdatingDirectory->GetNativeLeafName(updatingDirName);
   NS_ENSURE_SUCCESS(rv, rv);
 
   // Remove the destination directory first (just in case) the do the copy.
   mUpdatingDirectory->Remove(true);
-  rv = mRootStoreDirectory->CopyToNative(nullptr, updatingDirName);
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  // We moved some things to new places, so move the handles around, too.
-  rv = SetupPathNames();
+  if (!mRootStoreDirectoryForUpdate) {
+    LOG(("mRootStoreDirectoryForUpdate is null."));
+    return NS_ERROR_NULL_POINTER;
+  }
+  rv = mRootStoreDirectoryForUpdate->CopyToNative(nullptr, updatingDirName);
   NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
 
 nsresult
-Classifier::CopyInUseLookupCacheForUpdate()
-{
-  MOZ_ASSERT(mNewLookupCaches.IsEmpty(), "Update intermediaries is forgotten to "
-                                         "be removed in the previous update.");
-
-  // TODO: Bug 1341183 - Implement deep copy function or copy ctor.
-  // For now we just re-open every table we have opened.
-  // Another option is to NOT pre-open these LookupCaches and
-  // do a "merge" instead of "swap" after update.
-  for (auto c: mLookupCaches) {
-    if (!GetLookupCacheForUpdate(c->TableName())) {
-      return NS_ERROR_FAILURE;
-    }
-  }
-
-  return NS_OK;
-}
-
-nsresult
 Classifier::RecoverBackups()
 {
   bool backupExists;
   nsresult rv = mBackupDirectory->Exists(&backupExists);
   NS_ENSURE_SUCCESS(rv, rv);
 
   if (backupExists) {
     // Remove the safebrowsing dir if it exists
@@ -1315,16 +1404,21 @@ Classifier::UpdateCache(TableUpdate* aUp
 #endif
 
   return NS_OK;
 }
 
 LookupCache *
 Classifier::GetLookupCache(const nsACString& aTable, bool aForUpdate)
 {
+  if (aForUpdate) {
+    MOZ_ASSERT(NS_GetCurrentThread() == mUpdateThread,
+               "GetLookupCache(aForUpdate==true) can only be called on update thread.");
+  }
+
   nsTArray<LookupCache*>& lookupCaches = aForUpdate ? mNewLookupCaches
                                                     : mLookupCaches;
   auto& rootStoreDirectory = aForUpdate ? mUpdatingDirectory
                                         : mRootStoreDirectory;
 
   for (auto c: lookupCaches) {
     if (c->TableName().Equals(aTable)) {
       return c;
--- a/toolkit/components/url-classifier/Classifier.h
+++ b/toolkit/components/url-classifier/Classifier.h
@@ -11,16 +11,18 @@
 #include "ProtocolParser.h"
 #include "LookupCache.h"
 #include "nsCOMPtr.h"
 #include "nsString.h"
 #include "nsIFile.h"
 #include "nsICryptoHash.h"
 #include "nsDataHashtable.h"
 
+class nsIThread;
+
 namespace mozilla {
 namespace safebrowsing {
 
 /**
  * Maintains the stores and LookupCaches for the url classifier.
  */
 class Classifier {
 public:
@@ -60,42 +62,32 @@ public:
    * Check a URL against the specified tables.
    */
   nsresult Check(const nsACString& aSpec,
                  const nsACString& tables,
                  uint32_t aFreshnessGuarantee,
                  LookupResultArray& aResults);
 
   /**
-   * Apply the table updates in the array.  Takes ownership of
-   * the updates in the array and clears it.  Wacky!
+   * Asynchronously apply updates to the in-use databases. When the
+   * update is complete, the caller can be notified by |aCallback|, which
+   * will occur on the caller thread. Note that the ownership of
+   * |aUpdates| will be transferred. This design is inherited from the
+   * previous sync update function (ApplyUpdates) which has been removed.
    */
-  nsresult ApplyUpdates(nsTArray<TableUpdate*>* aUpdates);
+  using AsyncUpdateCallback = std::function<void(nsresult)>;
+  nsresult AsyncApplyUpdates(nsTArray<TableUpdate*>* aUpdates,
+                             AsyncUpdateCallback aCallback);
 
   /**
-   * The "background" part of ApplyUpdates. Once the background update
-   * is called, the foreground update has to be called along with the
-   * background result no matter whether the background update is
-   * successful or not.
+   * Wait until the ongoing async update is finished and callback
+   * is fired. Once this function returns, AsyncApplyUpdates is
+   * no longer available.
    */
-  nsresult ApplyUpdatesBackground(nsTArray<TableUpdate*>* aUpdates,
-                                  nsACString& aFailedTableName);
-
-  /**
-   * The "foreground" part of ApplyUpdates. The in-use data (in-memory and
-   * on-disk) will be touched so this MUST be mutually exclusive to other
-   * member functions.
-   *
-   * If |aBackgroundRv| is successful, the return value is the result of
-   * bringing stuff to the foreground. Otherwise, the foreground table may
-   * be reset according to the background update failed reason and
-   * |aBackgroundRv| will be returned to forward the background update result.
-   */
-  nsresult ApplyUpdatesForeground(nsresult aBackgroundRv,
-                                  const nsACString& aFailedTableName);
+  void FlushAndDisableAsyncUpdate();
 
   /**
    * Apply full hashes retrived from gethash to cache.
    */
   nsresult ApplyFullHashes(nsTArray<TableUpdate*>* aUpdates);
 
   void SetLastUpdateTime(const nsACString& aTableName, uint64_t updateTime);
   int64_t GetLastUpdateTime(const nsACString& aTableName);
@@ -142,20 +134,19 @@ private:
   void DropStores();
   void DeleteTables(nsIFile* aDirectory, const nsTArray<nsCString>& aTables);
 
   nsresult CreateStoreDirectory();
   nsresult SetupPathNames();
   nsresult RecoverBackups();
   nsresult CleanToDelete();
   nsresult CopyInUseDirForUpdate();
-  nsresult CopyInUseLookupCacheForUpdate();
   nsresult RegenActiveTables();
 
-
+  void MergeNewLookupCaches(); // Merge mNewLookupCaches into mLookupCaches.
 
   // Remove any intermediary for update, including in-memory
   // and on-disk data.
   void RemoveUpdateIntermediaries();
 
 #ifdef MOZ_SAFEBROWSING_DUMP_FAILED_UPDATES
   already_AddRefed<nsIFile> GetFailedUpdateDirectroy();
   nsresult DumpFailedUpdate();
@@ -182,16 +173,38 @@ private:
 
   bool CheckValidUpdate(nsTArray<TableUpdate*>* aUpdates,
                         const nsACString& aTable);
 
   nsresult LoadMetadata(nsIFile* aDirectory, nsACString& aResult);
 
   nsCString GetProvider(const nsACString& aTableName);
 
+  /**
+   * The "background" part of ApplyUpdates. Once the background update
+   * is called, the foreground update has to be called along with the
+   * background result no matter whether the background update is
+   * successful or not.
+   */
+  nsresult ApplyUpdatesBackground(nsTArray<TableUpdate*>* aUpdates,
+                                  nsACString& aFailedTableName);
+
+  /**
+   * The "foreground" part of ApplyUpdates. The in-use data (in-memory and
+   * on-disk) will be touched so this MUST be mutually exclusive to other
+   * member functions.
+   *
+   * If |aBackgroundRv| is successful, the return value is the result of
+   * bringing stuff to the foreground. Otherwise, the foreground table may
+   * be reset according to the background update failed reason and
+   * |aBackgroundRv| will be returned to forward the background update result.
+   */
+  nsresult ApplyUpdatesForeground(nsresult aBackgroundRv,
+                                  const nsACString& aFailedTableName);
+
   // Root dir of the Local profile.
   nsCOMPtr<nsIFile> mCacheDirectory;
   // Main directory where to store the databases.
   nsCOMPtr<nsIFile> mRootStoreDirectory;
   // Used for atomically updating the other dirs.
   nsCOMPtr<nsIFile> mBackupDirectory;
   nsCOMPtr<nsIFile> mUpdatingDirectory; // For update only.
   nsCOMPtr<nsIFile> mToDeleteDirectory;
@@ -207,14 +220,24 @@ private:
   nsCString mTableRequestResult;
 
   // Whether mTableRequestResult is outdated and needs to
   // be reloaded from disk.
   bool mIsTableRequestResultOutdated;
 
   // The copy of mLookupCaches for update only.
   nsTArray<LookupCache*> mNewLookupCaches;
+
+  bool mUpdateInterrupted;
+
+  nsCOMPtr<nsIThread> mUpdateThread; // For async update.
+
+  // Identical to mRootStoreDirectory but for update only because
+  // nsIFile is not thread safe and mRootStoreDirectory needs to
+  // be accessed in CopyInUseDirForUpdate().
+  // It will be initialized right before update on the worker thread.
+  nsCOMPtr<nsIFile> mRootStoreDirectoryForUpdate;
 };
 
 } // namespace safebrowsing
 } // namespace mozilla
 
 #endif
--- a/toolkit/components/url-classifier/HashStore.cpp
+++ b/toolkit/components/url-classifier/HashStore.cpp
@@ -96,32 +96,16 @@
 // Name of the SafeBrowsing store
 #define STORE_SUFFIX ".sbstore"
 
 // MOZ_LOG=UrlClassifierDbService:5
 extern mozilla::LazyLogModule gUrlClassifierDbServiceLog;
 #define LOG(args) MOZ_LOG(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug, args)
 #define LOG_ENABLED() MOZ_LOG_TEST(gUrlClassifierDbServiceLog, mozilla::LogLevel::Debug)
 
-// Either the return was successful or we call the Reset function (unless we
-// hit an OOM).  Used while reading in the store.
-#define SUCCESS_OR_RESET(res)                                             \
-  do {                                                                    \
-    nsresult __rv = res; /* Don't evaluate |res| more than once */        \
-    if (__rv == NS_ERROR_OUT_OF_MEMORY) {                                 \
-      NS_WARNING("SafeBrowsing OOM.");                                    \
-      return __rv;                                                        \
-    }                                                                     \
-    if (NS_FAILED(__rv)) {                                                \
-      NS_WARNING("SafeBrowsing store corrupted or out of date.");         \
-      Reset();                                                            \
-      return __rv;                                                        \
-    }                                                                     \
-  } while(0)
-
 namespace mozilla {
 namespace safebrowsing {
 
 const uint32_t STORE_MAGIC = 0x1231af3b;
 const uint32_t CURRENT_VERSION = 3;
 
 nsresult
 TableUpdateV2::NewAddPrefix(uint32_t aAddChunk, const Prefix& aHash)
@@ -301,34 +285,34 @@ HashStore::Open()
   nsCOMPtr<nsIInputStream> origStream;
   rv = NS_NewLocalFileInputStream(getter_AddRefs(origStream), storeFile,
                                   PR_RDONLY | nsIFile::OS_READAHEAD);
 
   if (rv == NS_ERROR_FILE_NOT_FOUND) {
     UpdateHeader();
     return NS_OK;
   }
-  SUCCESS_OR_RESET(rv);
+  NS_ENSURE_SUCCESS(rv, rv);
 
   int64_t fileSize;
   rv = storeFile->GetFileSize(&fileSize);
   NS_ENSURE_SUCCESS(rv, rv);
 
   if (fileSize < 0 || fileSize > UINT32_MAX) {
     return NS_ERROR_FAILURE;
   }
 
   mFileSize = static_cast<uint32_t>(fileSize);
   mInputStream = NS_BufferInputStream(origStream, mFileSize);
 
   rv = ReadHeader();
-  SUCCESS_OR_RESET(rv);
+  NS_ENSURE_SUCCESS(rv, rv);
 
   rv = SanityCheck();
-  SUCCESS_OR_RESET(rv);
+  NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
 
 nsresult
 HashStore::ReadHeader()
 {
   if (!mInputStream) {
@@ -507,23 +491,23 @@ HashStore::ReadCompletions()
 
   return NS_OK;
 }
 
 nsresult
 HashStore::PrepareForUpdate()
 {
   nsresult rv = CheckChecksum(mFileSize);
-  SUCCESS_OR_RESET(rv);
+  NS_ENSURE_SUCCESS(rv, rv);
 
   rv = ReadChunkNumbers();
-  SUCCESS_OR_RESET(rv);
+  NS_ENSURE_SUCCESS(rv, rv);
 
   rv = ReadHashes();
-  SUCCESS_OR_RESET(rv);
+  NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
 
 nsresult
 HashStore::BeginUpdate()
 {
   // Check wether the file is corrupted and read the rest of the store
--- a/toolkit/components/url-classifier/LookupCache.cpp
+++ b/toolkit/components/url-classifier/LookupCache.cpp
@@ -89,36 +89,16 @@ LookupCache::UpdateRootDirHandle(nsIFile
     LOG(("Private store directory for %s is %s", mTableName.get(),
                                                  NS_ConvertUTF16toUTF8(path).get()));
   }
 
   return rv;
 }
 
 nsresult
-LookupCache::Reset()
-{
-  LOG(("LookupCache resetting"));
-
-  nsCOMPtr<nsIFile> prefixsetFile;
-  nsresult rv = mStoreDirectory->Clone(getter_AddRefs(prefixsetFile));
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  rv = prefixsetFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  rv = prefixsetFile->Remove(false);
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  ClearAll();
-
-  return NS_OK;
-}
-
-nsresult
 LookupCache::AddCompletionsToCache(AddCompleteArray& aAddCompletes)
 {
   for (uint32_t i = 0; i < aAddCompletes.Length(); i++) {
     if (mGetHashCache.BinaryIndexOf(aAddCompletes[i].CompleteHash()) == mGetHashCache.NoIndex) {
       mGetHashCache.AppendElement(aAddCompletes[i].CompleteHash());
     }
   }
   mGetHashCache.Sort();
@@ -364,19 +344,16 @@ LookupCache::LoadPrefixSet()
   bool exists;
   rv = psFile->Exists(&exists);
   NS_ENSURE_SUCCESS(rv, rv);
 
   if (exists) {
     LOG(("stored PrefixSet exists, loading from disk"));
     rv = LoadFromFile(psFile);
     if (NS_FAILED(rv)) {
-      if (rv == NS_ERROR_FILE_CORRUPTED) {
-        Reset();
-      }
       return rv;
     }
     mPrimed = true;
   } else {
     LOG(("no (usable) stored PrefixSet found"));
   }
 
 #ifdef DEBUG
--- a/toolkit/components/url-classifier/LookupCache.h
+++ b/toolkit/components/url-classifier/LookupCache.h
@@ -177,17 +177,16 @@ public:
   virtual void ClearAll();
 
   template<typename T>
   static T* Cast(LookupCache* aThat) {
     return ((aThat && T::VER == aThat->Ver()) ? reinterpret_cast<T*>(aThat) : nullptr);
   }
 
 private:
-  nsresult Reset();
   nsresult LoadPrefixSet();
 
   virtual nsresult StoreToFile(nsIFile* aFile) = 0;
   virtual nsresult LoadFromFile(nsIFile* aFile) = 0;
   virtual size_t SizeOfPrefixSet() = 0;
 
   virtual int Ver() const = 0;
 
--- a/toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
+++ b/toolkit/components/url-classifier/nsUrlClassifierDBService.cpp
@@ -129,19 +129,16 @@ LazyLogModule gUrlClassifierDbServiceLog
 
 class nsUrlClassifierDBServiceWorker;
 
 // Singleton instance.
 static nsUrlClassifierDBService* sUrlClassifierDBService;
 
 nsIThread* nsUrlClassifierDBService::gDbBackgroundThread = nullptr;
 
-// For update only.
-static nsIThread* gDbUpdateThread = nullptr;
-
 // Once we've committed to shutting down, don't do work in the background
 // thread.
 static bool gShuttingDownThread = false;
 
 static mozilla::Atomic<int32_t> gFreshnessGuarantee(CONFIRM_AGE_DEFAULT_SEC);
 
 NS_IMPL_ISUPPORTS(nsUrlClassifierDBServiceWorker,
                   nsIUrlClassifierDBService)
@@ -605,16 +602,21 @@ nsUrlClassifierDBServiceWorker::FinishSt
   mProtocolParser = nullptr;
 
   return NS_OK;
 }
 
 NS_IMETHODIMP
 nsUrlClassifierDBServiceWorker::FinishUpdate()
 {
+  LOG(("nsUrlClassifierDBServiceWorker::FinishUpdate"));
+
+  MOZ_ASSERT(!NS_IsMainThread(), "nsUrlClassifierDBServiceWorker::FinishUpdate "
+                                 "NUST NOT be on the main thread.");
+
   if (gShuttingDownThread) {
     return NS_ERROR_NOT_INITIALIZED;
   }
 
   MOZ_ASSERT(!mProtocolParser, "Should have been nulled out in FinishStream() "
                                "or never created.");
 
   NS_ENSURE_STATE(mUpdateObserver);
@@ -626,55 +628,47 @@ nsUrlClassifierDBServiceWorker::FinishUp
   }
 
   if (mTableUpdates.IsEmpty()) {
     LOG(("Nothing to update. Just notify update observer."));
     return NotifyUpdateObserver(NS_OK);
   }
 
   RefPtr<nsUrlClassifierDBServiceWorker> self = this;
-
-  // TODO: Asynchronously dispatch |ApplyUpdatesBackground| to update thread.
-  //       See Bug 1339050. For now we *synchronously* run
-  //       ApplyUpdatesForeground() after ApplyUpdatesBackground().
-  nsresult backgroundRv;
-  nsCString failedTableName;
-  nsCOMPtr<nsIRunnable> r =
-    NS_NewRunnableFunction([self, &backgroundRv, &failedTableName] () -> void {
-      backgroundRv = self->ApplyUpdatesBackground(failedTableName);
-    });
+  nsresult rv = mClassifier->AsyncApplyUpdates(&mTableUpdates,
+                                               [=] (nsresult aRv) -> void {
+#ifdef MOZ_SAFEBROWSING_DUMP_FAILED_UPDATES
+    if (NS_FAILED(aRv) &&
+        NS_ERROR_OUT_OF_MEMORY != aRv &&
+        NS_ERROR_UC_UPDATE_SHUTDOWNING != aRv) {
+      self->mClassifier->DumpRawTableUpdates(mRawTableUpdates);
+    }
+    // Invalidate the raw table updates.
+    self->mRawTableUpdates = EmptyCString();
+#endif
 
-  LOG(("Bulding new tables..."));
-  mozilla::SyncRunnable::DispatchToThread(gDbUpdateThread, r);
-  LOG(("Done bulding new tables. Try to commit them."));
-
-  return ApplyUpdatesForeground(backgroundRv, failedTableName);
-}
+    self->NotifyUpdateObserver(aRv);
+  });
 
-nsresult
-nsUrlClassifierDBServiceWorker::ApplyUpdatesForeground(nsresult aBackgroundRv,
-                                                      const nsACString& aFailedTableName)
-{
-  if (gShuttingDownThread) {
-    return NS_ERROR_NOT_INITIALIZED;
+  if (NS_FAILED(rv)) {
+    LOG(("Failed to start async update. Notify immediately."));
+    NotifyUpdateObserver(rv);
   }
 
-  MOZ_ASSERT(NS_GetCurrentThread() == nsUrlClassifierDBService::BackgroundThread(),
-             "nsUrlClassifierDBServiceWorker::ApplyUpdatesForeground MUST "
-             "run on worker thread!!");
-
-  nsresult rv =
-    mClassifier->ApplyUpdatesForeground(aBackgroundRv, aFailedTableName);
-
-  return NotifyUpdateObserver(rv);
+  return rv;
 }
 
 nsresult
 nsUrlClassifierDBServiceWorker::NotifyUpdateObserver(nsresult aUpdateStatus)
 {
+  MOZ_ASSERT(!NS_IsMainThread(), "nsUrlClassifierDBServiceWorker::NotifyUpdateObserver "
+                                 "NUST NOT be on the main thread.");
+
+  LOG(("nsUrlClassifierDBServiceWorker::NotifyUpdateObserver"));
+
   // We've either
   //  1) failed starting a download stream
   //  2) succeeded in starting a download stream but failed to obtain
   //     table updates
   //  3) succeeded in obtaining table updates but failed to build new
   //     tables.
   //  4) succeeded in building new tables but failed to take them.
   //  5) succeeded in taking new tables.
@@ -697,75 +691,48 @@ nsUrlClassifierDBServiceWorker::NotifyUp
   // Do not record telemetry for testing tables.
   if (!provider.Equals(TESTING_TABLE_PROVIDER_NAME)) {
     Telemetry::Accumulate(Telemetry::URLCLASSIFIER_UPDATE_ERROR, provider,
                           NS_ERROR_GET_CODE(updateStatus));
   }
 
   mMissCache.Clear();
 
+  // Null out mUpdateObserver before notifying so that BeginUpdate()
+  // becomes available prior to callback.
+  nsCOMPtr<nsIUrlClassifierUpdateObserver> updateObserver = nullptr;
+  updateObserver.swap(mUpdateObserver);
+
   if (NS_SUCCEEDED(mUpdateStatus)) {
     LOG(("Notifying success: %d", mUpdateWaitSec));
-    mUpdateObserver->UpdateSuccess(mUpdateWaitSec);
+    updateObserver->UpdateSuccess(mUpdateWaitSec);
   } else if (NS_ERROR_NOT_IMPLEMENTED == mUpdateStatus) {
     LOG(("Treating NS_ERROR_NOT_IMPLEMENTED a successful update "
          "but still mark it spoiled."));
-    mUpdateObserver->UpdateSuccess(0);
+    updateObserver->UpdateSuccess(0);
     mClassifier->ResetTables(Classifier::Clear_Cache, mUpdateTables);
   } else {
     if (LOG_ENABLED()) {
       nsAutoCString errorName;
       mozilla::GetErrorName(mUpdateStatus, errorName);
       LOG(("Notifying error: %s (%" PRIu32 ")", errorName.get(),
            static_cast<uint32_t>(mUpdateStatus)));
     }
 
-    mUpdateObserver->UpdateError(mUpdateStatus);
+    updateObserver->UpdateError(mUpdateStatus);
     /*
      * mark the tables as spoiled(clear cache in LookupCache), we don't want to
      * block hosts longer than normal because our update failed
     */
     mClassifier->ResetTables(Classifier::Clear_Cache, mUpdateTables);
   }
-  mUpdateObserver = nullptr;
 
   return NS_OK;
 }
 
-nsresult
-nsUrlClassifierDBServiceWorker::ApplyUpdatesBackground(nsACString& aFailedTableName)
-{
-  // ********* Please be very careful while adding tasks. *********
-  // This function will run on the update thread (whereas most of the other
-  // functions will run on the worker thread) so any task that might mutate
-  // the in-use data is forbidden.
-
-  if (gShuttingDownThread) {
-    return NS_ERROR_NOT_INITIALIZED;
-  }
-
-  MOZ_ASSERT(NS_GetCurrentThread() == gDbUpdateThread,
-             "nsUrlClassifierDBServiceWorker::BuildNewTables MUST "
-             "run on update thread!!");
-
-  LOG(("nsUrlClassifierDBServiceWorker::BuildNewTables()"));
-  nsresult rv = mClassifier->ApplyUpdatesBackground(&mTableUpdates,
-                                                    aFailedTableName);
-
-#ifdef MOZ_SAFEBROWSING_DUMP_FAILED_UPDATES
-  if (NS_FAILED(rv) && NS_ERROR_OUT_OF_MEMORY != rv) {
-    mClassifier->DumpRawTableUpdates(mRawTableUpdates);
-  }
-  // Invalidate the raw table updates.
-  mRawTableUpdates = EmptyCString();
-#endif
-
-  return rv;
-}
-
 NS_IMETHODIMP
 nsUrlClassifierDBServiceWorker::ResetDatabase()
 {
   nsresult rv = OpenDb();
 
   if (NS_SUCCEEDED(rv)) {
     mClassifier->Reset();
   }
@@ -830,16 +797,26 @@ nsUrlClassifierDBServiceWorker::CancelUp
     ResetUpdate();
   } else {
     LOG(("No UpdateObserver, nothing to cancel"));
   }
 
   return NS_OK;
 }
 
+void
+nsUrlClassifierDBServiceWorker::FlushAndDisableAsyncUpdate()
+{
+  LOG(("nsUrlClassifierDBServiceWorker::FlushAndDisableAsyncUpdate()"));
+
+  if (mClassifier) {
+    mClassifier->FlushAndDisableAsyncUpdate();
+  }
+}
+
 // Allows the main thread to delete the connection which may be in
 // a background thread.
 // XXX This could be turned into a single shutdown event so the logic
 // is simpler in nsUrlClassifierDBService::Shutdown.
 nsresult
 nsUrlClassifierDBServiceWorker::CloseDb()
 {
   if (mClassifier) {
@@ -1586,23 +1563,16 @@ nsUrlClassifierDBService::Init()
     }
   }
 
   // Start the background thread.
   rv = NS_NewNamedThread("URL Classifier", &gDbBackgroundThread);
   if (NS_FAILED(rv))
     return rv;
 
-  // Start the update thread.
-  rv = NS_NewNamedThread(NS_LITERAL_CSTRING("SafeBrowsing DB Update"),
-                         &gDbUpdateThread);
-  if (NS_FAILED(rv)) {
-    return rv;
-  }
-
   mWorker = new nsUrlClassifierDBServiceWorker();
   if (!mWorker)
     return NS_ERROR_OUT_OF_MEMORY;
 
   rv = mWorker->Init(gethashNoise, cacheDir);
   if (NS_FAILED(rv)) {
     mWorker = nullptr;
     return rv;
@@ -2040,16 +2010,42 @@ nsUrlClassifierDBService::BeginUpdate(ns
                                       const nsACString &updateTables)
 {
   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
 
   if (mInUpdate) {
     LOG(("Already updating, not available"));
     return NS_ERROR_NOT_AVAILABLE;
   }
+  if (mWorker->IsBusyUpdating()) {
+    // |mInUpdate| used to work well because "notifying update observer"
+    // is synchronously done in Worker::FinishUpdate(). Even if the
+    // update observer hasn't been notified at this point, we can still
+    // dispatch BeginUpdate() since it will NOT be run until the
+    // previous Worker::FinishUpdate() returns.
+    //
+    // However, some tasks in Worker::FinishUpdate() have been moved to
+    // another thread. The update observer will NOT be notified when
+    // Worker::FinishUpdate() returns. If we only check |mInUpdate|,
+    // the following sequence might happen on worker thread:
+    //
+    // Worker::FinishUpdate() // for update 1
+    // Worker::BeginUpdate()  // for update 2
+    // Worker::NotifyUpdateObserver() // for update 1
+    //
+    // So, we have to find out a way to reject BeginUpdate() right here
+    // if the previous update observer hasn't been notified.
+    //
+    // Directly probing the worker's state is the most lightweight solution.
+    // No lock is required since Worker::BeginUpdate() and
+    // Worker::NotifyUpdateObserver() are by nature mutual exclusive.
+    // (both run on worker thread.)
+    LOG(("The previous update observer hasn't been notified."));
+    return NS_ERROR_NOT_AVAILABLE;
+  }
 
   mInUpdate = true;
 
   // The proxy observer uses the current thread
   nsCOMPtr<nsIUrlClassifierUpdateObserver> proxyObserver =
     new UrlClassifierUpdateObserverProxy(observer);
 
   return mWorkerProxy->BeginUpdate(proxyObserver, updateTables);
@@ -2100,24 +2096,34 @@ nsUrlClassifierDBService::CancelUpdate()
   return mWorkerProxy->CancelUpdate();
 }
 
 NS_IMETHODIMP
 nsUrlClassifierDBService::ResetDatabase()
 {
   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
 
+  if (mWorker->IsBusyUpdating()) {
+    LOG(("Failed to ResetDatabase because of the unfinished update."));
+    return NS_ERROR_FAILURE;
+  }
+
   return mWorkerProxy->ResetDatabase();
 }
 
 NS_IMETHODIMP
 nsUrlClassifierDBService::ReloadDatabase()
 {
   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
 
+  if (mWorker->IsBusyUpdating()) {
+    LOG(("Failed to ReloadDatabase because of the unfinished update."));
+    return NS_ERROR_FAILURE;
+  }
+
   return mWorkerProxy->ReloadDatabase();
 }
 
 nsresult
 nsUrlClassifierDBService::CacheCompletions(CacheResultArray *results)
 {
   NS_ENSURE_TRUE(gDbBackgroundThread, NS_ERROR_NOT_INITIALIZED);
 
@@ -2184,61 +2190,39 @@ nsUrlClassifierDBService::Observe(nsISup
       NS_LITERAL_STRING(DISALLOW_COMPLETION_TABLE_PREF).Equals(aData)) {
       // Just read everything again.
       ReadTablesFromPrefs();
     } else if (NS_LITERAL_STRING(CONFIRM_AGE_PREF).Equals(aData)) {
       gFreshnessGuarantee = Preferences::GetInt(CONFIRM_AGE_PREF,
         CONFIRM_AGE_DEFAULT_SEC);
     }
   } else if (!strcmp(aTopic, "quit-application")) {
-    Shutdown();
+    // Tell the update thread to finish as soon as possible.
+    gShuttingDownThread = true;
   } else if (!strcmp(aTopic, "profile-before-change")) {
-    // Unit test does not receive "quit-application",
-    // need call shutdown in this case
+    gShuttingDownThread = true;
     Shutdown();
-    LOG(("joining background thread"));
-    mWorkerProxy = nullptr;
-
-    if (gDbBackgroundThread) {
-      nsIThread *backgroundThread = gDbBackgroundThread;
-      // Nulling out gDbBackgroundThread MUST be done before Shutdown()
-      // since Shutdown() will lead to processing pending events.
-      gDbBackgroundThread = nullptr;
-      backgroundThread->Shutdown();
-      NS_RELEASE(backgroundThread);
-    }
-
-    if (gDbUpdateThread) {
-      nsIThread *updateThread = gDbUpdateThread;
-      gDbUpdateThread = nullptr;
-      updateThread->Shutdown();
-      NS_RELEASE(updateThread);
-    }
-
-    return NS_OK;
   } else {
     return NS_ERROR_UNEXPECTED;
   }
 
   return NS_OK;
 }
 
 // Join the background thread if it exists.
 nsresult
 nsUrlClassifierDBService::Shutdown()
 {
   LOG(("shutting down db service\n"));
   MOZ_ASSERT(XRE_IsParentProcess());
 
-  if ((!gDbBackgroundThread && !gDbUpdateThread) || gShuttingDownThread) {
+  if (!gDbBackgroundThread) {
     return NS_OK;
   }
 
-  gShuttingDownThread = true;
-
   Telemetry::AutoTimer<Telemetry::URLCLASSIFIER_SHUTDOWN_TIME> timer;
 
   mCompleters.Clear();
 
   nsCOMPtr<nsIPrefBranch> prefs = do_GetService(NS_PREFSERVICE_CONTRACTID);
   if (prefs) {
     prefs->RemoveObserver(CHECK_MALWARE_PREF, this);
     prefs->RemoveObserver(CHECK_PHISHING_PREF, this);
@@ -2249,24 +2233,55 @@ nsUrlClassifierDBService::Shutdown()
     prefs->RemoveObserver(TRACKING_WHITELIST_TABLE_PREF, this);
     prefs->RemoveObserver(BLOCKED_TABLE_PREF, this);
     prefs->RemoveObserver(DOWNLOAD_BLOCK_TABLE_PREF, this);
     prefs->RemoveObserver(DOWNLOAD_ALLOW_TABLE_PREF, this);
     prefs->RemoveObserver(DISALLOW_COMPLETION_TABLE_PREF, this);
     prefs->RemoveObserver(CONFIRM_AGE_PREF, this);
   }
 
+  // 1. Synchronize with worker thread and update thread by
+  //    *synchronously* dispatching an event to worker thread
+  //    for shutting down the update thread. The reason not
+  //    shutting down update thread directly from main thread
+  //    is to avoid racing for Classifier::mUpdateThread
+  //    between main thread and the worker thread. (Both threads
+  //    would access Classifier::mUpdateThread.)
+  using Worker = nsUrlClassifierDBServiceWorker;
+  RefPtr<nsIRunnable> r =
+    NewRunnableMethod(mWorker, &Worker::FlushAndDisableAsyncUpdate);
+  SyncRunnable::DispatchToThread(gDbBackgroundThread, r);
+
+  // At this point the update thread has been shut down and
+  // the worker thread should only have at most one event,
+  // which is the callback event.
+
+  // 2. Send CancelUpdate() event to notify the dangling update.
+  //    (i.e. BeginUpdate is called but FinishUpdate is not.)
+  //    and CloseDb() to clear mClassifier. They will be the last two
+  //    events on the worker thread in the shutdown process.
   DebugOnly<nsresult> rv;
-  // First close the db connection.
-  if (mWorker) {
-    rv = mWorkerProxy->CancelUpdate();
-    NS_ASSERTION(NS_SUCCEEDED(rv), "failed to post cancel update event");
+  rv = mWorkerProxy->CancelUpdate();
+  MOZ_ASSERT(NS_SUCCEEDED(rv), "failed to post 'cancel update' event");
+  rv = mWorkerProxy->CloseDb();
+  MOZ_ASSERT(NS_SUCCEEDED(rv), "failed to post 'close db' event");
+  mWorkerProxy = nullptr;
 
-    rv = mWorkerProxy->CloseDb();
-    NS_ASSERTION(NS_SUCCEEDED(rv), "failed to post close db event");
+  // 3. Invalidate XPCOM APIs by nulling out gDbBackgroundThread
+  //    since every API checks gDbBackgroundThread first. This has
+  //    to be done before calling nsIThread.shutdown because it
+  //    will cause the pending events on the joining thread to
+  //    be processed.
+  nsIThread *backgroundThread = nullptr;
+  Swap(backgroundThread, gDbBackgroundThread);
+
+  // 4. Wait until the worker thread is down.
+  if (backgroundThread) {
+    backgroundThread->Shutdown();
+    NS_RELEASE(backgroundThread);
   }
   return NS_OK;
 }
 
 nsIThread*
 nsUrlClassifierDBService::BackgroundThread()
 {
   return gDbBackgroundThread;
--- a/toolkit/components/url-classifier/nsUrlClassifierDBService.h
+++ b/toolkit/components/url-classifier/nsUrlClassifierDBService.h
@@ -177,30 +177,34 @@ public:
   nsresult GCC_MANGLING_WORKAROUND OpenDb();
 
   // Provide a way to forcibly close the db connection.
   nsresult GCC_MANGLING_WORKAROUND CloseDb();
 
   nsresult CacheCompletions(CacheResultArray * aEntries);
   nsresult CacheMisses(PrefixArray * aEntries);
 
+  // Used to probe the state of the worker thread. When the update begins,
+  // mUpdateObserver will be set. When the update finished, mUpdateObserver
+  // will be nulled out in NotifyUpdateObserver.
+  bool IsBusyUpdating() const { return !!mUpdateObserver; }
+
+  // Delegate Classifier to disable async update. If there is an
+  // ongoing update on the update thread, we will be blocked until
+  // the background update is done and callback is fired.
+  // Should be called on the worker thread.
+  void FlushAndDisableAsyncUpdate();
+
 private:
   // No subclassing
   ~nsUrlClassifierDBServiceWorker();
 
   // Disallow copy constructor
   nsUrlClassifierDBServiceWorker(nsUrlClassifierDBServiceWorker&);
 
-  // Perform update in the background. Can be run on/off the worker thread.
-  nsresult ApplyUpdatesBackground(nsACString& aFailedTableName);
-
-  // Perform update for the foreground part. MUST be run on the worker thread.
-  nsresult ApplyUpdatesForeground(nsresult aBackgroundRv,
-                                 const nsACString& aFailedTableName);
-
   nsresult NotifyUpdateObserver(nsresult aUpdateStatus);
 
   // Reset the in-progress update stream
   void ResetStream();
 
   // Reset the in-progress update
   void ResetUpdate();
 
--- a/toolkit/components/url-classifier/tests/gtest/Common.cpp
+++ b/toolkit/components/url-classifier/tests/gtest/Common.cpp
@@ -15,16 +15,61 @@ void RunTestInNewThread(Function&& aFunc
   nsCOMPtr<nsIRunnable> r = NS_NewRunnableFunction(mozilla::Forward<Function>(aFunction));
   nsCOMPtr<nsIThread> testingThread;
   nsresult rv =
     NS_NewNamedThread("Testing Thread", getter_AddRefs(testingThread), r);
   ASSERT_EQ(rv, NS_OK);
   testingThread->Shutdown();
 }
 
+nsresult SyncApplyUpdates(Classifier* aClassifier,
+                          nsTArray<TableUpdate*>* aUpdates)
+{
+  // We need to spin a new thread specifically because the callback
+  // will be on the caller thread. If we call Classifier::AsyncApplyUpdates
+  // and wait on the same thread, this function will never return.
+
+  nsresult ret;
+  bool done = false;
+  auto onUpdateComplete = [&done, &ret](nsresult rv) {
+    ret = rv;
+    done = true;
+  };
+
+  nsCOMPtr<nsIRunnable> r = NS_NewRunnableFunction([&]() {
+    nsresult rv = aClassifier->AsyncApplyUpdates(aUpdates, onUpdateComplete);
+    if (NS_FAILED(rv)) {
+      onUpdateComplete(rv);
+    }
+  });
+
+  nsCOMPtr<nsIThread> testingThread;
+  NS_NewNamedThread("ApplyUpdates", getter_AddRefs(testingThread));
+  if (!testingThread) {
+    return NS_ERROR_FAILURE;
+  }
+
+  testingThread->Dispatch(r, NS_DISPATCH_NORMAL);
+  while (!done) {
+    // NS_NewCheckSummedOutputStream in HashStore::WriteFile
+    // will synchronously init NS_CRYPTO_HASH_CONTRACTID on
+    // the main thread. As a result we have to keep processing
+    // pending event until |done| becomes true. If there's no
+    // more pending event, what we only can do is wait.
+    // Condition variable doesn't work here because instrusively
+    // notifying the from NS_NewCheckSummedOutputStream() or
+    // HashStore::WriteFile() is weird.
+    if (!NS_ProcessNextEvent(NS_GetCurrentThread(), false)) {
+      PR_Sleep(PR_MillisecondsToInterval(100));
+    }
+  }
+
+  return ret;
+}
+
 already_AddRefed<nsIFile>
 GetFile(const nsTArray<nsString>& path)
 {
   nsCOMPtr<nsIFile> file;
   nsresult rv = NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
   if (NS_WARN_IF(NS_FAILED(rv))) {
     return nullptr;
   }
@@ -48,19 +93,17 @@ void ApplyUpdate(nsTArray<TableUpdate*>&
     // because nsIUrlClassifierDBService will not run in advance
     // in gtest.
     nsresult rv;
     nsCOMPtr<nsIUrlClassifierUtils> dummy =
       do_GetService(NS_URLCLASSIFIERUTILS_CONTRACTID, &rv);
       ASSERT_TRUE(NS_SUCCEEDED(rv));
   }
 
-  RunTestInNewThread([&] () -> void {
-    classifier->ApplyUpdates(&updates);
-  });
+  SyncApplyUpdates(classifier.get(), &updates);
 }
 
 void ApplyUpdate(TableUpdate* update)
 {
   nsTArray<TableUpdate*> updates = { update };
   ApplyUpdate(updates);
 }
 
--- a/toolkit/components/url-classifier/tests/gtest/Common.h
+++ b/toolkit/components/url-classifier/tests/gtest/Common.h
@@ -1,19 +1,29 @@
 #include "HashStore.h"
 #include "nsIFile.h"
 #include "nsTArray.h"
 #include "gtest/gtest.h"
 
 using namespace mozilla;
 using namespace mozilla::safebrowsing;
 
+namespace mozilla {
+namespace safebrowsing {
+    class Classifier;
+}
+}
+
 template<typename Function>
 void RunTestInNewThread(Function&& aFunction);
 
+// Synchronously apply updates by calling Classifier::AsyncApplyUpdates.
+nsresult SyncApplyUpdates(Classifier* aClassifier,
+                          nsTArray<TableUpdate*>* aUpdates);
+
 // Return nsIFile with root directory - NS_APP_USER_PROFILE_50_DIR
 // Sub-directories are passed in path argument.
 already_AddRefed<nsIFile>
 GetFile(const nsTArray<nsString>& path);
 
 // ApplyUpdate will call |ApplyUpdates| of Classifier within a new thread
 void ApplyUpdate(nsTArray<TableUpdate*>& updates);
 
--- a/toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp
+++ b/toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp
@@ -193,38 +193,43 @@ static void
 testUpdateFail(nsTArray<TableUpdate*>& tableUpdates)
 {
   nsCOMPtr<nsIFile> file;
   NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
 
   UniquePtr<Classifier> classifier(new Classifier());
   classifier->Open(*file);
 
-  RunTestInNewThread([&] () -> void {
-    nsresult rv = classifier->ApplyUpdates(&tableUpdates);
-    ASSERT_TRUE(NS_FAILED(rv));
-  });
+  nsresult rv = SyncApplyUpdates(classifier.get(), &tableUpdates);
+  ASSERT_TRUE(NS_FAILED(rv));
 }
 
 static void
 testUpdate(nsTArray<TableUpdate*>& tableUpdates,
            PrefixStringMap& expected)
 {
   nsCOMPtr<nsIFile> file;
   NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(file));
 
+  {
+    // Force nsIUrlClassifierUtils loading on main thread
+    // because nsIUrlClassifierDBService will not run in advance
+    // in gtest.
+    nsresult rv;
+    nsCOMPtr<nsIUrlClassifierUtils> dummy =
+      do_GetService(NS_URLCLASSIFIERUTILS_CONTRACTID, &rv);
+      ASSERT_TRUE(NS_SUCCEEDED(rv));
+  }
+
   UniquePtr<Classifier> classifier(new Classifier());
   classifier->Open(*file);
 
-  RunTestInNewThread([&] () -> void {
-    nsresult rv = classifier->ApplyUpdates(&tableUpdates);
-    ASSERT_TRUE(rv == NS_OK);
-
-    VerifyPrefixSet(expected);
-  });
+  nsresult rv = SyncApplyUpdates(classifier.get(), &tableUpdates);
+  ASSERT_TRUE(rv == NS_OK);
+  VerifyPrefixSet(expected);
 }
 
 static void
 testFullUpdate(PrefixStringMap& add, nsCString* checksum)
 {
   nsTArray<TableUpdate*> tableUpdates;
 
   GenerateUpdateData(true, add, nullptr, checksum, tableUpdates);
--- a/toolkit/components/url-classifier/tests/unit/test_partial.js
+++ b/toolkit/components/url-classifier/tests/unit/test_partial.js
@@ -22,17 +22,17 @@ QueryInterface: function(iid)
 
 complete: function(partialHash, gethashUrl, tableName, cb)
 {
   this.queries.push(partialHash);
   var fragments = this.fragments;
   var self = this;
   var doCallback = function() {
       if (self.alwaysFail) {
-        cb.completionFinished(1);
+        cb.completionFinished(Cr.NS_ERROR_FAILURE);
         return;
       }
       var results;
       if (fragments[partialHash]) {
         for (var i = 0; i < fragments[partialHash].length; i++) {
           var chunkId = fragments[partialHash][i][0];
           var hash = fragments[partialHash][i][1];
           cb.completion(hash, self.tableName, chunkId);