Persist a Query's Sync State

Author: schmidt-sebastian

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

@@ -80,6 +80,7 @@ public void onViewSnapshot(ViewSnapshot newSnapshot) {
               documentChanges,
               newSnapshot.isFromCache(),
               newSnapshot.getMutatedKeys(),
+              newSnapshot.isSynced(),
               newSnapshot.didSyncStateChange(),
               /* excludesMetadataChanges= */ true);
     }
@@ -158,6 +159,7 @@ private void raiseInitialEvent(ViewSnapshot snapshot) {
             snapshot.getDocuments(),
             snapshot.getMutatedKeys(),
             snapshot.isFromCache(),
+            snapshot.isSynced(),
             snapshot.excludesMetadataChanges());
     raisedInitialEvent = true;
     listener.onEvent(snapshot, null);

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

@@ -488,6 +488,7 @@ private void emitNewSnapsAndNotifyLocalStore(
         // The query has a limit and some docs were removed/updated, so we need to re-run the query
         // against the local store to make sure we didn't lose any good docs that had been past the
         // limit.
+        localStore.markNeedsRefill(queryView.getTargetId());
         ImmutableSortedMap<DocumentKey, Document> docs =
             localStore.executeQuery(queryView.getQuery());
         viewDocChanges = view.computeDocChanges(docs, viewDocChanges);

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

@@ -311,6 +311,7 @@ public ViewChange applyChanges(DocumentChanges docChanges, TargetChange targetCh
               viewChanges,
               fromCache,
               docChanges.mutatedKeys,
+              synced,
               syncStatedChanged,
               /* excludesMetadataChanges= */ false);
     }

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

