Add Overlays and DocumentOverlayCache implementation (#5969)

* Overlay, DocumentOverlayCache, MemoryDocumentOverlayCache, IndexedDbDocumentOverlayCache.

* Fix for getOverlaysForCollectionGroup.

* Add Persistence APIs.

* Add document overlay cache tests.

* Lint fixes.

* Fix bug related to using SortedMap.

* Add schema update test.

* use custom mutation comparison function.

* Fix bugs in indexed db implementation.

* Address comments.

* Address comments.

* Create good-pugs-check.md

* Delete changeset.

* Address comments.

* Remove the index that is not necessary.

* Fix test code and add a new test.

* minor updates based on comments.

packages/firestore/src/local/document_overlay_cache.ts

@@ -0,0 +1,99 @@
+/**
+ * @license
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { DocumentKeySet } from '../model/collections';
+import { DocumentKey } from '../model/document_key';
+import { Mutation } from '../model/mutation';
+import { Overlay } from '../model/overlay';
+import { ResourcePath } from '../model/path';
+
+import { PersistencePromise } from './persistence_promise';
+import { PersistenceTransaction } from './persistence_transaction';
+
+/**
+ * Provides methods to read and write document overlays.
+ *
+ * An overlay is a saved mutation, that gives a local view of a document when
+ * applied to the remote version of the document.
+ *
+ * Each overlay stores the largest batch ID that is included in the overlay,
+ * which allows us to remove the overlay once all batches leading up to it have
+ * been acknowledged.
+ */
+export interface DocumentOverlayCache {
+  /**
+   * Gets the saved overlay mutation for the given document key.
+   * Returns null if there is no overlay for that key.
+   */
+  getOverlay(
+    transaction: PersistenceTransaction,
+    key: DocumentKey
+  ): PersistencePromise<Overlay | null>;
+
+  /**
+   * Saves the given document mutation map to persistence as overlays.
+   * All overlays will have their largest batch id set to `largestBatchId`.
+   */
+  saveOverlays(
+    transaction: PersistenceTransaction,
+    largestBatchId: number,
+    overlays: Map<DocumentKey, Mutation>
+  ): PersistencePromise<void>;
+
+  /** Removes overlays for the given document keys and batch ID. */
+  removeOverlaysForBatchId(
+    transaction: PersistenceTransaction,
+    documentKeys: DocumentKeySet,
+    batchId: number
+  ): PersistencePromise<void>;
+
+  /**
+   * Returns all saved overlays for the given collection.
+   *
+   * @param transaction - The persistence transaction to use for this operation.
+   * @param collection - The collection path to get the overlays for.
+   * @param sinceBatchId - The minimum batch ID to filter by (exclusive).
+   * Only overlays that contain a change past `sinceBatchId` are returned.
+   * @returns Mapping of each document key in the collection to its overlay.
+   */
+  getOverlaysForCollection(
+    transaction: PersistenceTransaction,
+    collection: ResourcePath,
+    sinceBatchId: number
+  ): PersistencePromise<Map<DocumentKey, Overlay>>;
+
+  /**
+   * Returns `count` overlays with a batch ID higher than `sinceBatchId` for the
+   * provided collection group, processed by ascending batch ID. The method
+   * always returns all overlays for a batch even if the last batch contains
+   * more documents than the remaining limit.
+   *
+   * @param transaction - The persistence transaction used for this operation.
+   * @param collectionGroup - The collection group to get the overlays for.
+   * @param sinceBatchId - The minimum batch ID to filter by (exclusive).
+   * Only overlays that contain a change past `sinceBatchId` are returned.
+   * @param count - The number of overlays to return. Can be exceeded if the last
+   * batch contains more entries.
+   * @return Mapping of each document key in the collection group to its overlay.
+   */
+  getOverlaysForCollectionGroup(
+    transaction: PersistenceTransaction,
+    collectionGroup: string,
+    sinceBatchId: number,
+    count: number
+  ): PersistencePromise<Map<DocumentKey, Overlay>>;
+}

packages/firestore/src/local/indexeddb_document_overlay_cache.ts

@@ -0,0 +1,203 @@
+/**
+ * @license
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { User } from '../auth/user';
+import { DocumentKeySet } from '../model/collections';
+import { DocumentKey } from '../model/document_key';
+import { Mutation } from '../model/mutation';
+import { Overlay } from '../model/overlay';
+import { ResourcePath } from '../model/path';
+
+import { DocumentOverlayCache } from './document_overlay_cache';
+import { encodeResourcePath } from './encoded_resource_path';
+import { DbDocumentOverlay, DbDocumentOverlayKey } from './indexeddb_schema';
+import { getStore } from './indexeddb_transaction';
+import {
+  fromDbDocumentOverlay,
+  LocalSerializer,
+  toDbDocumentOverlay,
+  toDbDocumentOverlayKey
+} from './local_serializer';
+import { PersistencePromise } from './persistence_promise';
+import { PersistenceTransaction } from './persistence_transaction';
+import { SimpleDbStore } from './simple_db';
+
+/**
+ * Implementation of DocumentOverlayCache using IndexedDb.
+ */
+export class IndexedDbDocumentOverlayCache implements DocumentOverlayCache {
+  /**
+   * @param serializer - The document serializer.
+   * @param userId - The userId for which we are accessing overlays.
+   */
+  constructor(
+    private readonly serializer: LocalSerializer,
+    private readonly userId: string
+  ) {}
+
+  static forUser(
+    serializer: LocalSerializer,
+    user: User
+  ): IndexedDbDocumentOverlayCache {
+    const userId = user.uid || '';
+    return new IndexedDbDocumentOverlayCache(serializer, userId);
+  }
+
+  getOverlay(
+    transaction: PersistenceTransaction,
+    key: DocumentKey
+  ): PersistencePromise<Overlay | null> {
+    return documentOverlayStore(transaction)
+      .get(toDbDocumentOverlayKey(this.userId, key))
+      .next(dbOverlay => {
+        if (dbOverlay) {
+          return fromDbDocumentOverlay(this.serializer, dbOverlay);
+        }
+        return null;
+      });
+  }
+
+  saveOverlays(
+    transaction: PersistenceTransaction,
+    largestBatchId: number,
+    overlays: Map<DocumentKey, Mutation>
+  ): PersistencePromise<void> {
+    const promises: Array<PersistencePromise<void>> = [];
+    overlays.forEach(mutation => {
+      const overlay = new Overlay(largestBatchId, mutation);
+      promises.push(this.saveOverlay(transaction, overlay));
+    });
+    return PersistencePromise.waitFor(promises);
+  }
+
+  removeOverlaysForBatchId(
+    transaction: PersistenceTransaction,
+    documentKeys: DocumentKeySet,
+    batchId: number
+  ): PersistencePromise<void> {
+    const collectionPaths = new Set<string>();
+
+    // Get the set of unique collection paths.
+    documentKeys.forEach(key =>
+      collectionPaths.add(encodeResourcePath(key.getCollectionPath()))
+    );
+
+    const promises: Array<PersistencePromise<void>> = [];
+    collectionPaths.forEach(collectionPath => {
+      const range = IDBKeyRange.bound(
+        [this.userId, collectionPath, batchId],
+        [this.userId, collectionPath, batchId + 1],
+        /*lowerOpen=*/ false,
+        /*upperOpen=*/ true
+      );
+      promises.push(
+        documentOverlayStore(transaction).deleteAll(
+          DbDocumentOverlay.collectionPathOverlayIndex,
+          range
+        )
+      );
+    });
+    return PersistencePromise.waitFor(promises);
+  }
+
+  getOverlaysForCollection(
+    transaction: PersistenceTransaction,
+    collection: ResourcePath,
+    sinceBatchId: number
+  ): PersistencePromise<Map<DocumentKey, Overlay>> {
+    const result = new Map<DocumentKey, Overlay>();
+    const collectionPath = encodeResourcePath(collection);
+    // We want batch IDs larger than `sinceBatchId`, and so the lower bound
+    // is not inclusive.
+    const range = IDBKeyRange.bound(
+      [this.userId, collectionPath, sinceBatchId],
+      [this.userId, collectionPath, Number.POSITIVE_INFINITY],
+      /*lowerOpen=*/ true
+    );
+    return documentOverlayStore(transaction)
+      .loadAll(DbDocumentOverlay.collectionPathOverlayIndex, range)
+      .next(dbOverlays => {
+        for (const dbOverlay of dbOverlays) {
+          const overlay = fromDbDocumentOverlay(this.serializer, dbOverlay);
+          result.set(overlay.getKey(), overlay);
+        }
+        return result;
+      });
+  }
+
+  getOverlaysForCollectionGroup(
+    transaction: PersistenceTransaction,
+    collectionGroup: string,
+    sinceBatchId: number,
+    count: number
+  ): PersistencePromise<Map<DocumentKey, Overlay>> {
+    const result = new Map<DocumentKey, Overlay>();
+    let currentBatchId: number | undefined = undefined;
+    // We want batch IDs larger than `sinceBatchId`, and so the lower bound
+    // is not inclusive.
+    const range = IDBKeyRange.bound(
+      [this.userId, collectionGroup, sinceBatchId],
+      [this.userId, collectionGroup, Number.POSITIVE_INFINITY],
+      /*lowerOpen=*/ true
+    );
+    return documentOverlayStore(transaction)
+      .iterate(
+        {
+          index: DbDocumentOverlay.collectionGroupOverlayIndex,
+          range
+        },
+        (_, dbOverlay, control) => {
+          // We do not want to return partial batch overlays, even if the size
+          // of the result set exceeds the given `count` argument. Therefore, we
+          // continue to aggregate results even after the result size exceeds
+          // `count` if there are more overlays from the `currentBatchId`.
+          const overlay = fromDbDocumentOverlay(this.serializer, dbOverlay);
+          if (
+            result.size < count ||
+            overlay.largestBatchId === currentBatchId
+          ) {
+            result.set(overlay.getKey(), overlay);
+            currentBatchId = overlay.largestBatchId;
+          } else {
+            control.done();
+          }
+        }
+      )
+      .next(() => result);
+  }
+
+  private saveOverlay(
+    transaction: PersistenceTransaction,
+    overlay: Overlay
+  ): PersistencePromise<void> {
+    return documentOverlayStore(transaction).put(
+      toDbDocumentOverlay(this.serializer, this.userId, overlay)
+    );
+  }
+}
+
+/**
+ * Helper to get a typed SimpleDbStore for the document overlay object store.
+ */
+function documentOverlayStore(
+  txn: PersistenceTransaction
+): SimpleDbStore<DbDocumentOverlayKey, DbDocumentOverlay> {
+  return getStore<DbDocumentOverlayKey, DbDocumentOverlay>(
+    txn,
+    DbDocumentOverlay.store
+  );
+}

packages/firestore/src/local/indexeddb_persistence.ts

@@ -28,8 +28,10 @@ import { logDebug, logError } from '../util/log';
 import { DocumentLike, WindowLike } from '../util/types';
 
 import { BundleCache } from './bundle_cache';
+import { DocumentOverlayCache } from './document_overlay_cache';
 import { IndexManager } from './index_manager';
 import { IndexedDbBundleCache } from './indexeddb_bundle_cache';
+import { IndexedDbDocumentOverlayCache } from './indexeddb_document_overlay_cache';
 import { IndexedDbIndexManager } from './indexeddb_index_manager';
 import { IndexedDbLruDelegateImpl } from './indexeddb_lru_delegate_impl';
 import { IndexedDbMutationQueue } from './indexeddb_mutation_queue';
@@ -744,6 +746,14 @@ export class IndexedDbPersistence implements Persistence {
     return new IndexedDbIndexManager(user);
   }
 
+  getDocumentOverlayCache(user: User): DocumentOverlayCache {
+    debugAssert(
+      this.started,
+      'Cannot initialize IndexedDbDocumentOverlayCache before persistence is started.'
+    );
+    return IndexedDbDocumentOverlayCache.forUser(this.serializer, user);
+  }
+
   getBundleCache(): BundleCache {
     debugAssert(
       this.started,