Store per-user state for indexes (#3153)

Author: schmidt-sebastian

firebase-firestore/src/main/java/com/google/firebase/firestore/FirebaseFirestore.java

@@ -44,7 +44,6 @@
 import com.google.firebase.firestore.model.FieldIndex;
 import com.google.firebase.firestore.model.FieldPath;
 import com.google.firebase.firestore.model.ResourcePath;
-import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.remote.FirestoreChannel;
 import com.google.firebase.firestore.remote.GrpcMetadataProvider;
 import com.google.firebase.firestore.util.AsyncQueue;
@@ -340,7 +339,9 @@ Task<Void> configureIndices(String json) {
             }
           }
 
-          parsedIndices.add(FieldIndex.create(-1, collectionGroup, segments, SnapshotVersion.NONE));
+          parsedIndices.add(
+              FieldIndex.create(
+                  FieldIndex.UNKNOWN_ID, collectionGroup, segments, FieldIndex.INITIAL_STATE));
         }
       }
     } catch (JSONException e) {

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.Timestamp;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.Document;
@@ -27,10 +26,13 @@
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.util.AsyncQueue;
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
-import java.util.PriorityQueue;
-import java.util.Queue;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
 /** Implements the steps for backfilling indexes. */
@@ -44,13 +46,15 @@ public class IndexBackfiller {
 
   private final Scheduler scheduler;
   private final Persistence persistence;
+  private final RemoteDocumentCache remoteDocumentCache;
   private LocalDocumentsView localDocumentsView;
   private IndexManager indexManager;
   private int maxDocumentsToProcess = MAX_DOCUMENTS_TO_PROCESS;
 
   public IndexBackfiller(Persistence persistence, AsyncQueue asyncQueue) {
     this.persistence = persistence;
     this.scheduler = new Scheduler(asyncQueue);
+    this.remoteDocumentCache = persistence.getRemoteDocumentCache();
   }
 
   public void setLocalDocumentsView(LocalDocumentsView localDocumentsView) {
@@ -131,71 +135,92 @@ public Results backfill() {
     return persistence.runTransaction(
         "Backfill Indexes",
         () -> {
-          // TODO(indexing): Handle field indexes that are removed by the user.
           int documentsProcessed = writeIndexEntries(localDocumentsView);
           return new Results(/* hasRun= */ true, documentsProcessed);
         });
   }
 
   /** Writes index entries until the cap is reached. Returns the number of documents processed. */
   private int writeIndexEntries(LocalDocumentsView localDocumentsView) {
-    int documentsProcessed = 0;
-    Timestamp startingTimestamp = Timestamp.now();
-
-    while (documentsProcessed < maxDocumentsToProcess) {
-      int documentsRemaining = maxDocumentsToProcess - documentsProcessed;
-      String collectionGroup = indexManager.getNextCollectionGroupToUpdate(startingTimestamp);
-      if (collectionGroup == null) {
+    Set<String> processedCollectionGroups = new HashSet<>();
+    int documentsRemaining = maxDocumentsToProcess;
+    while (documentsRemaining > 0) {
+      String collectionGroup = indexManager.getNextCollectionGroupToUpdate();
+      if (collectionGroup == null || processedCollectionGroups.contains(collectionGroup)) {
         break;
       }
-      documentsProcessed +=
+      documentsRemaining -=
           writeEntriesForCollectionGroup(localDocumentsView, collectionGroup, documentsRemaining);
+      processedCollectionGroups.add(collectionGroup);
     }
-
-    return documentsProcessed;
+    return maxDocumentsToProcess - documentsRemaining;
   }
 
   /** Writes entries for the fetched field indexes. */
   private int writeEntriesForCollectionGroup(
       LocalDocumentsView localDocumentsView, String collectionGroup, int entriesRemainingUnderCap) {
     Query query = new Query(ResourcePath.EMPTY, collectionGroup);
 
-    // Use the earliest updateTime of all field indexes as the base updateTime.
-    SnapshotVersion earliestUpdateTime =
-        getEarliestUpdateTime(indexManager.getFieldIndexes(collectionGroup));
+    // Use the earliest readTime of all field indexes as the base readtime.
+    SnapshotVersion earliestReadTime =
+        getEarliestReadTime(indexManager.getFieldIndexes(collectionGroup));
 
     // TODO(indexing): Use limit queries to only fetch the required number of entries.
     // TODO(indexing): Support mutation batch Ids when sorting and writing indexes.
     ImmutableSortedMap<DocumentKey, Document> documents =
-        localDocumentsView.getDocumentsMatchingQuery(query, earliestUpdateTime);
+        localDocumentsView.getDocumentsMatchingQuery(query, earliestReadTime);
 
-    Queue<Document> oldestDocuments = getOldestDocuments(documents, entriesRemainingUnderCap);
+    List<Document> oldestDocuments = getOldestDocuments(documents, entriesRemainingUnderCap);
     indexManager.updateIndexEntries(oldestDocuments);
+
+    SnapshotVersion latestReadTime = getPostUpdateReadTime(oldestDocuments, earliestReadTime);
+    indexManager.updateCollectionGroup(collectionGroup, latestReadTime);
     return oldestDocuments.size();
   }
 
+  /**
+   * Returns the new read time for the index.
+   *
+   * @param documents a list of documents sorted by read time (ascending)
+   * @param currentReadTime the current read time of the index
+   */
+  private SnapshotVersion getPostUpdateReadTime(
+      List<Document> documents, SnapshotVersion currentReadTime) {
+    SnapshotVersion latestReadTime =
+        documents.isEmpty()
+            ? remoteDocumentCache.getLatestReadTime()
+            : documents.get(documents.size() - 1).getReadTime();
+    // Make sure the index does not go back in time
+    latestReadTime =
+        latestReadTime.compareTo(currentReadTime) > 0 ? latestReadTime : currentReadTime;
+    return latestReadTime;
+  }
+
   /** Returns up to {@code count} documents sorted by read time. */
-  private Queue<Document> getOldestDocuments(
+  private List<Document> getOldestDocuments(
       ImmutableSortedMap<DocumentKey, Document> documents, int count) {
-    Queue<Document> oldestDocuments =
-        new PriorityQueue<>(count + 1, (l, r) -> r.getReadTime().compareTo(l.getReadTime()));
+    List<Document> oldestDocuments = new ArrayList<>();
     for (Map.Entry<DocumentKey, Document> entry : documents) {
       oldestDocuments.add(entry.getValue());
-      if (oldestDocuments.size() > count) {
-        oldestDocuments.poll();
-      }
     }
-    return oldestDocuments;
+    Collections.sort(
+        oldestDocuments,
+        (l, r) -> {
+          int cmp = l.getReadTime().compareTo(r.getReadTime());
+          if (cmp != 0) return cmp;
+          return l.getKey().compareTo(r.getKey());
+        });
+    return oldestDocuments.subList(0, Math.min(count, oldestDocuments.size()));
   }
 
-  private SnapshotVersion getEarliestUpdateTime(Collection<FieldIndex> fieldIndexes) {
+  private SnapshotVersion getEarliestReadTime(Collection<FieldIndex> fieldIndexes) {
     SnapshotVersion lowestVersion = null;
     for (FieldIndex fieldIndex : fieldIndexes) {
       lowestVersion =
           lowestVersion == null
-              ? fieldIndex.getUpdateTime()
-              : fieldIndex.getUpdateTime().compareTo(lowestVersion) < 0
-                  ? fieldIndex.getUpdateTime()
+              ? fieldIndex.getIndexState().getReadTime()
+              : fieldIndex.getIndexState().getReadTime().compareTo(lowestVersion) < 0
+                  ? fieldIndex.getIndexState().getReadTime()
                   : lowestVersion;
     }
     return lowestVersion;

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

@@ -15,12 +15,12 @@
 package com.google.firebase.firestore.local;
 
 import androidx.annotation.Nullable;
-import com.google.firebase.Timestamp;
 import com.google.firebase.firestore.core.Target;
 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.ResourcePath;
+import com.google.firebase.firestore.model.SnapshotVersion;
 import java.util.Collection;
 import java.util.List;
 import java.util.Set;
@@ -87,14 +87,22 @@ public interface IndexManager {
   /** Returns the documents that match the given target based on the provided index. */
   Set<DocumentKey> getDocumentsMatchingTarget(FieldIndex fieldIndex, Target target);
 
-  /** Returns the next collection group to update. */
+  /** Returns the next collection group to update. Returns {@code null} if no group exists. */
   @Nullable
-  String getNextCollectionGroupToUpdate(Timestamp lastUpdateTime);
+  String getNextCollectionGroupToUpdate();
+
+  /**
+   * Sets the collection group's latest read time.
+   *
+   * <p>This method updates the read time for all field indices for the collection group and
+   * increments their sequence number. Subsequent calls to {@link #getNextCollectionGroupToUpdate()}
+   * will return a different collection group (unless only one collection group is configured).
+   */
+  void updateCollectionGroup(String collectionGroup, SnapshotVersion readTime);
 
   /**
    * Updates the index entries for the provided documents and corresponding field indexes until the
-   * cap is reached. Updates the field indexes in persistence with the latest read time that was
-   * processed.
+   * cap is reached.
    */
   void updateIndexEntries(Collection<Document> documents);
 }

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

@@ -84,7 +84,7 @@ private ImmutableSortedMap<DocumentKey, Document> performCollectionQuery(Query q
       ImmutableSortedMap<DocumentKey, Document> indexedDocuments =
           localDocuments.getDocuments(keys);
       ImmutableSortedMap<DocumentKey, Document> additionalDocuments =
-          localDocuments.getDocumentsMatchingQuery(query, fieldIndex.getUpdateTime());
+          localDocuments.getDocumentsMatchingQuery(query, fieldIndex.getIndexState().getReadTime());
       for (Map.Entry<DocumentKey, Document> entry : additionalDocuments) {
         indexedDocuments = indexedDocuments.insert(entry.getKey(), entry.getValue());
       }

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

@@ -16,12 +16,12 @@
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
 import androidx.annotation.Nullable;
-import com.google.firebase.Timestamp;
 import com.google.firebase.firestore.core.Target;
 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.ResourcePath;
+import com.google.firebase.firestore.model.SnapshotVersion;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -79,12 +79,18 @@ public Set<DocumentKey> getDocumentsMatchingTarget(FieldIndex fieldIndex, Target
     return Collections.emptySet();
   }
 
+  @Nullable
   @Override
-  public String getNextCollectionGroupToUpdate(Timestamp startingTimestamp) {
+  public String getNextCollectionGroupToUpdate() {
     // Field indices are not supported with memory persistence.
     return null;
   }
 
+  @Override
+  public void updateCollectionGroup(String collectionGroup, SnapshotVersion readTime) {
+    // Field indices are not supported with memory persistence.
+  }
+
   @Override
   public Collection<FieldIndex> getFieldIndexes(String collectionGroup) {
     // Field indices are not supported with memory persistence.

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

@@ -36,9 +36,12 @@ final class MemoryRemoteDocumentCache implements RemoteDocumentCache {
   private ImmutableSortedMap<DocumentKey, Pair<MutableDocument, SnapshotVersion>> docs;
   /** Manages the collection group index. */
   private IndexManager indexManager;
+  /** The latest read time of any document in the cache. */
+  private SnapshotVersion latestReadTime;
 
   MemoryRemoteDocumentCache() {
     docs = ImmutableSortedMap.Builder.emptyMap(DocumentKey.comparator());
+    latestReadTime = SnapshotVersion.NONE;
   }
 
   @Override
@@ -53,6 +56,7 @@ 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(), new Pair<>(document.clone(), readTime));
+    latestReadTime = readTime.compareTo(latestReadTime) > 0 ? readTime : latestReadTime;
 
     indexManager.addToCollectionParentIndex(document.getKey().getPath().popLast());
   }
@@ -120,6 +124,11 @@ public ImmutableSortedMap<DocumentKey, MutableDocument> getAllDocumentsMatchingQ
     return result;
   }
 
+  @Override
+  public SnapshotVersion getLatestReadTime() {
+    return latestReadTime;
+  }
+
   Iterable<MutableDocument> getDocuments() {
     return new DocumentIterable();
   }

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

@@ -79,4 +79,7 @@ interface RemoteDocumentCache {
    */
   ImmutableSortedMap<DocumentKey, MutableDocument> getAllDocumentsMatchingQuery(
       Query query, SnapshotVersion sinceReadTime);
+
+  /** Returns the latest read time of any document in the cache. */
+  SnapshotVersion getLatestReadTime();
 }