Improve performance of LevelDbDocumentOverlayCache::RemoveOverlaysForBatchId() (#9386)

Author: dconeybe

Firestore/core/src/local/document_overlay_cache.h

@@ -25,12 +25,15 @@
 #include "Firestore/core/src/model/mutation.h"
 #include "Firestore/core/src/model/mutation/overlay.h"
 #include "Firestore/core/src/model/resource_path.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/optional.h"
 
 namespace firebase {
 namespace firestore {
 namespace local {
 
+class DocumentOverlayCacheTestHelper;
+
 /**
  * Provides methods to read and write document overlays.
  *
@@ -101,6 +104,13 @@ class DocumentOverlayCache {
       const std::string& collection_group,
       int since_batch_id,
       std::size_t count) const = 0;
+
+ private:
+  friend class DocumentOverlayCacheTestHelper;
+
+  // Returns the total number of overlays in the database.
+  // This method exists for unit testing only.
+  virtual int GetOverlayCount() const = 0;
 };
 
 }  // namespace local

Firestore/core/src/local/leveldb_document_overlay_cache.cc

@@ -53,24 +53,24 @@ LevelDbDocumentOverlayCache::LevelDbDocumentOverlayCache(
 }
 
 absl::optional<Overlay> LevelDbDocumentOverlayCache::GetOverlay(
-    const DocumentKey& key) const {
-  const std::string leveldb_key_prefix =
-      LevelDbDocumentOverlayKey::KeyPrefix(user_id_, key);
+    const DocumentKey& document_key) const {
+  const std::string key_prefix =
+      LevelDbDocumentOverlayKey::KeyPrefix(user_id_, document_key);
 
   auto it = db_->current_transaction()->NewIterator();
-  it->Seek(leveldb_key_prefix);
+  it->Seek(key_prefix);
 
-  if (!(it->Valid() && absl::StartsWith(it->key(), leveldb_key_prefix))) {
+  if (!(it->Valid() && absl::StartsWith(it->key(), key_prefix))) {
     return absl::nullopt;
   }
 
-  LevelDbDocumentOverlayKey decoded_key;
-  HARD_ASSERT(decoded_key.Decode(it->key()));
-  if (decoded_key.document_key() != key) {
+  LevelDbDocumentOverlayKey key;
+  HARD_ASSERT(key.Decode(it->key()));
+  if (key.document_key() != document_key) {
     return absl::nullopt;
   }
 
-  return ParseOverlay(decoded_key, it->value());
+  return ParseOverlay(key, it->value());
 }
 
 void LevelDbDocumentOverlayCache::SaveOverlays(
@@ -81,16 +81,8 @@ void LevelDbDocumentOverlayCache::SaveOverlays(
 }
 
 void LevelDbDocumentOverlayCache::RemoveOverlaysForBatchId(int batch_id) {
-  // TODO(dconeybe) Implement an index so that this query can be performed
-  // without requiring a full table scan.
-
-  ForEachOverlay([&](absl::string_view encoded_key,
-                     const LevelDbDocumentOverlayKey& decoded_key,
-                     absl::string_view) {
-    if (decoded_key.largest_batch_id() == batch_id) {
-      db_->current_transaction()->Delete(encoded_key);
-    }
-  });
+  ForEachKeyWithLargestBatchId(
+      batch_id, [&](LevelDbDocumentOverlayKey&& key) { DeleteOverlay(key); });
 }
 
 DocumentOverlayCache::OverlayByDocumentKeyMap
@@ -103,20 +95,18 @@ LevelDbDocumentOverlayCache::GetOverlays(const ResourcePath& collection,
 
   const size_t immediate_children_path_length{collection.size() + 1};
 
-  ForEachOverlay([&](absl::string_view,
-                     const LevelDbDocumentOverlayKey& decoded_key,
+  ForEachOverlay([&](LevelDbDocumentOverlayKey&& key,
                      absl::string_view encoded_mutation) {
-    const DocumentKey key = decoded_key.document_key();
-    if (!collection.IsPrefixOf(key.path())) {
+    if (!collection.IsPrefixOf(key.document_key().path())) {
       return;
     }
     // Documents from sub-collections
-    if (key.path().size() != immediate_children_path_length) {
+    if (key.document_key().path().size() != immediate_children_path_length) {
       return;
     }
 
-    if (decoded_key.largest_batch_id() > since_batch_id) {
-      result[key] = ParseOverlay(decoded_key, encoded_mutation);
+    if (key.largest_batch_id() > since_batch_id) {
+      result[key.document_key()] = ParseOverlay(key, encoded_mutation);
     }
   });
 
@@ -135,17 +125,16 @@ LevelDbDocumentOverlayCache::GetOverlays(const std::string& collection_group,
   // by largest_batch_id, the loop below can iterate over it ordered by
   // largest_batch_id.
   std::map<int, std::unordered_set<Overlay, OverlayHash>> overlays_by_batch_id;
-  ForEachOverlay([&](absl::string_view,
-                     const LevelDbDocumentOverlayKey& decoded_key,
-                     absl::string_view encoded_mutation) {
-    if (decoded_key.largest_batch_id() <= since_batch_id) {
-      return;
-    }
-    if (decoded_key.document_key().HasCollectionId(collection_group)) {
-      overlays_by_batch_id[decoded_key.largest_batch_id()].emplace(
-          ParseOverlay(decoded_key, encoded_mutation));
-    }
-  });
+  ForEachOverlay(
+      [&](LevelDbDocumentOverlayKey&& key, absl::string_view encoded_mutation) {
+        if (key.largest_batch_id() <= since_batch_id) {
+          return;
+        }
+        if (key.document_key().HasCollectionId(collection_group)) {
+          overlays_by_batch_id[key.largest_batch_id()].emplace(
+              ParseOverlay(key, encoded_mutation));
+        }
+      });
 
   // Trim down the overlays loaded above to respect the given `count`, and
   // return them.
@@ -157,8 +146,8 @@ LevelDbDocumentOverlayCache::GetOverlays(const std::string& collection_group,
   OverlayByDocumentKeyMap result;
   for (auto& overlays_by_batch_id_entry : overlays_by_batch_id) {
     for (auto& overlay : overlays_by_batch_id_entry.second) {
-      DocumentKey key = overlay.key();
-      result[key] = std::move(overlay);
+      DocumentKey document_key = overlay.key();
+      result[document_key] = std::move(overlay);
     }
     if (result.size() >= count) {
       break;
@@ -168,6 +157,27 @@ LevelDbDocumentOverlayCache::GetOverlays(const std::string& collection_group,
   return result;
 }
 
+int LevelDbDocumentOverlayCache::GetOverlayCount() const {
+  return CountEntriesWithKeyPrefix(
+      LevelDbDocumentOverlayKey::KeyPrefix(user_id_));
+}
+
+int LevelDbDocumentOverlayCache::GetLargestBatchIdIndexEntryCount() const {
+  return CountEntriesWithKeyPrefix(
+      LevelDbDocumentOverlayLargestBatchIdIndexKey::KeyPrefix(user_id_));
+}
+
+int LevelDbDocumentOverlayCache::CountEntriesWithKeyPrefix(
+    const std::string& key_prefix) const {
+  int count = 0;
+  auto it = db_->current_transaction()->NewIterator();
+  for (it->Seek(key_prefix);
+       it->Valid() && absl::StartsWith(it->key(), key_prefix); it->Next()) {
+    ++count;
+  }
+  return count;
+}
+
 Overlay LevelDbDocumentOverlayCache::ParseOverlay(
     const LevelDbDocumentOverlayKey& key,
     absl::string_view encoded_mutation) const {
@@ -181,42 +191,70 @@ Overlay LevelDbDocumentOverlayCache::ParseOverlay(
 }
 
 void LevelDbDocumentOverlayCache::SaveOverlay(int largest_batch_id,
-                                              const DocumentKey& key,
+                                              const DocumentKey& document_key,
                                               const Mutation& mutation) {
-  DeleteOverlay(key);
-  const std::string leveldb_key =
-      LevelDbDocumentOverlayKey::Key(user_id_, key, largest_batch_id);
-  auto serialized_mutation = serializer_->EncodeMutation(mutation);
-  db_->current_transaction()->Put(leveldb_key, serialized_mutation);
+  // Remove the existing overlay and any index entries pointing to it.
+  DeleteOverlay(document_key);
+
+  const LevelDbDocumentOverlayKey key(user_id_, document_key, largest_batch_id);
+
+  // Add the overlay to the database and index entries pointing to it.
+  auto* transaction = db_->current_transaction();
+  transaction->Put(key.Encode(), serializer_->EncodeMutation(mutation));
+  transaction->Put(LevelDbDocumentOverlayLargestBatchIdIndexKey::Key(key), "");
 }
 
-void LevelDbDocumentOverlayCache::DeleteOverlay(const model::DocumentKey& key) {
-  const std::string leveldb_key_prefix =
-      LevelDbDocumentOverlayKey::KeyPrefix(user_id_, key);
+void LevelDbDocumentOverlayCache::DeleteOverlay(
+    const model::DocumentKey& document_key) {
+  const std::string key_prefix =
+      LevelDbDocumentOverlayKey::KeyPrefix(user_id_, document_key);
   auto it = db_->current_transaction()->NewIterator();
-  for (it->Seek(leveldb_key_prefix);
-       it->Valid() && absl::StartsWith(it->key(), leveldb_key_prefix);
-       it->Next()) {
-    LevelDbDocumentOverlayKey decoded_key;
-    HARD_ASSERT(decoded_key.Decode(it->key()));
-    if (decoded_key.document_key() == key) {
-      db_->current_transaction()->Delete(it->key());
-    }
+  it->Seek(key_prefix);
+
+  if (!(it->Valid() && absl::StartsWith(it->key(), key_prefix))) {
+    return;
   }
+
+  LevelDbDocumentOverlayKey key;
+  HARD_ASSERT(key.Decode(it->key()));
+  if (key.document_key() == document_key) {
+    DeleteOverlay(key);
+  }
+}
+
+void LevelDbDocumentOverlayCache::DeleteOverlay(
+    const LevelDbDocumentOverlayKey& key) {
+  auto* transaction = db_->current_transaction();
+  transaction->Delete(key.Encode());
+  transaction->Delete(LevelDbDocumentOverlayLargestBatchIdIndexKey::Key(key));
 }
 
 void LevelDbDocumentOverlayCache::ForEachOverlay(
-    std::function<void(absl::string_view encoded_key,
-                       const LevelDbDocumentOverlayKey& decoded_key,
-                       absl::string_view encoded_mutation)> callback) const {
+    std::function<void((LevelDbDocumentOverlayKey && key,
+                        absl::string_view encoded_mutation))> callback) const {
   auto it = db_->current_transaction()->NewIterator();
   const std::string user_key = LevelDbDocumentOverlayKey::KeyPrefix(user_id_);
 
   for (it->Seek(user_key); it->Valid() && absl::StartsWith(it->key(), user_key);
        it->Next()) {
-    LevelDbDocumentOverlayKey decoded_key;
-    HARD_ASSERT(decoded_key.Decode(it->key()));
-    callback(it->key(), decoded_key, it->value());
+    LevelDbDocumentOverlayKey key;
+    HARD_ASSERT(key.Decode(it->key()));
+    callback(std::move(key), it->value());
+  }
+}
+
+void LevelDbDocumentOverlayCache::ForEachKeyWithLargestBatchId(
+    int largest_batch_id,
+    std::function<void(LevelDbDocumentOverlayKey&& key)> callback) const {
+  const std::string key_prefix =
+      LevelDbDocumentOverlayLargestBatchIdIndexKey::KeyPrefix(user_id_,
+                                                              largest_batch_id);
+  auto it = db_->current_transaction()->NewIterator();
+  for (it->Seek(key_prefix);
+       it->Valid() && absl::StartsWith(it->key(), key_prefix); it->Next()) {
+    LevelDbDocumentOverlayLargestBatchIdIndexKey key;
+    HARD_ASSERT(key.Decode(it->key()));
+    callback(std::move(key).ToLevelDbDocumentOverlayKey());
   }
 }
 

Firestore/core/src/local/leveldb_document_overlay_cache.h

@@ -33,6 +33,7 @@ class User;
 
 namespace local {
 
+class LevelDbDocumentOverlayCacheTestHelper;
 class LevelDbDocumentOverlayKey;
 class LevelDbPersistence;
 class LocalSerializer;
@@ -52,7 +53,7 @@ class LevelDbDocumentOverlayCache final : public DocumentOverlayCache {
       delete;
 
   absl::optional<model::mutation::Overlay> GetOverlay(
-      const model::DocumentKey& key) const override;
+      const model::DocumentKey&) const override;
 
   void SaveOverlays(int largest_batch_id,
                     const MutationByDocumentKeyMap& overlays) override;
@@ -67,21 +68,35 @@ class LevelDbDocumentOverlayCache final : public DocumentOverlayCache {
                                       std::size_t count) const override;
 
  private:
+  friend class LevelDbDocumentOverlayCacheTestHelper;
+
+  // Returns the number of index entries for the largest batch ID.
+  // This method exists for unit testing only.
+  int GetLargestBatchIdIndexEntryCount() const;
+
+  int GetOverlayCount() const override;
+  int CountEntriesWithKeyPrefix(const std::string& key_prefix) const;
+
   model::mutation::Overlay ParseOverlay(
       const LevelDbDocumentOverlayKey& key,
       absl::string_view encoded_mutation) const;
 
   void SaveOverlay(int largest_batch_id,
-                   const model::DocumentKey& key,
+                   const model::DocumentKey& document_key,
                    const model::Mutation& mutation);
 
-  void DeleteOverlay(const model::DocumentKey& key);
+  void DeleteOverlay(const model::DocumentKey&);
+
+  void DeleteOverlay(const LevelDbDocumentOverlayKey&);
 
   void ForEachOverlay(
-      std::function<void(absl::string_view encoded_key,
-                         const LevelDbDocumentOverlayKey& decoded_key,
+      std::function<void(LevelDbDocumentOverlayKey&&,
                          absl::string_view encoded_mutation)>) const;
 
+  void ForEachKeyWithLargestBatchId(
+      int largest_batch_id,
+      std::function<void(LevelDbDocumentOverlayKey&&)>) const;
+
   // The LevelDbDocumentOverlayCache instance is owned by LevelDbPersistence.
   LevelDbPersistence* db_;
 

Firestore/core/src/local/leveldb_key.cc

@@ -54,6 +54,8 @@ const char* kIndexConfigurationTable = "index_configuration";
 const char* kIndexStateTable = "index_state";
 const char* kIndexEntriesTable = "index_entries";
 const char* kDocumentOverlaysTable = "document_overlays";
+const char* kDocumentOverlaysLargestBatchIdIndexTable =
+    "document_overlays_largest_batch_id_index";
 
 /**
  * Labels for the components of keys. These serve to make keys self-describing.
@@ -1304,6 +1306,53 @@ bool LevelDbDocumentOverlayKey::Decode(absl::string_view key) {
   return reader.ok();
 }
 
+std::string LevelDbDocumentOverlayLargestBatchIdIndexKey::KeyPrefix() {
+  Writer writer;
+  writer.WriteTableName(kDocumentOverlaysLargestBatchIdIndexTable);
+  return writer.result();
+}
+
+std::string LevelDbDocumentOverlayLargestBatchIdIndexKey::KeyPrefix(
+    absl::string_view user_id) {
+  Writer writer;
+  writer.WriteTableName(kDocumentOverlaysLargestBatchIdIndexTable);
+  writer.WriteUserId(user_id);
+  return writer.result();
+}
+
+std::string LevelDbDocumentOverlayLargestBatchIdIndexKey::KeyPrefix(
+    absl::string_view user_id, model::BatchId largest_batch_id) {
+  Writer writer;
+  writer.WriteTableName(kDocumentOverlaysLargestBatchIdIndexTable);
+  writer.WriteUserId(user_id);
+  writer.WriteBatchId(largest_batch_id);
+  return writer.result();
+}
+
+std::string LevelDbDocumentOverlayLargestBatchIdIndexKey::Key(
+    absl::string_view user_id,
+    model::BatchId largest_batch_id,
+    const model::DocumentKey& document_key) {
+  Writer writer;
+  writer.WriteTableName(kDocumentOverlaysLargestBatchIdIndexTable);
+  writer.WriteUserId(user_id);
+  writer.WriteBatchId(largest_batch_id);
+  writer.WriteResourcePath(document_key.path());
+  writer.WriteTerminator();
+  return writer.result();
+}
+
+bool LevelDbDocumentOverlayLargestBatchIdIndexKey::Decode(
+    absl::string_view key) {
+  Reader reader{key};
+  reader.ReadTableNameMatching(kDocumentOverlaysLargestBatchIdIndexTable);
+  user_id_ = reader.ReadUserId();
+  largest_batch_id_ = reader.ReadBatchId();
+  document_key_ = reader.ReadDocumentKey();
+  reader.ReadTerminator();
+  return reader.ok();
+}
+
 }  // namespace local
 }  // namespace firestore
 }  // namespace firebase