Firestore: Improve efficiency of memory persistence when processing a large number of writes

firebase-firestore/CHANGELOG.md

@@ -1,5 +1,5 @@
 # Unreleased
-
+* [changed] Improve efficiency of memory persistence when processing a large number of writes. [#6233](//github.com/firebase/firebase-android-sdk/pull/6233)
 
 # 25.1.0
 * [feature] Add support for the VectorValue type. [#6154](//github.com/firebase/firebase-android-sdk/pull/6154)

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

@@ -208,17 +208,17 @@ public LocalDocumentsView getLocalDocumentsForCurrentUser() {
 
   public ImmutableSortedMap<DocumentKey, Document> handleUserChange(User user) {
     // Swap out the mutation queue, grabbing the pending mutation batches before and after.
-    List<MutationBatch> oldBatches = mutationQueue.getAllMutationBatches();
+    Collection<MutationBatch> oldBatches = mutationQueue.getAllMutationBatches();
 
     initializeUserComponents(user);
     startIndexManager();
     startMutationQueue();
 
-    List<MutationBatch> newBatches = mutationQueue.getAllMutationBatches();
+    Collection<MutationBatch> newBatches = mutationQueue.getAllMutationBatches();
 
     // Union the old/new changed keys.
     ImmutableSortedSet<DocumentKey> changedKeys = DocumentKey.emptyKeySet();
-    for (List<MutationBatch> batches : asList(oldBatches, newBatches)) {
+    for (Collection<MutationBatch> batches : asList(oldBatches, newBatches)) {
       for (MutationBatch batch : batches) {
         for (Mutation mutation : batch.getMutations()) {
           changedKeys = changedKeys.insert(mutation.getKey());

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

@@ -31,9 +31,12 @@
 import com.google.firebase.firestore.util.Util;
 import com.google.protobuf.ByteString;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
+import java.util.TreeMap;
 
 final class MemoryMutationQueue implements MutationQueue {
 
@@ -54,7 +57,7 @@ final class MemoryMutationQueue implements MutationQueue {
    * <p>Once the held write acknowledgements become visible they are removed from the head of the
    * queue along with any tombstones that follow.
    */
-  private final List<MutationBatch> queue;
+  private final TreeMap<Integer, MutationBatch> queue = new TreeMap<>();
 
   /** An ordered mapping between documents and the mutation batch IDs. */
   private ImmutableSortedSet<DocumentReference> batchesByDocumentKey;
@@ -74,7 +77,6 @@ final class MemoryMutationQueue implements MutationQueue {
 
   MemoryMutationQueue(MemoryPersistence persistence, User user) {
     this.persistence = persistence;
-    queue = new ArrayList<>();
 
     batchesByDocumentKey = new ImmutableSortedSet<>(emptyList(), DocumentReference.BY_KEY);
     nextBatchId = 1;
@@ -105,16 +107,20 @@ public boolean isEmpty() {
   @Override
   public void acknowledgeBatch(MutationBatch batch, ByteString streamToken) {
     int batchId = batch.getBatchId();
-    int batchIndex = indexOfExistingBatchId(batchId, "acknowledged");
-    hardAssert(batchIndex == 0, "Can only acknowledge the first batch in the mutation queue");
+    hardAssert(
+        queue.containsKey(batchId),
+        "Can not acknowledge unknown batch ID: %d (batch=%s)",
+        batchId,
+        batch);
 
-    // Verify that the batch in the queue is the one to be acknowledged.
-    MutationBatch check = queue.get(batchIndex);
+    int firstBatchId = queue.firstKey();
     hardAssert(
-        batchId == check.getBatchId(),
-        "Queue ordering failure: expected batch %d, got batch %d",
+        batchId == firstBatchId,
+        "Can only acknowledge the first batch in the mutation queue:"
+            + " got batchId %d, but expected %d (batch=%s)",
         batchId,
-        check.getBatchId());
+        firstBatchId,
+        batch);
 
     lastStreamToken = checkNotNull(streamToken);
   }
@@ -137,15 +143,17 @@ public MutationBatch addMutationBatch(
     int batchId = nextBatchId;
     nextBatchId += 1;
 
-    int size = queue.size();
-    if (size > 0) {
-      MutationBatch prior = queue.get(size - 1);
+    if (!queue.isEmpty()) {
+      int priorKey = queue.lastKey();
       hardAssert(
-          prior.getBatchId() < batchId, "Mutation batchIds must be monotonically increasing order");
+          priorKey < batchId,
+          "Mutation batchIds must be monotonically increasing order: priorKey=%d, batchId=%d",
+          priorKey,
+          batchId);
     }
 
     MutationBatch batch = new MutationBatch(batchId, localWriteTime, baseMutations, mutations);
-    queue.add(batch);
+    queue.put(batchId, batch);
 
     // Track references by document key and index collection parents.
     for (Mutation mutation : mutations) {
@@ -161,25 +169,14 @@ public MutationBatch addMutationBatch(
   @Nullable
   @Override
   public MutationBatch lookupMutationBatch(int batchId) {
-    int index = indexOfBatchId(batchId);
-    if (index < 0 || index >= queue.size()) {
-      return null;
-    }
-
-    MutationBatch batch = queue.get(index);
-    hardAssert(batch.getBatchId() == batchId, "If found batch must match");
-    return batch;
+    return queue.get(batchId);
   }
 
   @Nullable
   @Override
   public MutationBatch getNextMutationBatchAfterBatchId(int batchId) {
-    int nextBatchId = batchId + 1;
-
-    // The requested batchId may still be out of range so normalize it to the start of the queue.
-    int rawIndex = indexOfBatchId(nextBatchId);
-    int index = rawIndex < 0 ? 0 : rawIndex;
-    return queue.size() > index ? queue.get(index) : null;
+    Map.Entry<Integer, MutationBatch> nextEntry = queue.higherEntry(batchId);
+    return nextEntry == null ? null : nextEntry.getValue();
   }
 
   @Override
@@ -188,8 +185,8 @@ public int getHighestUnacknowledgedBatchId() {
   }
 
   @Override
-  public List<MutationBatch> getAllMutationBatches() {
-    return Collections.unmodifiableList(queue);
+  public Collection<MutationBatch> getAllMutationBatches() {
+    return Collections.unmodifiableCollection(queue.values());
   }
 
   @Override
@@ -292,12 +289,23 @@ private List<MutationBatch> lookupMutationBatches(ImmutableSortedSet<Integer> ba
 
   @Override
   public void removeMutationBatch(MutationBatch batch) {
-    // Find the position of the first batch for removal. This need not be the first entry in the
-    // queue.
-    int batchIndex = indexOfExistingBatchId(batch.getBatchId(), "removed");
-    hardAssert(batchIndex == 0, "Can only remove the first entry of the mutation queue");
+    int batchId = batch.getBatchId();
+    hardAssert(
+        queue.containsKey(batchId),
+        "Can not remove unknown batch ID: %d (batch=%s)",
+        batchId,
+        batch);
+
+    int firstBatchId = queue.firstKey();
+    hardAssert(
+        batchId == firstBatchId,
+        "Can only remove the first batch in the mutation queue:"
+            + " got batchId %d, but expected %d (batch=%s)",
+        batchId,
+        firstBatchId,
+        batch);
 
-    queue.remove(0);
+    queue.remove(batchId);
 
     // Remove entries from the index too.
     ImmutableSortedSet<DocumentReference> references = batchesByDocumentKey;
@@ -336,45 +344,9 @@ boolean containsKey(DocumentKey key) {
 
   // Helpers
 
-  /**
-   * Finds the index of the given batchId in the mutation queue. This operation is O(1).
-   *
-   * @return The computed index of the batch with the given batchId, based on the state of the
-   *     queue. Note this index can be negative if the requested batchId has already been removed
-   *     from the queue or past the end of the queue if the batchId is larger than the last added
-   *     batch.
-   */
-  private int indexOfBatchId(int batchId) {
-    if (queue.isEmpty()) {
-      // As an index this is past the end of the queue
-      return 0;
-    }
-
-    // Examine the front of the queue to figure out the difference between the batchId and indexes
-    // in the array. Note that since the queue is ordered by batchId, if the first batch has a
-    // larger batchId then the requested batchId doesn't exist in the queue.
-    MutationBatch firstBatch = queue.get(0);
-    int firstBatchId = firstBatch.getBatchId();
-    return batchId - firstBatchId;
-  }
-
-  /**
-   * Finds the index of the given batchId in the mutation queue and asserts that the resulting index
-   * is within the bounds of the queue.
-   *
-   * @param batchId The batchId to search for
-   * @param action A description of what the caller is doing, phrased in passive form (for example
-   *     "acknowledged" in a routine that acknowledges batches).
-   */
-  private int indexOfExistingBatchId(int batchId, String action) {
-    int index = indexOfBatchId(batchId);
-    hardAssert(index >= 0 && index < queue.size(), "Batches must exist to be %s", action);
-    return index;
-  }
-
   long getByteSize(LocalSerializer serializer) {
     long count = 0;
-    for (MutationBatch batch : queue) {
+    for (MutationBatch batch : queue.values()) {
       count += serializer.encodeMutationBatch(batch).getSerializedSize();
     }
     return count;

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

@@ -21,6 +21,7 @@
 import com.google.firebase.firestore.model.mutation.Mutation;
 import com.google.firebase.firestore.model.mutation.MutationBatch;
 import com.google.protobuf.ByteString;
+import java.util.Collection;
 import java.util.List;
 
 /** A queue of mutations to apply to the remote store. */
@@ -79,7 +80,7 @@ MutationBatch addMutationBatch(
   /** Returns all mutation batches in the mutation queue. */
   // TODO: PERF: Current consumer only needs mutated keys; if we can provide that
   // cheaply, we should replace this.
-  List<MutationBatch> getAllMutationBatches();
+  Collection<MutationBatch> getAllMutationBatches();
 
   /**
    * Finds all mutation batches that could @em possibly affect the given document key. Not all

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

@@ -20,8 +20,8 @@
 import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.mutation.MutationBatch;
+import java.util.Collection;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Set;
 
 /** Manages overlay migrations required to have overlay support . */
@@ -58,7 +58,7 @@ private void buildOverlays() {
 
             // Get all document keys that have local mutations
             Set<DocumentKey> allDocumentKeys = new HashSet<>();
-            List<MutationBatch> batches = mutationQueue.getAllMutationBatches();
+            Collection<MutationBatch> batches = mutationQueue.getAllMutationBatches();
             for (MutationBatch batch : batches) {
               allDocumentKeys.addAll(batch.getKeys());
             }

firebase-firestore/src/test/java/com/google/firebase/firestore/local/MutationQueueTestCase.java

@@ -40,6 +40,7 @@
 import com.google.firebase.firestore.remote.WriteStream;
 import com.google.protobuf.ByteString;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -343,10 +344,10 @@ public void testRemoveMutationBatches() {
     removeMutationBatches(batches.remove(0));
     assertEquals(9, batchCount());
 
-    List<MutationBatch> found;
+    Collection<MutationBatch> found;
 
     found = mutationQueue.getAllMutationBatches();
-    assertEquals(batches, found);
+    assertEquals(batches, new ArrayList<>(found));
     assertEquals(9, found.size());
 
     removeMutationBatches(batches.get(0), batches.get(1), batches.get(2));
@@ -356,14 +357,14 @@ public void testRemoveMutationBatches() {
     assertEquals(6, batchCount());
 
     found = mutationQueue.getAllMutationBatches();
-    assertEquals(batches, found);
+    assertEquals(batches, new ArrayList<>(found));
     assertEquals(6, found.size());
 
     removeMutationBatches(batches.remove(0));
     assertEquals(5, batchCount());
 
     found = mutationQueue.getAllMutationBatches();
-    assertEquals(batches, found);
+    assertEquals(batches, new ArrayList<>(found));
     assertEquals(5, found.size());
 
     removeMutationBatches(batches.remove(0));
@@ -373,12 +374,12 @@ public void testRemoveMutationBatches() {
     assertEquals(3, batchCount());
 
     found = mutationQueue.getAllMutationBatches();
-    assertEquals(batches, found);
+    assertEquals(batches, new ArrayList<>(found));
     assertEquals(3, found.size());
 
     removeMutationBatches(batches.toArray(new MutationBatch[0]));
     found = mutationQueue.getAllMutationBatches();
-    assertEquals(emptyList(), found);
+    assertEquals(emptyList(), new ArrayList<>(found));
     assertEquals(0, found.size());
     assertTrue(mutationQueue.isEmpty());
   }