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

* Rename

* First commit

* Create slimy-elephants-hear.md

Author: wu-hui

.changeset/slimy-elephants-hear.md

@@ -0,0 +1,5 @@
+---
+"@firebase/firestore": patch
+---
+
+Reduce memory usage by applying query check sooner in remote document cache.

packages/firestore/src/local/indexeddb_remote_document_cache.ts

@@ -15,12 +15,14 @@
  * limitations under the License.
  */
 
+import { Query, queryMatches } from '../core/query';
 import { SnapshotVersion } from '../core/snapshot_version';
 import {
   DocumentKeySet,
   DocumentSizeEntries,
   MutableDocumentMap,
-  mutableDocumentMap
+  mutableDocumentMap,
+  OverlayMap
 } from '../model/collections';
 import { MutableDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
@@ -273,11 +275,13 @@ class IndexedDbRemoteDocumentCacheImpl implements IndexedDbRemoteDocumentCache {
       });
   }
 
-  getAllFromCollection(
+  getDocumentsMatchingQuery(
     transaction: PersistenceTransaction,
-    collection: ResourcePath,
-    offset: IndexOffset
+    query: Query,
+    offset: IndexOffset,
+    mutatedDocs: OverlayMap
   ): PersistencePromise<MutableDocumentMap> {
+    const collection = query.path;
     const startKey = [
       collection.popLast().toArray(),
       collection.lastSegment(),
@@ -307,7 +311,13 @@ class IndexedDbRemoteDocumentCacheImpl implements IndexedDbRemoteDocumentCache {
             ),
             dbRemoteDoc
           );
-          results = results.insert(document.key, document);
+          if (
+            document.isFoundDocument() &&
+            (queryMatches(query, document) || mutatedDocs.has(document.key))
+          ) {
+            // Either the document matches the given query, or it is mutated.
+            results = results.insert(document.key, document);
+          }
         }
         return results;
       });

packages/firestore/src/local/local_documents_view.ts

