Add test cases for existence filter with updated/added documents (#7457)

milaGGL · web-flow · commit 74fc4cf9c0cd · 2023-09-05T10:13:18.000-04:00
diff --git a/packages/firestore/test/integration/api/query.test.ts b/packages/firestore/test/integration/api/query.test.ts
@@ -2203,6 +2203,175 @@ apiDescribe('Queries', persistence => {
     }
   ).timeout('90s');
 
+  // TODO(b/291365820): Stop skipping this test when running against the
+  // Firestore emulator once the emulator is improved to include a bloom filter
+  // in the existence filter messages that it sends.
+  // eslint-disable-next-line no-restricted-properties
+  (USE_EMULATOR ? it.skip : it)(
+    'bloom filter should avert a full re-query when documents were added, ' +
+      'deleted, removed, updated, and unchanged since the resume token',
+    async () => {
+      // Prepare the names and contents of the 20 documents to create.
+      const testDocs: { [key: string]: object } = {};
+      for (let i = 0; i < 20; i++) {
+        testDocs['doc' + (1000 + i)] = {
+          key: 42,
+          removed: false
+        };
+      }
+
+      // 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 => {
+          // Run a query to populate the local cache with the 20 documents
+          // and a resume token.
+          const snapshot1 = await getDocs(
+            query(coll, where('removed', '==', false))
+          );
+          expect(snapshot1.size, 'snapshot1.size').to.equal(20);
+          const createdDocuments = snapshot1.docs.map(snapshot => snapshot.ref);
+
+          // Out of the 20 existing documents, leave 5 docs untouched, delete 5 docs,
+          // remove 5 docs, update 5 docs, and add 15 new docs.
+          const deletedDocumentIds = new Set<string>();
+          const removedDocumentIds = new Set<string>();
+          const updatedDocumentIds = new Set<string>();
+          const addedDocumentIds: string[] = [];
+
+          // Use a different Firestore instance to avoid affecting the local cache.
+          await withTestDb(PERSISTENCE_MODE_UNSPECIFIED, async db2 => {
+            const batch = writeBatch(db2);
+
+            for (let i = 0; i < createdDocuments.length; i += 4) {
+              const documentToDelete = doc(db2, createdDocuments[i].path);
+              batch.delete(documentToDelete);
+              deletedDocumentIds.add(documentToDelete.id);
+            }
+            expect(deletedDocumentIds.size).to.equal(5);
+
+            // Update 5 documents to no longer match the query.
+            for (let i = 1; i < createdDocuments.length; i += 4) {
+              const documentToModify = doc(db2, createdDocuments[i].path);
+              batch.update(documentToModify, {
+                removed: true
+              });
+              removedDocumentIds.add(documentToModify.id);
+            }
+            expect(removedDocumentIds.size).to.equal(5);
+
+            // Update 5 documents, but ensure they still match the query.
+            for (let i = 2; i < createdDocuments.length; i += 4) {
+              const documentToModify = doc(db2, createdDocuments[i].path);
+              batch.update(documentToModify, {
+                key: 43
+              });
+              updatedDocumentIds.add(documentToModify.id);
+            }
+            expect(updatedDocumentIds.size).to.equal(5);
+
+            for (let i = 0; i < 15; i += 1) {
+              const documentToAdd = doc(
+                db2,
+                coll.path + '/newDoc' + (1000 + i)
+              );
+              batch.set(documentToAdd, {
+                key: 42,
+                removed: false
+              });
+              addedDocumentIds.push(documentToAdd.id);
+            }
+
+            // Ensure the sets above are disjoint.
+            const mergedSet = new Set<string>();
+            [
+              deletedDocumentIds,
+              removedDocumentIds,
+              updatedDocumentIds,
+              addedDocumentIds
+            ].forEach(set => {
+              set.forEach(documentId => mergedSet.add(documentId));
+            });
+            expect(mergedSet.size).to.equal(30);
+
+            await batch.commit();
+          });
+
+          // 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 to verify that Watch sent a bloom
+          // filter, and it was used to avert a full requery.
+          const [existenceFilterMismatches, snapshot2] =
+            await captureExistenceFilterMismatches(() =>
+              getDocs(query(coll, where('removed', '==', false)))
+            );
+
+          // Verify that the snapshot from the resumed query contains the
+          // expected documents; that is, 10 existing documents that still
+          // match the query, and 15 documents that are newly added.
+          const actualDocumentIds = snapshot2.docs
+            .map(documentSnapshot => documentSnapshot.ref.id)
+            .sort();
+          const expectedDocumentIds = createdDocuments
+            .map(documentRef => documentRef.id)
+            .filter(documentId => !deletedDocumentIds.has(documentId))
+            .filter(documentId => !removedDocumentIds.has(documentId))
+            .concat(addedDocumentIds)
+            .sort();
+
+          expect(actualDocumentIds, 'snapshot2.docs').to.deep.equal(
+            expectedDocumentIds
+          );
+          expect(actualDocumentIds.length).to.equal(25);
+
+          // Verify that Watch sent an existence filter with the correct
+          // counts when the query was resumed.
+          expect(
+            existenceFilterMismatches,
+            'existenceFilterMismatches'
+          ).to.have.length(1);
+          const { localCacheCount, existenceFilterCount, bloomFilter } =
+            existenceFilterMismatches[0];
+          expect(localCacheCount, 'localCacheCount').to.equal(35);
+          expect(existenceFilterCount, 'existenceFilterCount').to.equal(25);
+
+          // 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 then retry the entire
+          // test. 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) {
+            throw new RetryError();
+          }
+
+          expect(
+            bloomFilter.applied,
+            `bloomFilter.applied with attemptNumber=${attemptNumber}`
+          ).to.be.true;
+        });
+      });
+    }
+  ).timeout('90s');
+
   // TODO(b/291365820): Stop skipping this test when running against the
   // Firestore emulator once the emulator is improved to include a bloom filter
   // in the existence filter messages that it sends.
