WIP

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,16 +18,15 @@
 
 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;
 import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
+import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.util.AsyncQueue;
 import com.google.firebase.firestore.util.Logger;
 import java.util.Collection;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -52,8 +51,8 @@ public class IndexBackfiller {
 
   public IndexBackfiller(Persistence persistence, AsyncQueue asyncQueue) {
     this.persistence = persistence;
-    this.scheduler = new Scheduler(asyncQueue);
     this.remoteDocumentCache = persistence.getRemoteDocumentCache();
+    this.scheduler = new Scheduler(asyncQueue);
   }
 
   public void setLocalDocumentsView(LocalDocumentsView localDocumentsView) {
@@ -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,52 +122,43 @@ 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);
 
-    IndexOffset newOffset = getNewOffset(documents, existingOffset);
-    indexManager.updateCollectionGroup(collectionGroup, newOffset);
+    LocalDocumentsResult nextDocuments =
+        localDocumentsView.getNextDocuments(
+            collectionGroup, existingOffset, documentsRemainingUnderCap);
 
-    return documents.size();
-  }
+    indexManager.updateIndexEntries(nextDocuments.getDocuments());
+    IndexOffset newOffset = getNewOffset(existingOffset, nextDocuments);
+    indexManager.updateCollectionGroup(collectionGroup, newOffset);
 
-  /** 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();
-      }
-    }
-    return lowestOffset == null ? IndexOffset.NONE : lowestOffset;
+    return nextDocuments.getDocuments().size();
   }
 
-  /** 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());
+  private IndexOffset getNewOffset(IndexOffset existingOffset, LocalDocumentsResult lookupResult) {
+    if (lookupResult.getDocuments().isEmpty()) {
+      return IndexOffset.create(
+          remoteDocumentCache.getLatestReadTime(),
+          DocumentKey.empty(),
+          existingOffset.getLargestBatchId());
     } else {
-      IndexOffset latestOffset = currentOffset;
-      Iterator<Map.Entry<DocumentKey, Document>> it = documents.iterator();
-      while (it.hasNext()) {
-        IndexOffset newOffset = IndexOffset.fromDocument(it.next().getValue());
+      IndexOffset latestOffset = existingOffset;
+      for (Map.Entry<DocumentKey, Document> entry : lookupResult.getDocuments()) {
+        IndexOffset newOffset =
+            IndexOffset.fromDocument(entry.getValue(), lookupResult.getLargestBatchId());
         if (newOffset.compareTo(latestOffset) > 0) {
           latestOffset = newOffset;
         }
@@ -178,6 +167,26 @@ private IndexOffset getNewOffset(
     }
   }
 
+  /** Returns the lowest offset for the provided index group. */
+  private IndexOffset getExistingOffset(Collection<FieldIndex> fieldIndexes) {
+    SnapshotVersion readTime = SnapshotVersion.NONE;
+    DocumentKey key = DocumentKey.empty();
+    int batchId = 0;
+
+    for (FieldIndex fieldIndex : fieldIndexes) {
+      IndexOffset offset = fieldIndex.getIndexState().getOffset();
+      if (readTime.compareTo(offset.getReadTime()) > 0) {
+        readTime = offset.getReadTime();
+        key = offset.getDocumentKey();
+      } else if (readTime.compareTo(offset.getReadTime()) == 0) {
+        key = key.compareTo(offset.getDocumentKey()) < 0 ? key : offset.getDocumentKey();
+      }
+      batchId = Math.max(batchId, offset.getLargestBatchId());
+    }
+
+    return IndexOffset.create(readTime, key, batchId);
+  }
+
   @VisibleForTesting
   void setMaxDocumentsToProcess(int newMax) {
     maxDocumentsToProcess = newMax;

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

@@ -31,6 +31,8 @@
 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.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -112,31 +114,28 @@ ImmutableSortedMap<DocumentKey, Document> getDocuments(Iterable<DocumentKey> key
     return getLocalViewOfDocuments(docs, new HashSet<>());
   }
 
-  /**
-   * Similar to {@link #getDocuments}, but creates the local view from the given {@code baseDocs}
-   * without retrieving documents from the local store.
-   *
-   * @param docs The documents to apply local mutations to get the local views.
-   * @param existenceStateChanged The set of document keys whose existence state is changed. This is
-   *     useful to determine if some documents overlay needs to be recalculated.
-   */
-  ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
-      Map<DocumentKey, MutableDocument> docs, Set<DocumentKey> existenceStateChanged) {
+  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());
       }
     }
 
@@ -148,6 +147,19 @@ ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
     return results;
   }
 
+  /**
+   * Similar to {@link #getDocuments}, but creates the local view from the given {@code baseDocs}
+   * without retrieving documents from the local store.
+   *
+   * @param docs The documents to apply local mutations to get the local views.
+   * @param existenceStateChanged The set of document keys whose existence state is changed. This is
+   *     useful to determine if some documents overlay needs to be recalculated.
+   */
+  ImmutableSortedMap<DocumentKey, Document> getLocalViewOfDocuments(
+      Map<DocumentKey, MutableDocument> docs, Set<DocumentKey> existenceStateChanged) {
+    return computeView(docs, Collections.emptyMap(), existenceStateChanged);
+  }
+
   private void recalculateAndSaveOverlays(Map<DocumentKey, MutableDocument> docs) {
     List<MutationBatch> batches =
         mutationQueue.getAllMutationBatchesAffectingDocumentKeys(docs.keySet());
@@ -195,14 +207,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.
    *
@@ -255,12 +259,54 @@ 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 offset.
+   *
+   * <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;
+    List<DocumentKey> mutatedKeys = new ArrayList<>();
+    for (Overlay overlay : overlays.values()) {
+      if (!docs.containsKey(overlay.getKey())) {
+        mutatedKeys.add(overlay.getKey());
+      }
+      largestBatchId = Math.max(largestBatchId, overlay.getLargestBatchId());
+    }
+
+    if (!mutatedKeys.isEmpty()) {
+      docs.putAll(remoteDocumentCache.getAll(mutatedKeys));
+    }
+
+    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 =
         remoteDocumentCache.getAll(query.getPath(), offset);
-    Map<DocumentKey, Overlay> overlays = documentOverlayCache.getOverlays(query.getPath(), -1);
+    Map<DocumentKey, Overlay> overlays =
+        documentOverlayCache.getOverlays(query.getPath(), offset.getLargestBatchId());
 
     // As documents might match the query because of their overlay we need to include documents
     // for all overlays in the initial document set.
@@ -282,7 +328,6 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
         results = results.insert(docEntry.getKey(), docEntry.getValue());
       }
     }
-
     return results;
   }
 }

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

@@ -121,7 +121,7 @@ public Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOf
         continue;
       }
 
-      if (IndexOffset.fromDocument(doc).compareTo(offset) <= 0) {
+      if (IndexOffset.fromDocument(doc, 0).compareTo(offset) <= 0) {
         // The document sorts before the offset.
         continue;
       }