Wire DocumentOverlayCache into LocalStore and LocalDocumentsView (#9549)

dconeybe · web-flow · commit 020a5eeb32d2 · 2022-04-01T10:37:10.000-04:00
diff --git a/Firestore/core/src/local/local_documents_view.h b/Firestore/core/src/local/local_documents_view.h
@@ -19,6 +19,7 @@
 
 #include <vector>
 
+#include "Firestore/core/src/local/document_overlay_cache.h"
 #include "Firestore/core/src/local/index_manager.h"
 #include "Firestore/core/src/local/mutation_queue.h"
 #include "Firestore/core/src/local/remote_document_cache.h"
@@ -47,9 +48,11 @@ class LocalDocumentsView {
  public:
   LocalDocumentsView(RemoteDocumentCache* remote_document_cache,
                      MutationQueue* mutation_queue,
+                     DocumentOverlayCache* document_overlay_cache,
                      IndexManager* index_manager)
       : remote_document_cache_{remote_document_cache},
         mutation_queue_{mutation_queue},
+        document_overlay_cache_{document_overlay_cache},
         index_manager_{index_manager} {
   }
 
@@ -136,13 +139,18 @@ class LocalDocumentsView {
     return mutation_queue_;
   }
 
+  DocumentOverlayCache* document_overlay_cache() {
+    return document_overlay_cache_;
+  }
+
   IndexManager* index_manager() {
     return index_manager_;
   }
 
  private:
   RemoteDocumentCache* remote_document_cache_;
   MutationQueue* mutation_queue_;
+  DocumentOverlayCache* document_overlay_cache_;
   IndexManager* index_manager_;
 };
 
diff --git a/Firestore/core/src/local/local_store.cc b/Firestore/core/src/local/local_store.cc
@@ -88,8 +88,10 @@ LocalStore::LocalStore(Persistence* persistence,
       query_engine_(query_engine) {
   index_manager_ = persistence->GetIndexManager(initial_user);
   mutation_queue_ = persistence->GetMutationQueue(initial_user, index_manager_);
+  document_overlay_cache_ = persistence->GetDocumentOverlayCache(initial_user);
   local_documents_ = absl::make_unique<LocalDocumentsView>(
-      remote_document_cache_, mutation_queue_, index_manager_);
+      remote_document_cache_, mutation_queue_, document_overlay_cache_,
+      index_manager_);
   remote_document_cache_->SetIndexManager(index_manager_);
 
   persistence->reference_delegate()->AddInMemoryPins(&local_view_references_);
@@ -130,7 +132,8 @@ DocumentMap LocalStore::HandleUserChange(const User& user) {
 
     // Recreate our LocalDocumentsView using the new MutationQueue.
     local_documents_ = absl::make_unique<LocalDocumentsView>(
-        remote_document_cache_, mutation_queue_, index_manager_);
+        remote_document_cache_, mutation_queue_, document_overlay_cache_,
+        index_manager_);
     query_engine_->SetLocalDocumentsView(local_documents_.get());
 
     // Union the old/new changed keys.
diff --git a/Firestore/core/src/local/local_store.h b/Firestore/core/src/local/local_store.h
@@ -26,6 +26,7 @@
 #include "Firestore/core/src/bundle/bundle_metadata.h"
 #include "Firestore/core/src/bundle/named_query.h"
 #include "Firestore/core/src/core/target_id_generator.h"
+#include "Firestore/core/src/local/document_overlay_cache.h"
 #include "Firestore/core/src/local/reference_set.h"
 #include "Firestore/core/src/local/target_data.h"
 #include "Firestore/core/src/model/document.h"
@@ -339,6 +340,12 @@ class LocalStore : public bundle::BundleCallback {
    */
   MutationQueue* mutation_queue_ = nullptr;
 
+  /**
+   * The overlays that can be used to short circuit applying all mutations from
+   * mutation queue.
+   */
+  DocumentOverlayCache* document_overlay_cache_ = nullptr;
+
   /** The set of all cached remote documents. */
   RemoteDocumentCache* remote_document_cache_ = nullptr;
 
diff --git a/Firestore/core/test/unit/local/counting_query_engine.cc b/Firestore/core/test/unit/local/counting_query_engine.cc
@@ -22,6 +22,7 @@
 #include "Firestore/core/src/model/mutable_document.h"
 #include "Firestore/core/src/model/mutation_batch.h"
 #include "Firestore/core/src/nanopb/byte_string.h"
+#include "Firestore/core/src/util/hard_assert.h"
 
 namespace firebase {
 namespace firestore {
@@ -43,9 +44,11 @@ void CountingQueryEngine::SetLocalDocumentsView(
       local_documents->remote_document_cache(), this);
   mutation_queue_ = absl::make_unique<WrappedMutationQueue>(
       local_documents->mutation_queue(), this);
+  document_overlay_cache_ = absl::make_unique<WrappedDocumentOverlayCache>(
+      local_documents->document_overlay_cache(), this);
   local_documents_ = absl::make_unique<LocalDocumentsView>(
       remote_documents_.get(), mutation_queue_.get(),
-      local_documents->index_manager());
+      document_overlay_cache_.get(), local_documents->index_manager());
   QueryEngine::SetLocalDocumentsView(local_documents_.get());
 }
 
@@ -54,6 +57,9 @@ void CountingQueryEngine::ResetCounts() {
   mutations_read_by_key_ = 0;
   documents_read_by_query_ = 0;
   documents_read_by_key_ = 0;
+  overlays_read_by_key_ = 0;
+  overlays_read_by_collection_ = 0;
+  overlays_read_by_collection_group_ = 0;
 }
 
 // MARK: - WrappedMutationQueue
@@ -173,6 +179,44 @@ model::MutableDocumentMap WrappedRemoteDocumentCache::GetMatching(
   return result;
 }
 
+// MARK: - WrappedDocumentOverlayCache
+
+absl::optional<model::Overlay> WrappedDocumentOverlayCache::GetOverlay(
+    const model::DocumentKey& key) const {
+  ++query_engine_->overlays_read_by_key_;
+  return subject_->GetOverlay(key);
+}
+
+void WrappedDocumentOverlayCache::SaveOverlays(
+    int largest_batch_id, const MutationByDocumentKeyMap& overlays) {
+  subject_->SaveOverlays(largest_batch_id, overlays);
+}
+
+void WrappedDocumentOverlayCache::RemoveOverlaysForBatchId(int batch_id) {
+  subject_->RemoveOverlaysForBatchId(batch_id);
+}
+
+DocumentOverlayCache::OverlayByDocumentKeyMap
+WrappedDocumentOverlayCache::GetOverlays(const model::ResourcePath& collection,
+                                         int since_batch_id) const {
+  auto result = subject_->GetOverlays(collection, since_batch_id);
+  query_engine_->overlays_read_by_collection_ += result.size();
+  return result;
+}
+
+DocumentOverlayCache::OverlayByDocumentKeyMap
+WrappedDocumentOverlayCache::GetOverlays(absl::string_view collection_group,
+                                         int since_batch_id,
+                                         std::size_t count) const {
+  auto result = subject_->GetOverlays(collection_group, since_batch_id, count);
+  query_engine_->overlays_read_by_collection_group_ += result.size();
+  return result;
+}
+
+int WrappedDocumentOverlayCache::GetOverlayCount() const {
+  HARD_FAIL("WrappedDocumentOverlayCache::GetOverlayCount() not implemented");
+}
+
 }  // namespace local
 }  // namespace firestore
 }  // namespace firebase
