Port performance optimizations to speed up reading large collections from Android

packages/firestore/src/local/indexeddb_mutation_queue.ts

@@ -25,6 +25,7 @@ import { BATCHID_UNKNOWN, MutationBatch } from '../model/mutation_batch';
 import { ResourcePath } from '../model/path';
 import { assert, fail } from '../util/assert';
 import { primitiveComparator } from '../util/misc';
+import { SortedMap } from '../util/sorted_map';
 import { SortedSet } from '../util/sorted_set';
 
 import * as EncodedResourcePath from './encoded_resource_path';
@@ -46,6 +47,8 @@ import { PersistenceTransaction, ReferenceDelegate } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { SimpleDbStore, SimpleDbTransaction } from './simple_db';
 
+import { AnyJs } from '../../src/util/misc';
+
 /** A mutation queue for a specific user, backed by IndexedDB. */
 export class IndexedDbMutationQueue implements MutationQueue {
   /**
@@ -325,7 +328,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
 
   getAllMutationBatchesAffectingDocumentKeys(
     transaction: PersistenceTransaction,
-    documentKeys: DocumentKeySet
+    documentKeys: SortedMap<DocumentKey, AnyJs>
   ): PersistencePromise<MutationBatch[]> {
     let uniqueBatchIDs = new SortedSet<BatchId>(primitiveComparator);
 

packages/firestore/src/local/indexeddb_remote_document_cache.ts

@@ -16,15 +16,20 @@
 
 import { Query } from '../core/query';
 import {
+  DocumentKeySet,
   documentKeySet,
   DocumentMap,
   documentMap,
+  DocumentSizeEntries,
   DocumentSizeEntry,
   MaybeDocumentMap,
-  maybeDocumentMap
+  maybeDocumentMap,
+  nullableMaybeDocumentMap,
+  NullableMaybeDocumentMap
 } from '../model/collections';
 import { Document, MaybeDocument, NoDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
+import { SortedMap } from '../util/sorted_map';
 
 import { SnapshotVersion } from '../core/snapshot_version';
 import { assert, fail } from '../util/assert';
@@ -178,6 +183,110 @@ export class IndexedDbRemoteDocumentCache implements RemoteDocumentCache {
       });
   }
 
+  getEntries(
+    transaction: PersistenceTransaction,
+    documentKeys: DocumentKeySet
+  ): PersistencePromise<NullableMaybeDocumentMap> {
+    let results = nullableMaybeDocumentMap();
+    return this.forEachDbEntry(
+      transaction,
+      documentKeys,
+      (key, dbRemoteDoc) => {
+        if (dbRemoteDoc) {
+          results = results.insert(
+            key,
+            this.serializer.fromDbRemoteDocument(dbRemoteDoc)
+          );
+        } else {
+          results = results.insert(key, null);
+        }
+      }
+    ).next(() => results);
+  }
+
+  /**
+   * Looks up several entries in the cache.
+   *
+   * @param documentKeys The set of keys entries to look up.
+   * @return A map of MaybeDocuments indexed by key (if a document cannot be
+   *     found, the key will be mapped to null) and a map of sizes indexed by
+   *     key (zero if the key cannot be found).
+   */
+  getSizedEntries(
+    transaction: PersistenceTransaction,
+    documentKeys: DocumentKeySet
+  ): PersistencePromise<DocumentSizeEntries> {
+    let results = nullableMaybeDocumentMap();
+    let sizeMap = new SortedMap<DocumentKey, number>(DocumentKey.comparator);
+    return this.forEachDbEntry(
+      transaction,
+      documentKeys,
+      (key, dbRemoteDoc) => {
+        if (dbRemoteDoc) {
+          results = results.insert(
+            key,
+            this.serializer.fromDbRemoteDocument(dbRemoteDoc)
+          );
+          sizeMap = sizeMap.insert(key, dbDocumentSize(dbRemoteDoc));
+        } else {
+          results = results.insert(key, null);
+          sizeMap = sizeMap.insert(key, 0);
+        }
+      }
+    ).next(() => {
+      return { maybeDocuments: results, sizeMap };
+    });
+  }
+
+  private forEachDbEntry(
+    transaction: PersistenceTransaction,
+    documentKeys: DocumentKeySet,
+    callback: (key: DocumentKey, doc: DbRemoteDocument | null) => void
+  ): PersistencePromise<void> {
+    if (documentKeys.isEmpty()) {
+      return PersistencePromise.resolve();
+    }
+
+    const range = IDBKeyRange.bound(
+      documentKeys.first()!.path.toArray(),
+      documentKeys.last()!.path.toArray()
+    );
+    const keyIter = documentKeys.getIterator();
+    let nextKey: DocumentKey | null = keyIter.getNext();
+
+    return remoteDocumentsStore(transaction)
+      .iterate({ range }, (potentialKeyRaw, dbRemoteDoc, control) => {
+        const potentialKey = DocumentKey.fromSegments(potentialKeyRaw);
+
+        // Go through keys not found in cache.
+        while (nextKey && DocumentKey.comparator(nextKey!, potentialKey) < 0) {
+          callback(nextKey!, null);
+          nextKey = keyIter.getNext();
+        }
+
+        if (nextKey && nextKey!.isEqual(potentialKey)) {
+          // Key found in cache.
+          callback(nextKey!, dbRemoteDoc);
+          nextKey = keyIter.hasNext() ? keyIter.getNext() : null;
+        }
+
+        // Skip to the next key (if there is one).
+        if (nextKey) {
+          control.skip(nextKey!.path.toArray());
+        } else {
+          control.done();
+        }
+      })
+      .next(() => {
+        // The rest of the keys are not in the cache. One case where `iterate`
+        // above won't go through them is when the cache is empty.
+        while (nextKey) {
+          callback(nextKey!, null);
+          nextKey = keyIter.hasNext() ? keyIter.getNext() : null;
+        }
+      });
+  }
+
   getDocumentsMatchingQuery(
     transaction: PersistenceTransaction,
     query: Query
@@ -381,6 +490,13 @@ class IndexedDbRemoteDocumentChangeBuffer extends RemoteDocumentChangeBuffer {
   ): PersistencePromise<DocumentSizeEntry | null> {
     return this.documentCache.getSizedEntry(transaction, documentKey);
   }
+
+  protected getAllFromCache(
+    transaction: PersistenceTransaction,
+    documentKeys: DocumentKeySet
+  ): PersistencePromise<DocumentSizeEntries> {
+    return this.documentCache.getSizedEntries(transaction, documentKeys);
+  }
 }
 
 export function isDocumentChangeMissingError(err: FirestoreError): boolean {

packages/firestore/src/local/local_documents_view.ts

@@ -21,7 +21,9 @@ import {
   DocumentMap,
   documentMap,
   MaybeDocumentMap,
-  maybeDocumentMap
+  maybeDocumentMap,
+  NullableMaybeDocumentMap,
+  nullableMaybeDocumentMap
 } from '../model/collections';
 import { Document, MaybeDocument, NoDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
@@ -74,6 +76,23 @@ export class LocalDocumentsView {
     });
   }
 
+  // Returns the view of the given `docs` as they would appear after applying
+  // all mutations in the given `batches`.
+  private applyLocalMutationsToDocuments(
+    transaction: PersistenceTransaction,
+    docs: NullableMaybeDocumentMap,
+    batches: MutationBatch[]
+  ): NullableMaybeDocumentMap {
+    let results = nullableMaybeDocumentMap();
+    docs.forEach((key, localView) => {
+      for (const batch of batches) {
+        localView = batch.applyToLocalView(key, localView);
+      }
+      results = results.insert(key, localView);
+    });
+    return results;
+  }
+
   /**
    * Gets the local view of the documents identified by `keys`.
    *
@@ -83,29 +102,38 @@ export class LocalDocumentsView {
   getDocuments(
     transaction: PersistenceTransaction,
     keys: DocumentKeySet
+  ): PersistencePromise<MaybeDocumentMap> {
+    return this.remoteDocumentCache
+      .getEntries(transaction, keys)
+      .next(docs => this.getLocalViewOfDocuments(transaction, docs));
+  }
+
+  /**
+   * Similar to `getDocuments`, but creates the local view from the given
+   * `baseDocs` without retrieving documents from the local store.
+   */
+  getLocalViewOfDocuments(
+    transaction: PersistenceTransaction,
+    baseDocs: NullableMaybeDocumentMap
   ): PersistencePromise<MaybeDocumentMap> {
     return this.mutationQueue
-      .getAllMutationBatchesAffectingDocumentKeys(transaction, keys)
+      .getAllMutationBatchesAffectingDocumentKeys(transaction, baseDocs)
       .next(batches => {
-        const promises = [] as Array<PersistencePromise<void>>;
+        const docs = this.applyLocalMutationsToDocuments(
+          transaction,
+          baseDocs,
+          batches
+        );
         let results = maybeDocumentMap();
-        keys.forEach(key => {
-          promises.push(
-            this.getDocumentInternal(transaction, key, batches).next(
-              maybeDoc => {
-                // TODO(http://b/32275378): Don't conflate missing / deleted.
-                if (!maybeDoc) {
-                  maybeDoc = new NoDocument(
-                    key,
-                    SnapshotVersion.forDeletedDoc()
-                  );
-                }
-                results = results.insert(key, maybeDoc);
-              }
-            )
-          );
+        docs.forEach((key, maybeDoc) => {
+          // TODO(http://b/32275378): Don't conflate missing / deleted.
+          if (!maybeDoc) {
+            maybeDoc = new NoDocument(key, SnapshotVersion.forDeletedDoc());
+          }
+          results = results.insert(key, maybeDoc);
         });
-        return PersistencePromise.waitFor(promises).next(() => results);
+
+        return results;
       });
   }
 

packages/firestore/src/local/local_serializer.ts

@@ -71,7 +71,9 @@ export class LocalSerializer {
   /** Encodes a document for storage locally. */
   toDbRemoteDocument(maybeDoc: MaybeDocument): DbRemoteDocument {
     if (maybeDoc instanceof Document) {
-      const doc = this.remoteSerializer.toDocument(maybeDoc);
+      const doc = maybeDoc.proto
+        ? maybeDoc.proto
+        : this.remoteSerializer.toDocument(maybeDoc);
       const hasCommittedMutations = maybeDoc.hasCommittedMutations;
       return new DbRemoteDocument(
         /* unknownDocument= */ null,

packages/firestore/src/local/local_store.ts

@@ -23,6 +23,7 @@ import {
   DocumentKeySet,
   documentKeySet,
   DocumentMap,
+  maybeDocumentMap,
   MaybeDocumentMap
 } from '../model/collections';
 import { MaybeDocument } from '../model/document';
@@ -466,11 +467,18 @@ export class LocalStore {
           }
         );
 
-        let changedDocKeys = documentKeySet();
+        let changedDocs = maybeDocumentMap();
+        let updatedKeys = documentKeySet();
         remoteEvent.documentUpdates.forEach((key, doc) => {
-          changedDocKeys = changedDocKeys.add(key);
-          promises.push(
-            documentBuffer.getEntry(txn, key).next(existingDoc => {
+          updatedKeys = updatedKeys.add(key);
+        });
+
+        // Each loop iteration only affects its "own" doc, so it's safe to get all the remote
+        // documents in advance in a single call.
+        promises.push(
+          documentBuffer.getEntries(txn, updatedKeys).next(existingDocs => {
+            remoteEvent.documentUpdates.forEach((key, doc) => {
+              const existingDoc = existingDocs.get(key);
               // If a document update isn't authoritative, make sure we don't
               // apply an old document version to the remote cache. We make an
               // exception for SnapshotVersion.MIN which can happen for
@@ -484,6 +492,7 @@ export class LocalStore {
                 doc.version.compareTo(existingDoc.version) >= 0
               ) {
                 documentBuffer.addEntry(doc);
+                changedDocs = changedDocs.insert(key, doc);
               } else {
                 log.debug(
                   LOG_TAG,
@@ -495,14 +504,18 @@ export class LocalStore {
                   doc.version
                 );
               }
-            })
-          );
-          if (remoteEvent.resolvedLimboDocuments.has(key)) {
-            promises.push(
-              this.persistence.referenceDelegate.updateLimboDocument(txn, key)
-            );
-          }
-        });
+
+              if (remoteEvent.resolvedLimboDocuments.has(key)) {
+                promises.push(
+                  this.persistence.referenceDelegate.updateLimboDocument(
+                    txn,
+                    key
+                  )
+                );
+              }
+            });
+          })
+        );
 
         // HACK: The only reason we allow a null snapshot version is so that we
         // can synthesize remote events when we get permission denied errors while
@@ -532,7 +545,10 @@ export class LocalStore {
         return PersistencePromise.waitFor(promises)
           .next(() => documentBuffer.apply(txn))
           .next(() => {
-            return this.localDocuments.getDocuments(txn, changedDocKeys);
+            return this.localDocuments.getLocalViewOfDocuments(
+              txn,
+              changedDocs
+            );
           });
       }
     );

packages/firestore/src/local/memory_mutation_queue.ts

@@ -24,13 +24,16 @@ import { BATCHID_UNKNOWN, MutationBatch } from '../model/mutation_batch';
 import { emptyByteString } from '../platform/platform';
 import { assert } from '../util/assert';
 import { primitiveComparator } from '../util/misc';
+import { SortedMap } from '../util/sorted_map';
 import { SortedSet } from '../util/sorted_set';
 
 import { MutationQueue } from './mutation_queue';
 import { PersistenceTransaction, ReferenceDelegate } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { DocReference } from './reference_set';
 
+import { AnyJs } from '../../src/util/misc';
+
 export class MemoryMutationQueue implements MutationQueue {
   /**
    * The set of all mutations that have been sent but not yet been applied to
@@ -203,7 +206,7 @@ export class MemoryMutationQueue implements MutationQueue {
 
   getAllMutationBatchesAffectingDocumentKeys(
     transaction: PersistenceTransaction,
-    documentKeys: DocumentKeySet
+    documentKeys: SortedMap<DocumentKey, AnyJs>
   ): PersistencePromise<MutationBatch[]> {
     let uniqueBatchIDs = new SortedSet<number>(primitiveComparator);