Add expected count to target (#4574)

Author: milaGGL

firebase-firestore/src/main/java/com/google/firebase/firestore/core/SyncEngine.java

@@ -203,13 +203,14 @@ public int listen(Query query) {
     hardAssert(!queryViewsByQuery.containsKey(query), "We already listen to query: %s", query);
 
     TargetData targetData = localStore.allocateTarget(query.toTarget());
-    remoteStore.listen(targetData);
 
     ViewSnapshot viewSnapshot =
         initializeViewAndComputeSnapshot(
             query, targetData.getTargetId(), targetData.getResumeToken());
     syncEngineListener.onViewSnapshots(Collections.singletonList(viewSnapshot));
 
+    remoteStore.listen(targetData);
+
     return targetData.getTargetId();
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalSerializer.java

@@ -260,7 +260,8 @@ TargetData decodeTargetData(com.google.firebase.firestore.proto.Target targetPro
         QueryPurpose.LISTEN,
         version,
         lastLimboFreeSnapshotVersion,
-        resumeToken);
+        resumeToken,
+        null);
   }
 
   public com.google.firestore.bundle.BundledQuery encodeBundledQuery(BundledQuery bundledQuery) {

firebase-firestore/src/main/java/com/google/firebase/firestore/local/TargetData.java

@@ -16,10 +16,12 @@
 
 import static com.google.firebase.firestore.util.Preconditions.checkNotNull;
 
+import androidx.annotation.Nullable;
 import com.google.firebase.firestore.core.Target;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.remote.WatchStream;
 import com.google.protobuf.ByteString;
+import java.util.Objects;
 
 /** An immutable set of metadata that the store will need to keep track of for each target. */
 public final class TargetData {
@@ -30,6 +32,7 @@ public final class TargetData {
   private final SnapshotVersion snapshotVersion;
   private final SnapshotVersion lastLimboFreeSnapshotVersion;
   private final ByteString resumeToken;
+  private final @Nullable Integer expectedCount;
 
   /**
    * Creates a new TargetData with the given values.
@@ -45,6 +48,9 @@ public final class TargetData {
    * @param resumeToken An opaque, server-assigned token that allows watching a target to be resumed
    *     after disconnecting without retransmitting all the data that matches the target. The resume
    *     token essentially identifies a point in time from which the server should resume sending
+   * @param expectedCount The number of documents that last matched the query at the resume token or
+   *     read time. Documents are counted only when making a listen request with resume token or
+   *     read time, otherwise, keep it null.
    */
   TargetData(
       Target target,
@@ -53,14 +59,16 @@ public final class TargetData {
       QueryPurpose purpose,
       SnapshotVersion snapshotVersion,
       SnapshotVersion lastLimboFreeSnapshotVersion,
-      ByteString resumeToken) {
+      ByteString resumeToken,
+      @Nullable Integer expectedCount) {
     this.target = checkNotNull(target);
     this.targetId = targetId;
     this.sequenceNumber = sequenceNumber;
     this.lastLimboFreeSnapshotVersion = lastLimboFreeSnapshotVersion;
     this.purpose = purpose;
     this.snapshotVersion = checkNotNull(snapshotVersion);
     this.resumeToken = checkNotNull(resumeToken);
+    this.expectedCount = expectedCount;
   }
 
   /** Convenience constructor for use when creating a TargetData for the first time. */
@@ -72,7 +80,8 @@ public TargetData(Target target, int targetId, long sequenceNumber, QueryPurpose
         purpose,
         SnapshotVersion.NONE,
         SnapshotVersion.NONE,
-        WatchStream.EMPTY_RESUME_TOKEN);
+        WatchStream.EMPTY_RESUME_TOKEN,
+        null);
   }
 
   /** Creates a new target data instance with an updated sequence number. */
@@ -84,7 +93,8 @@ public TargetData withSequenceNumber(long sequenceNumber) {
         purpose,
         snapshotVersion,
         lastLimboFreeSnapshotVersion,
-        resumeToken);
+        resumeToken,
+        /* expectedCount= */ null);
   }
 
   /** Creates a new target data instance with an updated resume token and snapshot version. */
@@ -96,7 +106,21 @@ public TargetData withResumeToken(ByteString resumeToken, SnapshotVersion snapsh
         purpose,
         snapshotVersion,
         lastLimboFreeSnapshotVersion,
-        resumeToken);
+        resumeToken,
+        expectedCount);
+  }
+
+  /** Creates a new target data instance with an updated expected count. */
+  public TargetData withExpectedCount(@Nullable Integer expectedCount) {
+    return new TargetData(
+        target,
+        targetId,
+        sequenceNumber,
+        purpose,
+        snapshotVersion,
+        lastLimboFreeSnapshotVersion,
+        resumeToken,
+        expectedCount);
   }
 
   /** Creates a new target data instance with an updated last limbo free snapshot version number. */
