Removing the target state on stream error (#1420)

2 files changed, 6 insertions, 1 deletions

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];
     }
   }
 }