Fixing race in FSTWriteStream

2 files changed, 10 insertions, 7 deletions

diff --git a/Firestore/Source/Remote/FSTStream.m b/Firestore/Source/Remote/FSTStream.m
index 6cce921..bf54a6b 100644
--- a/Firestore/Source/Remote/FSTStream.m
+++ b/Firestore/Source/Remote/FSTStream.m
@@ -381,17 +381,17 @@ static const NSTimeInterval kIdleTimeout = 60.0;
   _messageReceived = NO;
   _rpc = nil;
 
-  // If the caller explicitly requested a stream stop, don't notify them of a closing stream (it
-  // could trigger undesirable recovery logic, etc.).
-  if (finalState != FSTStreamStateStopped) {
-    [self notifyStreamInterruptedWithError:error];
-  }
-
   // Clear the delegates to avoid any possible bleed through of events from GRPC.
   FSTAssert(_delegate,
             @"closeWithFinalState should only be called for a started stream that has an active "
             @"delegate.");
   _delegate = nil;
+
+  // If the caller explicitly requested a stream stop, don't notify them of a closing stream (it
+  // could trigger undesirable recovery logic, etc.).
+  if (finalState != FSTStreamStateStopped) {
+    [self notifyStreamInterruptedWithError:error];
+  }
 }
 
 - (void)stop {
diff --git a/Firestore/Source/Util/FSTDispatchQueue.m b/Firestore/Source/Util/FSTDispatchQueue.m
index 8c953fe..2c8ee55 100644
--- a/Firestore/Source/Util/FSTDispatchQueue.m
+++ b/Firestore/Source/Util/FSTDispatchQueue.m
@@ -58,7 +58,10 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)dispatchAfterDelay:(NSTimeInterval)delay block:(void (^)(void))block {
   dispatch_time_t delayNs = dispatch_time(DISPATCH_TIME_NOW, (int64_t)(delay * NSEC_PER_SEC));
-  dispatch_after(delayNs, self.queue, block);
+  dispatch_after(delayNs, self.queue, ^() {
+    // Make sure that we prioritize tasks that are already queued for immediate execution.
+    [self dispatchAsyncAllowingSameQueue:block];
+  });
 }
 
 #pragma mark - Private Methods