Update the integration test to verify that bloom filter averted full requery

.github/workflows/check-changeset.yml

@@ -2,8 +2,8 @@ name: Check Changeset
 
 on:
   pull_request:
-    branches-ignore:
-      - release
+    branches:
+      - master
 
 env:
   GITHUB_PULL_REQUEST_HEAD_SHA: ${{ github.event.pull_request.head.sha }}

integration/firestore/gulpfile.js

@@ -45,6 +45,7 @@ function copyTests() {
         testBase + '/integration/util/events_accumulator.ts',
         testBase + '/integration/util/helpers.ts',
         testBase + '/integration/util/settings.ts',
+        testBase + '/integration/util/testing_hooks_util.ts',
         testBase + '/util/equality_matcher.ts',
         testBase + '/util/promise.ts'
       ],

packages/firestore/src/api.ts

@@ -196,3 +196,4 @@ export type { ByteString as _ByteString } from './util/byte_string';
 export { logWarn as _logWarn } from './util/log';
 export { EmptyAuthCredentialsProvider as _EmptyAuthCredentialsProvider } from './api/credentials';
 export { EmptyAppCheckTokenProvider as _EmptyAppCheckTokenProvider } from './api/credentials';
+export { TestingHooks as _TestingHooks } from './util/testing_hooks';

packages/firestore/src/remote/watch_change.ts

@@ -37,6 +37,10 @@ import { logDebug, logWarn } from '../util/log';
 import { primitiveComparator } from '../util/misc';
 import { SortedMap } from '../util/sorted_map';
 import { SortedSet } from '../util/sorted_set';
+import {
+  ExistenceFilterMismatchInfo as TestingHooksExistenceFilterMismatchInfo,
+  TestingHooks
+} from '../util/testing_hooks';
 
 import { BloomFilter, BloomFilterError } from './bloom_filter';
 import { ExistenceFilter } from './existence_filter';
@@ -432,6 +436,13 @@ export class WatchChangeAggregator {
             this.resetTarget(targetId);
             this.pendingTargetResets = this.pendingTargetResets.add(targetId);
           }
+          TestingHooks.instance?.notifyOnExistenceFilterMismatch(
+            createExistenceFilterMismatchInfoForTestingHooks(
+              bloomFilterApplied,
+              currentSize,
+              watchChange.existenceFilter
+            )
+          );
         }
       }
     }
