Apply bloom filter on existence filter mismatch

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

@@ -27,26 +27,25 @@ public class BloomFilter {
   private final int hashCount;
   private final MessageDigest md5HashMessageDigest;
 
-  public BloomFilter(@NonNull byte[] bitmap, int padding, int hashCount) {
+  public BloomFilter(@NonNull byte[] bitmap, int padding, int hashCount)
+      throws BloomFilterException {
     if (bitmap == null) {
       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;

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 Exception {
+  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, 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

@@ -946,7 +946,8 @@ public WatchChange decodeWatchChange(ListenResponse protoChange) {
       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/WatchChangeAggregator.java

@@ -29,6 +29,7 @@
 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;
@@ -196,17 +197,69 @@ 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.getExistenceFilter(), targetId, 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(
+      ExistenceFilter existenceFilter, int targetId, int currentCount) {
+    int expectedCount = existenceFilter.getCount();
+    com.google.firestore.v1.BloomFilter unchangedNames = existenceFilter.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("Firestore", "BloomFilter error: %s", e);
+      return false;
+    }
+
+    int removedDocumentCount = this.filterRemovedDocuments(bloomFilter, targetId);
+
+    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) {
+      if (!bloomFilter.mightContain(
+          "projects/test-project/databases/(default)/documents/"
+              + key.getPath().canonicalString())) {
+        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

@@ -41,13 +41,13 @@ public class BloomFilterTest {
       "src/test/resources/bloom_filter_golden_test_data/";
 
   @Test
-  public void instantiateEmptyBloomFilter() {
+  public void instantiateEmptyBloomFilter() throws BloomFilterException {
     BloomFilter bloomFilter = new BloomFilter(new byte[0], 0, 0);
     assertEquals(bloomFilter.getBitCount(), 0);
   }
 
   @Test
-  public void instantiateNonEmptyBloomFilter() {
+  public void instantiateNonEmptyBloomFilter() throws BloomFilterException {
     {
       BloomFilter bloomFilter1 = new BloomFilter(new byte[] {1}, 0, 1);
       assertEquals(bloomFilter1.getBitCount(), 8);
@@ -73,75 +73,74 @@ 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");
   }
 
   @Test
-  public void mightContainCanProcessNonStandardCharacters() {
+  public void mightContainCanProcessNonStandardCharacters() throws BloomFilterException {
     // A non-empty BloomFilter object with 1 insertion : "ÀÒ∑"
     BloomFilter bloomFilter = new BloomFilter(new byte[] {(byte) 237, 5}, 5, 8);
     assertTrue(bloomFilter.mightContain("ÀÒ∑"));
     assertFalse(bloomFilter.mightContain("Ò∑À"));
   }
 
   @Test
-  public void mightContainOnEmptyBloomFilterShouldReturnFalse() {
+  public void mightContainOnEmptyBloomFilterShouldReturnFalse() throws BloomFilterException {
     BloomFilter bloomFilter = new BloomFilter(new byte[0], 0, 0);
     assertFalse(bloomFilter.mightContain(""));
     assertFalse(bloomFilter.mightContain("a"));
   }
 
   @Test
-  public void mightContainWithEmptyStringMightReturnFalsePositiveResult() {
+  public void mightContainWithEmptyStringMightReturnFalsePositiveResult()
+      throws BloomFilterException {
     {
       BloomFilter bloomFilter1 = new BloomFilter(new byte[] {1}, 1, 1);
       assertFalse(bloomFilter1.mightContain(""));
@@ -153,7 +152,7 @@ public void mightContainWithEmptyStringMightReturnFalsePositiveResult() {
   }
 
   @Test
-  public void bloomFilterToString() {
+  public void bloomFilterToString() throws BloomFilterException {
     {
       BloomFilter emptyBloomFilter = new BloomFilter(new byte[0], 0, 0);
       assertEquals(emptyBloomFilter.toString(), "BloomFilter{hashCount=0, size=0, bitmap=\"\"}");
@@ -179,7 +178,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);