Add expectedCount to Target in listen request

packages/firestore/src/core/sync_engine_impl.ts

@@ -316,9 +316,6 @@ export async function syncEngineListen(
       syncEngineImpl.localStore,
       queryToTarget(query)
     );
-    if (syncEngineImpl.isPrimaryClient) {
-      remoteStoreListen(syncEngineImpl.remoteStore, targetData);
-    }
 
     const status = syncEngineImpl.sharedClientState.addLocalQueryTarget(
       targetData.targetId
@@ -331,6 +328,10 @@ export async function syncEngineListen(
       status === 'current',
       targetData.resumeToken
     );
+
+    if (syncEngineImpl.isPrimaryClient) {
+      remoteStoreListen(syncEngineImpl.remoteStore, targetData);
+    }
   }
 
   return viewSnapshot;

packages/firestore/src/local/target_data.ts

@@ -66,7 +66,12 @@ export class TargetData {
      * matches the target. The resume token essentially identifies a point in
      * time from which the server should resume sending results.
      */
-    readonly resumeToken: ByteString = ByteString.EMPTY_BYTE_STRING
+    readonly resumeToken: ByteString = ByteString.EMPTY_BYTE_STRING,
+
+    /** The number of documents that last matched the query at the resume token or
+     * read time.
+     */
+    readonly expectedCount: number = 0
   ) {}
 
   /** Creates a new target data instance with an updated sequence number. */
@@ -78,17 +83,19 @@ export class TargetData {
       sequenceNumber,
       this.snapshotVersion,
       this.lastLimboFreeSnapshotVersion,
-      this.resumeToken
+      this.resumeToken,
+      this.expectedCount
     );
   }
 
   /**
-   * Creates a new target data instance with an updated resume token and
-   * snapshot version.
+   * Creates a new target data instance with an updated resume token,
+   * snapshot version, and expectedCount if presented.
    */
   withResumeToken(
     resumeToken: ByteString,
-    snapshotVersion: SnapshotVersion
+    snapshotVersion: SnapshotVersion,
+    expectedCount?: number
   ): TargetData {
     return new TargetData(
       this.target,
@@ -97,7 +104,24 @@ export class TargetData {
       this.sequenceNumber,
       snapshotVersion,
       this.lastLimboFreeSnapshotVersion,
-      resumeToken
+      resumeToken,
+      expectedCount || this.expectedCount
+    );
+  }
+
+  /**
+   * Creates a new target data instance with an updated expected count.
+   */
+  withExpectedCount(expectedCount: number): TargetData {
+    return new TargetData(
+      this.target,
+      this.targetId,
+      this.purpose,
+      this.sequenceNumber,
+      this.snapshotVersion,
+      this.lastLimboFreeSnapshotVersion,
+      this.resumeToken,
+      expectedCount
     );
   }
 
@@ -115,7 +139,8 @@ export class TargetData {
       this.sequenceNumber,
       this.snapshotVersion,
       lastLimboFreeSnapshotVersion,
-      this.resumeToken
+      this.resumeToken,
+      this.expectedCount
     );
   }
 }

packages/firestore/src/remote/remote_store.ts

@@ -334,6 +334,17 @@ function sendWatchRequest(
   remoteStoreImpl.watchChangeAggregator!.recordPendingTargetRequest(
     targetData.targetId
   );
+
+  if (
+    targetData.resumeToken.approximateByteSize() > 0 ||
+    targetData.snapshotVersion.compareTo(SnapshotVersion.min()) > 0
+  ) {
+    const expectedCount = remoteStoreImpl.remoteSyncer.getRemoteKeysForTarget!(
+      targetData.targetId
+    ).size;
+    targetData = targetData.withExpectedCount(expectedCount);
+  }
+
   ensureWatchStream(remoteStoreImpl).watch(targetData);
 }
 

packages/firestore/src/remote/serializer.ts