@@ -789,3 +800,26 @@ function documentTargetMap(): SortedMap<DocumentKey, SortedSet<TargetId>> {
 function snapshotChangesMap(): SortedMap<DocumentKey, ChangeType> {
   return new SortedMap<DocumentKey, ChangeType>(DocumentKey.comparator);
 }
+
+function createExistenceFilterMismatchInfoForTestingHooks(
+  bloomFilterApplied: boolean,
+  localCacheCount: number,
+  existenceFilter: ExistenceFilter
+): TestingHooksExistenceFilterMismatchInfo {
+  const result: TestingHooksExistenceFilterMismatchInfo = {
+    localCacheCount,
+    existenceFilterCount: existenceFilter.count
+  };
+
+  const unchangedNames = existenceFilter.unchangedNames;
+  if (unchangedNames) {
+    result.bloomFilter = {
+      applied: bloomFilterApplied,
+      hashCount: unchangedNames?.hashCount ?? 0,
+      bitmapLength: unchangedNames?.bits?.bitmap?.length ?? 0,
+      padding: unchangedNames?.bits?.padding ?? 0
+    };
+  }
+
+  return result;
+}

packages/firestore/src/util/testing_hooks.ts

@@ -0,0 +1,139 @@
+/**
+ * @license
+ * 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.
+ */
+
+/**
+ * Manages "testing hooks", hooks into the internals of the SDK to verify
+ * internal state and events during integration tests. Do not use this class
+ * except for testing purposes.
+ *
+ * There are two ways to retrieve the global singleton instance of this class:
+ * 1. The `instance` property, which returns null if the global singleton
+ *      instance has not been created. Use this property if the caller should
+ *      "do nothing" if there are no testing hooks registered, such as when
+ *      delivering an event to notify registered callbacks.
+ * 2. The `getOrCreateInstance()` method, which creates the global singleton
+ *      instance if it has not been created. Use this method if the instance is
+ *      needed to, for example, register a callback.
+ *
+ * @internal
+ */
+export class TestingHooks {
+  private readonly onExistenceFilterMismatchCallbacks = new Map<
+    Symbol,
+    ExistenceFilterMismatchCallback
+  >();
+
+  private constructor() {}
+
+  /**
+   * Returns the singleton instance of this class, or null if it has not been
+   * initialized.
+   */
+  static get instance(): TestingHooks | null {
+    return gTestingHooksSingletonInstance;
+  }
+
+  /**
+   * Returns the singleton instance of this class, creating it if is has never
+   * been created before.
+   */
+  static getOrCreateInstance(): TestingHooks {
+    if (gTestingHooksSingletonInstance === null) {
+      gTestingHooksSingletonInstance = new TestingHooks();
+    }
+    return gTestingHooksSingletonInstance;
+  }
+
+  /**
+   * Registers a callback to be notified when an existence filter mismatch
+   * occurs in the Watch listen stream.
+   *
+   * The relative order in which callbacks are notified is unspecified; do not
+   * rely on any particular ordering. If a given callback is registered multiple
+   * times then it will be notified multiple times, once per registration.
+   *
+   * @param callback the callback to invoke upon existence filter mismatch.
+   *
+   * @return a function that, when called, unregisters the given callback; only
+   * the first invocation of the returned function does anything; all subsequent
+   * invocations do nothing.
+   */
+  onExistenceFilterMismatch(
+    callback: ExistenceFilterMismatchCallback
+  ): () => void {
+    const key = Symbol();
+    this.onExistenceFilterMismatchCallbacks.set(key, callback);
+    return () => this.onExistenceFilterMismatchCallbacks.delete(key);
+  }
+
+  /**
+   * Invokes all currently-registered `onExistenceFilterMismatch` callbacks.
+   * @param info Information about the existence filter mismatch.
+   */
+  notifyOnExistenceFilterMismatch(info: ExistenceFilterMismatchInfo): void {
+    this.onExistenceFilterMismatchCallbacks.forEach(callback => callback(info));
+  }
+}
+
+/**
+ * Information about an existence filter mismatch, as specified to callbacks
+ * registered with `TestingUtils.onExistenceFilterMismatch()`.
+ */
+export interface ExistenceFilterMismatchInfo {
+  /** The number of documents that matched the query in the local cache. */
+  localCacheCount: number;
+
+  /**
+   * The number of documents that matched the query on the server, as specified
+   * in the ExistenceFilter message's `count` field.
+   */
+  existenceFilterCount: number;
+
+  /**
+   * Information about the bloom filter provided by Watch in the ExistenceFilter
+   * message's `unchangedNames` field. If this property is omitted or undefined
+   * then that means that Watch did _not_ provide a bloom filter.
+   */
+  bloomFilter?: {
+    /**
+     * Whether a full requery was averted by using the bloom filter. If false,
+     * then something happened, such as a false positive, to prevent using the
+     * bloom filter to avoid a full requery.
+     */
+    applied: boolean;
+
+    /** The number of hash functions used in the bloom filter. */
+    hashCount: number;
+
+    /** The number of bytes in the bloom filter's bitmask. */
+    bitmapLength: number;
+
+    /** The number of bits of padding in the last byte of the bloom filter. */
+    padding: number;
+  };
+}
+
+/**
+ * The signature of callbacks registered with
+ * `TestingUtils.onExistenceFilterMismatch()`.
+ */
+export type ExistenceFilterMismatchCallback = (
+  info: ExistenceFilterMismatchInfo
+) => void;
+
+/** The global singleton instance of `TestingHooks`. */
+let gTestingHooksSingletonInstance: TestingHooks | null = null;

packages/firestore/test/integration/api/query.test.ts

@@ -26,6 +26,7 @@ import {
   Bytes,
   collection,
   collectionGroup,
+  CollectionReference,
   deleteDoc,
   disableNetwork,
   doc,
@@ -35,6 +36,7 @@ import {
   enableNetwork,
   endAt,
   endBefore,
+  Firestore,
   GeoPoint,
   getDocs,
   getDocsFromCache,
@@ -65,7 +67,8 @@ import {
   withTestCollection,
   withTestDb
 } from '../util/helpers';
-import { USE_EMULATOR } from '../util/settings';
+import { USE_EMULATOR, TARGET_BACKEND } from '../util/settings';
+import { captureExistenceFilterMismatches } from '../util/testing_hooks_util';
 
 apiDescribe('Queries', (persistence: boolean) => {
   addEqualityMatcher();
@@ -2061,34 +2064,133 @@ apiDescribe('Queries', (persistence: boolean) => {
     });
   });
 
-  // TODO(Mila): Skip the test when using emulator as there is a bug related to 
-  // sending existence filter in response: b/270731363. Remove the condition  
-  // here once the bug is resolved.
+  // TODO(b/270731363): Re-enable this test to run against the Firestore
+  // emulator once the bug where an existence filter fails to be sent when a
+  // query is resumed is fixed.
   // eslint-disable-next-line no-restricted-properties
   (USE_EMULATOR ? it.skip : it)(
-    'resuming a query should remove deleted documents indicated by existence filter',
-    () => {
+    'resuming a query should use bloom filter to avoid full requery',
+    async () => {
+      // Create 100 documents in a new collection.
       const testDocs: { [key: string]: object } = {};
       for (let i = 1; i <= 100; i++) {
         testDocs['doc' + i] = { key: i };
       }
-      return withTestCollection(persistence, testDocs, async (coll, db) => {
+
+      // The function that runs a single iteration of the test.
+      // Below this definition, there is a "while" loop that calls this
+      // function potentially multiple times.
+      const runTestIteration = async (
+        coll: CollectionReference,
+        db: Firestore
+      ): Promise<'retry' | 'passed'> => {
+        // Run a query to populate the local cache with the 100 documents
+        // and a resume token.
         const snapshot1 = await getDocs(coll);
-        expect(snapshot1.size).to.equal(100);
-        // Delete 50 docs in transaction so that it doesn't affect local cache.
+        expect(snapshot1.size, 'snapshot1.size').to.equal(100);
+
+        // Delete 50 of the 100 documents. Do this in a transaction, rather
+        // than deleteDoc(), to avoid affecting the local cache.
         await runTransaction(db, async txn => {
           for (let i = 1; i <= 50; i++) {
             txn.delete(doc(coll, 'doc' + i));
           }
         });
-        // Wait 10 seconds, during which Watch will stop tracking the query
-        // and will send an existence filter rather than "delete" events.
+
+        // Wait for 10 seconds, during which Watch will stop tracking the
+        // query and will send an existence filter rather than "delete"
+        // events when the query is resumed.
         await new Promise(resolve => setTimeout(resolve, 10000));
-        const snapshot2 = await getDocs(coll);
-        expect(snapshot2.size).to.equal(50);
-      });
+
+        // Resume the query and expect to get a snapshot with the 50
+        // remaining documents. Use some internal testing hooks to "capture"
+        // the existence filter mismatches to later verify that Watch sent a
+        // bloom filter, and it was used to avert a full requery.
+        const existenceFilterMismatches =
+          await captureExistenceFilterMismatches(async () => {
+            const snapshot2 = await getDocs(coll);
+            expect(snapshot2.size, 'snapshot2.size').to.equal(50);
+          });
+
+        // Skip the verification of the existence filter mismatch when
+        // persistence is disabled because without persistence there is no
+        // resume token specified in the subsequent call to getDocs(), and,
+        // therefore, Watch will _not_ send an existence filter.
+        if (!persistence) {
+          return 'passed';
+        }
+
+        // Skip the verification of the existence filter mismatch when
+        // testing against the Firestore emulator because the Firestore
+        // emulator does not include the `unchanged_names` bloom filter when
+        // it sends ExistenceFilter messages. Some day the emulator _may_
+        // implement this logic, at which time this short-circuit can be
+        // removed.
+        if (USE_EMULATOR) {
+          return 'passed';
+        }
+
+        // Verify that upon resuming the query that Watch sent an existence
+        // filter that included a bloom filter, and that the bloom filter
+        // was successfully used to avoid a full requery.
+        // TODO(b/271949433) Remove this check for "nightly" once the bloom
+        // filter logic is deployed to production, circa May 2023.
+        if (TARGET_BACKEND === 'nightly') {
+          expect(
+            existenceFilterMismatches,
+            'existenceFilterMismatches'
+          ).to.have.length(1);
+          const { localCacheCount, existenceFilterCount, bloomFilter } =
+            existenceFilterMismatches[0];
+
+          expect(localCacheCount, 'localCacheCount').to.equal(100);
+          expect(existenceFilterCount, 'existenceFilterCount').to.equal(50);
+          if (!bloomFilter) {
+            expect.fail(
+              'The existence filter should have specified ' +
+                'a bloom filter in its `unchanged_names` field.'
+            );
+            throw new Error('should never get here');
+          }
+
+          expect(bloomFilter.hashCount, 'bloomFilter.hashCount').to.be.above(0);
+          expect(
+            bloomFilter.bitmapLength,
+            'bloomFilter.bitmapLength'
+          ).to.be.above(0);
+          expect(bloomFilter.padding, 'bloomFilterPadding').to.be.above(0);
+          expect(bloomFilter.padding, 'bloomFilterPadding').to.be.below(8);
+
+          // Retry the entire test if a bloom filter false positive occurs.
+          // Although statistically rare, false positives are expected to
+          // happen occasionally. When a false positive _does_ happen, just
+          // retry the test with a different set of documents. If that retry
+          // _also_ experiences a false positive, then fail the test because
+          // that is so improbable that something must have gone wrong.
+          if (attemptNumber > 1 && !bloomFilter.applied) {
+            return 'retry';
+          }
+          expect(bloomFilter.applied, 'bloomFilter.applied').to.be.true;
+        }
+
+        return 'passed';
+      };
+
+      // Run the test
+      let attemptNumber = 0;
+      while (true) {
+        attemptNumber++;
+        const iterationResult = await withTestCollection(
+          persistence,
+          testDocs,
+          runTestIteration
+        );
+        if (iterationResult === 'passed') {
+          break;
+        }
+      }
     }
-  ).timeout('20s');
+  ).timeout('90s');
 });
 
 function verifyDocumentChange<T>(