Handling local mutations, acknowledgement and rejections

Author: wu-hui

packages/firestore/src/core/sync_engine_impl.ts

@@ -400,8 +400,17 @@ export function toLocalChangesToApplyToView(
   results: DocumentMap
 ): LocalChangesToApplyToView {
   return isPipeline(query) && getPipelineFlavor(query) !== 'exact'
-    ? { changedDocs: documentMap(), augmentedResults: results }
-    : { changedDocs: results, augmentedResults: undefined };
+    ? {
+        changedDocs: documentMap(),
+        augmentedResults: { results, mutatedKeys: documentKeySet() }
+      }
+    : {
+        changedDocs: results,
+        augmentedResults: {
+          results: documentMap(),
+          mutatedKeys: documentKeySet()
+        }
+      };
 }
 
 /**
@@ -566,10 +575,10 @@ export async function syncEngineWrite(
     );
     syncEngineImpl.sharedClientState.addPendingMutation(result.batchId);
     addMutationCallback(syncEngineImpl, result.batchId, userCallback);
-    await syncEngineEmitNewSnapsAndNotifyLocalStore(syncEngineImpl, {
-      changedDocs: result.changes,
-      newAugmentedResults: new Map()
-    });
+    await syncEngineEmitNewSnapsAndNotifyLocalStore(
+      syncEngineImpl,
+      result.changes
+    );
     await fillWritePipeline(syncEngineImpl.remoteStore);
   } catch (e) {
     // If we can't persist the mutation, we reject the user callback and
@@ -798,10 +807,7 @@ export async function syncEngineApplySuccessfulWrite(
       batchId,
       'acknowledged'
     );
-    await syncEngineEmitNewSnapsAndNotifyLocalStore(syncEngineImpl, {
-      changedDocs: changes,
-      newAugmentedResults: new Map()
-    });
+    await syncEngineEmitNewSnapsAndNotifyLocalStore(syncEngineImpl, changes);
   } catch (error) {
     await ignoreIfPrimaryLeaseLoss(error as FirestoreError);
   }
@@ -832,10 +838,7 @@ export async function syncEngineRejectFailedWrite(
       'rejected',
       error
     );
-    await syncEngineEmitNewSnapsAndNotifyLocalStore(syncEngineImpl, {
-      changedDocs: changes,
-      newAugmentedResults: new Map()
-    });
+    await syncEngineEmitNewSnapsAndNotifyLocalStore(syncEngineImpl, changes);
   } catch (error) {
     await ignoreIfPrimaryLeaseLoss(error as FirestoreError);
   }

packages/firestore/src/core/view.ts

@@ -50,6 +50,7 @@ import {
   newPipelineComparator,
   queryOrPipelineMatches
 } from './pipeline_run';
+import { MergedPipelineResults } from '../local/local_documents_view';
 
 export type LimboDocumentChange = AddedLimboDocument | RemovedLimboDocument;
 export class AddedLimboDocument {
@@ -82,7 +83,7 @@ export interface ViewChange {
 
 export interface LocalChangesToApplyToView {
   changedDocs: DocumentMap;
-  augmentedResults: DocumentMap | undefined;
+  augmentedResults: MergedPipelineResults | undefined;
 }
 
 /**
@@ -152,14 +153,14 @@ export class View {
   }
 
   computeAugmentedResultChanges(
-    augmentedResults: DocumentMap | undefined
+    augmentedResults: MergedPipelineResults | undefined
   ): ViewDocumentChanges {
     const changeSet = new DocumentChangeSet();
     const oldDocumentSet = this.documentSet;
-    let newMutatedKeys = documentKeySet();
+    let newMutatedKeys = augmentedResults?.mutatedKeys ?? documentKeySet();
     let newDocumentSet = new DocumentSet(this.docComparator);
     let needsRefill = false;
-    augmentedResults?.inorderTraversal((key, entry) => {
+    augmentedResults?.results.inorderTraversal((key, entry) => {
       const oldDoc = oldDocumentSet.get(key);
       const newDoc = entry;
       newDocumentSet = newDocumentSet.add(newDoc);

packages/firestore/src/local/index_backfiller.ts

@@ -31,6 +31,7 @@ import { Persistence, Scheduler } from './persistence';
 import { PersistencePromise } from './persistence_promise';
 import { PersistenceTransaction } from './persistence_transaction';
 import { isIndexedDbTransactionError } from './simple_db';
+import { NextDocuments } from './local_documents_view';
 
 const LOG_TAG = 'IndexBackfiller';
 
@@ -203,7 +204,7 @@ export class IndexBackfiller {
   /** Returns the next offset based on the provided documents. */
   private getNewOffset(
     existingOffset: IndexOffset,
-    lookupResult: LocalWriteResult
+    lookupResult: NextDocuments
   ): IndexOffset {
     let maxOffset: IndexOffset = existingOffset;
     lookupResult.changes.forEach((key, document) => {

packages/firestore/src/local/local_documents_view.ts

@@ -56,7 +56,6 @@ import { SortedMap } from '../util/sorted_map';
 
 import { DocumentOverlayCache } from './document_overlay_cache';
 import { IndexManager } from './index_manager';
-import { LocalWriteResult } from './local_store_impl';
 import { MutationQueue } from './mutation_queue';
 import { OverlayedDocument } from './overlayed_document';
 import { PersistencePromise } from './persistence_promise';
@@ -73,22 +72,30 @@ import {
   isPipeline,
   QueryOrPipeline
 } from '../core/pipeline-util';
-import { Pipeline } from '../lite-api/pipeline';
 import { FirestoreError } from '../util/error';
 import {
   CorePipeline,
   pipelineEvaluate,
   runPipeline
 } from '../core/pipeline_run';
-import { SortedSet } from '../util/sorted_set';
 import {
   PipelineCachedResults,
   PipelineResultsCache
 } from './pipeline_results_cache';
-import { TargetId } from '../core/types';
+import { BatchId, TargetId } from '../core/types';
 import { TargetData } from './target_data';
 import { targetIsPipelineTarget } from '../core/target';
 
+export interface NextDocuments {
+  batchId: BatchId;
+  changes: DocumentMap;
+}
+
+export interface MergedPipelineResults {
+  results: DocumentMap;
+  mutatedKeys: DocumentKeySet;
+}
+
 /**
  * A readonly view of the local state of all documents we're tracking (i.e. we
  * have a cached version in remoteDocumentCache or local mutations for the
@@ -433,14 +440,14 @@ export class LocalDocumentsView {
    * @param collectionGroup The collection group for the documents.
    * @param offset The offset to index into.
    * @param count The number of documents to return
-   * @return A LocalWriteResult with the documents that follow the provided offset and the last processed batch id.
+   * @return A NextDocuments with the documents that follow the provided offset and the last processed batch id.
    */
   getNextDocuments(
     transaction: PersistenceTransaction,
     collectionGroup: string,
     offset: IndexOffset,
     count: number
-  ): PersistencePromise<LocalWriteResult> {
+  ): PersistencePromise<NextDocuments> {
     return this.remoteDocumentCache
       .getAllFromCollectionGroup(transaction, collectionGroup, offset, count)
       .next((originalDocs: MutableDocumentMap) => {
@@ -788,9 +795,12 @@ export class LocalDocumentsView {
 
   calculateMergedAugmentPipelineResults(
     targetMap: SortedMap<TargetId, TargetData>,
-    currentAugmentPipelineResults: Map<TargetId, PipelineCachedResults>,
-    changedDocs: SortedMap<DocumentKey, Document>
-  ): Map<TargetId, MutableDocumentMap> {
+    currentAugmentPipelineResults: Map<
+      TargetId,
+      PipelineCachedResults | undefined
+    >,
+    changedDocs: DocumentMap
+  ): Map<TargetId, MergedPipelineResults> {
     // We only care about documents with pending writes because changedDocs are results
     // of two scenarios:
     // 1. Global snapshot, which means currentAugmentPipelineResults should include all
@@ -803,7 +813,7 @@ export class LocalDocumentsView {
       }
     });
 
-    const mergedResults = new Map<TargetId, MutableDocumentMap>();
+    const mergedResults = new Map<TargetId, MergedPipelineResults>();
     currentAugmentPipelineResults.forEach((results, targetId) => {
       const pipeline = targetMap.get(targetId)?.target;
       debugAssert(!!pipeline, `Target Id ${targetId} not found`);
@@ -812,13 +822,21 @@ export class LocalDocumentsView {
         `Target Id ${targetId} is not a pipeline target`
       );
 
-      const merged = results.results;
+      let merged = results?.results ?? mutableDocumentMap();
+      let mutatedKeys = documentKeySet();
       // TODO(pipeline): We need to handle limit pipelines here!!
       const pipelineResult = runPipeline(pipeline, docsWithMutations);
       for (const result of pipelineResult) {
-        merged.insert(result.key, result);
+        mutatedKeys = mutatedKeys.add(result.key);
+        merged = merged.insert(result.key, result);
+      }
+      for (const doc of docsWithMutations) {
+        if (!mutatedKeys.has(doc.key) && !!merged.get(doc.key)) {
+          merged = merged.remove(doc.key);
+          mutatedKeys = mutatedKeys.add(doc.key);
+        }
       }
-      mergedResults.set(targetId, merged);
+      mergedResults.set(targetId, { results: merged, mutatedKeys });
     });
 
     return mergedResults;