@@ -509,18 +509,19 @@ export class LocalDocumentsView {
     offset: IndexOffset
   ): PersistencePromise<DocumentMap> {
     // Query the remote documents and overlay mutations.
-    let remoteDocuments: MutableDocumentMap;
-    return this.remoteDocumentCache
-      .getAllFromCollection(transaction, query.path, offset)
-      .next(queryResults => {
-        remoteDocuments = queryResults;
-        return this.documentOverlayCache.getOverlaysForCollection(
+    let overlays: OverlayMap;
+    return this.documentOverlayCache
+      .getOverlaysForCollection(transaction, query.path, offset.largestBatchId)
+      .next(result => {
+        overlays = result;
+        return this.remoteDocumentCache.getDocumentsMatchingQuery(
           transaction,
-          query.path,
-          offset.largestBatchId
+          query,
+          offset,
+          overlays
         );
       })
-      .next(overlays => {
+      .next(remoteDocuments => {
         // As documents might match the query because of their overlay we need to
         // include documents for all overlays in the initial document set.
         overlays.forEach((_, overlay) => {

packages/firestore/src/local/memory_remote_document_cache.ts

@@ -15,11 +15,13 @@
  * limitations under the License.
  */
 
+import { Query, queryMatches } from '../core/query';
 import { SnapshotVersion } from '../core/snapshot_version';
 import {
   DocumentKeySet,
   MutableDocumentMap,
-  mutableDocumentMap
+  mutableDocumentMap,
+  OverlayMap
 } from '../model/collections';
 import { Document, MutableDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
@@ -28,7 +30,6 @@ import {
   indexOffsetComparator,
   newIndexOffsetFromDocument
 } from '../model/field_index';
-import { ResourcePath } from '../model/path';
 import { debugAssert, fail } from '../util/assert';
 import { SortedMap } from '../util/sorted_map';
 
@@ -159,15 +160,17 @@ class MemoryRemoteDocumentCacheImpl implements MemoryRemoteDocumentCache {
     return PersistencePromise.resolve(results);
   }
 
-  getAllFromCollection(
+  getDocumentsMatchingQuery(
     transaction: PersistenceTransaction,
-    collectionPath: ResourcePath,
-    offset: IndexOffset
+    query: Query,
+    offset: IndexOffset,
+    mutatedDocs: OverlayMap
   ): PersistencePromise<MutableDocumentMap> {
     let results = mutableDocumentMap();
 
     // Documents are ordered by key, so we can use a prefix scan to narrow down
     // the documents we need to match the query against.
+    const collectionPath = query.path;
     const prefix = new DocumentKey(collectionPath.child(''));
     const iterator = this.docs.getIteratorFrom(prefix);
     while (iterator.hasNext()) {
@@ -188,6 +191,11 @@ class MemoryRemoteDocumentCacheImpl implements MemoryRemoteDocumentCache {
         // The document sorts before the offset.
         continue;
       }
+      if (!mutatedDocs.has(document.key) && !queryMatches(query, document)) {
+        // The document cannot possibly match the query.
+        continue;
+      }
+
       results = results.insert(document.key, document.mutableCopy());
     }
     return PersistencePromise.resolve(results);

packages/firestore/src/local/remote_document_cache.ts

@@ -15,11 +15,15 @@
  * limitations under the License.
  */
 
-import { DocumentKeySet, MutableDocumentMap } from '../model/collections';
+import { Query } from '../core/query';
+import {
+  DocumentKeySet,
+  MutableDocumentMap,
+  OverlayMap
+} from '../model/collections';
 import { MutableDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
 import { IndexOffset } from '../model/field_index';
-import { ResourcePath } from '../model/path';
 
 import { IndexManager } from './index_manager';
 import { PersistencePromise } from './persistence_promise';
@@ -62,16 +66,17 @@ export interface RemoteDocumentCache {
   ): PersistencePromise<MutableDocumentMap>;
 
   /**
-   * Returns the documents from the provided collection.
+   * Returns the documents matching the given query
    *
-   * @param collection - The collection to read.
+   * @param query - The query to match documents against.
    * @param offset - The offset to start the scan at (exclusive).
    * @returns The set of matching documents.
    */
-  getAllFromCollection(
+  getDocumentsMatchingQuery(
     transaction: PersistenceTransaction,
-    collection: ResourcePath,
-    offset: IndexOffset
+    query: Query,
+    offset: IndexOffset,
+    mutatedDocs: OverlayMap
   ): PersistencePromise<MutableDocumentMap>;
 
   /**

packages/firestore/test/unit/local/counting_query_engine.ts

@@ -101,9 +101,19 @@ export class CountingQueryEngine extends QueryEngine {
       setIndexManager: (indexManager: IndexManager) => {
         subject.setIndexManager(indexManager);
       },
-      getAllFromCollection: (transaction, collection, sinceReadTime) => {
+      getDocumentsMatchingQuery: (
+        transaction,
+        query,
+        sinceReadTime,
+        overlays
+      ) => {
         return subject
-          .getAllFromCollection(transaction, collection, sinceReadTime)
+          .getDocumentsMatchingQuery(
+            transaction,
+            query,
+            sinceReadTime,
+            overlays
+          )
           .next(result => {
             this.documentsReadByCollection += result.size;
             return result;

packages/firestore/test/unit/local/local_store.test.ts

@@ -2288,10 +2288,10 @@ function genericLocalStoreTests(
         .afterAcknowledgingMutation({ documentVersion: 10 })
         .afterAcknowledgingMutation({ documentVersion: 10 })
         .afterExecutingQuery(query1)
-        // Execute the query, but note that we read all existing documents
+        // Execute the query, but note that we read matching documents
         // from the RemoteDocumentCache since we do not yet have target
         // mapping.
-        .toHaveRead({ documentsByCollection: 3 })
+        .toHaveRead({ documentsByCollection: 2 })
         .after(
           docAddedRemoteEvent(
             [

packages/firestore/test/unit/local/remote_document_cache.test.ts

@@ -19,21 +19,27 @@ import { expect } from 'chai';
 
 import { User } from '../../../src/auth/user';
 import { IndexedDbPersistence } from '../../../src/local/indexeddb_persistence';
-import { documentKeySet, DocumentMap } from '../../../src/model/collections';
+import {
+  documentKeySet,
+  DocumentMap,
+  newOverlayMap
+} from '../../../src/model/collections';
 import { MutableDocument, Document } from '../../../src/model/document';
 import {
   IndexOffset,
   INITIAL_LARGEST_BATCH_ID,
   newIndexOffsetFromDocument,
   newIndexOffsetSuccessorFromReadTime
 } from '../../../src/model/field_index';
+import { Overlay } from '../../../src/model/overlay';
 import {
   deletedDoc,
   doc,
   expectEqual,
   field,
+  filter,
   key,
-  path,
+  query,
   version,
   wrap
 } from '../../util/helpers';
@@ -464,9 +470,10 @@ function genericRemoteDocumentCacheTests(
       doc('c/1', VERSION, DOC_DATA)
     ]);
 
-    const matchingDocs = await cache.getAllFromCollection(
-      path('b'),
-      IndexOffset.min()
+    const matchingDocs = await cache.getDocumentsMatchingQuery(
+      query('b'),
+      IndexOffset.min(),
+      newOverlayMap()
     );
     assertMatches(
       [doc('b/1', VERSION, DOC_DATA), doc('b/2', VERSION, DOC_DATA)],
@@ -480,9 +487,10 @@ function genericRemoteDocumentCacheTests(
       doc('a/1/b/1', VERSION, DOC_DATA)
     ]);
 
-    const matchingDocs = await cache.getAllFromCollection(
-      path('a'),
-      IndexOffset.min()
+    const matchingDocs = await cache.getDocumentsMatchingQuery(
+      query('a'),
+      IndexOffset.min(),
+      newOverlayMap()
     );
     assertMatches([doc('a/1', VERSION, DOC_DATA)], matchingDocs);
   });
@@ -498,20 +506,62 @@ function genericRemoteDocumentCacheTests(
       doc('b/new', 3, DOC_DATA).setReadTime(version(13))
     ]);
 
-    const matchingDocs = await cache.getAllFromCollection(
-      path('b'),
-      newIndexOffsetSuccessorFromReadTime(version(12), INITIAL_LARGEST_BATCH_ID)
+    const matchingDocs = await cache.getDocumentsMatchingQuery(
+      query('b'),
+      newIndexOffsetSuccessorFromReadTime(
+        version(12),
+        INITIAL_LARGEST_BATCH_ID
+      ),
+      newOverlayMap()
     );
     assertMatches([doc('b/new', 3, DOC_DATA)], matchingDocs);
   });
 
+  it('getDocumentsMatchingQuery() applies query check', async () => {
+    await cache.addEntries([
+      doc('a/1', 1, { matches: true }).setReadTime(version(1))
+    ]);
+    await cache.addEntries([
+      doc('a/2', 1, { matches: true }).setReadTime(version(2))
+    ]);
+    await cache.addEntries([
+      doc('a/3', 1, { matches: false }).setReadTime(version(3))
+    ]);
+
+    const matchingDocs = await cache.getDocumentsMatchingQuery(
+      query('a', filter('matches', '==', true)),
+      newIndexOffsetSuccessorFromReadTime(version(1), INITIAL_LARGEST_BATCH_ID),
+      newOverlayMap()
+    );
+    assertMatches([doc('a/2', 1, { matches: true })], matchingDocs);
+  });
+
+  it('getDocumentsMatchingQuery() respects mutated documents', async () => {
+    await cache.addEntries([
+      doc('a/1', 1, { matches: true }).setReadTime(version(1))
+    ]);
+    await cache.addEntries([
+      doc('a/2', 1, { matches: false }).setReadTime(version(2))
+    ]);
+
+    const mutatedDocs = newOverlayMap();
+    mutatedDocs.set(key('a/2'), {} as Overlay);
+    const matchingDocs = await cache.getDocumentsMatchingQuery(
+      query('a', filter('matches', '==', true)),
+      newIndexOffsetSuccessorFromReadTime(version(1), INITIAL_LARGEST_BATCH_ID),
+      mutatedDocs
+    );
+    assertMatches([doc('a/2', 1, { matches: false })], matchingDocs);
+  });
+
   it('getAll() uses read time rather than update time', async () => {
     await cache.addEntries([doc('b/old', 1, DOC_DATA).setReadTime(version(2))]);
     await cache.addEntries([doc('b/new', 2, DOC_DATA).setReadTime(version(1))]);
 
-    const matchingDocs = await cache.getAllFromCollection(
-      path('b'),
-      newIndexOffsetSuccessorFromReadTime(version(1), INITIAL_LARGEST_BATCH_ID)
+    const matchingDocs = await cache.getDocumentsMatchingQuery(
+      query('b'),
+      newIndexOffsetSuccessorFromReadTime(version(1), INITIAL_LARGEST_BATCH_ID),
+      newOverlayMap()
     );
     assertMatches([doc('b/old', 1, DOC_DATA)], matchingDocs);
   });
@@ -545,7 +595,11 @@ function genericRemoteDocumentCacheTests(
     document.data.set(field('state'), wrap('new'));
 
     document = await cache
-      .getAllFromCollection(path('coll'), IndexOffset.min())
+      .getDocumentsMatchingQuery(
+        query('coll'),
+        IndexOffset.min(),
+        newOverlayMap()
+      )
       .then(m => m.get(key('coll/doc'))!);
     verifyOldValue(document);
 

packages/firestore/test/unit/local/test_remote_document_cache.ts

@@ -15,6 +15,7 @@
  * limitations under the License.
  */
 
+import { Query } from '../../../src/core/query';
 import { SnapshotVersion } from '../../../src/core/snapshot_version';
 import { IndexManager } from '../../../src/local/index_manager';
 import { Persistence } from '../../../src/local/persistence';
@@ -23,12 +24,12 @@ import { RemoteDocumentCache } from '../../../src/local/remote_document_cache';
 import { RemoteDocumentChangeBuffer } from '../../../src/local/remote_document_change_buffer';
 import {
   DocumentKeySet,
-  MutableDocumentMap
+  MutableDocumentMap,
+  OverlayMap
 } from '../../../src/model/collections';
 import { Document, MutableDocument } from '../../../src/model/document';
 import { DocumentKey } from '../../../src/model/document_key';
 import { IndexOffset } from '../../../src/model/field_index';
-import { ResourcePath } from '../../../src/model/path';
 
 /**
  * A wrapper around a RemoteDocumentCache that automatically creates a
@@ -109,14 +110,16 @@ export class TestRemoteDocumentCache {
     });
   }
 
-  getAllFromCollection(
-    collection: ResourcePath,
-    offset: IndexOffset
+  getDocumentsMatchingQuery(
+    query: Query,
+    offset: IndexOffset,
+    mutatedDocs: OverlayMap
   ): Promise<MutableDocumentMap> {
     return this.persistence.runTransaction(
       'getAllFromCollection',
       'readonly',
-      txn => this.cache.getAllFromCollection(txn, collection, offset)
+      txn =>
+        this.cache.getDocumentsMatchingQuery(txn, query, offset, mutatedDocs)
     );
   }