Removing the target state on stream error (#1420)

3 files changed, 6 insertions, 3 deletions

diff --git a/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm b/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
index c6936f7..715f09b 100644
--- a/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
+++ b/Firestore/Example/Tests/Remote/FSTRemoteEventTests.mm
@@ -93,8 +93,6 @@ NS_ASSUME_NONNULL_BEGIN
  * Creates an aggregator initialized with the set of provided FSTWatchChanges. Tests can add further
  * changes via `handleDocumentChange`, `handleTargetChange` and `handleExistenceFilterChange`.
  *
- * @param snapshotVersion The version at which to create the remote event. This corresponds to the
- * snapshot version provided by a NO_CHANGE event.
  * @param targetMap A map of query data for all active targets. The map must include an entry for
  * every target referenced by any of the watch changes.
  * @param outstandingResponses The number of outstanding ACKs a target has to receive before it is
diff --git a/Firestore/Source/Remote/FSTRemoteEvent.h b/Firestore/Source/Remote/FSTRemoteEvent.h
index 9ea0f9c..02a5289 100644
--- a/Firestore/Source/Remote/FSTRemoteEvent.h
+++ b/Firestore/Source/Remote/FSTRemoteEvent.h
@@ -176,6 +176,9 @@ initWithSnapshotVersion:(firebase::firestore::model::SnapshotVersion)snapshotVer
 /** Processes and adds the WatchTargetChange to the current set of changes. */
 - (void)handleTargetChange:(FSTWatchTargetChange *)targetChange;
 
+/** Removes the in-memory state for the provided target. */
+- (void)removeTarget:(FSTTargetID)targetID;
+
 /**
  * Handles existence filters and synthesizes deletes for filter mismatches. Targets that are
  * invalidated by filter mismatches are added to `targetMismatches`.
diff --git a/Firestore/Source/Remote/FSTRemoteStore.mm b/Firestore/Source/Remote/FSTRemoteStore.mm
index 54d00c4..4309c74 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.mm
+++ b/Firestore/Source/Remote/FSTRemoteStore.mm
@@ -405,8 +405,10 @@ static const int kMaxPendingWrites = 10;
   // Ignore targets that have been removed already.
   for (FSTBoxedTargetID *targetID in change.targetIDs) {
     if (self.listenTargets[targetID]) {
+      int unboxedTargetId = targetID.intValue;
       [self.listenTargets removeObjectForKey:targetID];
-      [self.syncEngine rejectListenWithTargetID:[targetID intValue] error:change.cause];
+      [self.watchChangeAggregator removeTarget:unboxedTargetId];
+      [self.syncEngine rejectListenWithTargetID:unboxedTargetId error:change.cause];
     }
   }
 }