diff --git a/packages/firestore/test/unit/specs/existence_filter_spec.test.ts b/packages/firestore/test/unit/specs/existence_filter_spec.test.ts
@@ -273,7 +273,8 @@ describeSpec('Existence Filters:', [], () => {
   );
 
   /**
-   * Test existence filter with bloom filter.
+   * Test existence filter with bloom filter. Existence filters below is sent mid-stream for
+   * testing simplicity.
    */
   specTest(
     'Full re-query is skipped when bloom filter can identify documents deleted',
@@ -626,9 +627,185 @@ describeSpec('Existence Filters:', [], () => {
         // Doc0 to doc49 are deleted in the next sync.
         .watchFilters([query1], docKeys.slice(0, 50), bloomFilterProto)
         .watchSnapshots(2000)
-        // BloomFilter correctly identifies docs that deleted, skip full query.
+        // Bloom Filter correctly identifies docs that deleted, skips full query.
         .expectEvents(query1, { fromCache: true })
         .expectLimboDocs(...docKeys.slice(50))
     );
   });
+
+  specTest(
+    'Resume a query with bloom filter when there is no document changes',
+    [],
+    () => {
+      const query1 = query('collection');
+      const docA = doc('collection/a', 1000, { v: 1 });
+      const bloomFilterProto = generateBloomFilterProto({
+        contains: [docA],
+        notContains: []
+      });
+      return (
+        spec()
+          .userListens(query1)
+          .watchAcksFull(query1, 1000, docA)
+          .expectEvents(query1, { added: [docA] })
+          .disableNetwork()
+          .expectEvents(query1, { fromCache: true })
+          .enableNetwork()
+          .restoreListen(query1, 'resume-token-1000')
+          .watchAcks(query1)
+          // Nothing happened while this client was disconnected.
+          // Bloom Filter includes docA as there are no changes since the resume token.
+          .watchFilters([query1], [docA.key], bloomFilterProto)
+          // Expected count equals to documents in cache. Existence Filter matches.
+          .watchCurrents(query1, 'resume-token-2000')
+          .watchSnapshots(2000)
+          .expectEvents(query1, { fromCache: false })
+      );
+    }
+  );
+
+  specTest(
+    'Resume a query with bloom filter when new documents are added',
+    [],
+    () => {
+      const query1 = query('collection');
+      const docA = doc('collection/a', 1000, { v: 1 });
+      const docB = doc('collection/b', 1000, { v: 2 });
+      const bloomFilterProto = generateBloomFilterProto({
+        contains: [docA, docB],
+        notContains: []
+      });
+      return (
+        spec()
+          .userListens(query1)
+          .watchAcksFull(query1, 1000, docA)
+          .expectEvents(query1, { added: [docA] })
+          .disableNetwork()
+          .expectEvents(query1, { fromCache: true })
+          .enableNetwork()
+          .restoreListen(query1, 'resume-token-1000')
+          .watchAcks(query1)
+          // While this client was disconnected, another client added docB.
+          .watchSends({ affects: [query1] }, docB)
+          // Bloom Filter includes all the documents that match the query, both
+          // those that haven't changed since the resume token and those newly added.
+          .watchFilters([query1], [docA.key, docB.key], bloomFilterProto)
+          // Expected count equals to documents in cache. Existence Filter matches.
+          .watchCurrents(query1, 'resume-token-2000')
+          .watchSnapshots(2000)
+          .expectEvents(query1, { added: [docB], fromCache: false })
+      );
+    }
+  );
+
+  specTest(
+    'Resume a query with bloom filter when existing docs are updated',
+    [],
+    () => {
+      const query1 = query('collection', filter('v', '>=', 1));
+      const docA = doc('collection/a', 1000, { v: 1 });
+      const docB = doc('collection/b', 1000, { v: 1 });
+      const updatedDocB = doc('collection/b', 1000, { v: 2 });
+
+      const bloomFilterProto = generateBloomFilterProto({
+        contains: [docA, updatedDocB],
+        notContains: []
+      });
+      return (
+        spec()
+          .userListens(query1)
+          .watchAcksFull(query1, 1000, docA, docB)
+          .expectEvents(query1, { added: [docA, docB] })
+          .disableNetwork()
+          .expectEvents(query1, { fromCache: true })
+          .enableNetwork()
+          .restoreListen(query1, 'resume-token-1000')
+          .watchAcks(query1)
+          // While this client was disconnected, another client updated fields in docB.
+          .watchSends({ affects: [query1] }, updatedDocB)
+          // Bloom Filter includes all the documents that match the query, both
+          // those that have changed since the resume token and those that have not.
+          .watchFilters([query1], [docA.key, updatedDocB.key], bloomFilterProto)
+          // Expected count equals to documents in cache. Existence Filter matches.
+          .watchCurrents(query1, 'resume-token-2000')
+          .watchSnapshots(2000)
+          .expectEvents(query1, { fromCache: false })
+      );
+    }
+  );
+
+  specTest(
+    'Resume a query with bloom filter when documents are updated to no longer match the query',
+    [],
+    () => {
+      const query1 = query('collection', filter('v', '==', 1));
+      const docA = doc('collection/a', 1000, { v: 1 });
+      const docB = doc('collection/b', 1000, { v: 1 });
+      const updatedDocB = doc('collection/b', 2000, { v: 2 });
+
+      const bloomFilterProto = generateBloomFilterProto({
+        contains: [docA],
+        notContains: [docB]
+      });
+      return (
+        spec()
+          .userListens(query1)
+          .watchAcksFull(query1, 1000, docA, docB)
+          .expectEvents(query1, { added: [docA, docB] })
+          .disableNetwork()
+          .expectEvents(query1, { fromCache: true })
+          .enableNetwork()
+          .restoreListen(query1, 'resume-token-1000')
+          .watchAcks(query1)
+          // While this client was disconnected, another client modified docB to no longer match the
+          // query. Bloom Filter includes only docA that matches the query since the resume token.
+          .watchFilters([query1], [docA.key], bloomFilterProto)
+          .watchCurrents(query1, 'resume-token-2000')
+          .watchSnapshots(2000)
+          // Bloom Filter identifies that updatedDocB no longer matches the query, skips full query
+          // and puts updatedDocB into limbo directly.
+          .expectLimboDocs(updatedDocB.key) // updatedDocB is now in limbo.
+      );
+    }
+  );
+
+  specTest(
+    'Resume a query with bloom filter when documents are added, removed and deleted',
+    [],
+    () => {
+      const query1 = query('collection', filter('v', '==', 1));
+      const docA = doc('collection/a', 1000, { v: 1 });
+      const docB = doc('collection/b', 1000, { v: 1 });
+      const updatedDocB = doc('collection/b', 2000, { v: 2 });
+      const docC = doc('collection/c', 1000, { v: 1 });
+      const docD = doc('collection/d', 1000, { v: 1 });
+      const bloomFilterProto = generateBloomFilterProto({
+        contains: [docA, docD],
+        notContains: [docB, docC]
+      });
+
+      return (
+        spec()
+          .userListens(query1)
+          .watchAcksFull(query1, 1000, docA, docB, docC)
+          .expectEvents(query1, { added: [docA, docB, docC] })
+          .disableNetwork()
+          .expectEvents(query1, { fromCache: true })
+          .enableNetwork()
+          .restoreListen(query1, 'resume-token-1000')
+          .watchAcks(query1)
+          // While this client was disconnected, another client modified docB to no longer match the
+          // query, deleted docC and added docD.
+          .watchSends({ affects: [query1] }, docD)
+          // Bloom Filter includes all the documents that match the query.
+          .watchFilters([query1], [docA.key, docD.key], bloomFilterProto)
+          .watchCurrents(query1, 'resume-token-2000')
+          .watchSnapshots(2000)
+          .expectEvents(query1, { added: [docD], fromCache: true })
+          // Bloom Filter identifies that updatedDocB and docC no longer match the query, skips full
+          // query and puts them into limbo directly.
+          .expectLimboDocs(updatedDocB.key, docC.key) // updatedDocB and docC is now in limbo.
+      );
+    }
+  );
 });
diff --git a/packages/firestore/test/unit/specs/limbo_spec.test.ts b/packages/firestore/test/unit/specs/limbo_spec.test.ts
@@ -959,10 +959,10 @@ describeSpec('Limbo Documents:', [], () => {
           // While this client was disconnected, another client deleted all the
           // docAs replaced them with docBs. If Watch has to re-run the
           // underlying query when this client re-listens, Watch won't be able
-          // to tell that docAs were deleted and will only send us existing
-          // documents that changed since the resume token. This will cause it
-          // to just send the docBs with an existence filter with a count of 3.
+          // to tell that docAs were deleted and will only send us watch change
+          // for new docs added since the resume token.
           .watchSends({ affects: [query1] }, docB1, docB2, docB3)
+          // The existence filter will include the docBs with a count of 3.
           .watchFilters(
             [query1],
             [docB1.key, docB2.key, docB3.key],
