Bug 1305801 - Part 4: Store variable-length prefix to disk. r=francois, r=gcp
authorDimi Lee <dlee@mozilla.com>
Mon, 19 Sep 2016 11:51:01 +0800
changeset 316424 122db2c234f18b7e9ae8e6487ab673874d887a50
parent 316423 94d8f8d4e3d57a2bb17db6aaa911ba855ac2f69d
child 316425 2d4f6c10aacfb9a319624a9df50d7b0769cfafa9
push id30770
push userkwierso@gmail.com
push dateWed, 05 Oct 2016 00:00:48 +0000
treeherdermozilla-central@3470e326025c [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewersfrancois, gcp
bugs1305801
milestone52.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1305801 - Part 4: Store variable-length prefix to disk. r=francois, r=gcp MozReview-Commit-ID: BMTGtgMuQdg
toolkit/components/url-classifier/Classifier.cpp
toolkit/components/url-classifier/Classifier.h
toolkit/components/url-classifier/HashStore.h
toolkit/components/url-classifier/LookupCache.cpp
toolkit/components/url-classifier/LookupCache.h
toolkit/components/url-classifier/ProtocolParser.cpp
toolkit/components/url-classifier/tests/gtest/TestPerProviderDirectory.cpp
toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp
toolkit/components/url-classifier/tests/gtest/moz.build
--- a/toolkit/components/url-classifier/Classifier.cpp
+++ b/toolkit/components/url-classifier/Classifier.cpp
@@ -512,40 +512,35 @@ Classifier::ApplyUpdates(nsTArray<TableU
     clockStart = PR_IntervalNow();
   }
 
   nsresult rv;
 
   {
     ScopedUpdatesClearer scopedUpdatesClearer(aUpdates);
 
-    // In order to prevent any premature update code from being
-    // run against V4 updates, we bail out as early as possible
-    // if aUpdates is using V4.
-    for (auto update : *aUpdates) {
-      if (update && TableUpdate::Cast<TableUpdateV4>(update)) {
-        LOG(("V4 update is not supported yet."));
-        // TODO: Bug 1283009 - Supports applying table udpate V4.
-        return NS_ERROR_NOT_IMPLEMENTED;
-      }
-    }
-
     LOG(("Backup before update."));
 
     rv = BackupTables();
     NS_ENSURE_SUCCESS(rv, rv);
 
     LOG(("Applying %d 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());
-        rv = UpdateHashStore(aUpdates, updateTable);
+
+        if (TableUpdate::Cast<TableUpdateV2>((*aUpdates)[i])) {
+          rv = UpdateHashStore(aUpdates, updateTable);
+        } else {
+          rv = UpdateTableV4(aUpdates, updateTable);
+        }
+
         if (NS_FAILED(rv)) {
           if (rv != NS_ERROR_OUT_OF_MEMORY) {
             Reset();
           }
           return rv;
         }
       }
     }
@@ -852,17 +847,18 @@ Classifier::UpdateHashStore(nsTArray<Tab
   }
 
   nsresult rv = store.Open();
   NS_ENSURE_SUCCESS(rv, rv);
   rv = store.BeginUpdate();
   NS_ENSURE_SUCCESS(rv, rv);
 
   // Read the part of the store that is (only) in the cache
-  LookupCache *lookupCache = GetLookupCache(store.TableName());
+  LookupCacheV2* lookupCache =
+    LookupCache::Cast<LookupCacheV2>(GetLookupCache(store.TableName()));
   if (!lookupCache) {
     return NS_ERROR_FAILURE;
   }
 
   // Clear cache when update
   lookupCache->ClearCache();
 
   FallibleTArray<uint32_t> AddPrefixHashes;
@@ -917,40 +913,98 @@ Classifier::UpdateHashStore(nsTArray<Tab
   NS_ENSURE_SUCCESS(rv, rv);
 
   // At this point the store is updated and written out to disk, but
   // the data is still in memory.  Build our quick-lookup table here.
   rv = lookupCache->Build(store.AddPrefixes(), store.AddCompletes());
   NS_ENSURE_SUCCESS(rv, rv);
 
 #if defined(DEBUG)
-  lookupCache->Dump();
+  lookupCache->DumpCompletions();
 #endif
   rv = lookupCache->WriteFile();
   NS_ENSURE_SUCCESS(rv, rv);
 
   int64_t now = (PR_Now() / PR_USEC_PER_SEC);
   LOG(("Successfully updated %s", store.TableName().get()));
   mTableFreshness.Put(store.TableName(), now);
 
   return NS_OK;
 }
 
 nsresult
+Classifier::UpdateTableV4(nsTArray<TableUpdate*>* aUpdates,
+                          const nsACString& aTable)
+{
+  LOG(("Classifier::UpdateTableV4(%s)", PromiseFlatCString(aTable).get()));
+
+  if (!CheckValidUpdate(aUpdates, aTable)) {
+    return NS_OK;
+  }
+
+  LookupCacheV4* lookupCache =
+    LookupCache::Cast<LookupCacheV4>(GetLookupCache(aTable));
+  if (!lookupCache) {
+    return NS_ERROR_FAILURE;
+  }
+
+  PrefixStringMap prefixes;
+  for (uint32_t i = 0; i < aUpdates->Length(); i++) {
+    TableUpdate *update = aUpdates->ElementAt(i);
+    if (!update || !update->TableName().Equals(aTable)) {
+      continue;
+    }
+
+    auto updateV4 = TableUpdate::Cast<TableUpdateV4>(update);
+    NS_ENSURE_TRUE(updateV4, NS_ERROR_FAILURE);
+
+    if (updateV4->IsFullUpdate()) {
+      prefixes.Clear();
+      TableUpdateV4::PrefixesStringMap& map = updateV4->Prefixes();
+
+      for (auto iter = map.Iter(); !iter.Done(); iter.Next()) {
+        // prefixes is an nsClassHashtable object stores prefix string.
+        // It will take the ownership of the put object.
+        nsCString* prefix = new nsCString(iter.Data()->GetPrefixString());
+        prefixes.Put(iter.Key(), prefix);
+      }
+    } else {
+      // TODO: Bug 1287058, partial update
+    }
+
+    aUpdates->ElementAt(i) = nullptr;
+  }
+
+  nsresult rv = lookupCache->Build(prefixes);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  rv = lookupCache->WriteFile();
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  int64_t now = (PR_Now() / PR_USEC_PER_SEC);
+  LOG(("Successfully updated %s\n", PromiseFlatCString(aTable).get()));
+  mTableFreshness.Put(aTable, now);
+
+  return NS_OK;
+}
+
+nsresult
 Classifier::UpdateCache(TableUpdate* aUpdate)
 {
   if (!aUpdate) {
     return NS_OK;
   }
 
   nsAutoCString table(aUpdate->TableName());
   LOG(("Classifier::UpdateCache(%s)", table.get()));
 
   LookupCache *lookupCache = GetLookupCache(table);
-  NS_ENSURE_TRUE(lookupCache, NS_ERROR_FAILURE);
+  if (!lookupCache) {
+    return NS_ERROR_FAILURE;
+  }
 
   auto updateV2 = TableUpdate::Cast<TableUpdateV2>(aUpdate);
   lookupCache->AddCompletionsToCache(updateV2->AddCompletes());
 
 #if defined(DEBUG)
   lookupCache->DumpCache();
 #endif
 
@@ -961,17 +1015,26 @@ LookupCache *
 Classifier::GetLookupCache(const nsACString& aTable)
 {
   for (uint32_t i = 0; i < mLookupCaches.Length(); i++) {
     if (mLookupCaches[i]->TableName().Equals(aTable)) {
       return mLookupCaches[i];
     }
   }
 
-  UniquePtr<LookupCache> cache(new LookupCache(aTable, mRootStoreDirectory));
+  // TODO : Bug 1302600, It would be better if we have a more general non-main
+  //        thread method to convert table name to protocol version. Currently
+  //        we can only know this by checking if the table name ends with '-proto'.
+  UniquePtr<LookupCache> cache;
+  if (StringEndsWith(aTable, NS_LITERAL_CSTRING("-proto"))) {
+    cache = MakeUnique<LookupCacheV4>(aTable, mRootStoreDirectory);
+  } else {
+    cache = MakeUnique<LookupCacheV2>(aTable, mRootStoreDirectory);
+  }
+
   nsresult rv = cache->Init();
   if (NS_FAILED(rv)) {
     return nullptr;
   }
   rv = cache->Open();
   if (NS_FAILED(rv)) {
     if (rv == NS_ERROR_FILE_CORRUPTED) {
       Reset();
@@ -983,17 +1046,18 @@ Classifier::GetLookupCache(const nsACStr
 }
 
 nsresult
 Classifier::ReadNoiseEntries(const Prefix& aPrefix,
                              const nsACString& aTableName,
                              uint32_t aCount,
                              PrefixArray* aNoiseEntries)
 {
-  LookupCache *cache = GetLookupCache(aTableName);
+  // TODO : Bug 1297962, support adding noise for v4
+  LookupCacheV2 *cache = static_cast<LookupCacheV2*>(GetLookupCache(aTableName));
   if (!cache) {
     return NS_ERROR_FAILURE;
   }
 
   FallibleTArray<uint32_t> prefixes;
   nsresult rv = cache->GetPrefixes(prefixes);
   NS_ENSURE_SUCCESS(rv, rv);
 
--- a/toolkit/components/url-classifier/Classifier.h
+++ b/toolkit/components/url-classifier/Classifier.h
@@ -110,16 +110,19 @@ private:
   nsresult BackupTables();
   nsresult RemoveBackupTables();
   nsresult RegenActiveTables();
   nsresult ScanStoreDir(nsTArray<nsCString>& aTables);
 
   nsresult UpdateHashStore(nsTArray<TableUpdate*>* aUpdates,
                            const nsACString& aTable);
 
+  nsresult UpdateTableV4(nsTArray<TableUpdate*>* aUpdates,
+                         const nsACString& aTable);
+
   nsresult UpdateCache(TableUpdate* aUpdates);
 
   LookupCache *GetLookupCache(const nsACString& aTable);
 
   bool CheckValidUpdate(nsTArray<TableUpdate*>* aUpdates,
                         const nsACString& aTable);
 
   // Root dir of the Local profile.
--- a/toolkit/components/url-classifier/HashStore.h
+++ b/toolkit/components/url-classifier/HashStore.h
@@ -148,36 +148,40 @@ public:
   };
 
   typedef nsClassHashtable<nsUint32HashKey, PrefixString> PrefixesStringMap;
   typedef nsTArray<int32_t> RemovalIndiceArray;
 
 public:
   explicit TableUpdateV4(const nsACString& aTable)
     : TableUpdate(aTable)
+    , mFullUpdate(false)
   {
   }
 
   bool Empty() const override
   {
     return mPrefixesMap.IsEmpty() && mRemovalIndiceArray.IsEmpty();
   }
 
+  bool IsFullUpdate() const { return mFullUpdate; }
   PrefixesStringMap& Prefixes() { return mPrefixesMap; }
   RemovalIndiceArray& RemovalIndices() { return mRemovalIndiceArray; }
 
   // For downcasting.
   static const int TAG = 4;
 
+  void SetFullUpdate(bool aIsFullUpdate) { mFullUpdate = aIsFullUpdate; }
   void NewPrefixes(int32_t aSize, std::string& aPrefixes);
   void NewRemovalIndices(const uint32_t* aIndices, size_t aNumOfIndices);
 
 private:
   virtual int Tag() const override { return TAG; }
 
+  bool mFullUpdate;
   PrefixesStringMap mPrefixesMap;
   RemovalIndiceArray mRemovalIndiceArray;
 };
 
 // There is one hash store per table.
 class HashStore {
 public:
   HashStore(const nsACString& aTableName, nsIFile* aRootStoreFile);
--- a/toolkit/components/url-classifier/LookupCache.cpp
+++ b/toolkit/components/url-classifier/LookupCache.cpp
@@ -35,47 +35,32 @@
 // 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)
 
 namespace mozilla {
 namespace safebrowsing {
 
+const int LookupCacheV2::VER = 2;
+const int LookupCacheV4::VER = 4;
+
 LookupCache::LookupCache(const nsACString& aTableName, nsIFile* aRootStoreDir)
   : mPrimed(false)
   , mTableName(aTableName)
   , mRootStoreDirectory(aRootStoreDir)
 {
   UpdateRootDirHandle(mRootStoreDirectory);
 }
 
 nsresult
-LookupCache::Init()
-{
-  mPrefixSet = new nsUrlClassifierPrefixSet();
-  nsresult rv = mPrefixSet->Init(mTableName);
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  return NS_OK;
-}
-
-LookupCache::~LookupCache()
-{
-}
-
-nsresult
 LookupCache::Open()
 {
-  LOG(("Reading Completions"));
-  nsresult rv = ReadCompletions();
-  NS_ENSURE_SUCCESS(rv, rv);
-
   LOG(("Loading PrefixSet"));
-  rv = LoadPrefixSet();
+  nsresult rv = LoadPrefixSet();
   NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
 
 nsresult
 LookupCache::UpdateRootDirHandle(nsIFile* aNewRootStoreDirectory)
 {
@@ -120,42 +105,16 @@ LookupCache::Reset()
   rv = prefixsetFile->Remove(false);
   NS_ENSURE_SUCCESS(rv, rv);
 
   ClearAll();
 
   return NS_OK;
 }
 
-
-nsresult
-LookupCache::Build(AddPrefixArray& aAddPrefixes,
-                   AddCompleteArray& aAddCompletes)
-{
-  Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_COMPLETIONS,
-                        static_cast<uint32_t>(aAddCompletes.Length()));
-
-  mUpdateCompletions.Clear();
-  mUpdateCompletions.SetCapacity(aAddCompletes.Length());
-  for (uint32_t i = 0; i < aAddCompletes.Length(); i++) {
-    mUpdateCompletions.AppendElement(aAddCompletes[i].CompleteHash());
-  }
-  aAddCompletes.Clear();
-  mUpdateCompletions.Sort();
-
-  Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_PREFIXES,
-                        static_cast<uint32_t>(aAddPrefixes.Length()));
-
-  nsresult rv = ConstructPrefixSet(aAddPrefixes);
-  NS_ENSURE_SUCCESS(rv, rv);
-  mPrimed = true;
-
-  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());
     }
   }
@@ -172,115 +131,48 @@ LookupCache::DumpCache()
     return;
 
   for (uint32_t i = 0; i < mGetHashCache.Length(); i++) {
     nsAutoCString str;
     mGetHashCache[i].ToHexString(str);
     LOG(("Caches: %s", str.get()));
   }
 }
-
-void
-LookupCache::Dump()
-{
-  if (!LOG_ENABLED())
-    return;
-
-  for (uint32_t i = 0; i < mUpdateCompletions.Length(); i++) {
-    nsAutoCString str;
-    mUpdateCompletions[i].ToHexString(str);
-    LOG(("Update: %s", str.get()));
-  }
-}
 #endif
 
 nsresult
-LookupCache::Has(const Completion& aCompletion,
-                 bool* aHas, bool* aComplete)
-{
-  *aHas = *aComplete = false;
-
-  uint32_t prefix = aCompletion.ToUint32();
-
-  bool found;
-  nsresult rv = mPrefixSet->Contains(prefix, &found);
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  LOG(("Probe in %s: %X, found %d", mTableName.get(), prefix, found));
-
-  if (found) {
-    *aHas = true;
-  }
-
-  // TODO: We may need to distinguish completions found in cache or update in the future
-  if ((mGetHashCache.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex) ||
-      (mUpdateCompletions.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex)) {
-    LOG(("Complete in %s", mTableName.get()));
-    *aComplete = true;
-    *aHas = true;
-  }
-
-  return NS_OK;
-}
-
-nsresult
 LookupCache::WriteFile()
 {
   nsCOMPtr<nsIFile> psFile;
   nsresult rv = mStoreDirectory->Clone(getter_AddRefs(psFile));
   NS_ENSURE_SUCCESS(rv, rv);
 
   rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
   NS_ENSURE_SUCCESS(rv, rv);
 
-  rv = mPrefixSet->StoreToFile(psFile);
+  rv = StoreToFile(psFile);
   NS_WARNING_ASSERTION(NS_SUCCEEDED(rv), "failed to store the prefixset");
 
   return NS_OK;
 }
 
 void
 LookupCache::ClearAll()
 {
   ClearCache();
-  ClearUpdatedCompletions();
-  mPrefixSet->SetPrefixes(nullptr, 0);
+  ClearPrefixes();
   mPrimed = false;
 }
 
 void
-LookupCache::ClearUpdatedCompletions()
-{
-  mUpdateCompletions.Clear();
-}
-
-void
 LookupCache::ClearCache()
 {
   mGetHashCache.Clear();
 }
 
-nsresult
-LookupCache::ReadCompletions()
-{
-  HashStore store(mTableName, mRootStoreDirectory);
-
-  nsresult rv = store.Open();
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  mUpdateCompletions.Clear();
-
-  const AddCompleteArray& addComplete = store.AddCompletes();
-  for (uint32_t i = 0; i < addComplete.Length(); i++) {
-    mUpdateCompletions.AppendElement(addComplete[i].complete);
-  }
-
-  return NS_OK;
-}
-
 /* static */ bool
 LookupCache::IsCanonicalizedIP(const nsACString& aHost)
 {
   // The canonicalization process will have left IP addresses in dotted
   // decimal with no surprises.
   uint32_t i1, i2, i3, i4;
   char c;
   if (PR_sscanf(PromiseFlatCString(aHost).get(), "%u.%u.%u.%u%c",
@@ -447,19 +339,188 @@ LookupCache::GetHostKeys(const nsACStrin
     lookupHost2->Assign(hostComponents[last - 2]);
     lookupHost2->Append(".");
     lookupHost2->Append(*lookupHost);
   }
 
   return NS_OK;
 }
 
-bool LookupCache::IsPrimed()
+nsresult
+LookupCache::LoadPrefixSet()
+{
+  nsCOMPtr<nsIFile> psFile;
+  nsresult rv = mStoreDirectory->Clone(getter_AddRefs(psFile));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  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
+  if (mPrimed) {
+    uint32_t size = SizeOfPrefixSet();
+    LOG(("SB tree done, size = %d bytes\n", size));
+  }
+#endif
+
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV2::Init()
+{
+  mPrefixSet = new nsUrlClassifierPrefixSet();
+  nsresult rv = mPrefixSet->Init(mTableName);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV2::Open()
+{
+  nsresult rv = LookupCache::Open();
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  LOG(("Reading Completions"));
+  rv = ReadCompletions();
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  return NS_OK;
+}
+
+void
+LookupCacheV2::ClearAll()
+{
+  LookupCache::ClearAll();
+  mUpdateCompletions.Clear();
+}
+
+nsresult
+LookupCacheV2::Has(const Completion& aCompletion,
+                   bool* aHas, bool* aComplete)
 {
-  return mPrimed;
+  *aHas = *aComplete = false;
+
+  uint32_t prefix = aCompletion.ToUint32();
+
+  bool found;
+  nsresult rv = mPrefixSet->Contains(prefix, &found);
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  LOG(("Probe in %s: %X, found %d", mTableName.get(), prefix, found));
+
+  if (found) {
+    *aHas = true;
+  }
+
+  if ((mGetHashCache.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex) ||
+      (mUpdateCompletions.BinaryIndexOf(aCompletion) != nsTArray<Completion>::NoIndex)) {
+    LOG(("Complete in %s", mTableName.get()));
+    *aComplete = true;
+    *aHas = true;
+  }
+
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV2::Build(AddPrefixArray& aAddPrefixes,
+                     AddCompleteArray& aAddCompletes)
+{
+  Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_COMPLETIONS,
+                        static_cast<uint32_t>(aAddCompletes.Length()));
+
+  mUpdateCompletions.Clear();
+  mUpdateCompletions.SetCapacity(aAddCompletes.Length());
+  for (uint32_t i = 0; i < aAddCompletes.Length(); i++) {
+    mUpdateCompletions.AppendElement(aAddCompletes[i].CompleteHash());
+  }
+  aAddCompletes.Clear();
+  mUpdateCompletions.Sort();
+
+  Telemetry::Accumulate(Telemetry::URLCLASSIFIER_LC_PREFIXES,
+                        static_cast<uint32_t>(aAddPrefixes.Length()));
+
+  nsresult rv = ConstructPrefixSet(aAddPrefixes);
+  NS_ENSURE_SUCCESS(rv, rv);
+  mPrimed = true;
+
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV2::GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes)
+{
+  if (!mPrimed) {
+    // This can happen if its a new table, so no error.
+    LOG(("GetPrefixes from empty LookupCache"));
+    return NS_OK;
+  }
+  return mPrefixSet->GetPrefixesNative(aAddPrefixes);
+}
+
+nsresult
+LookupCacheV2::ReadCompletions()
+{
+  HashStore store(mTableName, mRootStoreDirectory);
+
+  nsresult rv = store.Open();
+  NS_ENSURE_SUCCESS(rv, rv);
+
+  mUpdateCompletions.Clear();
+
+  const AddCompleteArray& addComplete = store.AddCompletes();
+  for (uint32_t i = 0; i < addComplete.Length(); i++) {
+    mUpdateCompletions.AppendElement(addComplete[i].complete);
+  }
+
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV2::ClearPrefixes()
+{
+  return mPrefixSet->SetPrefixes(nullptr, 0);
+}
+
+nsresult
+LookupCacheV2::StoreToFile(nsIFile* aFile)
+{
+  return mPrefixSet->StoreToFile(aFile);
+}
+
+nsresult
+LookupCacheV2::LoadFromFile(nsIFile* aFile)
+{
+  return mPrefixSet->LoadFromFile(aFile);
+}
+
+size_t
+LookupCacheV2::SizeOfPrefixSet()
+{
+  return mPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
 }
 
 #ifdef DEBUG
 template <class T>
 static void EnsureSorted(T* aArray)
 {
   typename T::elem_type* start = aArray->Elements();
   typename T::elem_type* end = aArray->Elements() + aArray->Length();
@@ -473,17 +534,17 @@ static void EnsureSorted(T* aArray)
       MOZ_ASSERT(*previous <= *iter);
     }
   }
   return;
 }
 #endif
 
 nsresult
-LookupCache::ConstructPrefixSet(AddPrefixArray& aAddPrefixes)
+LookupCacheV2::ConstructPrefixSet(AddPrefixArray& aAddPrefixes)
 {
   Telemetry::AutoTimer<Telemetry::URLCLASSIFIER_PS_CONSTRUCT_TIME> timer;
 
   nsTArray<uint32_t> array;
   if (!array.SetCapacity(aAddPrefixes.Length(), fallible)) {
     return NS_ERROR_OUT_OF_MEMORY;
   }
 
@@ -507,60 +568,76 @@ LookupCache::ConstructPrefixSet(AddPrefi
   LOG(("SB tree done, size = %d bytes\n", size));
 #endif
 
   mPrimed = true;
 
   return NS_OK;
 }
 
-nsresult
-LookupCache::LoadPrefixSet()
+#if defined(DEBUG)
+void
+LookupCacheV2::DumpCompletions()
 {
-  nsCOMPtr<nsIFile> psFile;
-  nsresult rv = mStoreDirectory->Clone(getter_AddRefs(psFile));
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  rv = psFile->AppendNative(mTableName + NS_LITERAL_CSTRING(PREFIXSET_SUFFIX));
-  NS_ENSURE_SUCCESS(rv, rv);
-
-  bool exists;
-  rv = psFile->Exists(&exists);
-  NS_ENSURE_SUCCESS(rv, rv);
+  if (!LOG_ENABLED())
+    return;
 
-  if (exists) {
-    LOG(("stored PrefixSet exists, loading from disk"));
-    rv = mPrefixSet->LoadFromFile(psFile);
-    if (NS_FAILED(rv)) {
-      if (rv == NS_ERROR_FILE_CORRUPTED) {
-        Reset();
-      }
-      return rv;
-    }
-    mPrimed = true;
-  } else {
-    LOG(("no (usable) stored PrefixSet found"));
+  for (uint32_t i = 0; i < mUpdateCompletions.Length(); i++) {
+    nsAutoCString str;
+    mUpdateCompletions[i].ToHexString(str);
+    LOG(("Update: %s", str.get()));
   }
+}
+#endif
 
-#ifdef DEBUG
-  if (mPrimed) {
-    uint32_t size = mPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
-    LOG(("SB tree done, size = %d bytes\n", size));
-  }
-#endif
+nsresult
+LookupCacheV4::Init()
+{
+  mVLPrefixSet = new VariableLengthPrefixSet();
+  nsresult rv = mVLPrefixSet->Init(mTableName);
+  NS_ENSURE_SUCCESS(rv, rv);
 
   return NS_OK;
 }
 
+// TODO : Bug 1298257, Implement url matching for variable-length prefix set
 nsresult
-LookupCache::GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes)
+LookupCacheV4::Has(const Completion& aCompletion,
+                   bool* aHas, bool* aComplete)
 {
-  if (!mPrimed) {
-    // This can happen if its a new table, so no error.
-    LOG(("GetPrefixes from empty LookupCache"));
-    return NS_OK;
-  }
-  return mPrefixSet->GetPrefixesNative(aAddPrefixes);
+  *aHas = false;
+  return NS_OK;
+}
+
+nsresult
+LookupCacheV4::Build(PrefixStringMap& aPrefixMap)
+{
+  return mVLPrefixSet->SetPrefixes(aPrefixMap);
 }
 
+nsresult
+LookupCacheV4::ClearPrefixes()
+{
+  // Clear by seting a empty map
+  PrefixStringMap map;
+  return mVLPrefixSet->SetPrefixes(map);
+}
+
+nsresult
+LookupCacheV4::StoreToFile(nsIFile* aFile)
+{
+  return mVLPrefixSet->StoreToFile(aFile);
+}
+
+nsresult
+LookupCacheV4::LoadFromFile(nsIFile* aFile)
+{
+  return mVLPrefixSet->LoadFromFile(aFile);
+}
+
+size_t
+LookupCacheV4::SizeOfPrefixSet()
+{
+  return mVLPrefixSet->SizeOfIncludingThis(moz_malloc_size_of);
+}
 
 } // namespace safebrowsing
 } // namespace mozilla
--- a/toolkit/components/url-classifier/LookupCache.h
+++ b/toolkit/components/url-classifier/LookupCache.h
@@ -9,16 +9,17 @@
 #include "Entries.h"
 #include "nsString.h"
 #include "nsTArray.h"
 #include "nsCOMPtr.h"
 #include "nsIFile.h"
 #include "nsIFileStreams.h"
 #include "mozilla/RefPtr.h"
 #include "nsUrlClassifierPrefixSet.h"
+#include "VariableLengthPrefixSet.h"
 #include "mozilla/Logging.h"
 
 namespace mozilla {
 namespace safebrowsing {
 
 #define MAX_HOST_COMPONENTS 5
 #define MAX_PATH_COMPONENTS 4
 
@@ -91,63 +92,144 @@ public:
   // two keys will be returned:
   //  hostname.com/foo/bar -> [hostname.com]
   //  mail.hostname.com/foo/bar -> [hostname.com, mail.hostname.com]
   //  www.mail.hostname.com/foo/bar -> [hostname.com, mail.hostname.com]
   static nsresult GetHostKeys(const nsACString& aSpec,
                               nsTArray<nsCString>* aHostKeys);
 
   LookupCache(const nsACString& aTableName, nsIFile* aStoreFile);
-  ~LookupCache();
+  virtual ~LookupCache() {}
 
   const nsCString &TableName() const { return mTableName; }
 
-  nsresult Init();
-  nsresult Open();
   // The directory handle where we operate will
   // be moved away when a backup is made.
   nsresult UpdateRootDirHandle(nsIFile* aRootStoreDirectory);
+
   // This will Clear() the passed arrays when done.
-  nsresult Build(AddPrefixArray& aAddPrefixes,
-                 AddCompleteArray& aAddCompletes);
   nsresult AddCompletionsToCache(AddCompleteArray& aAddCompletes);
-  nsresult GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes);
-  void ClearUpdatedCompletions();
+
+  // Write data stored in lookup cache to disk.
+  nsresult WriteFile();
+
+  // Clear completions retrieved from gethash request.
   void ClearCache();
 
+  bool IsPrimed() const { return mPrimed; };
+
 #if DEBUG
   void DumpCache();
-  void Dump();
 #endif
-  nsresult WriteFile();
-  nsresult Has(const Completion& aCompletion,
-               bool* aHas, bool* aComplete);
-  bool IsPrimed();
+
+  virtual nsresult Open();
+  virtual nsresult Init() = 0;
+  virtual nsresult ClearPrefixes() = 0;
+  virtual nsresult Has(const Completion& aCompletion,
+                       bool* aHas, bool* aComplete) = 0;
+
+  template<typename T>
+  static T* Cast(LookupCache* aThat) {
+    return (T::VER == aThat->Ver() ? reinterpret_cast<T*>(aThat) : nullptr);
+  }
 
 private:
-  void ClearAll();
   nsresult Reset();
-  nsresult ReadCompletions();
   nsresult LoadPrefixSet();
-  nsresult LoadCompletions();
-  // Construct a Prefix Set with known prefixes.
-  // This will Clear() aAddPrefixes when done.
-  nsresult ConstructPrefixSet(AddPrefixArray& aAddPrefixes);
+
+  virtual nsresult StoreToFile(nsIFile* aFile) = 0;
+  virtual nsresult LoadFromFile(nsIFile* aFile) = 0;
+  virtual size_t SizeOfPrefixSet() = 0;
+
+  virtual int Ver() const = 0;
+
+protected:
+  virtual void ClearAll();
 
   bool mPrimed;
   nsCString mTableName;
   nsCOMPtr<nsIFile> mRootStoreDirectory;
   nsCOMPtr<nsIFile> mStoreDirectory;
-  // Set of prefixes known to be in the database
-  RefPtr<nsUrlClassifierPrefixSet> mPrefixSet;
-  // Full length hashes obtained in update request
-  CompletionArray mUpdateCompletions;
+
   // Full length hashes obtained in gethash request
   CompletionArray mGetHashCache;
 
   // For gtest to inspect private members.
   friend class PerProviderDirectoryTestUtils;
 };
 
+class LookupCacheV2 final : public LookupCache
+{
+public:
+  explicit LookupCacheV2(const nsACString& aTableName, nsIFile* aStoreFile)
+    : LookupCache(aTableName, aStoreFile) {}
+  ~LookupCacheV2() {}
+
+  virtual nsresult Init() override;
+  virtual nsresult Open() override;
+  virtual void ClearAll() override;
+  virtual nsresult Has(const Completion& aCompletion,
+                       bool* aHas, bool* aComplete) override;
+
+  nsresult Build(AddPrefixArray& aAddPrefixes,
+                 AddCompleteArray& aAddCompletes);
+
+  nsresult GetPrefixes(FallibleTArray<uint32_t>& aAddPrefixes);
+
+#if DEBUG
+  void DumpCompletions();
+#endif
+
+  static const int VER;
+
+protected:
+  nsresult ReadCompletions();
+
+  virtual nsresult ClearPrefixes() override;
+  virtual nsresult StoreToFile(nsIFile* aFile) override;
+  virtual nsresult LoadFromFile(nsIFile* aFile) override;
+  virtual size_t SizeOfPrefixSet() override;
+
+private:
+  virtual int Ver() const override { return VER; }
+
+  // Construct a Prefix Set with known prefixes.
+  // This will Clear() aAddPrefixes when done.
+  nsresult ConstructPrefixSet(AddPrefixArray& aAddPrefixes);
+
+  // Full length hashes obtained in update request
+  CompletionArray mUpdateCompletions;
+
+  // Set of prefixes known to be in the database
+  RefPtr<nsUrlClassifierPrefixSet> mPrefixSet;
+};
+
+class LookupCacheV4 final : public LookupCache
+{
+public:
+  explicit LookupCacheV4(const nsACString& aTableName, nsIFile* aStoreFile)
+    : LookupCache(aTableName, aStoreFile) {}
+  ~LookupCacheV4() {}
+
+  virtual nsresult Init() override;
+  virtual nsresult Has(const Completion& aCompletion,
+                       bool* aHas, bool* aComplete) override;
+
+  nsresult Build(PrefixStringMap& aPrefixMap);
+
+  static const int VER;
+
+protected:
+  virtual nsresult ClearPrefixes() override;
+  virtual nsresult StoreToFile(nsIFile* aFile) override;
+  virtual nsresult LoadFromFile(nsIFile* aFile) override;
+  virtual size_t SizeOfPrefixSet() override;
+
+private:
+  virtual int Ver() const override { return VER; }
+
+  RefPtr<VariableLengthPrefixSet> mVLPrefixSet;
+};
+
 } // namespace safebrowsing
 } // namespace mozilla
 
 #endif
--- a/toolkit/components/url-classifier/ProtocolParser.cpp
+++ b/toolkit/components/url-classifier/ProtocolParser.cpp
@@ -868,16 +868,18 @@ ProtocolParserProtobuf::ProcessOneRespon
     DebugOnly<nsresult> rv = SaveStateToPref(listName, state);
     NS_WARNING_ASSERTION(NS_SUCCEEDED(rv), "SaveStateToPref failed");
   }));
 
   PARSER_LOG(("==== Update for threat type '%d' ====", aResponse.threat_type()));
   PARSER_LOG(("* listName: %s\n", listName.get()));
   PARSER_LOG(("* newState: %s\n", aResponse.new_client_state().c_str()));
   PARSER_LOG(("* isFullUpdate: %s\n", (isFullUpdate ? "yes" : "no")));
+
+  tuV4->SetFullUpdate(isFullUpdate);
   ProcessAdditionOrRemoval(*tuV4, aResponse.additions(), true /*aIsAddition*/);
   ProcessAdditionOrRemoval(*tuV4, aResponse.removals(), false);
   PARSER_LOG(("\n\n"));
 
   return NS_OK;
 }
 
 nsresult