diff --git a/Firestore/core/test/unit/local/counting_query_engine.h b/Firestore/core/test/unit/local/counting_query_engine.h
@@ -21,6 +21,7 @@
 #include <utility>
 #include <vector>
 
+#include "Firestore/core/src/local/document_overlay_cache.h"
 #include "Firestore/core/src/local/mutation_queue.h"
 #include "Firestore/core/src/local/query_engine.h"
 #include "Firestore/core/src/local/remote_document_cache.h"
@@ -37,6 +38,7 @@ class Query;
 namespace local {
 
 class LocalDocumentsView;
+class WrappedDocumentOverlayCache;
 class WrappedMutationQueue;
 class WrappedRemoteDocumentCache;
 
@@ -90,17 +92,22 @@ class CountingQueryEngine : public QueryEngine {
   }
 
  private:
+  friend class WrappedDocumentOverlayCache;
   friend class WrappedMutationQueue;
   friend class WrappedRemoteDocumentCache;
 
   std::unique_ptr<LocalDocumentsView> local_documents_;
   std::unique_ptr<WrappedMutationQueue> mutation_queue_;
+  std::unique_ptr<WrappedDocumentOverlayCache> document_overlay_cache_;
   std::unique_ptr<WrappedRemoteDocumentCache> remote_documents_;
 
   size_t mutations_read_by_query_ = 0;
   size_t mutations_read_by_key_ = 0;
   size_t documents_read_by_query_ = 0;
   size_t documents_read_by_key_ = 0;
+  size_t overlays_read_by_key_ = 0;
+  size_t overlays_read_by_collection_ = 0;
+  size_t overlays_read_by_collection_group_ = 0;
 };
 
 /** A MutationQueue that counts document reads. */