@@ -1012,6 +1012,7 @@ export function toTarget(
 
   if (targetData.resumeToken.approximateByteSize() > 0) {
     result.resumeToken = toBytes(serializer, targetData.resumeToken);
+    result.expectedCount = targetData.expectedCount;
   } else if (targetData.snapshotVersion.compareTo(SnapshotVersion.min()) > 0) {
     // TODO(wuandy): Consider removing above check because it is most likely true.
     // Right now, many tests depend on this behaviour though (leaving min() out
@@ -1020,6 +1021,7 @@ export function toTarget(
       serializer,
       targetData.snapshotVersion.toTimestamp()
     );
+    result.expectedCount = targetData.expectedCount;
   }
 
   return result;

packages/firestore/test/unit/remote/serializer.helper.ts

@@ -1778,7 +1778,8 @@ export function serializerTest(
             }
           },
           resumeToken: new Uint8Array([1, 2, 3]),
-          targetId: 1
+          targetId: 1,
+          expectedCount: 0
         };
         expect(result).to.deep.equal(expected);
       });

packages/firestore/test/unit/specs/bundle_spec.test.ts

@@ -356,7 +356,7 @@ describeSpec('Bundles:', ['no-ios'], () => {
           // Read named query from loaded bundle by secondary.
           .client(0)
           .expectListen(query1, { readTime: 400 })
-          .expectActiveTargets({ query: query1, readTime: 400 })
+          .expectActiveTargets({ query: query1, resume: { readTime: 400 } })
           .userListens(query2, { readTime: 560 })
           .expectEvents(query2, {
             added: [doc('collection/a', 550, { value: 'c' })],

packages/firestore/test/unit/specs/existence_filter_spec.test.ts

@@ -116,7 +116,7 @@ describeSpec('Existence Filters:', [], () => {
         .watchSnapshots(2000)
         // query is now marked as "inconsistent" because of filter mismatch
         .expectEvents(query1, { fromCache: true })
-        .expectActiveTargets({ query: query1, resumeToken: '' })
+        .expectActiveTargets({ query: query1, resume: { resumeToken: '' } })
         .watchRemoves(query1) // Acks removal of query
         .watchAcksFull(query1, 2000, doc1)
         .expectLimboDocs(doc2.key) // doc2 is now in limbo
@@ -146,7 +146,7 @@ describeSpec('Existence Filters:', [], () => {
         .watchSnapshots(2000)
         // query is now marked as "inconsistent" because of filter mismatch
         .expectEvents(query1, { fromCache: true })
-        .expectActiveTargets({ query: query1, resumeToken: '' })
+        .expectActiveTargets({ query: query1, resume: { resumeToken: '' } })
         .watchRemoves(query1) // Acks removal of query
         .watchAcksFull(query1, 2000, doc1)
         .expectLimboDocs(doc2.key) // doc2 is now in limbo
@@ -173,14 +173,14 @@ describeSpec('Existence Filters:', [], () => {
         .watchStreamCloses(Code.UNAVAILABLE)
         .expectActiveTargets({
           query: query1,
-          resumeToken: 'existence-filter-resume-token'
+          resume: { resumeToken: 'existence-filter-resume-token' }
         })
         .watchAcks(query1)
         .watchFilters([query1], [doc1.key]) // in the next sync doc2 was deleted
         .watchSnapshots(2000)
         // query is now marked as "inconsistent" because of filter mismatch
         .expectEvents(query1, { fromCache: true })
-        .expectActiveTargets({ query: query1, resumeToken: '' })
+        .expectActiveTargets({ query: query1, resume: { resumeToken: '' } })
         .watchRemoves(query1) // Acks removal of query
         .watchAcksFull(query1, 2000, doc1)
         .expectLimboDocs(doc2.key) // doc2 is now in limbo
@@ -211,7 +211,7 @@ describeSpec('Existence Filters:', [], () => {
         // The query result includes doc3, but is marked as "inconsistent"
         // due to the filter mismatch
         .expectEvents(query1, { added: [doc3], fromCache: true })
-        .expectActiveTargets({ query: query1, resumeToken: '' })
+        .expectActiveTargets({ query: query1, resume: { resumeToken: '' } })
         .watchRemoves(query1) // Acks removal of query
         .watchAcksFull(query1, 3000, doc1, doc2, doc3)
         .expectEvents(query1, { added: [doc2] })
@@ -243,7 +243,7 @@ describeSpec('Existence Filters:', [], () => {
         .watchSnapshots(2000)
         // query is now marked as "inconsistent" because of filter mismatch
         .expectEvents(query1, { fromCache: true })
-        .expectActiveTargets({ query: query1, resumeToken: '' })
+        .expectActiveTargets({ query: query1, resume: { resumeToken: '' } })
         .watchRemoves(query1) // Acks removal of query
         .watchAcksFull(query1, 2000, doc1)
         .expectLimboDocs(doc2.key) // doc2 is now in limbo

packages/firestore/test/unit/specs/limbo_spec.test.ts

@@ -889,7 +889,7 @@ describeSpec('Limbo Documents:', [], () => {
           // The view now contains the docAs and the docBs (6 documents), but
           // the existence filter indicated only 3 should match. This causes
           // the client to re-listen without a resume token.
-          .expectActiveTargets({ query: query1, resumeToken: '' })
+          .expectActiveTargets({ query: query1, resume: { resumeToken: '' } })
           // When the existence filter mismatch was detected, the client removed
           // then re-added the target. Watch needs to acknowledge the removal.
           .watchRemoves(query1)

packages/firestore/test/unit/specs/limit_spec.test.ts

@@ -343,7 +343,10 @@ describeSpec('Limits:', [], () => {
           .watchSends({ affects: [limitQuery] }, secondDocument)
           .watchFilters([limitQuery], [secondDocument.key])
           .watchSnapshots(1004)
-          .expectActiveTargets({ query: limitQuery, resumeToken: '' })
+          .expectActiveTargets({
+            query: limitQuery,
+            resume: { resumeToken: '' }
+          })
           .watchRemoves(limitQuery)
           .watchAcksFull(limitQuery, 1005, secondDocument)
           // The snapshot after the existence filter mismatch triggers limbo

packages/firestore/test/unit/specs/listen_spec.test.ts

@@ -1809,4 +1809,85 @@ describeSpec('Listens:', [], () => {
       );
     }
   );
+
+  specTest(
+    'Query with resume target can pass in expectedCount to listen request',
+    [],
+    () => {
+      const query1 = query('collection');
+      const docA = doc('collection/a', 1000, { key: 'a' });
+      const docB = doc('collection/b', 1000, { key: 'b' });
+
+      return (
+        spec()
+          .withGCEnabled(false)
+          .userListens(query1)
+          .watchAcksFull(query1, 1000)
+          .expectEvents(query1, {})
+          .userUnlistens(query1)
+          .watchRemoves(query1)
+          // There is 0 remote document from previous listen.
+          .userListens(query1, { resumeToken: 'resume-token-1000' }, 0)
+          .expectEvents(query1, { fromCache: true })
+          .watchAcksFull(query1, 2000, docA, docB)
+          .expectEvents(query1, { added: [docA, docB] })
+          .userUnlistens(query1)
+          .userListens(query1, { resumeToken: 'resume-token-2000' }, 2)
+          .expectEvents(query1, { added: [docA, docB], fromCache: true })
+      );
+    }
+  );
+
+  specTest(
+    'ExpectedCount should equal to the number of documents that last matched the query at the resume token',
+    [],
+    () => {
+      const query1 = query('collection');
+      const docA = doc('collection/a', 1000, { key: 'a' });
+      const docBLocal = doc('collection/b', 1000, {
+        key: 'b'
+      }).setHasLocalMutations();
+
+      return spec()
+        .withGCEnabled(false)
+        .userListens(query1)
+        .watchAcksFull(query1, 1000, docA)
+        .expectEvents(query1, { added: [docA] })
+        .userUnlistens(query1)
+        .userSets('collection/b', { key: 'b' })
+        .userListens(query1, { resumeToken: 'resume-token-1000' }, 1)
+        .expectEvents(query1, {
+          added: [docA, docBLocal],
+          fromCache: true,
+          hasPendingWrites: true
+        });
+    }
+  );
+
+  specTest(
+    'ExpectedCount in listen request should work after coming back online',
+    [],
+    () => {
+      const query1 = query('collection');
+      const docA = doc('collection/a', 1000, { key: 'a' });
+      const docBLocal = doc('collection/b', 1000, {
+        key: 'b'
+      }).setHasLocalMutations();
+
+      return spec()
+        .withGCEnabled(false)
+        .userListens(query1)
+        .watchAcksFull(query1, 1000, docA)
+        .expectEvents(query1, { added: [docA] })
+        .userSets('collection/b', { key: 'b' })
+        .expectEvents(query1, {
+          hasPendingWrites: true,
+          added: [docBLocal]
+        })
+        .disableNetwork()
+        .expectEvents(query1, { hasPendingWrites: true, fromCache: true })
+        .enableNetwork()
+        .restoreListen(query1, 'resume-token-1000', 1);
+    }
+  );
 });

packages/firestore/test/unit/specs/persistence_spec.test.ts

@@ -166,7 +166,10 @@ describeSpec('Persistence:', [], () => {
         .expectEvents(query1, { added: [anonDoc], hasPendingWrites: true })
         .changeUser('user1')
         // A user change will re-send the query with the current resume token
-        .expectActiveTargets({ query: query1, resumeToken: 'resume-token-500' })
+        .expectActiveTargets({
+          query: query1,
+          resume: { resumeToken: 'resume-token-500' }
+        })
         .expectEvents(query1, { removed: [anonDoc] })
         .userSets('users/user1', { uid: 'user1' })
         .expectEvents(query1, { added: [user1Doc], hasPendingWrites: true })

packages/firestore/test/unit/specs/recovery_spec.test.ts

@@ -421,7 +421,7 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
         // Client is back online
         .expectActiveTargets({
           query: query1,
-          resumeToken: 'resume-token-1000'
+          resume: { resumeToken: 'resume-token-1000' }
         })
         .expectUserCallbacks({ acknowledged: ['collection/a'] })
         .watchAcksFull(query1, 1001, doc1)
@@ -549,7 +549,7 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
         .runTimer(TimerId.AsyncQueueRetry)
         .expectActiveTargets({
           query: query1,
-          resumeToken: 'resume-token-1000'
+          resume: { resumeToken: 'resume-token-1000' }
         })
         .watchAcksFull(query1, 2000, doc2)
         .expectEvents(query1, {
@@ -639,7 +639,7 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
         .expectActiveTargets(
           {
             query: filteredQuery,
-            resumeToken: 'resume-token-2000'
+            resume: { resumeToken: 'resume-token-2000' }
           },
           { query: limboQuery }
         )
@@ -681,7 +681,7 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
       .recoverDatabase()
       .runTimer(TimerId.AsyncQueueRetry)
       .expectActiveTargets(
-        { query: filteredQuery, resumeToken: 'resume-token-2000' },
+        { query: filteredQuery, resume: { resumeToken: 'resume-token-2000' } },
         { query: limboQuery }
       )
       .watchAcksFull(filteredQuery, 3000)

packages/firestore/test/unit/specs/remote_store_spec.test.ts

@@ -107,7 +107,7 @@ describeSpec('Remote store:', [], () => {
           // Change user (will shut down existing streams and start new ones).
           .changeUser('abc')
           // Our query should be sent to the new stream.
-          .expectActiveTargets({ query: query1, resumeToken: '' })
+          .expectActiveTargets({ query: query1, resume: { resumeToken: '' } })
 
           // Close the (newly-created) stream as if it too failed (should trigger
           // retry with backoff, potentially reproducing the crash in b/74749605).

packages/firestore/test/unit/specs/resume_token_spec.test.ts

@@ -35,7 +35,7 @@ describeSpec('Resume tokens:', [], () => {
       .watchStreamCloses(Code.UNAVAILABLE)
       .expectActiveTargets({
         query: query1,
-        resumeToken: 'custom-query-resume-token'
+        resume: { resumeToken: 'custom-query-resume-token' }
       });
   });