--- a/toolkit/components/url-classifier/tests/gtest/TestPerProviderDirectory.cpp
+++ b/toolkit/components/url-classifier/tests/gtest/TestPerProviderDirectory.cpp
@@ -76,22 +76,22 @@ void VerifyPrivateStorePath(const char* 
 TEST(PerProviderDirectory, LookupCache)
 {
   RunTestInNewThread([] () -> void {
     nsCOMPtr<nsIFile> rootDir;
     NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(rootDir));
 
     // For V2 tables (NOT ending with '-proto'), root directory should be
     // used as the private store.
-    VerifyPrivateStorePath<LookupCache>("goog-phish-shavar", "google", rootDir, false);
+    VerifyPrivateStorePath<LookupCacheV2>("goog-phish-shavar", "google", rootDir, false);
 
     // For V4 tables, if provider is found, use per-provider subdirectory;
     // If not found, use root directory.
-    VerifyPrivateStorePath<LookupCache>("goog-noprovider-proto", "", rootDir, false);
-    VerifyPrivateStorePath<LookupCache>("goog-phish-proto", "google4", rootDir, true);
+    VerifyPrivateStorePath<LookupCacheV4>("goog-noprovider-proto", "", rootDir, false);
+    VerifyPrivateStorePath<LookupCacheV4>("goog-phish-proto", "google4", rootDir, true);
   });
 }
 
 TEST(PerProviderDirectory, HashStore)
 {
   RunTestInNewThread([] () -> void {
     nsCOMPtr<nsIFile> rootDir;
     NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR, getter_AddRefs(rootDir));
new file mode 100644
--- /dev/null
+++ b/toolkit/components/url-classifier/tests/gtest/TestUrlClassifierTableUpdateV4.cpp
@@ -0,0 +1,190 @@
+#include "Classifier.h"
+#include "HashStore.h"
+#include "nsAppDirectoryServiceDefs.h"
+#include "nsIFile.h"
+#include "nsIThread.h"
+#include "string.h"
+#include "gtest/gtest.h"
+
+using namespace mozilla;
+using namespace mozilla::safebrowsing;
+
+typedef nsCString _Prefix;
+typedef nsTArray<_Prefix> _PrefixArray;
+
+static void
+PrefixArrayToPrefixStringMap(const _PrefixArray& prefixArray,
+                             PrefixStringMap& outMap)
+{
+  outMap.Clear();
+
+  for (uint32_t i = 0; i < prefixArray.Length(); i++) {
+    const _Prefix& prefix = prefixArray[i];
+    nsCString* prefixString = outMap.LookupOrAdd(prefix.Length());
+    prefixString->Append(prefix.BeginReading(), prefix.Length());
+  }
+}
+
+// N: Number of prefixes, MIN/MAX: minimum/maximum prefix size
+// This function will append generated prefixes to outArray.
+// All elements in the generated array will be different.
+static void
+CreateRandomSortedPrefixArray(uint32_t N,
+                              uint32_t MIN,
+                              uint32_t MAX,
+                              _PrefixArray& outArray)
+{
+  outArray.SetCapacity(outArray.Length() + N);
+
+  const uint32_t range = (MAX - MIN + 1);
+
+  for (uint32_t i = 0; i < N; i++) {
+    uint32_t prefixSize = (rand() % range) + MIN;
+    _Prefix prefix;
+    prefix.SetLength(prefixSize);
+
+    while (true) {
+      char* dst = prefix.BeginWriting();
+      for (uint32_t j = 0; j < prefixSize; j++) {
+        dst[j] = rand() % 256;
+      }
+
+      if (!outArray.Contains(prefix)) {
+        outArray.AppendElement(prefix);
+        break;
+      }
+    }
+  }
+
+  outArray.Sort();
+}
+
+// Function to generate TableUpdateV4.
+static void
+GenerateUpdateData(bool fullUpdate,
+                   PrefixStringMap& add,
+                   nsTArray<uint32_t>& removal,
+                   nsTArray<TableUpdate*>& tableUpdates)
+{
+  TableUpdateV4* tableUpdate = new TableUpdateV4(NS_LITERAL_CSTRING("gtest-malware-proto"));
+  tableUpdate->SetFullUpdate(fullUpdate);
+
+  for (auto iter = add.ConstIter(); !iter.Done(); iter.Next()) {
+    nsCString* pstring = iter.Data();
+    std::string str(pstring->BeginReading(), pstring->Length());
+
+    tableUpdate->NewPrefixes(iter.Key(), str);
+  }
+
+  tableUpdate->NewRemovalIndices(removal.Elements(), removal.Length());
+
+  tableUpdates.AppendElement(tableUpdate);
+}
+
+static void
+VerifyPrefixSet(PrefixStringMap& expected)
+{
+  // Verify the prefix set is written to disk.
+  nsCOMPtr<nsIFile> file;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
+                         getter_AddRefs(file));
+
+  file->AppendRelativeNativePath(NS_LITERAL_CSTRING("safebrowsing/gtest-malware-proto.pset"));
+
+  RefPtr<VariableLengthPrefixSet> load = new VariableLengthPrefixSet;
+  load->Init(NS_LITERAL_CSTRING("gtest-malware-proto"));
+
+  PrefixStringMap out;
+  load->LoadFromFile(file);
+  load->GetPrefixes(out);
+
+  for (auto iter = expected.ConstIter(); !iter.Done(); iter.Next()) {
+    nsCString* str1 = iter.Data();
+    nsCString* str2 = out.Get(iter.Key());
+
+    ASSERT_TRUE(*str1 == *str2);
+  }
+}
+
+static void
+Clear()
+{
+  nsCOMPtr<nsIFile> file;
+  NS_GetSpecialDirectory(NS_APP_USER_PROFILE_50_DIR,
+                         getter_AddRefs(file));
+  file->AppendRelativeNativePath(NS_LITERAL_CSTRING("safebrowsing/gtest-malware-proto.pset"));
+  file->Remove(false);
+}
+
+static void
+testUpdate(nsTArray<TableUpdate*>& tableUpdates,
+           PrefixStringMap& expected)
+{
+  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(rv == NS_OK);
+
+    VerifyPrefixSet(expected);
+  });
+}
+
+static void
+testFullUpdate(PrefixStringMap& add)
+{
+  nsTArray<uint32_t> empty;
+  nsTArray<TableUpdate*> tableUpdates;
+  GenerateUpdateData(true, add, empty, tableUpdates);
+
+  testUpdate(tableUpdates, add);
+}
+
+TEST(UrlClassifierTableUpdateV4, FixLenghtPSetFullUpdate)
+{
+  srand(time(NULL));
+
+   _PrefixArray array;
+  PrefixStringMap map;
+
+  CreateRandomSortedPrefixArray(5000, 4, 4, array);
+  PrefixArrayToPrefixStringMap(array, map);
+
+  testFullUpdate(map);
+
+  Clear();
+}
+
+
+TEST(UrlClassifierTableUpdateV4, VariableLenghtPSetFullUpdate)
+{
+   _PrefixArray array;
+  PrefixStringMap map;
+
+  CreateRandomSortedPrefixArray(5000, 5, 32, array);
+  PrefixArrayToPrefixStringMap(array, map);
+
+  testFullUpdate(map);
+
+  Clear();
+}
+
+// This test contain both variable length prefix set and fixed-length prefix set
+TEST(UrlClassifierTableUpdateV4, MixedPSetFullUpdate)
+{
+   _PrefixArray array;
+  PrefixStringMap map;
+
+  CreateRandomSortedPrefixArray(5000, 4, 4, array);
+  CreateRandomSortedPrefixArray(1000, 5, 32, array);
+  PrefixArrayToPrefixStringMap(array, map);
+
+  testFullUpdate(map);
+
+  Clear();
+}
--- a/toolkit/components/url-classifier/tests/gtest/moz.build
+++ b/toolkit/components/url-classifier/tests/gtest/moz.build
@@ -9,13 +9,14 @@ LOCAL_INCLUDES += [
 ]
 
 UNIFIED_SOURCES += [
     'TestChunkSet.cpp',
     'TestPerProviderDirectory.cpp',
     'TestSafebrowsingHash.cpp',
     'TestSafeBrowsingProtobuf.cpp',
     'TestTable.cpp',
+    'TestUrlClassifierTableUpdateV4.cpp',
     'TestUrlClassifierUtils.cpp',
     'TestVariableLengthPrefixSet.cpp',
 ]
 
 FINAL_LIBRARY = 'xul-gtest'