Port optimizations to LocalDocumentsView from iOS (#1055)

* add a method to find batches affecting a set of keys (port of [1479](firebase/firebase-ios-sdk#1479));
* use the newly-added method to avoid rereading batches when getting documents in `LocalDocumentsView` (port of [1505](firebase/firebase-ios-sdk#1505));
* avoid rereading batches when searching for documents in a collection (port of [1533](firebase/firebase-ios-sdk#1533)).

Speedup was measured by running tests in browser and checking time spent writing 10 batches of 500 mutations each, and then querying the resulting 5K docs collection from cache in offline mode. For this case, the writing speedup is about 3x, and querying speedup is about 6x (see PR for more details).

Author: var-const

packages/firestore/src/local/indexeddb_mutation_queue.ts

@@ -18,6 +18,7 @@ import { Timestamp } from '../api/timestamp';
 import { User } from '../auth/user';
 import { Query } from '../core/query';
 import { BatchId, ProtoByteString } from '../core/types';
+import { DocumentKeySet } from '../model/collections';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
 import { BATCHID_UNKNOWN, MutationBatch } from '../model/mutation_batch';
@@ -363,6 +364,50 @@ export class IndexedDbMutationQueue implements MutationQueue {
       .next(() => results);
   }
 
+  getAllMutationBatchesAffectingDocumentKeys(
+    transaction: PersistenceTransaction,
+    documentKeys: DocumentKeySet
+  ): PersistencePromise<MutationBatch[]> {
+    let uniqueBatchIDs = new SortedSet<BatchId>(primitiveComparator);
+
+    const promises: Array<PersistencePromise<void>> = [];
+    documentKeys.forEach(documentKey => {
+      const indexStart = DbDocumentMutation.prefixForPath(
+        this.userId,
+        documentKey.path
+      );
+      const range = IDBKeyRange.lowerBound(indexStart);
+
+      const promise = documentMutationsStore(transaction).iterate(
+        { range },
+        (indexKey, _, control) => {
+          const [userID, encodedPath, batchID] = indexKey;
+
+          // Only consider rows matching exactly the specific key of
+          // interest. Note that because we order by path first, and we
+          // order terminators before path separators, we'll encounter all
+          // the index rows for documentKey contiguously. In particular, all
+          // the rows for documentKey will occur before any rows for
+          // documents nested in a subcollection beneath documentKey so we
+          // can stop as soon as we hit any such row.
+          const path = EncodedResourcePath.decode(encodedPath);
+          if (userID !== this.userId || !documentKey.path.isEqual(path)) {
+            control.done();
+            return;
+          }
+
+          uniqueBatchIDs = uniqueBatchIDs.add(batchID);
+        }
+      );
+
+      promises.push(promise);
+    });
+
+    return PersistencePromise.waitFor(promises).next(() =>
+      this.lookupMutationBatches(transaction, uniqueBatchIDs)
+    );
+  }
+
   getAllMutationBatchesAffectingQuery(
     transaction: PersistenceTransaction,
     query: Query
@@ -414,29 +459,34 @@ export class IndexedDbMutationQueue implements MutationQueue {
         }
         uniqueBatchIDs = uniqueBatchIDs.add(batchID);
       })
-      .next(() => {
-        const results: MutationBatch[] = [];
-        const promises: Array<PersistencePromise<void>> = [];
-        // TODO(rockwood): Implement this using iterate.
-        uniqueBatchIDs.forEach(batchID => {
-          const mutationKey = this.keyForBatchId(batchID);
-          promises.push(
-            mutationsStore(transaction)
-              .get(mutationKey)
-              .next(mutation => {
-                if (mutation === null) {
-                  fail(
-                    'Dangling document-mutation reference found, ' +
-                      'which points to ' +
-                      mutationKey
-                  );
-                }
-                results.push(this.serializer.fromDbMutationBatch(mutation!));
-              })
-          );
-        });
-        return PersistencePromise.waitFor(promises).next(() => results);
-      });
+      .next(() => this.lookupMutationBatches(transaction, uniqueBatchIDs));
+  }
+
+  private lookupMutationBatches(
+    transaction: PersistenceTransaction,
+    batchIDs: SortedSet<BatchId>
+  ): PersistencePromise<MutationBatch[]> {
+    const results: MutationBatch[] = [];
+    const promises: Array<PersistencePromise<void>> = [];
+    // TODO(rockwood): Implement this using iterate.
+    batchIDs.forEach(batchID => {
+      const mutationKey = this.keyForBatchId(batchID);
+      promises.push(
+        mutationsStore(transaction)
+          .get(mutationKey)
+          .next(mutation => {
+            if (mutation === null) {
+              fail(
+                'Dangling document-mutation reference found, ' +
+                  'which points to ' +
+                  mutationKey
+              );
+            }
+            results.push(this.serializer.fromDbMutationBatch(mutation!));
+          })
+      );
+    });
+    return PersistencePromise.waitFor(promises).next(() => results);
   }
 
   removeMutationBatches(

packages/firestore/src/local/local_documents_view.ts

@@ -17,7 +17,6 @@
 import { Query } from '../core/query';
 import { SnapshotVersion } from '../core/snapshot_version';
 import {
-  documentKeySet,
   DocumentKeySet,
   DocumentMap,
   documentMap,
@@ -26,6 +25,7 @@ import {
 } from '../model/collections';
 import { Document, MaybeDocument, NoDocument } from '../model/document';
 import { DocumentKey } from '../model/document_key';
+import { MutationBatch } from '../model/mutation_batch';
 import { ResourcePath } from '../model/path';
 import { fail } from '../util/assert';
 
@@ -56,11 +56,23 @@ export class LocalDocumentsView {
     transaction: PersistenceTransaction,
     key: DocumentKey
   ): PersistencePromise<MaybeDocument | null> {
-    return this.remoteDocumentCache
-      .getEntry(transaction, key)
-      .next(remoteDoc => {
-        return this.computeLocalDocument(transaction, key, remoteDoc);
-      });
+    return this.mutationQueue
+      .getAllMutationBatchesAffectingDocumentKey(transaction, key)
+      .next(batches => this.getDocumentInternal(transaction, key, batches));
+  }
+
+  /** Internal version of `getDocument` that allows reusing batches. */
+  private getDocumentInternal(
+    transaction: PersistenceTransaction,
+    key: DocumentKey,
+    inBatches: MutationBatch[]
+  ): PersistencePromise<MaybeDocument | null> {
+    return this.remoteDocumentCache.getEntry(transaction, key).next(doc => {
+      for (const batch of inBatches) {
+        doc = batch.applyToLocalView(key, doc);
+      }
+      return doc;
+    });
   }
 
   /**
@@ -73,20 +85,29 @@ export class LocalDocumentsView {
     transaction: PersistenceTransaction,
     keys: DocumentKeySet
   ): PersistencePromise<MaybeDocumentMap> {
-    const promises = [] as Array<PersistencePromise<void>>;
-    let results = maybeDocumentMap();
-    keys.forEach(key => {
-      promises.push(
-        this.getDocument(transaction, key).next(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 this.mutationQueue
+      .getAllMutationBatchesAffectingDocumentKeys(transaction, keys)
+      .next(batches => {
+        const promises = [] as Array<PersistencePromise<void>>;
+        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);
+              }
+            )
+          );
+        });
+        return PersistencePromise.waitFor(promises).next(() => results);
+      });
   }
 
   /** Performs a query against the local view of all documents. */
@@ -122,48 +143,40 @@ export class LocalDocumentsView {
     query: Query
   ): PersistencePromise<DocumentMap> {
     // Query the remote documents and overlay mutations.
-    // TODO(mikelehen): There may be significant overlap between the mutations
-    // affecting these remote documents and the
-    // getAllMutationBatchesAffectingQuery() mutations. Consider optimizing.
     let results: DocumentMap;
     return this.remoteDocumentCache
       .getDocumentsMatchingQuery(transaction, query)
       .next(queryResults => {
-        return this.computeLocalDocuments(transaction, queryResults);
-      })
-      .next(promisedResults => {
-        results = promisedResults;
-        // Now use the mutation queue to discover any other documents that may
-        // match the query after applying mutations.
+        results = queryResults;
         return this.mutationQueue.getAllMutationBatchesAffectingQuery(
           transaction,
           query
         );
       })
       .next(matchingMutationBatches => {
-        let matchingKeys = documentKeySet();
         for (const batch of matchingMutationBatches) {
           for (const mutation of batch.mutations) {
-            // TODO(mikelehen): PERF: Check if this mutation actually
-            // affects the query to reduce work.
-            if (!results.get(mutation.key)) {
-              matchingKeys = matchingKeys.add(mutation.key);
+            const key = mutation.key;
+            // Only process documents belonging to the collection.
+            if (!query.path.isImmediateParentOf(key.path)) {
+              continue;
+            }
+
+            const baseDoc = results.get(key);
+            const mutatedDoc = mutation.applyToLocalView(
+              baseDoc,
+              baseDoc,
+              batch.localWriteTime
+            );
+            if (!mutatedDoc || mutatedDoc instanceof NoDocument) {
+              results = results.remove(key);
+            } else if (mutatedDoc instanceof Document) {
+              results = results.insert(key, mutatedDoc);
+            } else {
+              fail('Unknown MaybeDocument: ' + mutatedDoc);
             }
           }
         }
-
-        // Now add in the results for the matchingKeys.
-        const promises = [] as Array<PersistencePromise<void>>;
-        matchingKeys.forEach(key => {
-          promises.push(
-            this.getDocument(transaction, key).next(doc => {
-              if (doc instanceof Document) {
-                results = results.insert(doc.key, doc);
-              }
-            })
-          );
-        });
-        return PersistencePromise.waitFor(promises);
       })
       .next(() => {
         // Finally, filter out any documents that don't actually match
@@ -177,57 +190,4 @@ export class LocalDocumentsView {
         return results;
       });
   }
-
-  /**
-   * Takes a remote document and applies local mutations to generate the local
-   * view of the document.
-   * @param transaction The transaction in which to perform any persistence
-   *     operations.
-   * @param documentKey The key of the document (necessary when remoteDocument
-   *     is null).
-   * @param document The base remote document to apply mutations to or null.
-   */
-  private computeLocalDocument(
-    transaction: PersistenceTransaction,
-    documentKey: DocumentKey,
-    document: MaybeDocument | null
-  ): PersistencePromise<MaybeDocument | null> {
-    return this.mutationQueue
-      .getAllMutationBatchesAffectingDocumentKey(transaction, documentKey)
-      .next(batches => {
-        for (const batch of batches) {
-          document = batch.applyToLocalView(documentKey, document);
-        }
-        return document;
-      });
-  }
-
-  /**
-   * Takes a set of remote documents and applies local mutations to generate the
-   * local view of the documents.
-   * @param transaction The transaction in which to perform any persistence
-   *     operations.
-   * @param documents The base remote documents to apply mutations to.
-   * @return The local view of the documents.
-   */
-  private computeLocalDocuments(
-    transaction: PersistenceTransaction,
-    documents: DocumentMap
-  ): PersistencePromise<DocumentMap> {
-    const promises = [] as Array<PersistencePromise<void>>;
-    documents.forEach((key, doc) => {
-      promises.push(
-        this.computeLocalDocument(transaction, key, doc).next(mutatedDoc => {
-          if (mutatedDoc instanceof Document) {
-            documents = documents.insert(mutatedDoc.key, mutatedDoc);
-          } else if (mutatedDoc instanceof NoDocument) {
-            documents = documents.remove(mutatedDoc.key);
-          } else {
-            fail('Unknown MaybeDocument: ' + mutatedDoc);
-          }
-        })
-      );
-    });
-    return PersistencePromise.waitFor(promises).next(() => documents);
-  }
 }

packages/firestore/src/local/memory_mutation_queue.ts

@@ -17,6 +17,7 @@
 import { Timestamp } from '../api/timestamp';
 import { Query } from '../core/query';
 import { BatchId, ProtoByteString } from '../core/types';
+import { DocumentKeySet } from '../model/collections';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
 import { BATCHID_UNKNOWN, MutationBatch } from '../model/mutation_batch';
@@ -252,6 +253,28 @@ export class MemoryMutationQueue implements MutationQueue {
     return PersistencePromise.resolve(result);
   }
 
+  getAllMutationBatchesAffectingDocumentKeys(
+    transaction: PersistenceTransaction,
+    documentKeys: DocumentKeySet
+  ): PersistencePromise<MutationBatch[]> {
+    let uniqueBatchIDs = new SortedSet<number>(primitiveComparator);
+
+    documentKeys.forEach(documentKey => {
+      const start = new DocReference(documentKey, 0);
+      const end = new DocReference(documentKey, Number.POSITIVE_INFINITY);
+      this.batchesByDocumentKey.forEachInRange([start, end], ref => {
+        assert(
+          documentKey.isEqual(ref.key),
+          "For each key, should only iterate over a single key's batches"
+        );
+
+        uniqueBatchIDs = uniqueBatchIDs.add(ref.targetOrBatchId);
+      });
+    });
+
+    return PersistencePromise.resolve(this.findMutationBatches(uniqueBatchIDs));
+  }
+
   getAllMutationBatchesAffectingQuery(
     transaction: PersistenceTransaction,
     query: Query
@@ -293,16 +316,20 @@ export class MemoryMutationQueue implements MutationQueue {
       }
     }, start);
 
+    return PersistencePromise.resolve(this.findMutationBatches(uniqueBatchIDs));
+  }
+
+  private findMutationBatches(batchIDs: SortedSet<number>): MutationBatch[] {
     // Construct an array of matching batches, sorted by batchID to ensure that
     // multiple mutations affecting the same document key are applied in order.
     const result: MutationBatch[] = [];
-    uniqueBatchIDs.forEach(batchId => {
+    batchIDs.forEach(batchId => {
       const batch = this.findMutationBatch(batchId);
       if (batch !== null) {
         result.push(batch);
       }
     });
-    return PersistencePromise.resolve(result);
+    return result;
   }
 
   removeMutationBatches(