firestore::local::LevelDbDocumentOverlayCache added. (#9345)

Author: dconeybe

Firestore/core/src/local/leveldb_document_overlay_cache.cc

@@ -16,52 +16,200 @@
 
 #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"
+#include "Firestore/core/src/local/leveldb_key.h"
+#include "Firestore/core/src/local/leveldb_persistence.h"
+#include "Firestore/core/src/local/local_serializer.h"
+#include "Firestore/core/src/nanopb/message.h"
+#include "Firestore/core/src/nanopb/reader.h"
+#include "Firestore/core/src/util/hard_assert.h"
+#include "absl/strings/match.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
+
 namespace firebase {
 namespace firestore {
 namespace local {
 
+using credentials::User;
 using model::DocumentKey;
 using model::Mutation;
 using model::ResourcePath;
 using model::mutation::Overlay;
+using model::mutation::OverlayHash;
+using nanopb::Message;
+using nanopb::StringReader;
 
-// TODO(dconeybe) Implement these methods.
-
-LevelDbDocumentOverlayCache::LevelDbDocumentOverlayCache() {
+LevelDbDocumentOverlayCache::LevelDbDocumentOverlayCache(
+    const User& user, LevelDbPersistence* db, LocalSerializer* serializer)
+    : db_(NOT_NULL(db)),
+      serializer_(NOT_NULL(serializer)),
+      user_id_(user.is_authenticated() ? user.uid() : "") {
 }
 
 absl::optional<Overlay> LevelDbDocumentOverlayCache::GetOverlay(
     const DocumentKey& key) const {
-  (void)key;
-  return absl::nullopt;
+  const std::string leveldb_key_prefix =
+      LevelDbDocumentOverlayKey::KeyPrefix(user_id_, key);
+
+  auto it = db_->current_transaction()->NewIterator();
+  it->Seek(leveldb_key_prefix);
+
+  if (!(it->Valid() && absl::StartsWith(it->key(), leveldb_key_prefix))) {
+    return absl::nullopt;
+  }
+
+  LevelDbDocumentOverlayKey decoded_key;
+  HARD_ASSERT(decoded_key.Decode(it->key()));
+  return ParseOverlay(decoded_key, it->value());
 }
 
 void LevelDbDocumentOverlayCache::SaveOverlays(
     int largest_batch_id, const MutationByDocumentKeyMap& overlays) {
-  (void)largest_batch_id;
-  (void)overlays;
+  for (const auto& overlays_entry : overlays) {
+    SaveOverlay(largest_batch_id, overlays_entry.first, overlays_entry.second);
+  }
 }
 
 void LevelDbDocumentOverlayCache::RemoveOverlaysForBatchId(int batch_id) {
-  (void)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);
+    }
+  });
 }
 
 DocumentOverlayCache::OverlayByDocumentKeyMap
 LevelDbDocumentOverlayCache::GetOverlays(const ResourcePath& collection,
                                          int since_batch_id) const {
-  (void)collection;
-  (void)since_batch_id;
-  return {};
+  // TODO(dconeybe) Implement an index so that this query can be performed
+  // without requiring a full table scan.
+
+  OverlayByDocumentKeyMap result;
+
+  const size_t immediate_children_path_length{collection.size() + 1};
+
+  ForEachOverlay([&](absl::string_view,
+                     const LevelDbDocumentOverlayKey& decoded_key,
+                     absl::string_view encoded_mutation) {
+    const DocumentKey key = decoded_key.document_key();
+    if (!collection.IsPrefixOf(key.path())) {
+      return;
+    }
+    // Documents from sub-collections
+    if (key.path().size() != immediate_children_path_length) {
+      return;
+    }
+
+    if (decoded_key.largest_batch_id() > since_batch_id) {
+      result[key] = ParseOverlay(decoded_key, encoded_mutation);
+    }
+  });
+
+  return result;
 }
 
 DocumentOverlayCache::OverlayByDocumentKeyMap
 LevelDbDocumentOverlayCache::GetOverlays(const std::string& collection_group,
                                          int since_batch_id,
                                          std::size_t count) const {
-  (void)collection_group;
-  (void)since_batch_id;
-  (void)count;
-  return {};
+  // 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([&](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));
+    }
+  });
+
+  // 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`.
+  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);
+    }
+    if (result.size() >= count) {
+      break;
+    }
+  }
+
+  return result;
+}
+
+Overlay LevelDbDocumentOverlayCache::ParseOverlay(
+    const LevelDbDocumentOverlayKey& key,
+    absl::string_view encoded_mutation) const {
+  StringReader reader{encoded_mutation};
+  auto maybe_message = Message<google_firestore_v1_Write>::TryParse(&reader);
+  Mutation mutation = serializer_->DecodeMutation(&reader, *maybe_message);
+  if (!reader.ok()) {
+    HARD_FAIL("Mutation proto failed to parse: %s", reader.status().ToString());
+  }
+  return Overlay(key.largest_batch_id(), std::move(mutation));
+}
+
+void LevelDbDocumentOverlayCache::SaveOverlay(int largest_batch_id,
+                                              const DocumentKey& 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);
+}
+
+void LevelDbDocumentOverlayCache::DeleteOverlay(const model::DocumentKey& key) {
+  const std::string leveldb_key_prefix =
+      LevelDbDocumentOverlayKey::KeyPrefix(user_id_, key);
+  auto it = db_->current_transaction()->NewIterator();
+  for (it->Seek(leveldb_key_prefix);
+       it->Valid() && absl::StartsWith(it->key(), leveldb_key_prefix);
+       it->Next()) {
+    db_->current_transaction()->Delete(it->key());
+  }
+}
+
+void LevelDbDocumentOverlayCache::ForEachOverlay(
+    std::function<void(absl::string_view encoded_key,
+                       const LevelDbDocumentOverlayKey& decoded_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());
+  }
 }
 
 }  // namespace local

