Backfill Mutations

Author: schmidt-sebastian

firebase-firestore/src/main/java/com/google/firebase/firestore/core/SyncEngine.java

@@ -32,9 +32,9 @@
 import com.google.firebase.firestore.bundle.BundleMetadata;
 import com.google.firebase.firestore.bundle.BundleReader;
 import com.google.firebase.firestore.core.ViewSnapshot.SyncState;
+import com.google.firebase.firestore.local.LocalDocumentsResult;
 import com.google.firebase.firestore.local.LocalStore;
 import com.google.firebase.firestore.local.LocalViewChanges;
-import com.google.firebase.firestore.local.LocalWriteResult;
 import com.google.firebase.firestore.local.QueryPurpose;
 import com.google.firebase.firestore.local.QueryResult;
 import com.google.firebase.firestore.local.ReferenceSet;
@@ -274,10 +274,10 @@ void stopListening(Query query) {
   public void writeMutations(List<Mutation> mutations, TaskCompletionSource<Void> userTask) {
     assertCallback("writeMutations");
 
-    LocalWriteResult result = localStore.writeLocally(mutations);
-    addUserCallback(result.getBatchId(), userTask);
+    LocalDocumentsResult result = localStore.writeLocally(mutations);
+    addUserCallback(result.getLargestBatchId(), userTask);
 
-    emitNewSnapsAndNotifyLocalStore(result.getChanges(), /*remoteEvent=*/ null);
+    emitNewSnapsAndNotifyLocalStore(result.getDocuments(), /*remoteEvent=*/ null);
     remoteStore.fillWritePipeline();
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/local/IndexBackfiller.java

@@ -18,7 +18,6 @@
 
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
-import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex;
@@ -110,12 +109,11 @@ public Scheduler getScheduler() {
   public int backfill() {
     hardAssert(localDocumentsView != null, "setLocalDocumentsView() not called");
     hardAssert(indexManager != null, "setIndexManager() not called");
-    return persistence.runTransaction(
-        "Backfill Indexes", () -> writeIndexEntries(localDocumentsView));
+    return persistence.runTransaction("Backfill Indexes", this::writeIndexEntries);
   }
 
   /** Writes index entries until the cap is reached. Returns the number of documents processed. */
-  private int writeIndexEntries(LocalDocumentsView localDocumentsView) {
+  private int writeIndexEntries() {
     Set<String> processedCollectionGroups = new HashSet<>();
     int documentsRemaining = maxDocumentsToProcess;
     while (documentsRemaining > 0) {
@@ -124,58 +122,65 @@ private int writeIndexEntries(LocalDocumentsView localDocumentsView) {
         break;
       }
       Logger.debug(LOG_TAG, "Processing collection: %s", collectionGroup);
-      documentsRemaining -=
-          writeEntriesForCollectionGroup(localDocumentsView, collectionGroup, documentsRemaining);
+      documentsRemaining -= writeEntriesForCollectionGroup(collectionGroup, documentsRemaining);
       processedCollectionGroups.add(collectionGroup);
     }
     return maxDocumentsToProcess - documentsRemaining;
   }
 
-  /** Writes entries for the fetched field indexes. */
+  /**
+   * Writes entries for the provided collection group. Returns the number of documents processed.
+   */
   private int writeEntriesForCollectionGroup(
-      LocalDocumentsView localDocumentsView, String collectionGroup, int entriesRemainingUnderCap) {
-    // TODO(indexing): Support mutation batch Ids when sorting and writing indexes.
-
+      String collectionGroup, int documentsRemainingUnderCap) {
     // Use the earliest offset of all field indexes to query the local cache.
-    IndexOffset existingOffset = getExistingOffset(indexManager.getFieldIndexes(collectionGroup));
-    ImmutableSortedMap<DocumentKey, Document> documents =
-        localDocumentsView.getDocuments(collectionGroup, existingOffset, entriesRemainingUnderCap);
-    indexManager.updateIndexEntries(documents);
+    Collection<FieldIndex> fieldIndexes = indexManager.getFieldIndexes(collectionGroup);
+    IndexOffset existingOffset = getExistingOffset(fieldIndexes);
+
+    LocalDocumentsResult nextBatch =
+        localDocumentsView.getNextDocuments(
+            collectionGroup, existingOffset, documentsRemainingUnderCap);
+    indexManager.updateIndexEntries(nextBatch.getDocuments());
 
-    IndexOffset newOffset = getNewOffset(documents, existingOffset);
+    IndexOffset newOffset = getNewOffset(existingOffset, nextBatch);
     indexManager.updateCollectionGroup(collectionGroup, newOffset);
 
-    return documents.size();
+    return nextBatch.getDocuments().size();
   }
 
-  /** Returns the lowest offset for the provided index group. */
-  private IndexOffset getExistingOffset(Collection<FieldIndex> fieldIndexes) {
-    IndexOffset lowestOffset = null;
-    for (FieldIndex fieldIndex : fieldIndexes) {
-      if (lowestOffset == null
-          || fieldIndex.getIndexState().getOffset().compareTo(lowestOffset) < 0) {
-        lowestOffset = fieldIndex.getIndexState().getOffset();
+  /** Returns the next offset based on the provided documents. */
+  private IndexOffset getNewOffset(IndexOffset existingOffset, LocalDocumentsResult lookupResult) {
+    IndexOffset maxOffset = existingOffset;
+    for (Map.Entry<DocumentKey, Document> entry : lookupResult.getDocuments()) {
+      IndexOffset newOffset = IndexOffset.fromDocument(entry.getValue());
+      if (newOffset.compareTo(maxOffset) > 0) {
+        maxOffset = newOffset;
       }
     }
-    return lowestOffset == null ? IndexOffset.NONE : lowestOffset;
+    return IndexOffset.create(
+        maxOffset.getReadTime(),
+        maxOffset.getDocumentKey(),
+        Math.max(lookupResult.getLargestBatchId(), existingOffset.getLargestBatchId()));
   }
 
-  /** Returns the offset for the index based on the newly indexed documents. */
-  private IndexOffset getNewOffset(
-      ImmutableSortedMap<DocumentKey, Document> documents, IndexOffset currentOffset) {
-    if (documents.isEmpty()) {
-      return IndexOffset.create(remoteDocumentCache.getLatestReadTime());
-    } else {
-      IndexOffset latestOffset = currentOffset;
-      Iterator<Map.Entry<DocumentKey, Document>> it = documents.iterator();
-      while (it.hasNext()) {
-        IndexOffset newOffset = IndexOffset.fromDocument(it.next().getValue());
-        if (newOffset.compareTo(latestOffset) > 0) {
-          latestOffset = newOffset;
-        }
+  /** Returns the lowest offset for the provided index group. */
+  private IndexOffset getExistingOffset(Collection<FieldIndex> fieldIndexes) {
+    if (fieldIndexes.isEmpty()) {
+      return IndexOffset.NONE;
+    }
+
+    Iterator<FieldIndex> it = fieldIndexes.iterator();
+    IndexOffset minOffset = it.next().getIndexState().getOffset();
+    int minBatchId = minOffset.getLargestBatchId();
+    while (it.hasNext()) {
+      IndexOffset newOffset = it.next().getIndexState().getOffset();
+      if (newOffset.compareTo(minOffset) < 0) {
+        minOffset = newOffset;
       }
-      return latestOffset;
+      minBatchId = Math.max(newOffset.getLargestBatchId(), minBatchId);
     }
+
+    return IndexOffset.create(minOffset.getReadTime(), minOffset.getDocumentKey(), minBatchId);
   }
 
   @VisibleForTesting

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalWriteResult.java renamed to firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalDocumentsResult.java

@@ -18,22 +18,21 @@
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
 
-/** The result of a write to the local store. */
-public final class LocalWriteResult {
-  private final int batchId;
+/** The result of a applying local mutations. */
+public final class LocalDocumentsResult {
+  private final int largestBatchId;
+  private final ImmutableSortedMap<DocumentKey, Document> documents;
 
-  private final ImmutableSortedMap<DocumentKey, Document> changes;
-
-  LocalWriteResult(int batchId, ImmutableSortedMap<DocumentKey, Document> changes) {
-    this.batchId = batchId;
-    this.changes = changes;
+  LocalDocumentsResult(int largestBatchId, ImmutableSortedMap<DocumentKey, Document> documents) {
+    this.largestBatchId = largestBatchId;
+    this.documents = documents;
   }
 
-  public int getBatchId() {
-    return batchId;
+  public int getLargestBatchId() {
+    return largestBatchId;
   }
 
-  public ImmutableSortedMap<DocumentKey, Document> getChanges() {
-    return changes;
+  public ImmutableSortedMap<DocumentKey, Document> getDocuments() {
+    return documents;
   }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalDocumentsView.java

@@ -17,6 +17,7 @@
 import static com.google.firebase.firestore.model.DocumentCollections.emptyDocumentMap;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.firebase.Timestamp;
 import com.google.firebase.database.collection.ImmutableSortedMap;
@@ -31,6 +32,7 @@
 import com.google.firebase.firestore.model.mutation.MutationBatch;
 import com.google.firebase.firestore.model.mutation.Overlay;
 import com.google.firebase.firestore.model.mutation.PatchMutation;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -84,15 +86,10 @@ DocumentOverlayCache getDocumentOverlayCache() {
    */
   Document getDocument(DocumentKey key) {
     Overlay overlay = documentOverlayCache.getOverlay(key);
-    // Only read from remote document cache if overlay is a patch.
-    MutableDocument document =
-        (overlay == null || overlay.getMutation() instanceof PatchMutation)
-            ? remoteDocumentCache.get(key)
-            : MutableDocument.newInvalidDocument(key);
+    MutableDocument document = createBaseDocument(key, overlay);
     if (overlay != null) {
       overlay.getMutation().applyToLocalView(document, null, Timestamp.now());
     }
-
     return document;
   }
 
@@ -117,21 +114,35 @@ ImmutableSortedMap<DocumentKey, Document> getDocuments(Iterable<DocumentKey> key
    */
   ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
       Map<DocumentKey, MutableDocument> docs, Set<DocumentKey> existenceStateChanged) {
+    return computeView(docs, Collections.emptyMap(), existenceStateChanged);
+  }
+
+  /**
+   * Computes the local view for doc, applying overlays from both {@code memoizedOverlays} and the
+   * overlay cache.
+   */
+  private ImmutableSortedMap<DocumentKey, Document> computeView(
+      Map<DocumentKey, MutableDocument> docs,
+      Map<DocumentKey, Overlay> memoizedOverlays,
+      Set<DocumentKey> existenceStateChanged) {
     ImmutableSortedMap<DocumentKey, Document> results = emptyDocumentMap();
     Map<DocumentKey, MutableDocument> recalculateDocuments = new HashMap<>();
-    for (Map.Entry<DocumentKey, MutableDocument> entry : docs.entrySet()) {
-      Overlay overlay = documentOverlayCache.getOverlay(entry.getKey());
+    for (MutableDocument doc : docs.values()) {
+      Overlay overlay =
+          memoizedOverlays.containsKey(doc.getKey())
+              ? memoizedOverlays.get(doc.getKey())
+              : documentOverlayCache.getOverlay(doc.getKey());
       // 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.contains(entry.getKey())
+      if (existenceStateChanged.contains(doc.getKey())
           && (overlay == null || overlay.getMutation() instanceof PatchMutation)) {
-        recalculateDocuments.put(entry.getKey(), docs.get(entry.getKey()));
+        recalculateDocuments.put(doc.getKey(), doc);
       } else if (overlay != null) {
-        overlay.getMutation().applyToLocalView(entry.getValue(), null, Timestamp.now());
+        overlay.getMutation().applyToLocalView(doc, null, Timestamp.now());
       }
     }
 
@@ -190,14 +201,6 @@ void recalculateAndSaveOverlays(Set<DocumentKey> documentKeys) {
     recalculateAndSaveOverlays(docs);
   }
 
-  /** Gets the local view of the next {@code count} documents based on their read time. */
-  ImmutableSortedMap<DocumentKey, Document> getDocuments(
-      String collectionGroup, IndexOffset offset, int count) {
-    Map<DocumentKey, MutableDocument> docs =
-        remoteDocumentCache.getAll(collectionGroup, offset, count);
-    return getLocalViewOfDocuments(docs, new HashSet<>());
-  }
-
   /**
    * Performs a query against the local view of all documents.
    *
@@ -250,7 +253,43 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
     return results;
   }
 
-  /** Queries the remote documents and overlays by doing a full collection scan. */
+  /**
+   * Given a collection group, returns the next documents that follow the provided offset, along
+   * with an updated batch ID.
+   *
+   * <p>The documents returned by this method are ordered by remote version from the provided
+   * offset. If there are no more remote documents after the provided offset, documents with
+   * mutations in order of batch id from the offset are returned. Since all documents in a batch are
+   * returned together, the total number of documents returned can exceed {@code count}.
+   *
+   * @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 LocalDocumentsRResult with the documents that follow the provided offset and the last
+   *     processed batch id.
+   */
+  LocalDocumentsResult getNextDocuments(String collectionGroup, IndexOffset offset, int count) {
+    Map<DocumentKey, MutableDocument> docs =
+        remoteDocumentCache.getAll(collectionGroup, offset, count);
+    Map<DocumentKey, Overlay> overlays =
+        count - docs.size() > 0
+            ? documentOverlayCache.getOverlays(
+                collectionGroup, offset.getLargestBatchId(), count - docs.size())
+            : Collections.emptyMap();
+
+    int largestBatchId = -1;
+    for (Overlay overlay : overlays.values()) {
+      if (!docs.containsKey(overlay.getKey())) {
+        docs.put(overlay.getKey(), createBaseDocument(overlay.getKey(), overlay));
+      }
+      largestBatchId = Math.max(largestBatchId, overlay.getLargestBatchId());
+    }
+
+    ImmutableSortedMap<DocumentKey, Document> localDocs =
+        computeView(docs, overlays, Collections.emptySet());
+    return new LocalDocumentsResult(largestBatchId, localDocs);
+  }
+
   private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollectionQuery(
       Query query, IndexOffset offset) {
     Map<DocumentKey, MutableDocument> remoteDocuments =
@@ -281,4 +320,11 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
 
     return results;
   }
+
+  /** Returns a base document that can be used to apply `overlay`. */
+  private MutableDocument createBaseDocument(DocumentKey key, @Nullable Overlay overlay) {
+    return (overlay == null || overlay.getMutation() instanceof PatchMutation)
+        ? remoteDocumentCache.get(key)
+        : MutableDocument.newInvalidDocument(key);
+  }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalStore.java

@@ -236,7 +236,7 @@ public ImmutableSortedMap<DocumentKey, Document> handleUserChange(User user) {
   }
 
   /** Accepts locally generated Mutations and commits them to storage. */
-  public LocalWriteResult writeLocally(List<Mutation> mutations) {
+  public LocalDocumentsResult writeLocally(List<Mutation> mutations) {
     Timestamp localWriteTime = Timestamp.now();
 
     // TODO: Call queryEngine.handleDocumentChange() appropriately.
@@ -277,7 +277,7 @@ public LocalWriteResult writeLocally(List<Mutation> mutations) {
               mutationQueue.addMutationBatch(localWriteTime, baseMutations, mutations);
           Map<DocumentKey, Mutation> overlays = batch.applyToLocalDocumentSet(documents);
           documentOverlayCache.saveOverlays(batch.getBatchId(), overlays);
-          return new LocalWriteResult(batch.getBatchId(), documents);
+          return new LocalDocumentsResult(batch.getBatchId(), documents);
         });
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryRemoteDocumentCache.java

@@ -37,12 +37,9 @@ final class MemoryRemoteDocumentCache implements RemoteDocumentCache {
   private ImmutableSortedMap<DocumentKey, Document> docs;
   /** Manages the collection group index. */
   private IndexManager indexManager;
-  /** The latest read time of any document in the cache. */
-  private SnapshotVersion latestReadTime;
 
   MemoryRemoteDocumentCache() {
     docs = emptyDocumentMap();
-    latestReadTime = SnapshotVersion.NONE;
   }
 
   @Override
@@ -57,7 +54,6 @@ public void add(MutableDocument document, SnapshotVersion readTime) {
         !readTime.equals(SnapshotVersion.NONE),
         "Cannot add document to the RemoteDocumentCache with a read time of zero");
     docs = docs.insert(document.getKey(), document.mutableCopy().withReadTime(readTime));
-    latestReadTime = readTime.compareTo(latestReadTime) > 0 ? readTime : latestReadTime;
 
     indexManager.addToCollectionParentIndex(document.getKey().getCollectionPath());
   }
@@ -132,11 +128,6 @@ public Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOf
     return result;
   }
 
-  @Override
-  public SnapshotVersion getLatestReadTime() {
-    return latestReadTime;
-  }
-
   Iterable<Document> getDocuments() {
     return new DocumentIterable();
   }

firebase-firestore/src/main/java/com/google/firebase/firestore/local/QueryEngine.java

@@ -160,7 +160,9 @@ && needsRefill(
     }
 
     return appendRemainingResults(
-        previousResults, query, IndexOffset.create(lastLimboFreeSnapshotVersion));
+        previousResults,
+        query,
+        IndexOffset.create(lastLimboFreeSnapshotVersion, FieldIndex.INITIAL_LARGEST_BATCH_ID));
   }
 
   /** Applies the query filter and sorting to the provided documents. */

firebase-firestore/src/main/java/com/google/firebase/firestore/local/RemoteDocumentCache.java

@@ -84,7 +84,4 @@ interface RemoteDocumentCache {
    * @return A newly created map with the set of documents in the collection.
    */
   Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOffset offset);
-
-  /** Returns the latest read time of any document in the cache. */
-  SnapshotVersion getLatestReadTime();
 }