Port changes from Android SDK PR#3420.

Note that we are not going to do any processing in the background.

Author: ehsannas

packages/firestore/src/local/document_overlay_cache.ts

@@ -43,6 +43,15 @@ export interface DocumentOverlayCache {
     key: DocumentKey
   ): PersistencePromise<Overlay | null>;
 
+  /**
+   * Gets the saved overlay mutation for the given document keys. Skips keys for
+   * which there are no overlays.
+   */
+  getOverlays(
+    transaction: PersistenceTransaction,
+    keys: DocumentKeySet
+  ): PersistencePromise<OverlayMap>;
+
   /**
    * Saves the given document mutation map to persistence as overlays.
    * All overlays will have their largest batch id set to `largestBatchId`.

packages/firestore/src/local/indexeddb_document_overlay_cache.ts

@@ -81,6 +81,24 @@ export class IndexedDbDocumentOverlayCache implements DocumentOverlayCache {
       });
   }
 
+  getOverlays(
+    transaction: PersistenceTransaction,
+    keys: DocumentKeySet
+  ): PersistencePromise<OverlayMap> {
+    const result = newOverlayMap();
+    const promises: Array<PersistencePromise<void>> = [];
+    keys.forEach(key => {
+      promises.push(
+        this.getOverlay(transaction, key).next(overlay => {
+          if (overlay !== null) {
+            result.set(key, overlay);
+          }
+        })
+      );
+    });
+    return PersistencePromise.waitFor(promises).next(() => result);
+  }
+
   saveOverlays(
     transaction: PersistenceTransaction,
     largestBatchId: number,

packages/firestore/src/local/local_documents_view.ts

@@ -134,12 +134,39 @@ export class LocalDocumentsView {
     docs: MutableDocumentMap,
     existenceStateChanged: DocumentKeySet
   ): PersistencePromise<DocumentMap> {
-    return this.computeViews(
-      transaction,
-      docs,
-      newOverlayMap(),
-      existenceStateChanged
-    );
+    const overlays = newOverlayMap();
+    return this.populateOverlays(transaction, overlays, docs).next(() => {
+      return this.computeViews(
+        transaction,
+        docs,
+        overlays,
+        existenceStateChanged
+      );
+    });
+  }
+
+  /**
+   * Fetches the overlays for {@code docs} and adds them to provided overlay map
+   * if the map does not already contain an entry for the given document key.
+   */
+  private populateOverlays(
+    transaction: PersistenceTransaction,
+    overlays: OverlayMap,
+    docs: MutableDocumentMap
+  ): PersistencePromise<void> {
+    let missingOverlays = documentKeySet();
+    docs.forEach(key => {
+      if (!overlays.has(key)) {
+        missingOverlays = missingOverlays.add(key);
+      }
+    });
+    return this.documentOverlayCache
+      .getOverlays(transaction, missingOverlays)
+      .next(result => {
+        result.forEach((key, val) => {
+          overlays.set(key, val);
+        });
+      });
   }
 
   /**
@@ -149,53 +176,39 @@ export class LocalDocumentsView {
   computeViews(
     transaction: PersistenceTransaction,
     docs: MutableDocumentMap,
-    memoizedOverlays: OverlayMap,
+    overlays: OverlayMap,
     existenceStateChanged: DocumentKeySet
   ): PersistencePromise<DocumentMap> {
     let results = documentMap();
     let recalculateDocuments = mutableDocumentMap();
-    const promises: Array<PersistencePromise<void>> = [];
     docs.forEach((_, doc) => {
-      const overlayPromise = memoizedOverlays.has(doc.key)
-        ? PersistencePromise.resolve(memoizedOverlays.get(doc.key)!)
-        : this.documentOverlayCache.getOverlay(transaction, doc.key);
-
-      promises.push(
-        overlayPromise.next((overlay: Overlay | null) => {
-          // Recalculate an overlay if the document's existence state is changed
-          // due to a remote event *and* the overlay is a PatchMutation. This is
-          // because document existence state can change if some patch mutation's
-          // preconditions are met.
-          // NOTE: we recalculate when `overlay` is null as well, because there
-          // might be a patch mutation whose precondition does not match before
-          // the change (hence overlay==null), but would now match.
-          if (
-            existenceStateChanged.has(doc.key) &&
-            (overlay === null || overlay.mutation instanceof PatchMutation)
-          ) {
-            recalculateDocuments = recalculateDocuments.insert(doc.key, doc);
-          } else if (overlay !== null) {
-            mutationApplyToLocalView(
-              overlay.mutation,
-              doc,
-              null,
-              Timestamp.now()
-            );
-          }
-        })
-      );
+      const overlay = overlays.get(doc.key);
+      // Recalculate an overlay if the document's existence state is changed due
+      // to a remote event *and* the overlay is a PatchMutation. This is because
+      // document existence state can change if some patch mutation's
+      // preconditions are met.
+      // NOTE: we recalculate when `overlay` is undefined as well, because there
+      // might be a patch mutation whose precondition does not match before the
+      // change (hence overlay is undefined), but would now match.
+      if (
+        existenceStateChanged.has(doc.key) &&
+        (overlay === undefined || overlay.mutation instanceof PatchMutation)
+      ) {
+        recalculateDocuments = recalculateDocuments.insert(doc.key, doc);
+      } else if (overlay !== undefined) {
+        mutationApplyToLocalView(overlay.mutation, doc, null, Timestamp.now());
+      }
     });
 
-    return PersistencePromise.waitFor(promises)
-      .next(() =>
-        this.recalculateAndSaveOverlays(transaction, recalculateDocuments)
-      )
-      .next(() => {
-        docs.forEach((key, value) => {
-          results = results.insert(key, value);
-        });
-        return results;
+    return this.recalculateAndSaveOverlays(
+      transaction,
+      recalculateDocuments
+    ).next(() => {
+      docs.forEach((key, value) => {
+        results = results.insert(key, value);
       });
+      return results;
+    });
   }
 
   private recalculateAndSaveOverlays(

packages/firestore/src/local/memory_document_overlay_cache.ts

@@ -50,6 +50,24 @@ export class MemoryDocumentOverlayCache implements DocumentOverlayCache {
     return PersistencePromise.resolve(this.overlays.get(key));
   }
 
+  getOverlays(
+    transaction: PersistenceTransaction,
+    keys: DocumentKeySet
+  ): PersistencePromise<OverlayMap> {
+    const result = newOverlayMap();
+    const promises: Array<PersistencePromise<void>> = [];
+    keys.forEach(key => {
+      promises.push(
+        this.getOverlay(transaction, key).next(overlay => {
+          if (overlay !== null) {
+            result.set(key, overlay);
+          }
+        })
+      );
+    });
+    return PersistencePromise.waitFor(promises).next(() => result);
+  }
+
   saveOverlays(
     transaction: PersistenceTransaction,
     largestBatchId: number,

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

@@ -327,4 +327,29 @@ function genericDocumentOverlayCacheTests(): void {
     );
     expect(await overlayCache.getOverlay(key('coll/doc'))).to.equal(null);
   });
+
+  it('skips non-existing overlay in batch lookup', async () => {
+    const result = await overlayCache.getOverlays(
+      documentKeySet(key('coll/doc1'))
+    );
+    expect(result.isEmpty()).to.equal(true);
+  });
+
+  it('supports empty batch in batch lookup', async () => {
+    const result = await overlayCache.getOverlays(documentKeySet());
+    expect(result.isEmpty()).to.equal(true);
+  });
+
+  it('can read saved overlays in batches', async () => {
+    const m1 = setMutation('coll/a', { 'a': 1 });
+    const m2 = setMutation('coll/b', { 'b': 2 });
+    const m3 = setMutation('coll/c', { 'c': 3 });
+    await saveOverlaysForMutations(3, m1, m2, m3);
+    const overlays = await overlayCache.getOverlays(
+      documentKeySet(key('coll/a'), key('coll/b'), key('coll/c'))
+    );
+    verifyEqualMutations(overlays.get(key('coll/a'))!.mutation, m1);
+    verifyEqualMutations(overlays.get(key('coll/b'))!.mutation, m2);
+    verifyEqualMutations(overlays.get(key('coll/c'))!.mutation, m3);
+  });
 }

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

@@ -50,6 +50,12 @@ export class TestDocumentOverlayCache {
     });
   }
 
+  getOverlays(keys: DocumentKeySet): Promise<OverlayMap> {
+    return this.persistence.runTransaction('getOverlays', 'readonly', txn => {
+      return this.cache.getOverlays(txn, keys);
+    });
+  }
+
   getOverlayMutation(docKey: string): Promise<Mutation | null> {
     return this.getOverlay(key(docKey)).then(value => {
       if (!value) {