Firestore/core/src/local/leveldb_document_overlay_cache.h

@@ -18,17 +18,30 @@
 #define FIRESTORE_CORE_SRC_LOCAL_LEVELDB_DOCUMENT_OVERLAY_CACHE_H_
 
 #include <cstdlib>
+#include <functional>
 #include <string>
 
 #include "Firestore/core/src/local/document_overlay_cache.h"
+#include "absl/strings/string_view.h"
 
 namespace firebase {
 namespace firestore {
+
+namespace credentials {
+class User;
+}  // namespace credentials
+
 namespace local {
 
+class LevelDbDocumentOverlayKey;
+class LevelDbPersistence;
+class LocalSerializer;
+
 class LevelDbDocumentOverlayCache final : public DocumentOverlayCache {
  public:
-  LevelDbDocumentOverlayCache();
+  LevelDbDocumentOverlayCache(const credentials::User& user,
+                              LevelDbPersistence* db,
+                              LocalSerializer* serializer);
 
   LevelDbDocumentOverlayCache(const LevelDbDocumentOverlayCache&) = delete;
   LevelDbDocumentOverlayCache& operator=(const LevelDbDocumentOverlayCache&) =
@@ -52,6 +65,34 @@ class LevelDbDocumentOverlayCache final : public DocumentOverlayCache {
   OverlayByDocumentKeyMap GetOverlays(const std::string& collection_group,
                                       int since_batch_id,
                                       std::size_t count) const override;
+
+ private:
+  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::Mutation& mutation);
+
+  void DeleteOverlay(const model::DocumentKey& key);
+
+  void ForEachOverlay(
+      std::function<void(absl::string_view encoded_key,
+                         const LevelDbDocumentOverlayKey& decoded_key,
+                         absl::string_view encoded_mutation)>) const;
+
+  // The LevelDbDocumentOverlayCache instance is owned by LevelDbPersistence.
+  LevelDbPersistence* db_;
+
+  // Owned by LevelDbPersistence.
+  LocalSerializer* serializer_ = nullptr;
+
+  /**
+   * The normalized user_id (i.e. after converting null to empty) as used in our
+   * LevelDB keys.
+   */
+  std::string user_id_;
 };
 
 }  // namespace local

Firestore/core/src/local/leveldb_key.cc

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