Remove local references from reference delegate (#42)

Author: schmidt-sebastian

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

@@ -558,7 +558,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());
 
@@ -622,7 +630,7 @@ private Set<DocumentKey> releaseBatchResults(List<MutationBatchResult> batchResu
     ArrayList<MutationBatch> batches = new ArrayList<>(batchResults.size());
     // TODO: Call queryEngine.handleDocumentChange() as appropriate.
     for (MutationBatchResult batchResult : batchResults) {
-      applyBatchResult(batchResult);
+      applyWriteToRemoteDocuments(batchResult);
       batches.add(batchResult.getBatch());
     }
 
@@ -647,7 +655,7 @@ private Set<DocumentKey> removeMutationBatches(List<MutationBatch> batches) {
     return affectedDocs;
   }
 
-  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

@@ -239,6 +239,31 @@ public void testHandlesSetMutationThenDocument() {
     assertContains(doc("foo/bar", 2, map("foo", "bar"), true));
   }
 
+  @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"), true));
+    assertContains(doc("foo/bar", 0, map("foo", "bar"), true));
+
+    acknowledgeMutation(1);
+    notifyLocalViewChanges(viewChanges(2, asList("foo/bar"), emptyList()));
+
+    assertChanged();
+    assertContains(doc("foo/bar", 0, map("foo", "bar"), true));
+
+    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"), false));
+    }
+  }
+
   @Test
   public void testHandlesAckThenRejectThenRemoteEvent() {
     // Start a query that requires acks to be held.