@@ -186,6 +193,36 @@ class WrappedRemoteDocumentCache : public RemoteDocumentCache {
   CountingQueryEngine* query_engine_ = nullptr;
 };
 
+/** A DocumentOverlayCache that counts document reads. */
+class WrappedDocumentOverlayCache final : public DocumentOverlayCache {
+ public:
+  WrappedDocumentOverlayCache(DocumentOverlayCache* subject,
+                              CountingQueryEngine* query_engine)
+      : subject_(subject), query_engine_(query_engine) {
+  }
+
+  absl::optional<model::Overlay> GetOverlay(
+      const model::DocumentKey& key) const override;
+
+  void SaveOverlays(int largest_batch_id,
+                    const MutationByDocumentKeyMap& overlays) override;
+
+  void RemoveOverlaysForBatchId(int batch_id) override;
+
+  OverlayByDocumentKeyMap GetOverlays(const model::ResourcePath& collection,
+                                      int since_batch_id) const override;
+
+  OverlayByDocumentKeyMap GetOverlays(absl::string_view collection_group,
+                                      int since_batch_id,
+                                      std::size_t count) const override;
+
+ private:
+  int GetOverlayCount() const override;
+
+  DocumentOverlayCache* subject_ = nullptr;
+  CountingQueryEngine* query_engine_ = nullptr;
+};
+
 }  // namespace local
 }  // namespace firestore
 }  // namespace firebase
diff --git a/Firestore/core/test/unit/local/query_engine_test.cc b/Firestore/core/test/unit/local/query_engine_test.cc
@@ -120,10 +120,12 @@ class QueryEngineTest : public ::testing::Test {
         target_cache_(persistence_->target_cache()),
         index_manager_(dynamic_cast<MemoryIndexManager*>(
             persistence_->GetIndexManager(User::Unauthenticated()))),
-        local_documents_view_(remote_document_cache_,
-                              persistence_->GetMutationQueue(
-                                  User::Unauthenticated(), index_manager_),
-                              index_manager_) {
+        local_documents_view_(
+            remote_document_cache_,
+            persistence_->GetMutationQueue(User::Unauthenticated(),
+                                           index_manager_),
+            persistence_->GetDocumentOverlayCache(User::Unauthenticated()),
+            index_manager_) {
     remote_document_cache_->SetIndexManager(index_manager_);
     query_engine_.SetLocalDocumentsView(&local_documents_view_);
   }
