Apply bloom filter on existence filter mismatch (#4601)

Author: milaGGL

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/BloomFilter.java

@@ -32,21 +32,19 @@ public BloomFilter(@NonNull byte[] bitmap, int padding, int hashCount) {
       throw new NullPointerException("Bitmap cannot be null.");
     }
     if (padding < 0 || padding >= 8) {
-      throw new IllegalArgumentException("Invalid padding: " + padding);
+      throw new BloomFilterException("Invalid padding: " + padding);
     }
     if (hashCount < 0) {
-      throw new IllegalArgumentException("Invalid hash count: " + hashCount);
+      throw new BloomFilterException("Invalid hash count: " + hashCount);
     }
     if (bitmap.length > 0 && hashCount == 0) {
       // Only empty bloom filter can have 0 hash count.
-      throw new IllegalArgumentException("Invalid hash count: " + hashCount);
+      throw new BloomFilterException("Invalid hash count: " + hashCount);
     }
-    if (bitmap.length == 0) {
+    if (bitmap.length == 0 && padding != 0) {
       // Empty bloom filter should have 0 padding.
-      if (padding != 0) {
-        throw new IllegalArgumentException(
-            "Expected padding of 0 when bitmap length is 0, but got " + padding);
-      }
+      throw new BloomFilterException(
+          "Expected padding of 0 when bitmap length is 0, but got " + padding);
     }
 
     this.bitmap = bitmap;
@@ -133,8 +131,8 @@ private int getBitIndex(long hash1, long hash2, int hashIndex) {
   /**
    * Calculate modulo, where the dividend and divisor are treated as unsigned 64-bit longs.
    *
-   * <p>The implementation is taken from <a
-   * href="https://github.com/google/guava/blob/553037486901cc60820ab7dcb38a25b6f34eba43/android/guava/src/com/google/common/primitives/UnsignedLongs.java">Guava</a>,
+   * <p>The implementation is taken from <a href=
+   * "https://github.com/google/guava/blob/553037486901cc60820ab7dcb38a25b6f34eba43/android/guava/src/com/google/common/primitives/UnsignedLongs.java">Guava</a>,
    * simplified to our needs.
    *
    * <p>

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/BloomFilterException.java

@@ -0,0 +1,23 @@
+// Copyright 2023 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.firestore.remote;
+
+import androidx.annotation.NonNull;
+
+public class BloomFilterException extends RuntimeException {
+  public BloomFilterException(@NonNull String detailMessage) {
+    super(detailMessage);
+  }
+}

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/ExistenceFilter.java

@@ -14,20 +14,34 @@
 
 package com.google.firebase.firestore.remote;
 
+import androidx.annotation.Nullable;
+import com.google.firestore.v1.BloomFilter;
+
 /** Simplest form of existence filter */
 public final class ExistenceFilter {
   private final int count;
+  private BloomFilter unchangedNames;
 
   public ExistenceFilter(int count) {
     this.count = count;
   }
 
+  public ExistenceFilter(int count, @Nullable BloomFilter unchangedNames) {
+    this.count = count;
+    this.unchangedNames = unchangedNames;
+  }
+
   public int getCount() {
     return count;
   }
 
+  @Nullable
+  public BloomFilter getUnchangedNames() {
+    return unchangedNames;
+  }
+
   @Override
   public String toString() {
-    return "ExistenceFilter{count=" + count + '}';
+    return "ExistenceFilter{count=" + count + ", unchangedNames=" + unchangedNames + '}';
   }
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/RemoteSerializer.java

@@ -945,8 +945,8 @@ public WatchChange decodeWatchChange(ListenResponse protoChange) {
         break;
       case FILTER:
         com.google.firestore.v1.ExistenceFilter protoFilter = protoChange.getFilter();
-        // TODO: implement existence filter parsing (see b/33076578)
-        ExistenceFilter filter = new ExistenceFilter(protoFilter.getCount());
+        ExistenceFilter filter =
+            new ExistenceFilter(protoFilter.getCount(), protoFilter.getUnchangedNames());
         int targetId = protoFilter.getTargetId();
         watchChange = new ExistenceFilterWatchChange(targetId, filter);
         break;

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/RemoteStore.java

@@ -28,6 +28,7 @@
 import com.google.firebase.firestore.local.LocalStore;
 import com.google.firebase.firestore.local.QueryPurpose;
 import com.google.firebase.firestore.local.TargetData;
+import com.google.firebase.firestore.model.DatabaseId;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.model.mutation.MutationBatch;
@@ -758,6 +759,11 @@ public TargetData getTargetDataForTarget(int targetId) {
     return this.listenTargets.get(targetId);
   }
 
+  @Override
+  public DatabaseId getDatabaseId() {
+    return this.datastore.getDatabaseInfo().getDatabaseId();
+  }
+
   public Task<Long> runCountQuery(Query query) {
     if (canUseNetwork()) {
       return datastore.runCountQuery(query);

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/WatchChangeAggregator.java

@@ -23,12 +23,14 @@
 import com.google.firebase.firestore.core.Target;
 import com.google.firebase.firestore.local.QueryPurpose;
 import com.google.firebase.firestore.local.TargetData;
+import com.google.firebase.firestore.model.DatabaseId;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.MutableDocument;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.remote.WatchChange.DocumentChange;
 import com.google.firebase.firestore.remote.WatchChange.ExistenceFilterWatchChange;
 import com.google.firebase.firestore.remote.WatchChange.WatchTargetChange;
+import com.google.firebase.firestore.util.Logger;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -56,6 +58,9 @@ public interface TargetMetadataProvider {
      */
     @Nullable
     TargetData getTargetDataForTarget(int targetId);
+
+    /** Returns the database ID of the Firestore instance. */
+    DatabaseId getDatabaseId();
   }
 
   private final TargetMetadataProvider targetMetadataProvider;
@@ -75,6 +80,9 @@ public interface TargetMetadataProvider {
    */
   private Set<Integer> pendingTargetResets = new HashSet<>();
 
+  /** The log tag to use for this class. */
+  private static final String LOG_TAG = "WatchChangeAggregator";
+
   public WatchChangeAggregator(TargetMetadataProvider targetMetadataProvider) {
     this.targetMetadataProvider = targetMetadataProvider;
   }
@@ -196,17 +204,78 @@ public void handleExistenceFilter(ExistenceFilterWatchChange watchChange) {
               expectedCount == 1, "Single document existence filter with count: %d", expectedCount);
         }
       } else {
-        long currentSize = getCurrentDocumentCountForTarget(targetId);
+        int currentSize = getCurrentDocumentCountForTarget(targetId);
         if (currentSize != expectedCount) {
-          // Existence filter mismatch: We reset the mapping and raise a new snapshot with
-          // `isFromCache:true`.
-          resetTarget(targetId);
-          pendingTargetResets.add(targetId);
+
+          // Apply bloom filter to identify and mark removed documents.
+          boolean bloomFilterApplied = this.applyBloomFilter(watchChange, currentSize);
+
+          if (!bloomFilterApplied) {
+            // If bloom filter application fails, we reset the mapping and
+            // trigger re-run of the query.
+            resetTarget(targetId);
+            pendingTargetResets.add(targetId);
+          }
         }
       }
     }
   }
 
+  /** Returns whether a bloom filter removed the deleted documents successfully. */
+  private boolean applyBloomFilter(ExistenceFilterWatchChange watchChange, int currentCount) {
+    int expectedCount = watchChange.getExistenceFilter().getCount();
+    com.google.firestore.v1.BloomFilter unchangedNames =
+        watchChange.getExistenceFilter().getUnchangedNames();
+
+    if (unchangedNames == null || !unchangedNames.hasBits()) {
+      return false;
+    }
+
+    byte[] bitmap = unchangedNames.getBits().getBitmap().toByteArray();
+    BloomFilter bloomFilter;
+
+    try {
+      bloomFilter =
+          new BloomFilter(
+              bitmap, unchangedNames.getBits().getPadding(), unchangedNames.getHashCount());
+    } catch (BloomFilterException e) {
+      Logger.warn(
+          LOG_TAG,
+          "Decoding the base64 bloom filter in existence filter failed ("
+              + e.getMessage()
+              + "); ignoring the bloom filter and falling back to full re-query.");
+      return false;
+    }
+
+    int removedDocumentCount = this.filterRemovedDocuments(bloomFilter, watchChange.getTargetId());
+
+    return expectedCount == (currentCount - removedDocumentCount);
+  }
+
+  /**
+   * Filter out removed documents based on bloom filter membership result and return number of
+   * documents removed.
+   */
+  private int filterRemovedDocuments(BloomFilter bloomFilter, int targetId) {
+    ImmutableSortedSet<DocumentKey> existingKeys =
+        targetMetadataProvider.getRemoteKeysForTarget(targetId);
+    int removalCount = 0;
+    for (DocumentKey key : existingKeys) {
+      DatabaseId databaseId = targetMetadataProvider.getDatabaseId();
+      String documentPath =
+          "projects/"
+              + databaseId.getProjectId()
+              + "/databases/"
+              + databaseId.getDatabaseId()
+              + "/documents/"
+              + key.getPath().canonicalString();
+      if (!bloomFilter.mightContain(documentPath)) {
+        this.removeDocumentFromTarget(targetId, key, /*updatedDocument=*/ null);
+        removalCount++;
+      }
+    }
+    return removalCount;
+  }
   /**
    * Converts the currently accumulated state into a remote event at the provided snapshot version.
    * Resets the accumulated changes before returning.

firebase-firestore/src/test/java/com/google/firebase/firestore/remote/BloomFilterTest.java

@@ -73,55 +73,53 @@ public void constructorShouldThrowNPEOnNullBitmap() {
   }
 
   @Test
-  public void constructorShouldThrowIAEOnEmptyBloomFilterWithNonZeroPadding() {
-    IllegalArgumentException exception =
-        assertThrows(IllegalArgumentException.class, () -> new BloomFilter(new byte[0], 1, 0));
+  public void constructorShouldThrowBFEOnEmptyBloomFilterWithNonZeroPadding() {
+    BloomFilterException exception =
+        assertThrows(BloomFilterException.class, () -> new BloomFilter(new byte[0], 1, 0));
     assertThat(exception)
         .hasMessageThat()
         .contains("Expected padding of 0 when bitmap length is 0, but got 1");
   }
 
   @Test
-  public void constructorShouldThrowIAEOnNonEmptyBloomFilterWithZeroHashCount() {
-    IllegalArgumentException zeroHashCountException =
-        assertThrows(IllegalArgumentException.class, () -> new BloomFilter(new byte[] {1}, 1, 0));
+  public void constructorShouldThrowBFEOnNonEmptyBloomFilterWithZeroHashCount() {
+    BloomFilterException zeroHashCountException =
+        assertThrows(BloomFilterException.class, () -> new BloomFilter(new byte[] {1}, 1, 0));
     assertThat(zeroHashCountException).hasMessageThat().contains("Invalid hash count: 0");
   }
 
   @Test
-  public void constructorShouldThrowIAEOnNegativePadding() {
+  public void constructorShouldThrowBFEOnNegativePadding() {
     {
-      IllegalArgumentException emptyBloomFilterException =
-          assertThrows(IllegalArgumentException.class, () -> new BloomFilter(new byte[0], -1, 0));
+      BloomFilterException emptyBloomFilterException =
+          assertThrows(BloomFilterException.class, () -> new BloomFilter(new byte[0], -1, 0));
       assertThat(emptyBloomFilterException).hasMessageThat().contains("Invalid padding: -1");
     }
     {
-      IllegalArgumentException nonEmptyBloomFilterException =
-          assertThrows(
-              IllegalArgumentException.class, () -> new BloomFilter(new byte[] {1}, -1, 1));
+      BloomFilterException nonEmptyBloomFilterException =
+          assertThrows(BloomFilterException.class, () -> new BloomFilter(new byte[] {1}, -1, 1));
       assertThat(nonEmptyBloomFilterException).hasMessageThat().contains("Invalid padding: -1");
     }
   }
 
   @Test
-  public void constructorShouldThrowIAEOnNegativeHashCount() {
+  public void constructorShouldThrowBFEOnNegativeHashCount() {
     {
-      IllegalArgumentException emptyBloomFilterException =
-          assertThrows(IllegalArgumentException.class, () -> new BloomFilter(new byte[0], 0, -1));
+      BloomFilterException emptyBloomFilterException =
+          assertThrows(BloomFilterException.class, () -> new BloomFilter(new byte[0], 0, -1));
       assertThat(emptyBloomFilterException).hasMessageThat().contains("Invalid hash count: -1");
     }
     {
-      IllegalArgumentException nonEmptyBloomFilterException =
-          assertThrows(
-              IllegalArgumentException.class, () -> new BloomFilter(new byte[] {1}, 1, -1));
+      BloomFilterException nonEmptyBloomFilterException =
+          assertThrows(BloomFilterException.class, () -> new BloomFilter(new byte[] {1}, 1, -1));
       assertThat(nonEmptyBloomFilterException).hasMessageThat().contains("Invalid hash count: -1");
     }
   }
 
   @Test
-  public void constructorShouldThrowIAEIfPaddingIsTooLarge() {
-    IllegalArgumentException exception =
-        assertThrows(IllegalArgumentException.class, () -> new BloomFilter(new byte[] {1}, 8, 1));
+  public void constructorShouldThrowBFEIfPaddingIsTooLarge() {
+    BloomFilterException exception =
+        assertThrows(BloomFilterException.class, () -> new BloomFilter(new byte[] {1}, 8, 1));
     assertThat(exception).hasMessageThat().contains("Invalid padding: 8");
   }
 
@@ -179,7 +177,7 @@ public void bloomFilterToString() {
   private static void runGoldenTest(String testFile) throws Exception {
     String resultFile = testFile.replace("bloom_filter_proto", "membership_test_result");
     if (resultFile.equals(testFile)) {
-      throw new IllegalArgumentException("Cannot find corresponding result file for " + testFile);
+      throw new BloomFilterException("Cannot find corresponding result file for " + testFile);
     }
 
     JSONObject testJson = readJsonFile(testFile);