Make notifyLocalViewChanges idempotent

Author: schmidt-sebastian

packages/firestore/src/core/sync_engine.ts

@@ -682,8 +682,7 @@ export class SyncEngine implements RemoteSyncer {
     this.queriesByTarget.delete(targetId);
 
     if (this.isPrimaryClient) {
-      const limboKeys = this.limboDocumentRefs.referencesForId(targetId);
-      this.limboDocumentRefs.removeReferencesForId(targetId);
+      const limboKeys = this.limboDocumentRefs.removeReferencesForId(targetId);
       limboKeys.forEach(limboKey => {
         const isReferenced = this.limboDocumentRefs.containsKey(limboKey);
         if (!isReferenced) {

packages/firestore/src/local/indexeddb_mutation_queue.ts

@@ -519,7 +519,7 @@ export class IndexedDbMutationQueue implements MutationQueue {
       return PersistencePromise.forEach(
         removedDocuments,
         (key: DocumentKey) => {
-          return this.referenceDelegate.removeMutationReference(
+          return this.referenceDelegate.markPotentiallyOrphaned(
             transaction,
             key
           );

packages/firestore/src/local/indexeddb_persistence.ts

@@ -1070,8 +1070,6 @@ function clientMetadataStore(
 
 /** Provides LRU functionality for IndexedDB persistence. */
 export class IndexedDbLruDelegate implements ReferenceDelegate, LruDelegate {
-  private inMemoryPins: ReferenceSet | null = null;
-
   readonly garbageCollector: LruGarbageCollector;
 
   constructor(private readonly db: IndexedDbPersistence, params: LruParams) {
@@ -1081,14 +1079,14 @@ export class IndexedDbLruDelegate implements ReferenceDelegate, LruDelegate {
   getSequenceNumberCount(
     txn: PersistenceTransaction
   ): PersistencePromise<number> {
-    const docCountPromise = this.orphanedDocmentCount(txn);
+    const docCountPromise = this.orphanedDocumentCount(txn);
     const targetCountPromise = this.db.getTargetCache().getTargetCount(txn);
     return targetCountPromise.next(targetCount =>
       docCountPromise.next(docCount => targetCount + docCount)
     );
   }
 
-  private orphanedDocmentCount(
+  private orphanedDocumentCount(
     txn: PersistenceTransaction
   ): PersistencePromise<number> {
     let orphanedCount = 0;
@@ -1113,19 +1111,17 @@ export class IndexedDbLruDelegate implements ReferenceDelegate, LruDelegate {
     );
   }
 
-  setInMemoryPins(inMemoryPins: ReferenceSet): void {
-    this.inMemoryPins = inMemoryPins;
-  }
-
   addReference(
     txn: PersistenceTransaction,
+    targetId: TargetId,
     key: DocumentKey
   ): PersistencePromise<void> {
     return writeSentinelKey(txn, key);
   }
 
   removeReference(
     txn: PersistenceTransaction,
+    targetId: TargetId,
     key: DocumentKey
   ): PersistencePromise<void> {
     return writeSentinelKey(txn, key);
@@ -1141,7 +1137,7 @@ export class IndexedDbLruDelegate implements ReferenceDelegate, LruDelegate {
       .removeTargets(txn, upperBound, activeTargetIds);
   }
 
-  removeMutationReference(
+  markPotentiallyOrphaned(
     txn: PersistenceTransaction,
     key: DocumentKey
   ): PersistencePromise<void> {
@@ -1158,11 +1154,7 @@ export class IndexedDbLruDelegate implements ReferenceDelegate, LruDelegate {
     txn: PersistenceTransaction,
     docKey: DocumentKey
   ): PersistencePromise<boolean> {
-    if (this.inMemoryPins!.containsKey(docKey)) {
-      return PersistencePromise.resolve<boolean>(true);
-    } else {
-      return mutationQueuesContainKey(txn, docKey);
-    }
+    return mutationQueuesContainKey(txn, docKey);
   }
 
   removeOrphanedDocuments(

packages/firestore/src/local/indexeddb_target_cache.ts

@@ -284,7 +284,7 @@ export class IndexedDbTargetCache implements TargetCache {
     keys.forEach(key => {
       const path = encodeResourcePath(key.path);
       promises.push(store.put(new DbTargetDocument(targetId, path)));
-      promises.push(this.referenceDelegate.addReference(txn, key));
+      promises.push(this.referenceDelegate.addReference(txn, targetId, key));
     });
     return PersistencePromise.waitFor(promises);
   }
@@ -301,7 +301,7 @@ export class IndexedDbTargetCache implements TargetCache {
       const path = encodeResourcePath(key.path);
       return PersistencePromise.waitFor([
         store.delete([targetId, path]),
-        this.referenceDelegate.removeReference(txn, key)
+        this.referenceDelegate.removeReference(txn, targetId, key)
       ]);
     });
   }

packages/firestore/src/local/local_store.ts

@@ -165,11 +165,6 @@ export class LocalStore {
    */
   protected localDocuments: LocalDocumentsView;
 
-  /**
-   * The set of document references maintained by any local views.
-   */
-  private localViewReferences = new ReferenceSet();
-
   /** Maps a target to its `TargetData`. */
   protected targetCache: TargetCache;
 
@@ -206,9 +201,6 @@ export class LocalStore {
       persistence.started,
       'LocalStore was passed an unstarted persistence implementation'
     );
-    this.persistence.referenceDelegate.setInMemoryPins(
-      this.localViewReferences
-    );
     this.mutationQueue = persistence.getMutationQueue(initialUser);
     this.remoteDocuments = persistence.getRemoteDocumentCache();
     this.targetCache = persistence.getTargetCache();
@@ -704,49 +696,56 @@ export class LocalStore {
    * Notify local store of the changed views to locally pin documents.
    */
   notifyLocalViewChanges(viewChanges: LocalViewChanges[]): Promise<void> {
-    for (const viewChange of viewChanges) {
-      const targetId = viewChange.targetId;
-
-      this.localViewReferences.addReferences(viewChange.addedKeys, targetId);
-      this.localViewReferences.removeReferences(
-        viewChange.removedKeys,
-        targetId
-      );
-
-      if (!viewChange.fromCache) {
-        const targetData = this.targetDataByTarget.get(targetId);
-        debugAssert(
-          targetData !== null,
-          `Can't set limbo-free snapshot version for unknown target: ${targetId}`
-        );
-
-        // Advance the last limbo free snapshot version
-        const lastLimboFreeSnapshotVersion = targetData.snapshotVersion;
-        const updatedTargetData = targetData.withLastLimboFreeSnapshotVersion(
-          lastLimboFreeSnapshotVersion
-        );
-        this.targetDataByTarget = this.targetDataByTarget.insert(
-          targetId,
-          updatedTargetData
-        );
-      }
-    }
-    return this.persistence.runTransaction(
-      'notifyLocalViewChanges',
-      'readwrite',
-      txn => {
+    return this.persistence
+      .runTransaction('notifyLocalViewChanges', 'readwrite', txn => {
         return PersistencePromise.forEach(
           viewChanges,
           (viewChange: LocalViewChanges) => {
             return PersistencePromise.forEach(
-              viewChange.removedKeys,
+              viewChange.addedKeys,
               (key: DocumentKey) =>
-                this.persistence.referenceDelegate.removeReference(txn, key)
+                this.persistence.referenceDelegate.addReference(
+                  txn,
+                  viewChange.targetId,
+                  key
+                )
+            ).next(() =>
+              PersistencePromise.forEach(
+                viewChange.removedKeys,
+                (key: DocumentKey) =>
+                  this.persistence.referenceDelegate.removeReference(
+                    txn,
+                    viewChange.targetId,
+                    key
+                  )
+              )
             );
           }
         );
-      }
-    );
+      })
+      .then(() => {
+        for (const viewChange of viewChanges) {
+          const targetId = viewChange.targetId;
+
+          if (!viewChange.fromCache) {
+            const targetData = this.targetDataByTarget.get(targetId);
+            debugAssert(
+              targetData !== null,
+              `Can't set limbo-free snapshot version for unknown target: ${targetId}`
+            );
+
+            // Advance the last limbo free snapshot version
+            const lastLimboFreeSnapshotVersion = targetData.snapshotVersion;
+            const updatedTargetData = targetData.withLastLimboFreeSnapshotVersion(
+              lastLimboFreeSnapshotVersion
+            );
+            this.targetDataByTarget = this.targetDataByTarget.insert(
+              targetId,
+              updatedTargetData
+            );
+          }
+        }
+      });
   }
 
   /**
@@ -869,26 +868,11 @@ export class LocalStore {
     const mode = keepPersistedTargetData ? 'readwrite' : 'readwrite-primary';
     return this.persistence
       .runTransaction('Release target', mode, txn => {
-        // References for documents sent via Watch are automatically removed
-        // when we delete a target's data from the reference delegate.
-        // Since this does not remove references for locally mutated documents,
-        // we have to remove the target associations for these documents
-        // manually.
-        // This operation needs to be run inside the transaction since EagerGC
-        // uses the local view references during the transaction's commit.
-        // Fortunately, the operation is safe to be re-run in case the
-        // transaction fails since there are no side effects if the target has
-        // already been removed.
-        const removed = this.localViewReferences.removeReferencesForId(
-          targetId
-        );
-
         if (!keepPersistedTargetData) {
-          return PersistencePromise.forEach(removed, (key: DocumentKey) =>
-            this.persistence.referenceDelegate.removeReference(txn, key)
-          ).next(() => {
-            this.persistence.referenceDelegate.removeTarget(txn, targetData!);
-          });
+          return this.persistence.referenceDelegate.removeTarget(
+            txn,
+            targetData!
+          );
         } else {
           return PersistencePromise.resolve();
         }

packages/firestore/src/local/memory_mutation_queue.ts

@@ -299,7 +299,7 @@ export class MemoryMutationQueue implements MutationQueue {
     return PersistencePromise.forEach(batch.mutations, (mutation: Mutation) => {
       const ref = new DocReference(mutation.key, batch.batchId);
       references = references.delete(ref);
-      return this.referenceDelegate.removeMutationReference(
+      return this.referenceDelegate.markPotentiallyOrphaned(
         transaction,
         mutation.key
       );

packages/firestore/src/local/memory_persistence.ts

@@ -29,7 +29,7 @@ import {
   LruParams
 } from './lru_garbage_collector';
 import { ListenSequence } from '../core/listen_sequence';
-import { ListenSequenceNumber } from '../core/types';
+import { ListenSequenceNumber, TargetId } from '../core/types';
 import { estimateByteSize } from '../model/values';
 import { MemoryIndexManager } from './memory_index_manager';
 import { MemoryMutationQueue } from './memory_mutation_queue';
@@ -184,7 +184,9 @@ export interface MemoryReferenceDelegate extends ReferenceDelegate {
 }
 
 export class MemoryEagerDelegate implements MemoryReferenceDelegate {
-  private inMemoryPins: ReferenceSet | null = null;
+  /** Manages all documents that are active in Query views. */
+  private localViewReferences: ReferenceSet = new ReferenceSet();
+  /** The list of documents that are potentially GCed after each transaction. */
   private _orphanedDocuments: Set<DocumentKey> | null = null;
 
   private constructor(private readonly persistence: MemoryPersistence) {}
@@ -201,27 +203,27 @@ export class MemoryEagerDelegate implements MemoryReferenceDelegate {
     }
   }
 
-  setInMemoryPins(inMemoryPins: ReferenceSet): void {
-    this.inMemoryPins = inMemoryPins;
-  }
-
   addReference(
     txn: PersistenceTransaction,
+    targetId: TargetId,
     key: DocumentKey
   ): PersistencePromise<void> {
+    this.localViewReferences.addReference(key, targetId);
     this.orphanedDocuments.delete(key);
     return PersistencePromise.resolve();
   }
 
   removeReference(
     txn: PersistenceTransaction,
+    targetId: TargetId,
     key: DocumentKey
   ): PersistencePromise<void> {
+    this.localViewReferences.removeReference(key, targetId);
     this.orphanedDocuments.add(key);
     return PersistencePromise.resolve();
   }
 
-  removeMutationReference(
+  markPotentiallyOrphaned(
     txn: PersistenceTransaction,
     key: DocumentKey
   ): PersistencePromise<void> {
@@ -233,6 +235,10 @@ export class MemoryEagerDelegate implements MemoryReferenceDelegate {
     txn: PersistenceTransaction,
     targetData: TargetData
   ): PersistencePromise<void> {
+    const orphaned = this.localViewReferences.removeReferencesForId(
+      targetData.targetId
+    );
+    orphaned.forEach(key => this.orphanedDocuments.add(key));
     const cache = this.persistence.getTargetCache();
     return cache
       .getMatchingKeysForTargetId(txn, targetData.targetId)
@@ -290,15 +296,15 @@ export class MemoryEagerDelegate implements MemoryReferenceDelegate {
     key: DocumentKey
   ): PersistencePromise<boolean> {
     return PersistencePromise.or([
+      () =>
+        PersistencePromise.resolve(this.localViewReferences.containsKey(key)),
       () => this.persistence.getTargetCache().containsKey(txn, key),
-      () => this.persistence.mutationQueuesContainKey(txn, key),
-      () => PersistencePromise.resolve(this.inMemoryPins!.containsKey(key))
+      () => this.persistence.mutationQueuesContainKey(txn, key)
     ]);
   }
 }
 
 export class MemoryLruDelegate implements ReferenceDelegate, LruDelegate {
-  private inMemoryPins: ReferenceSet | null = null;
   private orphanedSequenceNumbers: ObjectMap<
     DocumentKey,
     ListenSequenceNumber
@@ -371,10 +377,6 @@ export class MemoryLruDelegate implements ReferenceDelegate, LruDelegate {
     );
   }
 
-  setInMemoryPins(inMemoryPins: ReferenceSet): void {
-    this.inMemoryPins = inMemoryPins;
-  }
-
   removeTargets(
     txn: PersistenceTransaction,
     upperBound: ListenSequenceNumber,
@@ -403,7 +405,7 @@ export class MemoryLruDelegate implements ReferenceDelegate, LruDelegate {
     return p.next(() => changeBuffer.apply(txn)).next(() => count);
   }
 
-  removeMutationReference(
+  markPotentiallyOrphaned(
     txn: PersistenceTransaction,
     key: DocumentKey
   ): PersistencePromise<void> {
@@ -421,6 +423,7 @@ export class MemoryLruDelegate implements ReferenceDelegate, LruDelegate {
 
   addReference(
     txn: PersistenceTransaction,
+    targetId: TargetId,
     key: DocumentKey
   ): PersistencePromise<void> {
     this.orphanedSequenceNumbers.set(key, txn.currentSequenceNumber);
@@ -429,6 +432,7 @@ export class MemoryLruDelegate implements ReferenceDelegate, LruDelegate {
 
   removeReference(
     txn: PersistenceTransaction,
+    targetId: TargetId,
     key: DocumentKey
   ): PersistencePromise<void> {
     this.orphanedSequenceNumbers.set(key, txn.currentSequenceNumber);
@@ -458,7 +462,6 @@ export class MemoryLruDelegate implements ReferenceDelegate, LruDelegate {
   ): PersistencePromise<boolean> {
     return PersistencePromise.or([
       () => this.persistence.mutationQueuesContainKey(txn, key),
-      () => PersistencePromise.resolve(this.inMemoryPins!.containsKey(key)),
       () => this.persistence.getTargetCache().containsKey(txn, key),
       () => {
         const orphanedAt = this.orphanedSequenceNumbers.get(key);