Merge

Author: schmidt-sebastian

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

@@ -208,7 +208,7 @@ public ImmutableSortedMap<DocumentKey, MaybeDocument> acknowledgeBatch(
         () -> {
           MutationBatch batch = batchResult.getBatch();
           mutationQueue.acknowledgeBatch(batch, batchResult.getStreamToken());
-          applyBatchResult(batchResult);
+          applyWriteToRemoteDocuments(batchResult);
           mutationQueue.performConsistencyCheck();
           return localDocuments.getDocuments(batch.getKeys());
         });
@@ -503,7 +503,15 @@ public void releaseQuery(Query query) {
             queryCache.updateQueryData(queryData);
           }
 
-          localViewReferences.removeReferencesForId(queryData.getTargetId());
+          // References for documents sent via Watch are automatically removed when we delete a
+          // query's target data from the reference delegate. Since this does not remove references
+          // for locally mutated documents, we have to remove the target associations for these
+          // documents manually.
+          ImmutableSortedSet<DocumentKey> removedReferences =
+              localViewReferences.removeReferencesForId(queryData.getTargetId());
+          for (DocumentKey key : removedReferences) {
+            persistence.getReferenceDelegate().removeReference(key);
+          }
           persistence.getReferenceDelegate().removeTarget(queryData);
           targetIds.remove(queryData.getTargetId());
         });
@@ -522,7 +530,7 @@ public ImmutableSortedSet<DocumentKey> getRemoteDocumentKeys(int targetId) {
     return queryCache.getMatchingKeysForTargetId(targetId);
   }
 
-  private void applyBatchResult(MutationBatchResult batchResult) {
+  private void applyWriteToRemoteDocuments(MutationBatchResult batchResult) {
     MutationBatch batch = batchResult.getBatch();
     Set<DocumentKey> docKeys = batch.getKeys();
     for (DocumentKey docKey : docKeys) {

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

@@ -77,19 +77,27 @@ public void removeReferences(ImmutableSortedSet<DocumentKey> keys, int targetOrB
     }
   }
 
-  /** Clears all references with a given ID. Calls removeReference() for each key removed. */
-  public void removeReferencesForId(int targetId) {
+  /**
+   * Clears all references with a given ID. Calls removeReference() for each key removed.
+   *
+   * @return The keys of the documents that were removed.
+   */
+  public ImmutableSortedSet<DocumentKey> removeReferencesForId(int targetId) {
     DocumentKey emptyKey = DocumentKey.empty();
     DocumentReference startRef = new DocumentReference(emptyKey, targetId);
     Iterator<DocumentReference> it = referencesByTarget.iteratorFrom(startRef);
+    ImmutableSortedSet<DocumentKey> keys = DocumentKey.emptyKeySet();
     while (it.hasNext()) {
       DocumentReference ref = it.next();
       if (ref.getId() == targetId) {
+        keys = keys.insert(ref.getKey());
         removeReference(ref);
       } else {
         break;
       }
     }
+
+    return keys;
   }
 
   /** Clears all references for all IDs. */

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

@@ -243,6 +243,31 @@ public void testHandlesSetMutationThenDocument() {
     assertContains(doc("foo/bar", 2, map("foo", "bar"), Document.DocumentState.LOCAL_MUTATIONS));
   }
 
+  @Test
+  public void testHandlesSetMutationThenAckThenRelease() {
+    Query query = Query.atPath(ResourcePath.fromSegments(asList("foo")));
+    allocateQuery(query);
+
+    writeMutation(setMutation("foo/bar", map("foo", "bar")));
+    assertChanged(doc("foo/bar", 0, map("foo", "bar"), Document.DocumentState.LOCAL_MUTATIONS));
+    assertContains(doc("foo/bar", 0, map("foo", "bar"), Document.DocumentState.LOCAL_MUTATIONS));
+
+    acknowledgeMutation(1);
+    notifyLocalViewChanges(viewChanges(2, asList("foo/bar"), emptyList()));
+
+    assertChanged(doc("foo/bar", 1, map("foo", "bar"), Document.DocumentState.COMMITTED_MUTATIONS));
+    assertContains(doc("foo/bar", 1, map("foo", "bar"), Document.DocumentState.COMMITTED_MUTATIONS));
+
+    releaseQuery(query);
+
+    // It has been acknowledged, and should no longer be retained as there is no target and mutation
+    if (garbageCollectorIsEager()) {
+      assertNotContains("foo/bar");
+    } else {
+      assertContains(doc("foo/bar", 0, map("foo", "bar"), Document.DocumentState.COMMITTED_MUTATIONS));
+    }
+  }
+
   @Test
   public void testHandlesAckThenRejectThenRemoteEvent() {
     // Start a query that requires acks to be held.