@@ -108,7 +132,8 @@ public TargetData withLastLimboFreeSnapshotVersion(SnapshotVersion lastLimboFree
         purpose,
         snapshotVersion,
         lastLimboFreeSnapshotVersion,
-        resumeToken);
+        resumeToken,
+        expectedCount);
   }
 
   public Target getTarget() {
@@ -135,6 +160,11 @@ public ByteString getResumeToken() {
     return resumeToken;
   }
 
+  @Nullable
+  public Integer getExpectedCount() {
+    return expectedCount;
+  }
+
   /**
    * Returns the last snapshot version for which the associated view contained no limbo documents.
    */
@@ -158,7 +188,8 @@ public boolean equals(Object o) {
         && purpose.equals(targetData.purpose)
         && snapshotVersion.equals(targetData.snapshotVersion)
         && lastLimboFreeSnapshotVersion.equals(targetData.lastLimboFreeSnapshotVersion)
-        && resumeToken.equals(targetData.resumeToken);
+        && resumeToken.equals(targetData.resumeToken)
+        && Objects.equals(expectedCount, targetData.expectedCount);
   }
 
   @Override
@@ -170,6 +201,7 @@ public int hashCode() {
     result = 31 * result + snapshotVersion.hashCode();
     result = 31 * result + lastLimboFreeSnapshotVersion.hashCode();
     result = 31 * result + resumeToken.hashCode();
+    result = 31 * result + Objects.hashCode(expectedCount);
     return result;
   }
 
@@ -190,6 +222,8 @@ public String toString() {
         + lastLimboFreeSnapshotVersion
         + ", resumeToken="
         + resumeToken
+        + ", expectedCount="
+        + expectedCount
         + '}';
   }
 }

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

@@ -499,6 +499,13 @@ public Target encodeTarget(TargetData targetData) {
       builder.setResumeToken(targetData.getResumeToken());
     }
 
+    // TODO(Mila) Incorporate this into the if statement above.
+    if (targetData.getExpectedCount() != null
+        && (!targetData.getResumeToken().isEmpty()
+            || targetData.getSnapshotVersion().compareTo(SnapshotVersion.NONE) > 0)) {
+      builder.setExpectedCount(Int32Value.newBuilder().setValue(targetData.getExpectedCount()));
+    }
+
     return builder.build();
   }
 

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

@@ -369,6 +369,12 @@ public void listen(TargetData targetData) {
 
   private void sendWatchRequest(TargetData targetData) {
     watchChangeAggregator.recordPendingTargetRequest(targetData.getTargetId());
+    if (!targetData.getResumeToken().isEmpty()
+        || targetData.getSnapshotVersion().compareTo(SnapshotVersion.NONE) > 0) {
+      int expectedCount = this.getRemoteKeysForTarget(targetData.getTargetId()).size();
+      targetData = targetData.withExpectedCount(expectedCount);
+    }
+
     watchStream.watchQuery(targetData);
   }
 

firebase-firestore/src/test/java/com/google/firebase/firestore/local/LocalSerializerTest.java

@@ -378,7 +378,8 @@ public void testEncodesTargetData() {
             QueryPurpose.LISTEN,
             snapshotVersion,
             limboFreeVersion,
-            resumeToken);
+            resumeToken,
+            null);
 
     // Let the RPC serializer test various permutations of query serialization.
     com.google.firestore.v1.Target.QueryTarget queryTarget =

firebase-firestore/src/test/java/com/google/firebase/firestore/local/TargetCacheTestCase.java

@@ -331,7 +331,8 @@ private TargetData newTargetData(Query query, int targetId, long version) {
         QueryPurpose.LISTEN,
         version(version),
         version(version),
-        resumeToken(version));
+        resumeToken(version),
+        null);
   }
 
   /** Adds the given query data to the targetCache under test, committing immediately. */

firebase-firestore/src/test/java/com/google/firebase/firestore/spec/SpecTestCase.java

@@ -1006,6 +1006,9 @@ private void validateExpectedState(@Nullable JSONObject expectedState) throws JS
                   targetData.withResumeToken(
                       ByteString.EMPTY, version(queryDataJson.getInt("readTime")));
             }
+            if (queryDataJson.has("expectedCount")) {
+              targetData = targetData.withExpectedCount(queryDataJson.getInt("expectedCount"));
+            }
 
             expectedActiveTargets.get(targetId).add(targetData);
           }
@@ -1144,6 +1147,10 @@ private void validateActiveTargets() {
           expectedTarget.getResumeToken().toStringUtf8(),
           actualTarget.getResumeToken().toStringUtf8());
 
+      if (expectedTarget.getExpectedCount() != null) {
+        assertEquals(expectedTarget.getExpectedCount(), actualTarget.getExpectedCount());
+      }
+
       actualTargets.remove(expected.getKey());
     }