Defend against users making API calls during deinit (#863)

Track operation in progress FSTDispatchQueue

... and allow dispatchAsync onto the worker queue even from the same
queue if an operation is not in progress. Fixes #753.

This fixes a corner case that happens when users pass us an object and
we retain that object in a block submitted to the worker queue for
processing. Users are not obligated to retain that object themselves.
This can lead to the object's destructor (or dealloc or deinit) being
called on the Firestore worker thread. If that destructor makes
Firestore API calls (most likely ListenerRegistration.remove but any are
possible) this will trigger the assertion.

5 files changed, 233 insertions, 34 deletions

diff --git a/Firestore/Example/Tests/Integration/FSTStreamTests.mm b/Firestore/Example/Tests/Integration/FSTStreamTests.mm
index 6550368..7e37913 100644
--- a/Firestore/Example/Tests/Integration/FSTStreamTests.mm
+++ b/Firestore/Example/Tests/Integration/FSTStreamTests.mm
@@ -243,9 +243,9 @@ using firebase::firestore::model::DatabaseId;
   }];
 
   // Writing before the handshake should throw
-  dispatch_sync(_testQueue, ^{
+  [_workerDispatchQueue dispatchSync:^{
     XCTAssertThrows([writeStream writeMutations:_mutations]);
-  });
+  }];
 
   [_delegate awaitNotificationFromBlock:^{
     [writeStream writeHandshake];
@@ -285,9 +285,9 @@ using firebase::firestore::model::DatabaseId;
 
   [_workerDispatchQueue runDelayedCallbacksUntil:FSTTimerIDWriteStreamIdle];
 
-  dispatch_sync(_testQueue, ^{
+  [_workerDispatchQueue dispatchSync:^{
     XCTAssertFalse([writeStream isOpen]);
-  });
+  }];
 
   [self verifyDelegateObservedStates:@[
     @"writeStreamDidOpen", @"writeStreamDidCompleteHandshake", @"writeStreamWasInterrupted"
@@ -315,9 +315,9 @@ using firebase::firestore::model::DatabaseId;
         [_workerDispatchQueue containsDelayedCallbackWithTimerID:FSTTimerIDWriteStreamIdle]);
   }];
 
-  dispatch_sync(_testQueue, ^{
+  [_workerDispatchQueue dispatchSync:^{
     XCTAssertTrue([writeStream isOpen]);
-  });
+  }];
 
   [self verifyDelegateObservedStates:@[
     @"writeStreamDidOpen", @"writeStreamDidCompleteHandshake",
diff --git a/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm b/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm
index 5ef860c..60b1705 100644
--- a/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm
+++ b/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm
@@ -29,6 +29,7 @@ static const FSTTimerID timerID3 = FSTTimerIDWriteStreamConnectionBackoff;
 @end
 
 @implementation FSTDispatchQueueTests {
+  dispatch_queue_t _underlyingQueue;
   FSTDispatchQueue *_queue;
   NSMutableArray *_completedSteps;
   NSArray *_expectedSteps;
@@ -37,13 +38,167 @@ static const FSTTimerID timerID3 = FSTTimerIDWriteStreamConnectionBackoff;
 
 - (void)setUp {
   [super setUp];
-  dispatch_queue_t dispatch_queue =
-      dispatch_queue_create("FSTDispatchQueueTests", DISPATCH_QUEUE_SERIAL);
-  _queue = [[FSTDispatchQueue alloc] initWithQueue:dispatch_queue];
+  _underlyingQueue = dispatch_queue_create("FSTDispatchQueueTests", DISPATCH_QUEUE_SERIAL);
+  _queue = [[FSTDispatchQueue alloc] initWithQueue:_underlyingQueue];
   _completedSteps = [NSMutableArray array];
   _expectedSteps = nil;
 }
 
+- (void)testDispatchAsyncBlocksSubmissionFromTasksOnTheQueue {
+  XCTestExpectation *expectation = [self expectationWithDescription:@"completion"];
+  __block NSException *caught = nil;
+  __block NSString *problem = nil;
+
+  [_queue dispatchAsync:^{
+    @try {
+      [self->_queue dispatchAsync:^{
+      }];
+      problem = @"Should have disallowed submission into the queue while running";
+      [expectation fulfill];
+    } @catch (NSException *ex) {
+      caught = ex;
+      [expectation fulfill];
+    }
+  }];
+
+  [self awaitExpectations];
+  XCTAssertNil(problem);
+  XCTAssertNotNil(caught);
+
+  XCTAssertEqualObjects(caught.name, NSInternalInconsistencyException);
+  XCTAssertTrue(
+      [caught.reason hasPrefix:@"FIRESTORE INTERNAL ASSERTION FAILED: "
+                               @"dispatchAsync called when we are already running on target"]);
+}
+
+- (void)testDispatchAsyncAllowingSameQueueActuallyAllowsSameQueue {
+  XCTestExpectation *expectation = [self expectationWithDescription:@"completion"];
+  __block NSException *caught = nil;
+
+  [_queue dispatchAsync:^{
+    @try {
+      [self->_queue dispatchAsyncAllowingSameQueue:^{
+        [expectation fulfill];
+      }];
+    } @catch (NSException *ex) {
+      caught = ex;
+      [expectation fulfill];
+    }
+  }];
+
+  [self awaitExpectations];
+  XCTAssertNil(caught);
+}
+
+- (void)testDispatchAsyncAllowsSameQueueForUnownedActions {
+  XCTestExpectation *expectation = [self expectationWithDescription:@"completion"];
+  __block NSException *caught = nil;
+
+  // Simulate the case of an action that runs on our queue because e.g. it's run by a user-owned
+  // deinitializer that happened to be last held in one of our API methods.
+  dispatch_async(_underlyingQueue, ^{
+    @try {
+      [self->_queue dispatchAsync:^{
+        [expectation fulfill];
+      }];
+    } @catch (NSException *ex) {
+      caught = ex;
+      [expectation fulfill];
+    }
+  });
+
+  [self awaitExpectations];
+  XCTAssertNil(caught);
+}
+
+- (void)testDispatchSyncBlocksSubmissionFromTasksOnTheQueue {
+  XCTestExpectation *expectation = [self expectationWithDescription:@"completion"];
+  __block NSException *caught = nil;
+  __block NSString *problem = nil;
+
+  [_queue dispatchSync:^{
+    @try {
+      [self->_queue dispatchSync:^{
+      }];
+      problem = @"Should have disallowed submission into the queue while running";
+      [expectation fulfill];
+    } @catch (NSException *ex) {
+      caught = ex;
+      [expectation fulfill];
+    }
+  }];
+
+  [self awaitExpectations];
+  XCTAssertNil(problem);
+  XCTAssertNotNil(caught);
+
+  XCTAssertEqualObjects(caught.name, NSInternalInconsistencyException);
+  XCTAssertTrue(
+      [caught.reason hasPrefix:@"FIRESTORE INTERNAL ASSERTION FAILED: "
+                               @"dispatchSync called when we are already running on target"]);
+}
+
+- (void)testVerifyIsCurrentQueueActuallyRequiresCurrentQueue {
+  XCTAssertNotEqualObjects(_underlyingQueue, dispatch_get_main_queue());
+
+  __block NSException *caught = nil;
+  @try {
+    // Run on the main queue not the FSTDispatchQueue's queue
+    [_queue verifyIsCurrentQueue];
+  } @catch (NSException *ex) {
+    caught = ex;
+  }
+  XCTAssertNotNil(caught);
+  XCTAssertTrue([caught.reason hasPrefix:@"FIRESTORE INTERNAL ASSERTION FAILED: "
+                                         @"We are running on the wrong dispatch queue"]);
+}
+
+- (void)testVerifyIsCurrentQueueRequiresOperationIsInProgress {
+  __block NSException *caught = nil;
+  dispatch_sync(_underlyingQueue, ^{
+    @try {
+      [_queue verifyIsCurrentQueue];
+    } @catch (NSException *ex) {
+      caught = ex;
+    }
+  });
+  XCTAssertNotNil(caught);
+  XCTAssertTrue(
+      [caught.reason hasPrefix:@"FIRESTORE INTERNAL ASSERTION FAILED: "
+                               @"verifyIsCurrentQueue called outside enterCheckedOperation"]);
+}
+
+- (void)testVerifyIsCurrentQueueWorksWithOperationIsInProgress {
+  __block NSException *caught = nil;
+  [_queue dispatchSync:^{
+    @try {
+      [_queue verifyIsCurrentQueue];
+    } @catch (NSException *ex) {
+      caught = ex;
+    }
+  }];
+  XCTAssertNil(caught);
+}
+
+- (void)testEnterCheckedOperationDisallowsNesting {
+  __block NSException *caught = nil;
+  __block NSString *problem = nil;
+  [_queue dispatchSync:^{
+    @try {
+      [_queue enterCheckedOperation:^{
+      }];
+      problem = @"Should not have been able to enter nested enterCheckedOperation";
+    } @catch (NSException *ex) {
+      caught = ex;
+    }
+  }];
+  XCTAssertNil(problem);
+  XCTAssertNotNil(caught);
+  XCTAssertTrue([caught.reason
+      hasPrefix:@"FIRESTORE INTERNAL ASSERTION FAILED: "
+                @"enterCheckedOperation may not be called when an operation is in progress"]);
+}
+
 /**
  * Helper to return a block that adds @(n) to _completedSteps when run and fulfils _expectation if
  * the _completedSteps match the _expectedSteps.
diff --git a/Firestore/Source/Remote/FSTStream.mm b/Firestore/Source/Remote/FSTStream.mm
index 44e3ef0..019b0bc 100644
--- a/Firestore/Source/Remote/FSTStream.mm
+++ b/Firestore/Source/Remote/FSTStream.mm
@@ -564,24 +564,25 @@ static const NSTimeInterval kIdleTimeout = 60.0;
  * the RPC.
  */
 - (void)writeValue:(id)value {
-  [self.workerDispatchQueue verifyIsCurrentQueue];
-  FSTAssert([self isStarted], @"writeValue: called for stopped stream.");
-
-  if (!self.messageReceived) {
-    self.messageReceived = YES;
-    if ([FIRFirestore isLoggingEnabled]) {
-      FSTLog(@"%@ %p headers (whitelisted): %@", NSStringFromClass([self class]),
-             (__bridge void *)self,
-             [FSTDatastore extractWhiteListedHeaders:self.rpc.responseHeaders]);
+  [self.workerDispatchQueue enterCheckedOperation:^{
+    FSTAssert([self isStarted], @"writeValue: called for stopped stream.");
+
+    if (!self.messageReceived) {
+      self.messageReceived = YES;
+      if ([FIRFirestore isLoggingEnabled]) {
+        FSTLog(@"%@ %p headers (whitelisted): %@", NSStringFromClass([self class]),
+               (__bridge void *)self,
+               [FSTDatastore extractWhiteListedHeaders:self.rpc.responseHeaders]);
+      }
     }
-  }
-  NSError *error;
-  id proto = [self parseProto:self.responseMessageClass data:value error:&error];
-  if (proto) {
-    [self handleStreamMessage:proto];
-  } else {
-    [_rpc finishWithError:error];
-  }
+    NSError *error;
+    id proto = [self parseProto:self.responseMessageClass data:value error:&error];
+    if (proto) {
+      [self handleStreamMessage:proto];
+    } else {
+      [self.rpc finishWithError:error];
+    }
+  }];
 }
 
 /**
@@ -597,10 +598,11 @@ static const NSTimeInterval kIdleTimeout = 60.0;
  */
 - (void)writesFinishedWithError:(nullable NSError *)error __used {
   error = [FSTDatastore firestoreErrorForError:error];
-  [self.workerDispatchQueue verifyIsCurrentQueue];
-  FSTAssert([self isStarted], @"writesFinishedWithError: called for stopped stream.");
+  [self.workerDispatchQueue enterCheckedOperation:^{
+    FSTAssert([self isStarted], @"writesFinishedWithError: called for stopped stream.");
 
-  [self handleStreamClose:error];
+    [self handleStreamClose:error];
+  }];
 }
 
 @end
diff --git a/Firestore/Source/Util/FSTDispatchQueue.h b/Firestore/Source/Util/FSTDispatchQueue.h
index 7922600..8e9273c 100644
--- a/Firestore/Source/Util/FSTDispatchQueue.h
+++ b/Firestore/Source/Util/FSTDispatchQueue.h
@@ -75,6 +75,14 @@ typedef NS_ENUM(NSInteger, FSTTimerID) {
 - (void)verifyIsCurrentQueue;
 
 /**
+ * Declares that we are already executing on the correct dispatch_queue_t and would like to
+ * officially execute code on behalf of this FSTDispatchQueue. To be used only when called  back
+ * by some other API directly onto our queue. This allows us to safely dispatch directly onto the
+ * worker queue without destroying the invariants this class helps us maintain.
+ */
+- (void)enterCheckedOperation:(void (^)(void))block;
+
+/**
  * Same as dispatch_async() except it asserts that we're not already on the queue, since this
  * generally indicates a bug (and can lead to re-ordering of operations, etc).
  *
diff --git a/Firestore/Source/Util/FSTDispatchQueue.mm b/Firestore/Source/Util/FSTDispatchQueue.mm
index 3184d29..15d6e7b 100644
--- a/Firestore/Source/Util/FSTDispatchQueue.mm
+++ b/Firestore/Source/Util/FSTDispatchQueue.mm
@@ -104,7 +104,9 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)startWithDelay:(NSTimeInterval)delay {
   dispatch_time_t delayNs = dispatch_time(DISPATCH_TIME_NOW, (int64_t)(delay * NSEC_PER_SEC));
   dispatch_after(delayNs, self.queue.queue, ^{
-    [self delayDidElapse];
+    [self.queue enterCheckedOperation:^{
+      [self delayDidElapse];
+    }];
   });
 }
 
@@ -151,6 +153,12 @@ NS_ASSUME_NONNULL_BEGIN
  */
 @property(nonatomic, strong, readonly) NSMutableArray<FSTDelayedCallback *> *delayedCallbacks;
 
+/**
+ * Flag set while an FSTDispatchQueue operation is currently executing. Used for assertion
+ * sanity-checks.
+ */
+@property(nonatomic, assign) BOOL operationInProgress;
+
 - (instancetype)initWithQueue:(dispatch_queue_t)queue NS_DESIGNATED_INITIALIZER;
 
 @end
@@ -165,6 +173,7 @@ NS_ASSUME_NONNULL_BEGIN
   if (self = [super init]) {
     _queue = queue;
     _delayedCallbacks = [NSMutableArray array];
+    _operationInProgress = NO;
   }
   return self;
 }
@@ -173,22 +182,47 @@ NS_ASSUME_NONNULL_BEGIN
   FSTAssert([self onTargetQueue],
             @"We are running on the wrong dispatch queue. Expected '%@' Actual: '%@'",
             [self targetQueueLabel], [self currentQueueLabel]);
+  FSTAssert(_operationInProgress,
+            @"verifyIsCurrentQueue called outside enterCheckedOperation on queue '%@'",
+            [self currentQueueLabel]);
+}
+
+- (void)enterCheckedOperation:(void (^)(void))block {
+  FSTAssert(!_operationInProgress,
+            @"enterCheckedOperation may not be called when an operation is in progress");
+  @try {
+    _operationInProgress = YES;
+    [self verifyIsCurrentQueue];
+    block();
+  } @finally {
+    _operationInProgress = NO;
+  }
 }
 
 - (void)dispatchAsync:(void (^)(void))block {
-  FSTAssert(![self onTargetQueue],
+  FSTAssert(!_operationInProgress || ![self onTargetQueue],
             @"dispatchAsync called when we are already running on target dispatch queue '%@'",
             [self targetQueueLabel]);
 
-  dispatch_async(self.queue, block);
+  dispatch_async(self.queue, ^{
+    [self enterCheckedOperation:block];
+  });
 }
 
 - (void)dispatchAsyncAllowingSameQueue:(void (^)(void))block {
-  dispatch_async(self.queue, block);
+  dispatch_async(self.queue, ^{
+    [self enterCheckedOperation:block];
+  });
 }
 
 - (void)dispatchSync:(void (^)(void))block {
-  dispatch_sync(self.queue, block);
+  FSTAssert(!_operationInProgress || ![self onTargetQueue],
+            @"dispatchSync called when we are already running on target dispatch queue '%@'",
+            [self targetQueueLabel]);
+
+  dispatch_sync(self.queue, ^{
+    [self enterCheckedOperation:block];
+  });
 }
 
 - (FSTDelayedCallback *)dispatchAfterDelay:(NSTimeInterval)delay