Reduce memory usage by applying query check sooner (#10745)

Author: wu-hui

Firestore/CHANGELOG.md

@@ -1,3 +1,7 @@
+
+# Unreleased
+- [fixed] Fix a potential high memory usage issue.
+
 # 10.5.0
 - [fixed] Add @discardableResult to addDocument API for easy handling unused return value. (#10640)
 

Firestore/core/src/local/leveldb_remote_document_cache.cc

@@ -26,7 +26,9 @@
 #include "Firestore/core/src/local/leveldb_persistence.h"
 #include "Firestore/core/src/local/local_serializer.h"
 #include "Firestore/core/src/model/document_key_set.h"
+#include "Firestore/core/src/model/model_fwd.h"
 #include "Firestore/core/src/model/mutable_document.h"
+#include "Firestore/core/src/model/overlay.h"
 #include "Firestore/core/src/nanopb/message.h"
 #include "Firestore/core/src/nanopb/reader.h"
 #include "Firestore/core/src/util/background_queue.h"
@@ -171,13 +173,18 @@ MutableDocumentMap LevelDbRemoteDocumentCache::GetAll(
 }
 
 MutableDocumentMap LevelDbRemoteDocumentCache::GetAllExisting(
-    DocumentVersionMap&& remote_map) const {
+    DocumentVersionMap&& remote_map,
+    const core::Query& query,
+    const model::OverlayByDocumentKeyMap& mutated_docs) const {
   BackgroundQueue tasks(executor_.get());
   AsyncResults<std::pair<DocumentKey, MutableDocument>> results;
   for (const auto& key_version : remote_map) {
-    tasks.Execute([this, &results, key_version] {
+    tasks.Execute([this, &results, &key_version, query, &mutated_docs] {
       auto document = Get(key_version.first).WithReadTime(key_version.second);
-      if (document.is_found_document()) {
+      if (document.is_found_document() &&
+          // Either the document matches the given query, or it is mutated.
+          (query.Matches(document) ||
+           mutated_docs.find(key_version.first) != mutated_docs.end())) {
         results.Insert(std::make_pair(key_version.first, std::move(document)));
       }
     });
@@ -206,24 +213,27 @@ MutableDocumentMap LevelDbRemoteDocumentCache::GetAll(
   MutableDocumentMap result;
   for (auto path = collections.cbegin();
        path != collections.cend() && result.size() < limit; path++) {
-    const auto remote_docs = GetAll(*path, offset, limit - result.size());
+    const auto remote_docs =
+        GetDocumentsMatchingQuery(Query(*path), offset, limit - result.size());
     for (const auto& doc : remote_docs) {
       result = result.insert(doc.first, doc.second);
     }
   }
   return result;
 }
 
-MutableDocumentMap LevelDbRemoteDocumentCache::GetAll(
-    const model::ResourcePath& path,
+MutableDocumentMap LevelDbRemoteDocumentCache::GetDocumentsMatchingQuery(
+    const core::Query& query,
     const model::IndexOffset& offset,
-    const absl::optional<size_t> limit) const {
+    absl::optional<size_t> limit,
+    const model::OverlayByDocumentKeyMap& mutated_docs) const {
   // Use the query path as a prefix for testing if a document matches the query.
 
   // Execute an index-free query and filter by read time. This is safe since
   // all document changes to queries that have a
   // last_limbo_free_snapshot_version (`since_read_time`) have a read time
   // set.
+  auto path = query.path();
   std::string start_key =
       LevelDbRemoteDocumentReadTimeKey::KeyPrefix(path, offset.read_time());
   auto it = db_->current_transaction()->NewIterator();
@@ -252,7 +262,8 @@ MutableDocumentMap LevelDbRemoteDocumentCache::GetAll(
     }
   }
 
-  return LevelDbRemoteDocumentCache::GetAllExisting(std::move(remote_map));
+  return LevelDbRemoteDocumentCache::GetAllExisting(std::move(remote_map),
+                                                    query, mutated_docs);
 }
 
 MutableDocument LevelDbRemoteDocumentCache::DecodeMaybeDocument(

Firestore/core/src/local/leveldb_remote_document_cache.h

@@ -22,11 +22,14 @@
 #include <thread>  // NOLINT(build/c++11)
 #include <vector>
 
+#include "Firestore/core/src/core/query.h"
 #include "Firestore/core/src/local/leveldb_index_manager.h"
 #include "Firestore/core/src/local/remote_document_cache.h"
 #include "Firestore/core/src/model/model_fwd.h"
+#include "Firestore/core/src/model/overlay.h"
 #include "Firestore/core/src/model/types.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/optional.h"
 
 namespace firebase {
 namespace firestore {
@@ -62,10 +65,11 @@ class LevelDbRemoteDocumentCache : public RemoteDocumentCache {
   model::MutableDocumentMap GetAll(const std::string& collection_group,
                                    const model::IndexOffset& offset,
                                    size_t limit) const override;
-  model::MutableDocumentMap GetAll(
-      const model::ResourcePath& path,
+  model::MutableDocumentMap GetDocumentsMatchingQuery(
+      const core::Query& query,
       const model::IndexOffset& offset,
-      absl::optional<size_t> limit = absl::nullopt) const override;
+      absl::optional<size_t> limit = absl::nullopt,
+      const model::OverlayByDocumentKeyMap& mutated_docs = {}) const override;
 
   void SetIndexManager(IndexManager* manager) override;
 
@@ -75,7 +79,9 @@ class LevelDbRemoteDocumentCache : public RemoteDocumentCache {
    * Type::Document together with its SnapshotVersion.
    */
   model::MutableDocumentMap GetAllExisting(
-      model::DocumentVersionMap&& remote_map) const;
+      model::DocumentVersionMap&& remote_map,
+      const core::Query& query,
+      const model::OverlayByDocumentKeyMap& mutated_docs = {}) const;
 
   model::MutableDocument DecodeMaybeDocument(
       absl::string_view encoded, const model::DocumentKey& key) const;

Firestore/core/src/local/local_documents_view.cc

@@ -64,9 +64,6 @@ using model::OverlayByDocumentKeyMap;
 using model::ResourcePath;
 using model::SnapshotVersion;
 
-using OverlayByDocumentKeyMap = std::
-    unordered_map<model::DocumentKey, model::Overlay, model::DocumentKeyHash>;
-
 Document LocalDocumentsView::GetDocument(
     const DocumentKey& key, const std::vector<MutationBatch>& batches) {
   MutableDocument document = remote_document_cache_->Get(key);
@@ -157,11 +154,12 @@ LocalWriteResult LocalDocumentsView::GetNextDocuments(
 
 DocumentMap LocalDocumentsView::GetDocumentsMatchingCollectionQuery(
     const Query& query, const IndexOffset& offset) {
-  MutableDocumentMap remote_documents =
-      remote_document_cache_->GetAll(query.path(), offset);
-  // Get locally persisted mutation batches.
+  // Get locally mutated documents
   OverlayByDocumentKeyMap overlays = document_overlay_cache_->GetOverlays(
       query.path(), offset.largest_batch_id());
+  MutableDocumentMap remote_documents =
+      remote_document_cache_->GetDocumentsMatchingQuery(
+          query, offset, absl::nullopt, overlays);
 
   // As documents might match the query because of their overlay we need to
   // include documents for all overlays in the initial document set.

Firestore/core/src/local/memory_remote_document_cache.cc

@@ -21,6 +21,7 @@
 #include "Firestore/core/src/local/memory_persistence.h"
 #include "Firestore/core/src/local/sizer.h"
 #include "Firestore/core/src/model/document.h"
+#include "Firestore/core/src/model/overlay.h"
 #include "Firestore/core/src/util/hard_assert.h"
 
 namespace firebase {
@@ -83,14 +84,16 @@ MutableDocumentMap MemoryRemoteDocumentCache::GetAll(const std::string&,
       "getAll(String, IndexOffset, int) is not supported.");
 }
 
-MutableDocumentMap MemoryRemoteDocumentCache::GetAll(
-    const model::ResourcePath& path,
+MutableDocumentMap MemoryRemoteDocumentCache::GetDocumentsMatchingQuery(
+    const core::Query& query,
     const model::IndexOffset& offset,
-    const absl::optional<size_t>) const {
+    absl::optional<size_t>,
+    const model::OverlayByDocumentKeyMap& mutated_docs) const {
   MutableDocumentMap results;
 
   // Documents are ordered by key, so we can use a prefix scan to narrow down
   // the documents we need to match the query against.
+  auto path = query.path();
   DocumentKey prefix{path.Append("")};
   size_t immediate_children_path_length = path.size() + 1;
   for (auto it = docs_.lower_bound(prefix); it != docs_.end(); ++it) {
@@ -110,6 +113,11 @@ MutableDocumentMap MemoryRemoteDocumentCache::GetAll(
       continue;
     }
 
+    if (mutated_docs.find(document.key()) == mutated_docs.end() &&
+        !query.Matches(document)) {
+      continue;
+    }
+
     // Note: We create an explicit copy to prevent modifications on the backing
     // data.
     results = results.insert(key, document.Clone());

Firestore/core/src/local/memory_remote_document_cache.h

@@ -27,6 +27,7 @@
 #include "Firestore/core/src/model/document_key.h"
 #include "Firestore/core/src/model/model_fwd.h"
 #include "Firestore/core/src/model/mutable_document.h"
+#include "Firestore/core/src/model/overlay.h"
 #include "Firestore/core/src/model/types.h"
 
 namespace firebase {
@@ -51,9 +52,12 @@ class MemoryRemoteDocumentCache : public RemoteDocumentCache {
   model::MutableDocumentMap GetAll(const std::string&,
                                    const model::IndexOffset&,
                                    size_t) const override;
-  model::MutableDocumentMap GetAll(const model::ResourcePath& path,
-                                   const model::IndexOffset& offset,
-                                   absl::optional<size_t>) const override;
+  model::MutableDocumentMap GetDocumentsMatchingQuery(
+      const core::Query& query,
+      const model::IndexOffset& offset,
+      absl::optional<size_t> limit = absl::nullopt,
+      const model::OverlayByDocumentKeyMap& mutated_docs = {}) const override;
+
   void SetIndexManager(IndexManager* manager) override;
 
   std::vector<model::DocumentKey> RemoveOrphanedDocuments(

Firestore/core/src/local/remote_document_cache.h

@@ -19,7 +19,9 @@
 
 #include <string>
 
+#include "Firestore/core/src/model/document_key.h"
 #include "Firestore/core/src/model/model_fwd.h"
+#include "Firestore/core/src/model/overlay.h"
 
 namespace firebase {
 namespace firestore {
@@ -100,17 +102,20 @@ class RemoteDocumentCache {
    *
    * Cached DeletedDocument entries have no bearing on query results.
    *
-   * @param path The collection path to match documents against.
+   * @param query The query to match documents against.
    * @param offset The read time and document key to start scanning at
    * (exclusive).
    * @param limit The maximum number of results to return.
    * If the limit is not defined, returns all matching documents.
+   * @param mutated_docs The documents with local mutations, they are read
+   * regardless if the remote version matches the given query.
    * @return The set of matching documents.
    */
-  virtual model::MutableDocumentMap GetAll(
-      const model::ResourcePath& path,
+  virtual model::MutableDocumentMap GetDocumentsMatchingQuery(
+      const core::Query& query,
       const model::IndexOffset& offset,
-      absl::optional<size_t> limit = absl::nullopt) const = 0;
+      absl::optional<size_t> limit = absl::nullopt,
+      const model::OverlayByDocumentKeyMap& mutated_docs = {}) const = 0;
 
   /**
    * Sets the index manager used by remote document cache.

Firestore/core/test/unit/local/counting_query_engine.cc

@@ -185,11 +185,13 @@ model::MutableDocumentMap WrappedRemoteDocumentCache::GetAll(
   return result;
 }
 
-model::MutableDocumentMap WrappedRemoteDocumentCache::GetAll(
-    const model::ResourcePath& path,
+model::MutableDocumentMap WrappedRemoteDocumentCache::GetDocumentsMatchingQuery(
+    const core::Query& query,
     const model::IndexOffset& offset,
-    absl::optional<size_t>) const {
-  auto result = subject_->GetAll(path, offset);
+    absl::optional<size_t> limit,
+    const model::OverlayByDocumentKeyMap& mutated_docs) const {
+  auto result =
+      subject_->GetDocumentsMatchingQuery(query, offset, limit, mutated_docs);
   query_engine_->documents_read_by_query_ += result.size();
   return result;
 }

Firestore/core/test/unit/local/counting_query_engine.h

@@ -196,9 +196,11 @@ class WrappedRemoteDocumentCache : public RemoteDocumentCache {
                                    const model::IndexOffset& offset,
                                    size_t limit) const override;
 
-  model::MutableDocumentMap GetAll(const model::ResourcePath& path,
-                                   const model::IndexOffset& offset,
-                                   absl::optional<size_t>) const override;
+  model::MutableDocumentMap GetDocumentsMatchingQuery(
+      const core::Query& query,
+      const model::IndexOffset& offset,
+      absl::optional<size_t>,
+      const model::OverlayByDocumentKeyMap& mutated_docs) const override;
 
   void SetIndexManager(IndexManager* manager) override {
     index_manager_ = NOT_NULL(manager);

Firestore/core/test/unit/local/local_store_test.cc

@@ -1036,10 +1036,10 @@ TEST_P(LocalStoreTest, UsesTargetMappingToExecuteQueries) {
   AcknowledgeMutationWithVersion(10);
   AcknowledgeMutationWithVersion(10);
 
-  // Execute the query, but note that we read all existing documents from the
+  // Execute the query, but note that we read matching documents from the
   // RemoteDocumentCache since we do not yet have target mapping.
   ExecuteQuery(query);
-  FSTAssertRemoteDocumentsRead(/* by_key */ 0, /* by_query= */ 3);
+  FSTAssertRemoteDocumentsRead(/* by_key */ 0, /* by_query= */ 2);
 
   // Issue a RemoteEvent to persist the target mapping.
   ApplyRemoteEvent(AddedRemoteEvent({Doc("foo/a", 10, Map("matches", true)),