Delete index entries for deleted documents (#3209)

Author: schmidt-sebastian

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

@@ -494,6 +494,7 @@ private DocumentChangeResult populateDocumentChanges(
       @Nullable Map<DocumentKey, SnapshotVersion> documentVersions,
       SnapshotVersion globalVersion) {
     Map<DocumentKey, MutableDocument> changedDocs = new HashMap<>();
+    List<DocumentKey> removedDocs = new ArrayList<>();
     Set<DocumentKey> conditionChanged = new HashSet<>();
 
     // Each loop iteration only affects its "own" doc, so it's safe to get all the remote
@@ -517,7 +518,7 @@ private DocumentChangeResult populateDocumentChanges(
       if (doc.isNoDocument() && doc.getVersion().equals(SnapshotVersion.NONE)) {
         // NoDocuments with SnapshotVersion.NONE are used in manufactured events. We remove
         // these documents from cache since we lost access.
-        remoteDocuments.remove(doc.getKey());
+        removedDocs.add(doc.getKey());
         changedDocs.put(key, doc);
       } else if (!existingDoc.isValidDocument()
           || doc.getVersion().compareTo(existingDoc.getVersion()) > 0
@@ -537,6 +538,7 @@ private DocumentChangeResult populateDocumentChanges(
             doc.getVersion());
       }
     }
+    remoteDocuments.removeAll(removedDocs);
     return new DocumentChangeResult(changedDocs, conditionChanged);
   }
 

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

@@ -16,7 +16,9 @@
 
 import com.google.firebase.firestore.core.ListenSequence;
 import com.google.firebase.firestore.model.DocumentKey;
+import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Set;
 
 /** Provides eager garbage collection for MemoryPersistence. */
@@ -72,11 +74,13 @@ public void onTransactionStarted() {
   @Override
   public void onTransactionCommitted() {
     MemoryRemoteDocumentCache remoteDocuments = persistence.getRemoteDocumentCache();
+    List<DocumentKey> docsToRemove = new ArrayList<>();
     for (DocumentKey key : orphanedDocuments) {
       if (!isReferenced(key)) {
-        remoteDocuments.remove(key);
+        docsToRemove.add(key);
       }
     }
+    remoteDocuments.removeAll(docsToRemove);
     orphanedDocuments = null;
   }
 

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

@@ -18,10 +18,12 @@
 
 import android.util.SparseArray;
 import com.google.firebase.firestore.core.ListenSequence;
+import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
-import com.google.firebase.firestore.model.MutableDocument;
 import com.google.firebase.firestore.util.Consumer;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 /** Provides LRU garbage collection functionality for MemoryPersistence. */
@@ -115,17 +117,17 @@ public int removeTargets(long upperBound, SparseArray<?> activeTargetIds) {
 
   @Override
   public int removeOrphanedDocuments(long upperBound) {
-    int count = 0;
     MemoryRemoteDocumentCache cache = persistence.getRemoteDocumentCache();
-    for (MutableDocument doc : cache.getDocuments()) {
+    List<DocumentKey> docsToRemove = new ArrayList<>();
+    for (Document doc : cache.getDocuments()) {
       DocumentKey key = doc.getKey();
       if (!isPinned(key, upperBound)) {
-        cache.remove(key);
+        docsToRemove.add(key);
         orphanedSequenceNumbers.remove(key);
-        count++;
       }
     }
-    return count;
+    cache.removeAll(docsToRemove);
+    return docsToRemove.size();
   }
 
   @Override

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

@@ -14,15 +14,18 @@
 
 package com.google.firebase.firestore.local;
 
+import static com.google.firebase.firestore.model.DocumentCollections.emptyDocumentMap;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
 import androidx.annotation.NonNull;
 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.IndexOffset;
 import com.google.firebase.firestore.model.MutableDocument;
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.SnapshotVersion;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.Map;
@@ -60,8 +63,16 @@ public void add(MutableDocument document, SnapshotVersion readTime) {
   }
 
   @Override
-  public void remove(DocumentKey key) {
-    docs = docs.remove(key);
+  public void removeAll(Collection<DocumentKey> keys) {
+    hardAssert(indexManager != null, "setIndexManager() not called");
+
+    ImmutableSortedMap<DocumentKey, Document> deletedDocs = emptyDocumentMap();
+    for (DocumentKey key : keys) {
+      docs = docs.remove(key);
+      deletedDocs =
+          deletedDocs.insert(key, MutableDocument.newNoDocument(key, SnapshotVersion.NONE));
+    }
+    indexManager.updateIndexEntries(deletedDocs);
   }
 
   @Override

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

@@ -19,6 +19,7 @@
 import com.google.firebase.firestore.model.MutableDocument;
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.SnapshotVersion;
+import java.util.Collection;
 import java.util.Map;
 
 /**
@@ -44,8 +45,8 @@ interface RemoteDocumentCache {
    */
   void add(MutableDocument document, SnapshotVersion readTime);
 
-  /** Removes the cached entry for the given key (no-op if no entry exists). */
-  void remove(DocumentKey documentKey);
+  /** Removes the cached entries for the given keys (no-op if no entry exists). */
+  void removeAll(Collection<DocumentKey> keys);
 
   /**
    * Looks up an entry in the cache.

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

@@ -225,6 +225,8 @@ public void deleteFieldIndex(FieldIndex index) {
   @Override
   public void updateIndexEntries(ImmutableSortedMap<DocumentKey, Document> documents) {
     hardAssert(started, "IndexManager not started");
+    if (!Persistence.INDEXING_SUPPORT_ENABLED) return;
+
     for (Map.Entry<DocumentKey, Document> entry : documents) {
       Collection<FieldIndex> fieldIndexes = getFieldIndexes(entry.getKey().getCollectionGroup());
       for (FieldIndex fieldIndex : fieldIndexes) {

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

@@ -21,6 +21,8 @@
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.util.Consumer;
+import java.util.ArrayList;
+import java.util.List;
 
 /** Provides LRU functionality for SQLite persistence. */
 class SQLiteLruReferenceDelegate implements ReferenceDelegate, LruDelegate {
@@ -159,6 +161,7 @@ public int removeOrphanedDocuments(long upperBound) {
 
     boolean resultsRemaining = true;
 
+    List<DocumentKey> docsToRemove = new ArrayList<>();
     while (resultsRemaining) {
       int rowsProccessed =
           persistence
@@ -171,14 +174,15 @@ public int removeOrphanedDocuments(long upperBound) {
                     DocumentKey key = DocumentKey.fromPath(path);
                     if (!isPinned(key)) {
                       count[0]++;
-                      persistence.getRemoteDocumentCache().remove(key);
+                      docsToRemove.add(key);
                       removeSentinel(key);
                     }
                   });
 
       resultsRemaining = (rowsProccessed == REMOVE_ORPHANED_DOCUMENTS_BATCH_SIZE);
     }
 
+    persistence.getRemoteDocumentCache().removeAll(docsToRemove);
     return count[0];
   }
 

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

@@ -16,6 +16,7 @@
 
 import static com.google.firebase.firestore.util.Assert.fail;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
+import static com.google.firebase.firestore.util.Util.repeatSequence;
 
 import android.content.Context;
 import android.database.Cursor;
@@ -675,23 +676,27 @@ boolean hasMoreSubqueries() {
       return argsIter.hasNext();
     }
 
-    /** Performs the next subquery and returns a {@link Query} object for method chaining. */
-    Query performNextSubquery() {
-      ++subqueriesPerformed;
-
+    private Object[] getNextSubqueryArgs() {
       List<Object> subqueryArgs = new ArrayList<>(argsHead);
-      StringBuilder placeholdersBuilder = new StringBuilder();
       for (int i = 0; argsIter.hasNext() && i < LIMIT - argsHead.size(); i++) {
-        if (i > 0) {
-          placeholdersBuilder.append(", ");
-        }
-        placeholdersBuilder.append("?");
-
         subqueryArgs.add(argsIter.next());
       }
-      String placeholders = placeholdersBuilder.toString();
+      return subqueryArgs.toArray();
+    }
 
-      return db.query(head + placeholders + tail).binding(subqueryArgs.toArray());
+    /** Performs the next subquery and returns a {@link Query} object for method chaining. */
+    Query performNextSubquery() {
+      ++subqueriesPerformed;
+      Object[] subqueryArgs = getNextSubqueryArgs();
+      return db.query(head + repeatSequence("?", subqueryArgs.length, ", ") + tail)
+          .binding(subqueryArgs);
+    }
+
+    /** Executes the next subquery. */
+    void executeNextSubquery() {
+      ++subqueriesPerformed;
+      Object[] subqueryArgs = getNextSubqueryArgs();
+      db.execute(head + repeatSequence("?", subqueryArgs.length, ", ") + tail, subqueryArgs);
     }
 
     /** How many subqueries were performed. */

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

@@ -14,13 +14,16 @@
 
 package com.google.firebase.firestore.local;
 
+import static com.google.firebase.firestore.model.DocumentCollections.emptyDocumentMap;
 import static com.google.firebase.firestore.util.Assert.fail;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 import static com.google.firebase.firestore.util.Util.firstNEntries;
 import static com.google.firebase.firestore.util.Util.repeatSequence;
 
 import androidx.annotation.VisibleForTesting;
 import com.google.firebase.Timestamp;
+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.IndexOffset;
 import com.google.firebase.firestore.model.MutableDocument;
@@ -31,6 +34,7 @@
 import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.MessageLite;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -79,10 +83,26 @@ public void add(MutableDocument document, SnapshotVersion readTime) {
   }
 
   @Override
-  public void remove(DocumentKey documentKey) {
-    String path = EncodedPath.encode(documentKey.getPath());
+  public void removeAll(Collection<DocumentKey> keys) {
+    if (keys.isEmpty()) return;
+
+    List<Object> encodedPaths = new ArrayList<>();
+    ImmutableSortedMap<DocumentKey, Document> deletedDocs = emptyDocumentMap();
+
+    for (DocumentKey key : keys) {
+      encodedPaths.add(EncodedPath.encode(key.getPath()));
+      deletedDocs =
+          deletedDocs.insert(key, MutableDocument.newNoDocument(key, SnapshotVersion.NONE));
+    }
+
+    SQLitePersistence.LongQuery longQuery =
+        new SQLitePersistence.LongQuery(
+            db, "DELETE FROM remote_documents WHERE path IN (", encodedPaths, ")");
+    while (longQuery.hasMoreSubqueries()) {
+      longQuery.executeNextSubquery();
+    }
 
-    db.execute("DELETE FROM remote_documents WHERE path = ?", path);
+    indexManager.updateIndexEntries(deletedDocs);
   }
 
   @Override

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

@@ -28,6 +28,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;
 import java.util.Map;
 
@@ -124,8 +125,8 @@ public void add(MutableDocument document, SnapshotVersion readTime) {
       }
 
       @Override
-      public void remove(DocumentKey documentKey) {
-        subject.remove(documentKey);
+      public void removeAll(Collection<DocumentKey> keys) {
+        subject.removeAll(keys);
       }
 
       @Override

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

@@ -16,11 +16,14 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.firebase.firestore.testutil.TestUtil.addedRemoteEvent;
+import static com.google.firebase.firestore.testutil.TestUtil.deletedDoc;
 import static com.google.firebase.firestore.testutil.TestUtil.doc;
 import static com.google.firebase.firestore.testutil.TestUtil.fieldIndex;
 import static com.google.firebase.firestore.testutil.TestUtil.filter;
 import static com.google.firebase.firestore.testutil.TestUtil.map;
 import static com.google.firebase.firestore.testutil.TestUtil.query;
+import static com.google.firebase.firestore.testutil.TestUtil.updateRemoteEvent;
+import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 
 import com.google.firebase.firestore.core.Query;
@@ -100,4 +103,30 @@ public void testUsesPartialIndexesWhenAvailable() {
     assertRemoteDocumentsRead(/* byKey= */ 1, /* byQuery= */ 1);
     assertQueryReturned("coll/a", "coll/b");
   }
+
+  @Test
+  public void testDeletedDocumentRemovesIndex() {
+    FieldIndex index = fieldIndex("coll", 0, FieldIndex.INITIAL_STATE, "matches", Kind.ASCENDING);
+    configureFieldIndexes(singletonList(index));
+
+    Query query = query("coll").filter(filter("matches", "==", true));
+    int targetId = allocateQuery(query);
+
+    applyRemoteEvent(addedRemoteEvent(doc("coll/a", 10, map("matches", true)), targetId));
+
+    // Add the document to the index
+    backfillIndexes();
+
+    executeQuery(query);
+    assertRemoteDocumentsRead(/* byKey= */ 1, /* byQuery= */ 0);
+    assertQueryReturned("coll/a");
+
+    applyRemoteEvent(
+        updateRemoteEvent(deletedDoc("coll/a", 0), singletonList(targetId), emptyList()));
+
+    // No backfill needed for deleted document.
+    executeQuery(query);
+    assertRemoteDocumentsRead(/* byKey= */ 0, /* byQuery= */ 0);
+    assertQueryReturned();
+  }
 }

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

@@ -72,6 +72,7 @@ public void setUp() {
     persistence.getReferenceDelegate().setInMemoryPins(new ReferenceSet());
     IndexManager indexManager = persistence.getIndexManager(User.UNAUTHENTICATED);
     indexManager.start();
+    persistence.getRemoteDocumentCache().setIndexManager(indexManager);
     mutationQueue = persistence.getMutationQueue(User.UNAUTHENTICATED, indexManager);
     mutationQueue.start();
   }