Reduce memory usage by applying query check sooner

firebase-firestore/CHANGELOG.md

@@ -1,4 +1,5 @@
 # Unreleased
+* [fixed] Fix a potential high-memory usage issue.
 * [fixed] Fix an issue that stops some performance optimization being applied.
 
 # 24.4.1

firebase-firestore/firebase-firestore.gradle

@@ -146,6 +146,7 @@ dependencies {
     androidTestAnnotationProcessor 'com.google.auto.value:auto-value:1.6.5'
     annotationProcessor 'com.google.auto.value:auto-value:1.6.5'
 
+    testImplementation project(':firebase-firestore')
     testImplementation 'junit:junit:4.13.2'
     testImplementation 'androidx.test:core:1.2.0'
     testImplementation "org.hamcrest:hamcrest-junit:2.0.0.0"
@@ -155,6 +156,7 @@ dependencies {
     testImplementation 'com.fasterxml.jackson.core:jackson-databind:2.9.8'
     testImplementation 'com.google.guava:guava-testlib:12.0-rc2'
 
+    androidTestImplementation project(':firebase-firestore')
     androidTestImplementation 'junit:junit:4.13.2'
     androidTestImplementation("com.google.truth:truth:$googleTruthVersion"){
         exclude group: "org.codehaus.mojo", module: "animal-sniffer-annotations"

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalDocumentsView.java

@@ -363,10 +363,10 @@ private void populateOverlays(Map<DocumentKey, Overlay> overlays, Set<DocumentKe
 
   private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollectionQuery(
       Query query, IndexOffset offset) {
-    Map<DocumentKey, MutableDocument> remoteDocuments =
-        remoteDocumentCache.getAll(query.getPath(), offset);
     Map<DocumentKey, Overlay> overlays =
         documentOverlayCache.getOverlays(query.getPath(), offset.getLargestBatchId());
+    Map<DocumentKey, MutableDocument> remoteDocuments =
+        remoteDocumentCache.getDocumentsMatchingQuery(query, offset, overlays.keySet());
 
     // As documents might match the query because of their overlay we need to include documents
     // for all overlays in the initial document set.

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryRemoteDocumentCache.java

@@ -19,16 +19,18 @@
 
 import androidx.annotation.NonNull;
 import com.google.firebase.database.collection.ImmutableSortedMap;
+import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
 import com.google.firebase.firestore.model.MutableDocument;
-import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Set;
+import javax.annotation.Nonnull;
 
 /** In-memory cache of remote documents. */
 final class MemoryRemoteDocumentCache implements RemoteDocumentCache {
@@ -94,25 +96,26 @@ public Map<DocumentKey, MutableDocument> getAll(
   }
 
   @Override
-  public Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOffset offset) {
+  public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
+      Query query, IndexOffset offset, @Nonnull Set<DocumentKey> mutatedKeys) {
     Map<DocumentKey, MutableDocument> result = new HashMap<>();
 
     // Documents are ordered by key, so we can use a prefix scan to narrow down the documents
     // we need to match the query against.
-    DocumentKey prefix = DocumentKey.fromPath(collection.append(""));
+    DocumentKey prefix = DocumentKey.fromPath(query.getPath().append(""));
     Iterator<Map.Entry<DocumentKey, Document>> iterator = docs.iteratorFrom(prefix);
 
     while (iterator.hasNext()) {
       Map.Entry<DocumentKey, Document> entry = iterator.next();
       Document doc = entry.getValue();
 
       DocumentKey key = entry.getKey();
-      if (!collection.isPrefixOf(key.getPath())) {
+      if (!query.getPath().isPrefixOf(key.getPath())) {
         // We are now scanning the next collection. Abort.
         break;
       }
 
-      if (key.getPath().length() > collection.length() + 1) {
+      if (key.getPath().length() > query.getPath().length() + 1) {
         // Exclude entries from subcollections.
         continue;
       }
@@ -122,6 +125,10 @@ public Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOf
         continue;
       }
 
+      if (!mutatedKeys.contains(doc.getKey()) && !query.matches(doc)) {
+        continue;
+      }
+
       result.put(doc.getKey(), doc.mutableCopy());
     }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/local/RemoteDocumentCache.java

@@ -14,13 +14,15 @@
 
 package com.google.firebase.firestore.local;
 
+import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
 import com.google.firebase.firestore.model.MutableDocument;
-import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import java.util.Collection;
 import java.util.Map;
+import java.util.Set;
+import javax.annotation.Nonnull;
 
 /**
  * Represents cached documents received from the remote backend.
@@ -77,11 +79,14 @@ interface RemoteDocumentCache {
   Map<DocumentKey, MutableDocument> getAll(String collectionGroup, IndexOffset offset, int limit);
 
   /**
-   * Returns the documents from the provided collection.
+   * Returns the documents that match the given query.
    *
-   * @param collection The collection to read.
+   * @param query The query to match against remote documents.
    * @param offset The read time and document key to start scanning at (exclusive).
+   * @param mutatedKeys The keys of documents who have mutations attached, they should be read
+   *     regardless whether they match the given query.
    * @return A newly created map with the set of documents in the collection.
    */
-  Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOffset offset);
+  Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
+      Query query, IndexOffset offset, @Nonnull Set<DocumentKey> mutatedKeys);
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLiteRemoteDocumentCache.java

@@ -24,6 +24,7 @@
 import androidx.annotation.VisibleForTesting;
 import com.google.firebase.Timestamp;
 import com.google.firebase.database.collection.ImmutableSortedMap;
+import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
@@ -32,6 +33,7 @@
 import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.util.BackgroundQueue;
 import com.google.firebase.firestore.util.Executors;
+import com.google.firebase.firestore.util.Function;
 import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.MessageLite;
 import java.util.ArrayList;
@@ -40,7 +42,10 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.Executor;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 final class SQLiteRemoteDocumentCache implements RemoteDocumentCache {
   /** The number of bind args per collection group in {@link #getAll(String, IndexOffset, int)} */
@@ -135,7 +140,7 @@ public Map<DocumentKey, MutableDocument> getAll(Iterable<DocumentKey> documentKe
     while (longQuery.hasMoreSubqueries()) {
       longQuery
           .performNextSubquery()
-          .forEach(row -> processRowInBackground(backgroundQueue, results, row));
+          .forEach(row -> processRowInBackground(backgroundQueue, results, row, /*filter*/ null));
     }
     backgroundQueue.drain();
     return results;
@@ -153,15 +158,18 @@ public Map<DocumentKey, MutableDocument> getAll(
     if (collections.isEmpty()) {
       return Collections.emptyMap();
     } else if (BINDS_PER_STATEMENT * collections.size() < SQLitePersistence.MAX_ARGS) {
-      return getAll(collections, offset, limit);
+      return getAll(collections, offset, limit, /*filter*/ null);
     } else {
       // We need to fan out our collection scan since SQLite only supports 999 binds per statement.
       Map<DocumentKey, MutableDocument> results = new HashMap<>();
       int pageSize = SQLitePersistence.MAX_ARGS / BINDS_PER_STATEMENT;
       for (int i = 0; i < collections.size(); i += pageSize) {
         results.putAll(
             getAll(
-                collections.subList(i, Math.min(collections.size(), i + pageSize)), offset, limit));
+                collections.subList(i, Math.min(collections.size(), i + pageSize)),
+                offset,
+                limit,
+                /*filter*/ null));
       }
       return firstNEntries(results, limit, IndexOffset.DOCUMENT_COMPARATOR);
     }
@@ -171,7 +179,10 @@ public Map<DocumentKey, MutableDocument> getAll(
    * Returns the next {@code count} documents from the provided collections, ordered by read time.
    */
   private Map<DocumentKey, MutableDocument> getAll(
-      List<ResourcePath> collections, IndexOffset offset, int count) {
+      List<ResourcePath> collections,
+      IndexOffset offset,
+      int count,
+      @Nullable Function<MutableDocument, Boolean> filter) {
     Timestamp readTime = offset.getReadTime().getTimestamp();
     DocumentKey documentKey = offset.getDocumentKey();
 
@@ -207,13 +218,16 @@ private Map<DocumentKey, MutableDocument> getAll(
     Map<DocumentKey, MutableDocument> results = new HashMap<>();
     db.query(sql.toString())
         .binding(bindVars)
-        .forEach(row -> processRowInBackground(backgroundQueue, results, row));
+        .forEach(row -> processRowInBackground(backgroundQueue, results, row, filter));
     backgroundQueue.drain();
     return results;
   }
 
   private void processRowInBackground(
-      BackgroundQueue backgroundQueue, Map<DocumentKey, MutableDocument> results, Cursor row) {
+      BackgroundQueue backgroundQueue,
+      Map<DocumentKey, MutableDocument> results,
+      Cursor row,
+      @Nullable Function<MutableDocument, Boolean> filter) {
     byte[] rawDocument = row.getBlob(0);
     int readTimeSeconds = row.getInt(1);
     int readTimeNanos = row.getInt(2);
@@ -225,15 +239,22 @@ private void processRowInBackground(
         () -> {
           MutableDocument document =
               decodeMaybeDocument(rawDocument, readTimeSeconds, readTimeNanos);
-          synchronized (results) {
-            results.put(document.getKey(), document);
+          if (filter == null || filter.apply(document)) {
+            synchronized (results) {
+              results.put(document.getKey(), document);
+            }
           }
         });
   }
 
   @Override
-  public Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOffset offset) {
-    return getAll(Collections.singletonList(collection), offset, Integer.MAX_VALUE);
+  public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
+      Query query, IndexOffset offset, @Nonnull Set<DocumentKey> mutatedKeys) {
+    return getAll(
+        Collections.singletonList(query.getPath()),
+        offset,
+        Integer.MAX_VALUE,
+        (MutableDocument doc) -> query.matches(doc) || mutatedKeys.contains(doc.getKey()));
   }
 
   private MutableDocument decodeMaybeDocument(

firebase-firestore/src/test/java/com/google/firebase/firestore/local/CountingQueryEngine.java

@@ -33,6 +33,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.SortedSet;
 
 /**
@@ -167,8 +168,10 @@ public Map<DocumentKey, MutableDocument> getAll(
       }
 
       @Override
-      public Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOffset offset) {
-        Map<DocumentKey, MutableDocument> result = subject.getAll(collection, offset);
+      public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
+          Query query, IndexOffset offset, Set<DocumentKey> mutatedKeys) {
+        Map<DocumentKey, MutableDocument> result =
+            subject.getDocumentsMatchingQuery(query, offset, mutatedKeys);
         documentsReadByCollection[0] += result.size();
         return result;
       }

firebase-firestore/src/test/java/com/google/firebase/firestore/local/LocalStoreTestCase.java

@@ -1105,7 +1105,7 @@ public void testUsesTargetMappingToExecuteQueries() {
     // Execute the query, but note that we read all existing documents from the RemoteDocumentCache
     // since we do not yet have target mapping.
     executeQuery(query);
-    assertRemoteDocumentsRead(/* byKey= */ 0, /* byCollection= */ 3);
+    assertRemoteDocumentsRead(/* byKey= */ 0, /* byCollection= */ 2);
 
     // Issue a RemoteEvent to persist the target mapping.
     applyRemoteEvent(