firebase · schmidt-sebastian · Jan 31, 2022 · Jan 28, 2022
diff --git a/firebase-firestore/src/main/java/com/google/firebase/firestore/bundle/BundleLoader.java b/firebase-firestore/src/main/java/com/google/firebase/firestore/bundle/BundleLoader.java
@@ -41,7 +41,7 @@ public class BundleLoader {
 
   private ImmutableSortedMap<DocumentKey, MutableDocument> documents;
   private long bytesLoaded;
-  @Nullable private DocumentKey currentDocument;
+  @Nullable private BundledDocumentMetadata currentMetadata;
 
   public BundleLoader(BundleCallback bundleCallback, BundleMetadata bundleMetadata) {
     this.bundleCallback = bundleCallback;
@@ -68,23 +68,27 @@ public BundleLoader(BundleCallback bundleCallback, BundleMetadata bundleMetadata
     } else if (bundleElement instanceof BundledDocumentMetadata) {
       BundledDocumentMetadata bundledDocumentMetadata = (BundledDocumentMetadata) bundleElement;
       documentsMetadata.put(bundledDocumentMetadata.getKey(), bundledDocumentMetadata);
-      currentDocument = bundledDocumentMetadata.getKey();
+      currentMetadata = bundledDocumentMetadata;
       if (!((BundledDocumentMetadata) bundleElement).exists()) {
         documents =
             documents.insert(
                 bundledDocumentMetadata.getKey(),
                 MutableDocument.newNoDocument(
-                    bundledDocumentMetadata.getKey(), bundledDocumentMetadata.getReadTime()));
-        currentDocument = null;
+                        bundledDocumentMetadata.getKey(), bundledDocumentMetadata.getReadTime())
+                    .setReadTime(bundledDocumentMetadata.getReadTime()));
+        currentMetadata = null;
       }
     } else if (bundleElement instanceof BundleDocument) {
       BundleDocument bundleDocument = (BundleDocument) bundleElement;
-      if (!bundleDocument.getKey().equals(currentDocument)) {
+      if (currentMetadata == null || !bundleDocument.getKey().equals(currentMetadata.getKey())) {
         throw new IllegalArgumentException(
             "The document being added does not match the stored metadata.");
       }
-      documents = documents.insert(bundleDocument.getKey(), bundleDocument.getDocument());
-      currentDocument = null;
+      documents =
+          documents.insert(
+              bundleDocument.getKey(),
+              bundleDocument.getDocument().setReadTime(currentMetadata.getReadTime()));
+      currentMetadata = null;
     }
 
     bytesLoaded += byteSize;
@@ -103,7 +107,7 @@ public BundleLoader(BundleCallback bundleCallback, BundleMetadata bundleMetadata
   /** Applies the loaded documents and queries to local store. Returns the document view changes. */
   public ImmutableSortedMap<DocumentKey, Document> applyChanges() {
     Preconditions.checkArgument(
-        currentDocument == null,
+        currentMetadata == null,
         "Bundled documents end with a document metadata element instead of a document.");
     Preconditions.checkArgument(bundleMetadata.getBundleId() != null, "Bundle ID must be set");
     Preconditions.checkArgument(

diff --git a/firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalStore.java b/firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalStore.java
@@ -442,8 +442,7 @@ public ImmutableSortedMap<DocumentKey, Document> applyRemoteEvent(RemoteEvent re
             }
           }
 
-          DocumentChangeResult result =
-              populateDocumentChanges(documentUpdates, null, remoteEvent.getSnapshotVersion());
+          DocumentChangeResult result = populateDocumentChanges(documentUpdates);
           Map<DocumentKey, MutableDocument> changedDocs = result.changedDocuments;
 
           // HACK: The only reason we allow snapshot version NONE is so that we can synthesize
@@ -483,15 +482,9 @@ private DocumentChangeResult(
    * resorts to `globalVersion`.
    *
    * @param documents Documents to be applied.
-   * @param documentVersions A DocumentKey-to-SnapshotVersion map if documents have their own read
-   *     time.
-   * @param globalVersion A SnapshotVersion representing the read time if all documents have the
-   *     same read time.
    */
   private DocumentChangeResult populateDocumentChanges(
-      Map<DocumentKey, MutableDocument> documents,
-      @Nullable Map<DocumentKey, SnapshotVersion> documentVersions,
-      SnapshotVersion globalVersion) {
+      Map<DocumentKey, MutableDocument> documents) {
     Map<DocumentKey, MutableDocument> changedDocs = new HashMap<>();
     List<DocumentKey> removedDocs = new ArrayList<>();
     Set<DocumentKey> conditionChanged = new HashSet<>();
@@ -504,8 +497,6 @@ private DocumentChangeResult populateDocumentChanges(
       DocumentKey key = entry.getKey();
       MutableDocument doc = entry.getValue();
       MutableDocument existingDoc = existingDocs.get(key);
-      SnapshotVersion readTime =
-          documentVersions != null ? documentVersions.get(key) : globalVersion;
       // Check if see if there is a existence state change for this document.
       if (doc.isFoundDocument() != existingDoc.isFoundDocument()) {
         conditionChanged.add(key);
@@ -524,9 +515,9 @@ private DocumentChangeResult populateDocumentChanges(
           || (doc.getVersion().compareTo(existingDoc.getVersion()) == 0
               && existingDoc.hasPendingWrites())) {
         hardAssert(
-            !SnapshotVersion.NONE.equals(readTime),
+            !SnapshotVersion.NONE.equals(doc.getReadTime()),
             "Cannot add a document when the remote version is zero");
-        remoteDocuments.add(doc, readTime);
+        remoteDocuments.add(doc, doc.getReadTime());
         changedDocs.put(key, doc);
       } else {
         Logger.debug(
@@ -713,7 +704,6 @@ public ImmutableSortedMap<DocumentKey, Document> applyBundledDocuments(
         () -> {
           ImmutableSortedSet<DocumentKey> documentKeys = DocumentKey.emptyKeySet();
           Map<DocumentKey, MutableDocument> documentMap = new HashMap<>();
-          Map<DocumentKey, SnapshotVersion> versionMap = new HashMap<>();
 
           for (Entry<DocumentKey, MutableDocument> entry : documents) {
             DocumentKey documentKey = entry.getKey();
@@ -723,14 +713,12 @@ public ImmutableSortedMap<DocumentKey, Document> applyBundledDocuments(
               documentKeys = documentKeys.insert(documentKey);
             }
             documentMap.put(documentKey, document);
-            versionMap.put(documentKey, document.getVersion());
           }
 
           targetCache.removeMatchingKeysForTargetId(umbrellaTargetData.getTargetId());
           targetCache.addMatchingKeys(documentKeys, umbrellaTargetData.getTargetId());
 
-          DocumentChangeResult result =
-              populateDocumentChanges(documentMap, versionMap, SnapshotVersion.NONE);
+          DocumentChangeResult result = populateDocumentChanges(documentMap);
           Map<DocumentKey, MutableDocument> changedDocs = result.changedDocuments;
           return localDocuments.getLocalViewOfDocuments(changedDocs, result.existenceChangedKeys);
         });

diff --git a/...irestore/src/main/java/com/google/firebase/firestore/local/MemoryRemoteDocumentCache.java b/...irestore/src/main/java/com/google/firebase/firestore/local/MemoryRemoteDocumentCache.java
@@ -53,7 +53,7 @@ public void add(MutableDocument document, SnapshotVersion readTime) {
     hardAssert(
         !readTime.equals(SnapshotVersion.NONE),
         "Cannot add document to the RemoteDocumentCache with a read time of zero");
-    docs = docs.insert(document.getKey(), document.mutableCopy().withReadTime(readTime));
+    docs = docs.insert(document.getKey(), document.mutableCopy().setReadTime(readTime));
 
     indexManager.addToCollectionParentIndex(document.getKey().getCollectionPath());
   }

diff --git a/...irestore/src/main/java/com/google/firebase/firestore/local/SQLiteRemoteDocumentCache.java b/...irestore/src/main/java/com/google/firebase/firestore/local/SQLiteRemoteDocumentCache.java
@@ -260,7 +260,7 @@ private MutableDocument decodeMaybeDocument(
     try {
       return serializer
           .decodeMaybeDocument(com.google.firebase.firestore.proto.MaybeDocument.parseFrom(bytes))
-          .withReadTime(new SnapshotVersion(new Timestamp(readTimeSeconds, readTimeNanos)));
+          .setReadTime(new SnapshotVersion(new Timestamp(readTimeSeconds, readTimeNanos)));
     } catch (InvalidProtocolBufferException e) {
       throw fail("MaybeDocument failed to parse: %s", e);
     }

diff --git a/firebase-firestore/src/main/java/com/google/firebase/firestore/model/MutableDocument.java b/firebase-firestore/src/main/java/com/google/firebase/firestore/model/MutableDocument.java
@@ -163,7 +163,7 @@ public MutableDocument setHasLocalMutations() {
     return this;
   }
 
-  public MutableDocument withReadTime(SnapshotVersion readTime) {
+  public MutableDocument setReadTime(SnapshotVersion readTime) {
     this.readTime = readTime;
     return this;
   }

diff --git a/...e-firestore/src/main/java/com/google/firebase/firestore/remote/WatchChangeAggregator.java b/...e-firestore/src/main/java/com/google/firebase/firestore/remote/WatchChangeAggregator.java
@@ -264,6 +264,10 @@ public RemoteEvent createRemoteEvent(SnapshotVersion snapshotVersion) {
       }
     }
 
+    for (MutableDocument document : pendingDocumentUpdates.values()) {
+      document.setReadTime(snapshotVersion);
+    }
+
     RemoteEvent remoteEvent =
         new RemoteEvent(
             snapshotVersion,

diff --git a/firebase-firestore/src/test/java/com/google/firebase/firestore/local/LocalStoreTestCase.java b/firebase-firestore/src/test/java/com/google/firebase/firestore/local/LocalStoreTestCase.java
@@ -1358,7 +1358,7 @@ public void testHoldsBackTransforms() {
     acknowledgeMutationWithTransformResults(3, 1338, asList("bar", "foo"));
     assertChanged(
         doc("foo/bar", 3, map("sum", 1338, "array_union", asList("bar", "foo")))
-            .withReadTime(new SnapshotVersion(new Timestamp(0, 3000)))
+            .setReadTime(new SnapshotVersion(new Timestamp(0, 3000)))
             .setHasCommittedMutations());
   }
 

diff --git a/.../src/test/java/com/google/firebase/firestore/local/SQLiteOverlayMigrationManagerTest.java b/.../src/test/java/com/google/firebase/firestore/local/SQLiteOverlayMigrationManagerTest.java
@@ -65,7 +65,7 @@ public void tearDown() {
 
   private void writeRemoteDocument(MutableDocument document) {
     // Set read time to update time.
-    document = document.withReadTime(document.getVersion());
+    document.setReadTime(document.getVersion());
     persistence.getRemoteDocumentCache().add(document, document.getReadTime());
   }
 

diff --git a/firebase-firestore/src/testUtil/java/com/google/firebase/firestore/testutil/TestUtil.java b/firebase-firestore/src/testUtil/java/com/google/firebase/firestore/testutil/TestUtil.java
@@ -215,11 +215,12 @@ public static MutableDocument doc(String key, long version, ObjectValue data) {
   }
 
   public static MutableDocument doc(DocumentKey key, long version, ObjectValue data) {
-    return MutableDocument.newFoundDocument(key, version(version), data);
+    return MutableDocument.newFoundDocument(key, version(version), data)
+        .setReadTime(version(version));
   }
 
   public static MutableDocument deletedDoc(String key, long version) {
-    return MutableDocument.newNoDocument(key(key), version(version));
+    return MutableDocument.newNoDocument(key(key), version(version)).setReadTime(version(version));
   }
 
   public static MutableDocument unknownDoc(String key, long version) {