Update test 'bloom filter should correctly encode special unicode characters' in query.test.ts (#7413)

Author: dconeybe

packages/firestore/src/remote/watch_change.ts

@@ -433,15 +433,15 @@ export class WatchChangeAggregator {
         // raise a snapshot with `isFromCache:true`.
         if (currentSize !== expectedCount) {
           // Apply bloom filter to identify and mark removed documents.
-          const status = this.applyBloomFilter(watchChange, currentSize);
+          const applyResult = this.applyBloomFilter(watchChange, currentSize);
 
-          if (status !== BloomFilterApplicationStatus.Success) {
+          if (applyResult.status !== BloomFilterApplicationStatus.Success) {
             // If bloom filter application fails, we reset the mapping and
             // trigger re-run of the query.
             this.resetTarget(targetId);
 
             const purpose: TargetPurpose =
-              status === BloomFilterApplicationStatus.FalsePositive
+              applyResult.status === BloomFilterApplicationStatus.FalsePositive
                 ? TargetPurpose.ExistenceFilterMismatchBloom
                 : TargetPurpose.ExistenceFilterMismatch;
             this.pendingTargetResets = this.pendingTargetResets.insert(
@@ -451,7 +451,8 @@ export class WatchChangeAggregator {
           }
           TestingHooks.instance?.notifyOnExistenceFilterMismatch(
             createExistenceFilterMismatchInfoForTestingHooks(
-              status,
+              applyResult.status,
+              applyResult.bloomFilterMightContain ?? null,
               currentSize,
               watchChange.existenceFilter
             )
@@ -468,12 +469,15 @@ export class WatchChangeAggregator {
   private applyBloomFilter(
     watchChange: ExistenceFilterChange,
     currentCount: number
-  ): BloomFilterApplicationStatus {
+  ): {
+    status: BloomFilterApplicationStatus;
+    bloomFilterMightContain?: (documentPath: string) => boolean;
+  } {
     const { unchangedNames, count: expectedCount } =
       watchChange.existenceFilter;
 
     if (!unchangedNames || !unchangedNames.bits) {
-      return BloomFilterApplicationStatus.Skipped;
+      return { status: BloomFilterApplicationStatus.Skipped };
     }
 
     const {
@@ -491,7 +495,7 @@ export class WatchChangeAggregator {
             err.message +
             '); ignoring the bloom filter and falling back to full re-query.'
         );
-        return BloomFilterApplicationStatus.Skipped;
+        return { status: BloomFilterApplicationStatus.Skipped };
       } else {
         throw err;
       }
@@ -507,23 +511,31 @@ export class WatchChangeAggregator {
       } else {
         logWarn('Applying bloom filter failed: ', err);
       }
-      return BloomFilterApplicationStatus.Skipped;
+      return { status: BloomFilterApplicationStatus.Skipped };
     }
 
+    const bloomFilterMightContain = (documentPath: string): boolean => {
+      const databaseId = this.metadataProvider.getDatabaseId();
+      return bloomFilter.mightContain(
+        `projects/${databaseId.projectId}/databases/${databaseId.database}` +
+          `/documents/${documentPath}`
+      );
+    };
+
     if (bloomFilter.bitCount === 0) {
-      return BloomFilterApplicationStatus.Skipped;
+      return { status: BloomFilterApplicationStatus.Skipped };
     }
 
     const removedDocumentCount = this.filterRemovedDocuments(
       watchChange.targetId,
-      bloomFilter
+      bloomFilterMightContain
     );
 
-    if (expectedCount !== currentCount - removedDocumentCount) {
-      return BloomFilterApplicationStatus.FalsePositive;
-    }
-
-    return BloomFilterApplicationStatus.Success;
+    const status =
+      expectedCount === currentCount - removedDocumentCount
+        ? BloomFilterApplicationStatus.Success
+        : BloomFilterApplicationStatus.FalsePositive;
+    return { status, bloomFilterMightContain };
   }
 
   /**
@@ -532,18 +544,13 @@ export class WatchChangeAggregator {
    */
   private filterRemovedDocuments(
     targetId: number,
-    bloomFilter: BloomFilter
+    bloomFilterMightContain: (documentPath: string) => boolean
   ): number {
     const existingKeys = this.metadataProvider.getRemoteKeysForTarget(targetId);
     let removalCount = 0;
 
     existingKeys.forEach(key => {
-      const databaseId = this.metadataProvider.getDatabaseId();
-      const documentPath = `projects/${databaseId.projectId}/databases/${
-        databaseId.database
-      }/documents/${key.path.canonicalString()}`;
-
-      if (!bloomFilter.mightContain(documentPath)) {
+      if (!bloomFilterMightContain(key.path.canonicalString())) {
         this.removeDocumentFromTarget(targetId, key, /*updatedDocument=*/ null);
         removalCount++;
       }
@@ -829,6 +836,7 @@ function snapshotChangesMap(): SortedMap<DocumentKey, ChangeType> {
 
 function createExistenceFilterMismatchInfoForTestingHooks(
   status: BloomFilterApplicationStatus,
+  bloomFilterMightContain: null | ((documentPath: string) => boolean),
   localCacheCount: number,
   existenceFilter: ExistenceFilter
 ): TestingHooksExistenceFilterMismatchInfo {
@@ -845,6 +853,9 @@ function createExistenceFilterMismatchInfoForTestingHooks(
       bitmapLength: unchangedNames?.bits?.bitmap?.length ?? 0,
       padding: unchangedNames?.bits?.padding ?? 0
     };
+    if (bloomFilterMightContain) {
+      result.bloomFilter.mightContain = bloomFilterMightContain;
+    }
   }
 
   return result;

packages/firestore/src/util/testing_hooks.ts

@@ -124,6 +124,21 @@ export interface ExistenceFilterMismatchInfo {
 
     /** The number of bits of padding in the last byte of the bloom filter. */
     padding: number;
+
+    /**
+     * Check if the given document path is contained in the bloom filter.
+     *
+     * The "path" of a document can be retrieved from the
+     * `DocumentReference.path` property.
+     *
+     * Note that due to the probabilistic nature of a bloom filter, it is
+     * possible that false positives may occur; that is, this function may
+     * return `true` even though the given string is not in the bloom filter.
+     *
+     * This property is "optional"; if it is undefined then parsing the bloom
+     * filter failed.
+     */
+    mightContain?(documentPath: string): boolean;
   };
 }
 

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

@@ -2243,99 +2243,98 @@ apiDescribe('Queries', persistence => {
         return map;
       }, {} as { [key: string]: DocumentData });
 
-      // Ensure that the local cache is configured to use LRU garbage
-      // collection (rather than eager garbage collection) so that the resume
-      // token and document data does not get prematurely evicted.
+      // Ensure that the local cache is configured to use LRU garbage collection
+      // (rather than eager garbage collection) so that the resume token and
+      // document data does not get prematurely evicted.
       const lruPersistence = persistence.toLruGc();
 
-      return withRetry(async attemptNumber => {
-        return withTestCollection(
-          lruPersistence,
-          testDocs,
-          async (coll, db) => {
-            // Run a query to populate the local cache with documents that have
-            // names with complex Unicode characters.
-            const snapshot1 = await getDocs(coll);
-            const snapshot1DocumentIds = snapshot1.docs.map(
-              documentSnapshot => documentSnapshot.id
-            );
-            expect(
-              snapshot1DocumentIds,
-              'snapshot1DocumentIds'
-            ).to.have.members(testDocIds);
-
-            // Delete one of the documents so that the next call to getDocs() will
-            // experience an existence filter mismatch. Do this deletion in a
-            // transaction, rather than using deleteDoc(), to avoid affecting the
-            // local cache.
-            await runTransaction(db, async txn => {
-              const snapshotOfDocumentToDelete = await txn.get(
-                doc(coll, 'DocumentToDelete')
-              );
-              expect(
-                snapshotOfDocumentToDelete.exists(),
-                'snapshotOfDocumentToDelete.exists()'
-              ).to.be.true;
-              txn.delete(snapshotOfDocumentToDelete.ref);
-            });
-
-            // 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));
-
-            // Resume the query and save the resulting snapshot for verification.
-            // Use some internal testing hooks to "capture" the existence filter
-            // mismatches.
-            const [existenceFilterMismatches, snapshot2] =
-              await captureExistenceFilterMismatches(() => getDocs(coll));
-            const snapshot2DocumentIds = snapshot2.docs.map(
-              documentSnapshot => documentSnapshot.id
-            );
-            const testDocIdsMinusDeletedDocId = testDocIds.filter(
-              documentId => documentId !== 'DocumentToDelete'
-            );
-            expect(
-              snapshot2DocumentIds,
-              'snapshot2DocumentIds'
-            ).to.have.members(testDocIdsMinusDeletedDocId);
-
-            // Verify that Watch sent an existence filter with the correct counts.
-            expect(
-              existenceFilterMismatches,
-              'existenceFilterMismatches'
-            ).to.have.length(1);
-            const { localCacheCount, existenceFilterCount, bloomFilter } =
-              existenceFilterMismatches[0];
-            expect(localCacheCount, 'localCacheCount').to.equal(
-              testDocIds.length
-            );
-            expect(existenceFilterCount, 'existenceFilterCount').to.equal(
-              testDocIds.length - 1
-            );
+      return withTestCollection(lruPersistence, testDocs, async (coll, db) => {
+        // Run a query to populate the local cache with documents that have
+        // names with complex Unicode characters.
+        const snapshot1 = await getDocs(coll);
+        const snapshot1DocumentIds = snapshot1.docs.map(
+          documentSnapshot => documentSnapshot.id
+        );
+        expect(snapshot1DocumentIds, 'snapshot1DocumentIds').to.have.members(
+          testDocIds
+        );
 
-            // Verify that Watch sent a valid bloom filter.
-            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');
-            }
-
-            // Verify that the bloom filter was successfully used to avert a full
-            // requery. If a false positive occurred, which is statistically rare,
-            // but technically possible, then retry the entire test.
-            if (attemptNumber === 1 && !bloomFilter.applied) {
-              throw new RetryError();
-            }
-
-            expect(
-              bloomFilter.applied,
-              `bloomFilter.applied with attemptNumber=${attemptNumber}`
-            ).to.be.true;
-          }
+        // Delete one of the documents so that the next call to getDocs() will
+        // experience an existence filter mismatch. Do this deletion in a
+        // transaction, rather than using deleteDoc(), to avoid affecting the
+        // local cache.
+        await runTransaction(db, async txn => {
+          const snapshotOfDocumentToDelete = await txn.get(
+            doc(coll, 'DocumentToDelete')
+          );
+          expect(
+            snapshotOfDocumentToDelete.exists(),
+            'snapshotOfDocumentToDelete.exists()'
+          ).to.be.true;
+          txn.delete(snapshotOfDocumentToDelete.ref);
+        });
+
+        // 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));
+
+        // Resume the query and save the resulting snapshot for verification.
+        // Use some internal testing hooks to "capture" the existence filter
+        // mismatches.
+        const [existenceFilterMismatches, snapshot2] =
+          await captureExistenceFilterMismatches(() => getDocs(coll));
+        const snapshot2DocumentIds = snapshot2.docs.map(
+          documentSnapshot => documentSnapshot.id
+        );
+        const testDocIdsMinusDeletedDocId = testDocIds.filter(
+          documentId => documentId !== 'DocumentToDelete'
+        );
+        expect(snapshot2DocumentIds, 'snapshot2DocumentIds').to.have.members(
+          testDocIdsMinusDeletedDocId
+        );
+
+        // Verify that Watch sent an existence filter with the correct counts.
+        expect(
+          existenceFilterMismatches,
+          'existenceFilterMismatches'
+        ).to.have.length(1);
+        const existenceFilterMismatch = existenceFilterMismatches[0];
+        expect(
+          existenceFilterMismatch.localCacheCount,
+          'localCacheCount'
+        ).to.equal(testDocIds.length);
+        expect(
+          existenceFilterMismatch.existenceFilterCount,
+          'existenceFilterCount'
+        ).to.equal(testDocIds.length - 1);
+
+        // Verify that we got a bloom filter from Watch.
+        const bloomFilter = existenceFilterMismatch.bloomFilter!;
+        expect(bloomFilter?.mightContain, 'bloomFilter.mightContain').to.not.be
+          .undefined;
+
+        // The bloom filter application should statistically be successful
+        // almost every time; the _only_ time when it would _not_ be successful
+        // is if there is a false positive when testing for 'DocumentToDelete'
+        // in the bloom filter. So verify that the bloom filter application is
+        // successful, unless there was a false positive.
+        const isFalsePositive = bloomFilter.mightContain!(
+          `${coll.path}/DocumentToDelete`
+        );
+        expect(bloomFilter.applied, 'bloomFilter.applied').to.equal(
+          !isFalsePositive
         );
+
+        // Verify that the bloom filter contains the document paths with complex
+        // Unicode characters.
+        for (const testDocId of testDocIdsMinusDeletedDocId) {
+          const testDocPath = `${coll.path}/${testDocId}`;
+          expect(
+            bloomFilter.mightContain!(testDocPath),
+            `bloomFilter.mightContain('${testDocPath}')`
+          ).to.be.true;
+        }
       });
     }
   ).timeout('90s');

packages/firestore/test/integration/util/testing_hooks_util.ts

@@ -69,5 +69,6 @@ export interface ExistenceFilterMismatchInfo {
     hashCount: number;
     bitmapLength: number;
     padding: number;
+    mightContain?(documentPath: string): boolean;
   };
 }