Improve performance of LevelDbDocumentOverlayCache::GetOverlays(absl::string_view) (#9420)

Author: dconeybe

Firestore/core/src/local/document_overlay_cache.h

@@ -18,7 +18,6 @@
 #define FIRESTORE_CORE_SRC_LOCAL_DOCUMENT_OVERLAY_CACHE_H_
 
 #include <cstdlib>
-#include <string>
 #include <unordered_map>
 
 #include "Firestore/core/src/model/document_key.h"
@@ -101,7 +100,7 @@ class DocumentOverlayCache {
    * overlay.
    */
   virtual OverlayByDocumentKeyMap GetOverlays(
-      const std::string& collection_group,
+      absl::string_view collection_group,
       int since_batch_id,
       std::size_t count) const = 0;
 

Firestore/core/src/local/leveldb_document_overlay_cache.cc

@@ -16,9 +16,7 @@
 
 #include "Firestore/core/src/local/leveldb_document_overlay_cache.h"
 
-#include <map>
 #include <string>
-#include <unordered_set>
 #include <utility>
 
 #include "Firestore/core/src/credentials/user.h"
@@ -99,46 +97,28 @@ LevelDbDocumentOverlayCache::GetOverlays(const ResourcePath& collection,
 }
 
 DocumentOverlayCache::OverlayByDocumentKeyMap
-LevelDbDocumentOverlayCache::GetOverlays(const std::string& collection_group,
+LevelDbDocumentOverlayCache::GetOverlays(absl::string_view collection_group,
                                          int since_batch_id,
                                          std::size_t count) const {
-  // TODO(dconeybe) Implement an index so that this query can be performed
-  // without requiring a full table scan.
-
-  // Load ALL overlays for the given `collection_group` whose largest_batch_id
-  // are greater than the given `since_batch_id`. By using a `std::map` keyed
-  // 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(
-      [&](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.
-  //
-  // Note that, as documented, all overlays for the largest_batch_id that pushes
-  // the size of the result set above the given `count` will be returned, even
-  // though this likely means that the size of the result set will be strictly
-  // greater than the given `count`.
+  absl::optional<int> current_batch_id;
   OverlayByDocumentKeyMap result;
-  for (auto& overlays_by_batch_id_entry : overlays_by_batch_id) {
-    for (auto& overlay : overlays_by_batch_id_entry.second) {
-      DocumentKey document_key = overlay.key();
-      result[document_key] = std::move(overlay);
-    }
-    if (result.size() >= count) {
-      break;
-    }
-  }
+  ForEachKeyInCollectionGroup(
+      collection_group, since_batch_id,
+      [&](LevelDbDocumentOverlayKey&& key) -> ForEachKeyAction {
+        if (!current_batch_id.has_value()) {
+          current_batch_id = key.largest_batch_id();
+        } else if (current_batch_id.value() != key.largest_batch_id()) {
+          if (result.size() >= count) {
+            return ForEachKeyAction::kStop;
+          }
+          current_batch_id = key.largest_batch_id();
+        }
 
+        absl::optional<Overlay> overlay = GetOverlay(key);
+        HARD_ASSERT(overlay.has_value());
+        result[std::move(key).document_key()] = std::move(overlay).value();
+        return ForEachKeyAction::kKeepGoing;
+      });
   return result;
 }
 
@@ -157,6 +137,11 @@ int LevelDbDocumentOverlayCache::GetCollectionIndexEntryCount() const {
       LevelDbDocumentOverlayCollectionIndexKey::KeyPrefix(user_id_));
 }
 
+int LevelDbDocumentOverlayCache::GetCollectionGroupIndexEntryCount() const {
+  return CountEntriesWithKeyPrefix(
+      LevelDbDocumentOverlayCollectionGroupIndexKey::KeyPrefix(user_id_));
+}
+
 int LevelDbDocumentOverlayCache::CountEntriesWithKeyPrefix(
     const std::string& key_prefix) const {
   int count = 0;
@@ -193,6 +178,12 @@ void LevelDbDocumentOverlayCache::SaveOverlay(int largest_batch_id,
   transaction->Put(key.Encode(), serializer_->EncodeMutation(mutation));
   transaction->Put(LevelDbDocumentOverlayLargestBatchIdIndexKey::Key(key), "");
   transaction->Put(LevelDbDocumentOverlayCollectionIndexKey::Key(key), "");
+
+  absl::optional<std::string> collection_group_index_key =
+      LevelDbDocumentOverlayCollectionGroupIndexKey::Key(key);
+  if (collection_group_index_key.has_value()) {
+    transaction->Put(std::move(collection_group_index_key).value(), "");
+  }
 }
 
 void LevelDbDocumentOverlayCache::DeleteOverlay(
@@ -219,19 +210,11 @@ void LevelDbDocumentOverlayCache::DeleteOverlay(
   transaction->Delete(key.Encode());
   transaction->Delete(LevelDbDocumentOverlayLargestBatchIdIndexKey::Key(key));
   transaction->Delete(LevelDbDocumentOverlayCollectionIndexKey::Key(key));
-}
 
-void LevelDbDocumentOverlayCache::ForEachOverlay(
-    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 key;
-    HARD_ASSERT(key.Decode(it->key()));
-    callback(std::move(key), it->value());
+  absl::optional<std::string> collection_group_index_key =
+      LevelDbDocumentOverlayCollectionGroupIndexKey::Key(key);
+  if (collection_group_index_key.has_value()) {
+    transaction->Delete(std::move(collection_group_index_key).value());
   }
 }
 
@@ -273,6 +256,36 @@ void LevelDbDocumentOverlayCache::ForEachKeyInCollection(
   }
 }
 
+void LevelDbDocumentOverlayCache::ForEachKeyInCollectionGroup(
+    absl::string_view collection_group,
+    int since_batch_id,
+    std::function<ForEachKeyAction(LevelDbDocumentOverlayKey&&)> callback)
+    const {
+  const std::string index_start_key =
+      LevelDbDocumentOverlayCollectionGroupIndexKey::KeyPrefix(
+          user_id_, collection_group, since_batch_id + 1);
+  const std::string index_key_prefix =
+      LevelDbDocumentOverlayCollectionGroupIndexKey::KeyPrefix(
+          user_id_, collection_group);
+
+  auto it = db_->current_transaction()->NewIterator();
+  for (it->Seek(index_start_key);
+       it->Valid() && absl::StartsWith(it->key(), index_key_prefix);
+       it->Next()) {
+    LevelDbDocumentOverlayCollectionGroupIndexKey key;
+    HARD_ASSERT(key.Decode(it->key()));
+    if (key.collection_group() != collection_group) {
+      break;
+    }
+    const ForEachKeyAction action =
+        callback(std::move(key).ToLevelDbDocumentOverlayKey());
+    if (action == ForEachKeyAction::kStop) {
+      break;
+    }
+    HARD_ASSERT(action == ForEachKeyAction::kKeepGoing);
+  }
+}
+
 absl::optional<Overlay> LevelDbDocumentOverlayCache::GetOverlay(
     const LevelDbDocumentOverlayKey& key) const {
   auto it = db_->current_transaction()->NewIterator();

Firestore/core/src/local/leveldb_document_overlay_cache.h

@@ -63,7 +63,7 @@ class LevelDbDocumentOverlayCache final : public DocumentOverlayCache {
   OverlayByDocumentKeyMap GetOverlays(const model::ResourcePath& collection,
                                       int since_batch_id) const override;
 
-  OverlayByDocumentKeyMap GetOverlays(const std::string& collection_group,
+  OverlayByDocumentKeyMap GetOverlays(absl::string_view collection_group,
                                       int since_batch_id,
                                       std::size_t count) const override;
 
@@ -74,10 +74,16 @@ class LevelDbDocumentOverlayCache final : public DocumentOverlayCache {
   // These methods exist for unit testing only.
   int GetLargestBatchIdIndexEntryCount() const;
   int GetCollectionIndexEntryCount() const;
+  int GetCollectionGroupIndexEntryCount() const;
 
   int GetOverlayCount() const override;
   int CountEntriesWithKeyPrefix(const std::string& key_prefix) const;
 
+  enum class ForEachKeyAction {
+    kKeepGoing,
+    kStop,
+  };
+
   model::mutation::Overlay ParseOverlay(
       const LevelDbDocumentOverlayKey& key,
       absl::string_view encoded_mutation) const;
@@ -90,10 +96,6 @@ class LevelDbDocumentOverlayCache final : public DocumentOverlayCache {
 
   void DeleteOverlay(const LevelDbDocumentOverlayKey&);
 
-  void ForEachOverlay(
-      std::function<void(LevelDbDocumentOverlayKey&&,
-                         absl::string_view encoded_mutation)>) const;
-
   void ForEachKeyWithLargestBatchId(
       int largest_batch_id,
       std::function<void(LevelDbDocumentOverlayKey&&)>) const;
@@ -103,6 +105,11 @@ class LevelDbDocumentOverlayCache final : public DocumentOverlayCache {
       int since_batch_id,
       std::function<void(LevelDbDocumentOverlayKey&&)>) const;
 
+  void ForEachKeyInCollectionGroup(
+      absl::string_view collection_group,
+      int since_batch_id,
+      std::function<ForEachKeyAction(LevelDbDocumentOverlayKey&&)>) const;
+
   absl::optional<model::mutation::Overlay> GetOverlay(
       const LevelDbDocumentOverlayKey& decoded_key) const;
 

Firestore/core/src/local/leveldb_key.cc

@@ -21,10 +21,12 @@
 
 #include "Firestore/core/src/local/leveldb_util.h"
 #include "Firestore/core/src/model/mutation_batch.h"
+#include "Firestore/core/src/util/hard_assert.h"
 #include "Firestore/core/src/util/ordered_code.h"
 #include "absl/base/attributes.h"
 #include "absl/strings/escaping.h"
 #include "absl/strings/str_cat.h"
+#include "absl/types/optional.h"
 
 using firebase::firestore::model::DocumentKey;
 using firebase::firestore::model::ResourcePath;
@@ -58,6 +60,8 @@ const char* kDocumentOverlaysLargestBatchIdIndexTable =
     "document_overlays_largest_batch_id_index";
 const char* kDocumentOverlaysCollectionIndexTable =
     "document_overlays_collection_index";
+const char* kDocumentOverlaysCollectionGroupIndexTable =
+    "document_overlays_collection_group_index";
 
 /**
  * Labels for the components of keys. These serve to make keys self-describing.
@@ -1401,6 +1405,72 @@ bool LevelDbDocumentOverlayCollectionIndexKey::Decode(absl::string_view key) {
   return reader.ok();
 }
 
+std::string LevelDbDocumentOverlayCollectionGroupIndexKey::KeyPrefix(
+    absl::string_view user_id) {
+  Writer writer;
+  writer.WriteTableName(kDocumentOverlaysCollectionGroupIndexTable);
+  writer.WriteUserId(user_id);
+  return writer.result();
+}
+
+std::string LevelDbDocumentOverlayCollectionGroupIndexKey::KeyPrefix(
+    absl::string_view user_id, absl::string_view collection_group) {
+  Writer writer;
+  writer.WriteTableName(kDocumentOverlaysCollectionGroupIndexTable);
+  writer.WriteUserId(user_id);
+  writer.WriteCollectionGroup(collection_group);
+  return writer.result();
+}
+
+std::string LevelDbDocumentOverlayCollectionGroupIndexKey::KeyPrefix(
+    absl::string_view user_id,
+    absl::string_view collection_group,
+    model::BatchId largest_batch_id) {
+  Writer writer;
+  writer.WriteTableName(kDocumentOverlaysCollectionGroupIndexTable);
+  writer.WriteUserId(user_id);
+  writer.WriteCollectionGroup(collection_group);
+  writer.WriteBatchId(largest_batch_id);
+  return writer.result();
+}
+
+std::string LevelDbDocumentOverlayCollectionGroupIndexKey::Key(
+    absl::string_view user_id,
+    absl::string_view collection_group,
+    model::BatchId largest_batch_id,
+    const model::DocumentKey& document_key) {
+  Writer writer;
+  writer.WriteTableName(kDocumentOverlaysCollectionGroupIndexTable);
+  writer.WriteUserId(user_id);
+  writer.WriteCollectionGroup(collection_group);
+  writer.WriteBatchId(largest_batch_id);
+  writer.WriteResourcePath(document_key.path());
+  writer.WriteTerminator();
+  return writer.result();
+}
+
+std::string LevelDbDocumentOverlayCollectionGroupIndexKey::Key(
+    const LevelDbDocumentOverlayKey& key) {
+  const absl::optional<std::string> collection_group =
+      key.document_key().GetCollectionId();
+  HARD_ASSERT(collection_group.has_value());
+  return Key(key.user_id(), collection_group.value(), key.largest_batch_id(),
+             key.document_key());
+}
+
+bool LevelDbDocumentOverlayCollectionGroupIndexKey::Decode(
+    absl::string_view key) {
+  Reader reader{key};
+  reader.ReadTableNameMatching(kDocumentOverlaysCollectionGroupIndexTable);
+  auto user_id = reader.ReadUserId();
+  collection_group_ = reader.ReadCollectionGroup();
+  auto largest_batch_id = reader.ReadBatchId();
+  auto document_key = reader.ReadDocumentKey();
+  reader.ReadTerminator();
+  Reset(std::move(user_id), largest_batch_id, std::move(document_key));
+  return reader.ok();
+}
+
 }  // namespace local
 }  // namespace firestore
 }  // namespace firebase