@@ -37,6 +37,7 @@ public enum SyncState {
   private final List<DocumentViewChange> changes;
   private final boolean isFromCache;
   private final ImmutableSortedSet<DocumentKey> mutatedKeys;
+  private final boolean synced;
   private final boolean didSyncStateChange;
   private boolean excludesMetadataChanges;
 
@@ -47,6 +48,7 @@ public ViewSnapshot(
       List<DocumentViewChange> changes,
       boolean isFromCache,
       ImmutableSortedSet<DocumentKey> mutatedKeys,
+      boolean synced,
       boolean didSyncStateChange,
       boolean excludesMetadataChanges) {
     this.query = query;
@@ -55,6 +57,7 @@ public ViewSnapshot(
     this.changes = changes;
     this.isFromCache = isFromCache;
     this.mutatedKeys = mutatedKeys;
+    this.synced = synced;
     this.didSyncStateChange = didSyncStateChange;
     this.excludesMetadataChanges = excludesMetadataChanges;
   }
@@ -65,6 +68,7 @@ public static ViewSnapshot fromInitialDocuments(
       DocumentSet documents,
       ImmutableSortedSet<DocumentKey> mutatedKeys,
       boolean fromCache,
+      boolean synced,
       boolean excludesMetadataChanges) {
     List<DocumentViewChange> viewChanges = new ArrayList<>();
     for (Document doc : documents) {
@@ -77,6 +81,7 @@ public static ViewSnapshot fromInitialDocuments(
         viewChanges,
         fromCache,
         mutatedKeys,
+        synced,
         /* didSyncStateChange= */ true,
         excludesMetadataChanges);
   }
@@ -85,6 +90,14 @@ public Query getQuery() {
     return query;
   }
 
+  /**
+   * Returns whether the view was synced with the backend at the time the snapshot was raised. A
+   * synced view is marked CURRENT and contains no Limbo documents.
+   */
+  public boolean isSynced() {
+    return synced;
+  }
+
   public DocumentSet getDocuments() {
     return documents;
   }
@@ -131,6 +144,9 @@ public final boolean equals(Object o) {
     if (isFromCache != that.isFromCache) {
       return false;
     }
+    if (synced != that.synced) {
+      return false;
+    }
     if (didSyncStateChange != that.didSyncStateChange) {
       return false;
     }
@@ -160,6 +176,7 @@ public int hashCode() {
     result = 31 * result + changes.hashCode();
     result = 31 * result + mutatedKeys.hashCode();
     result = 31 * result + (isFromCache ? 1 : 0);
+    result = 31 * result + (synced ? 1 : 0);
     result = 31 * result + (didSyncStateChange ? 1 : 0);
     result = 31 * result + (excludesMetadataChanges ? 1 : 0);
     return result;
@@ -179,6 +196,8 @@ public String toString() {
         + isFromCache
         + ", mutatedKeys="
         + mutatedKeys.size()
+        + ", synced="
+        + synced
         + ", didSyncStateChange="
         + didSyncStateChange
         + ", excludesMetadataChanges="

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

@@ -205,6 +205,7 @@ com.google.firebase.firestore.proto.Target encodeQueryData(QueryData queryData)
     result
         .setTargetId(queryData.getTargetId())
         .setLastListenSequenceNumber(queryData.getSequenceNumber())
+        .setSynced(queryData.isSynced())
         .setSnapshotVersion(rpcSerializer.encodeVersion(queryData.getSnapshotVersion()))
         .setResumeToken(queryData.getResumeToken());
 
@@ -223,6 +224,7 @@ QueryData decodeQueryData(com.google.firebase.firestore.proto.Target target) {
     SnapshotVersion version = rpcSerializer.decodeVersion(target.getSnapshotVersion());
     ByteString resumeToken = target.getResumeToken();
     long sequenceNumber = target.getLastListenSequenceNumber();
+    boolean synced = target.getSynced();
 
     Query query;
     switch (target.getTargetTypeCase()) {
@@ -239,6 +241,6 @@ QueryData decodeQueryData(com.google.firebase.firestore.proto.Target target) {
     }
 
     return new QueryData(
-        query, targetId, sequenceNumber, QueryPurpose.LISTEN, version, resumeToken);
+        query, targetId, sequenceNumber, synced, QueryPurpose.LISTEN, version, resumeToken);
   }
 }

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

@@ -359,7 +359,11 @@ public ImmutableSortedMap<DocumentKey, MaybeDocument> applyRemoteEvent(RemoteEve
             if (!resumeToken.isEmpty()) {
               QueryData oldQueryData = queryData;
               queryData =
-                  queryData.copy(remoteEvent.getSnapshotVersion(), resumeToken, sequenceNumber);
+                  queryData.copy(
+                      remoteEvent.getSnapshotVersion(),
+                      resumeToken,
+                      sequenceNumber,
+                      oldQueryData.isSynced());
               targetIds.put(boxedTargetId, queryData);
 
               if (shouldPersistQueryData(oldQueryData, queryData, change)) {
@@ -469,12 +473,26 @@ public void notifyLocalViewChanges(List<LocalViewChanges> viewChanges) {
         "notifyLocalViewChanges",
         () -> {
           for (LocalViewChanges viewChange : viewChanges) {
-            localViewReferences.addReferences(viewChange.getAdded(), viewChange.getTargetId());
+            int targetId = viewChange.getTargetId();
+
+            // Update the local view references for LRU garbage collection
+            localViewReferences.addReferences(viewChange.getAdded(), targetId);
             ImmutableSortedSet<DocumentKey> removed = viewChange.getRemoved();
             for (DocumentKey key : removed) {
               persistence.getReferenceDelegate().removeReference(key);
             }
-            localViewReferences.removeReferences(removed, viewChange.getTargetId());
+            localViewReferences.removeReferences(removed, targetId);
+
+            // Mark whether the local query view is in sync with the backend.
+            QueryData queryData = targetIds.get(targetId);
+            hardAssert(queryData != null, "Can't mark inactive target synchronized: %s", targetId);
+            QueryData updatedQueryData =
+                queryData.copy(
+                    queryData.getSnapshotVersion(),
+                    queryData.getResumeToken(),
+                    queryData.getSequenceNumber(),
+                    viewChange.isSynced());
+            targetIds.put(targetId, updatedQueryData);
           }
         });
   }
@@ -548,10 +566,18 @@ public void releaseQuery(Query query) {
 
           int targetId = queryData.getTargetId();
           QueryData cachedQueryData = targetIds.get(targetId);
+
+          boolean needsUpdate = false;
           if (cachedQueryData.getSnapshotVersion().compareTo(queryData.getSnapshotVersion()) > 0) {
             // If we've been avoiding persisting the resumeToken (see shouldPersistQueryData for
             // conditions and rationale) we need to persist the token now because there will no
             // longer be an in-memory version to fall back on.
+            needsUpdate = true;
+          } else if (cachedQueryData.isSynced() != queryData.isSynced()) {
+            needsUpdate = true;
+          }
+
+          if (needsUpdate) {
             queryData = cachedQueryData;
             queryCache.updateQueryData(queryData);
           }
@@ -570,6 +596,22 @@ public void releaseQuery(Query query) {
         });
   }
 
+  /**
+   * Mark the query as not-consistent with the backend, indicating that we have to re-run the query
+   * against the full set of local documents.
+   */
+  public void markNeedsRefill(int targetId) {
+    QueryData queryData = targetIds.get(targetId);
+    hardAssert(queryData != null, "Can't mark inactive target: %s", targetId);
+    QueryData updatedQueryData =
+        queryData.copy(
+            queryData.getSnapshotVersion(),
+            queryData.getResumeToken(),
+            queryData.getSequenceNumber(),
+            /* synced= */ false);
+    targetIds.put(targetId, updatedQueryData);
+  }
+
   /** Runs the given query against all the documents in the local store and returns the results. */
   public ImmutableSortedMap<DocumentKey, Document> executeQuery(Query query) {
     return queryEngine.getDocumentsMatchingQuery(query);

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

@@ -49,19 +49,22 @@ public static LocalViewChanges fromViewSnapshot(int targetId, ViewSnapshot snaps
       }
     }
 
-    return new LocalViewChanges(targetId, addedKeys, removedKeys);
+    return new LocalViewChanges(targetId, snapshot.isSynced(), addedKeys, removedKeys);
   }
 
   private final int targetId;
+  private final boolean synced;
 
   private final ImmutableSortedSet<DocumentKey> added;
   private final ImmutableSortedSet<DocumentKey> removed;
 
   public LocalViewChanges(
       int targetId,
+      boolean synced,
       ImmutableSortedSet<DocumentKey> added,
       ImmutableSortedSet<DocumentKey> removed) {
     this.targetId = targetId;
+    this.synced = synced;
     this.added = added;
     this.removed = removed;
   }
@@ -70,6 +73,11 @@ public int getTargetId() {
     return targetId;
   }
 
+  /** Returns whether the local query view is in sync with the backend. * */
+  public boolean isSynced() {
+    return synced;
+  }
+
   public ImmutableSortedSet<DocumentKey> getAdded() {
     return added;
   }

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

@@ -137,7 +137,10 @@ public void removeMutationReference(DocumentKey key) {
   public void removeTarget(QueryData queryData) {
     QueryData updated =
         queryData.copy(
-            queryData.getSnapshotVersion(), queryData.getResumeToken(), getCurrentSequenceNumber());
+            queryData.getSnapshotVersion(),
+            queryData.getResumeToken(),
+            getCurrentSequenceNumber(),
+            queryData.isSynced());
     persistence.getQueryCache().updateQueryData(updated);
   }
 

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

@@ -26,6 +26,7 @@ public final class QueryData {
   private final Query query;
   private final int targetId;
   private final long sequenceNumber;
+  private final boolean synced;
   private final QueryPurpose purpose;
   private final SnapshotVersion snapshotVersion;
   private final ByteString resumeToken;
@@ -36,6 +37,8 @@ public final class QueryData {
    * @param query The query being listened to.
    * @param targetId The target to which the query corresponds, assigned by the LocalStore for user
    *     queries or the SyncEngine for limbo queries.
+   * @param sequenceNumber The sequence number, denoting the last time this query was used.
+   * @param synced Whether the query's local state is in sync with the backend.
    * @param purpose The purpose of the query.
    * @param snapshotVersion The latest snapshot version seen for this target.
    * @param resumeToken An opaque, server-assigned token that allows watching a query to be resumed
@@ -47,12 +50,14 @@ public QueryData(
       Query query,
       int targetId,
       long sequenceNumber,
+      boolean synced,
       QueryPurpose purpose,
       SnapshotVersion snapshotVersion,
       ByteString resumeToken) {
     this.query = checkNotNull(query);
     this.targetId = targetId;
     this.sequenceNumber = sequenceNumber;
+    this.synced = synced;
     this.purpose = purpose;
     this.snapshotVersion = checkNotNull(snapshotVersion);
     this.resumeToken = checkNotNull(resumeToken);
@@ -64,6 +69,7 @@ public QueryData(Query query, int targetId, long sequenceNumber, QueryPurpose pu
         query,
         targetId,
         sequenceNumber,
+        /* synced= */ false,
         purpose,
         SnapshotVersion.NONE,
         WatchStream.EMPTY_RESUME_TOKEN);
@@ -93,6 +99,11 @@ public ByteString getResumeToken() {
     return resumeToken;
   }
 
+  /** Whether the query's local state is in sync with the backend. */
+  public boolean isSynced() {
+    return synced;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {
@@ -106,6 +117,7 @@ public boolean equals(Object o) {
     return query.equals(queryData.query)
         && targetId == queryData.targetId
         && sequenceNumber == queryData.sequenceNumber
+        && synced == queryData.synced
         && purpose.equals(queryData.purpose)
         && snapshotVersion.equals(queryData.snapshotVersion)
         && resumeToken.equals(queryData.resumeToken);
@@ -116,6 +128,7 @@ public int hashCode() {
     int result = query.hashCode();
     result = 31 * result + targetId;
     result = 31 * result + (int) sequenceNumber;
+    result = 31 * result + (synced ? 1 : 0);
     result = 31 * result + purpose.hashCode();
     result = 31 * result + snapshotVersion.hashCode();
     result = 31 * result + resumeToken.hashCode();
@@ -131,6 +144,8 @@ public String toString() {
         + targetId
         + ", sequenceNumber="
         + sequenceNumber
+        + ", isSynced="
+        + synced
         + ", purpose="
         + purpose
         + ", snapshotVersion="
@@ -140,9 +155,16 @@ public String toString() {
         + '}';
   }
 
-  /** Creates a new query data instance with an updated snapshot version and resume token. */
+  /**
+   * Creates a new query data instance with an updated snapshot version, sequence number, sync
+   * status and resume token.
+   */
   public QueryData copy(
-      SnapshotVersion snapshotVersion, ByteString resumeToken, long sequenceNumber) {
-    return new QueryData(query, targetId, sequenceNumber, purpose, snapshotVersion, resumeToken);
+      SnapshotVersion snapshotVersion,
+      ByteString resumeToken,
+      long sequenceNumber,
+      boolean synced) {
+    return new QueryData(
+        query, targetId, sequenceNumber, synced, purpose, snapshotVersion, resumeToken);
   }
 }

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

@@ -169,7 +169,10 @@ public int removeOrphanedDocuments(long upperBound) {
   public void removeTarget(QueryData queryData) {
     QueryData updated =
         queryData.copy(
-            queryData.getSnapshotVersion(), queryData.getResumeToken(), getCurrentSequenceNumber());
+            queryData.getSnapshotVersion(),
+            queryData.getResumeToken(),
+            getCurrentSequenceNumber(),
+            queryData.isSynced());
     persistence.getQueryCache().updateQueryData(updated);
   }