DispatchQueue delayed callback improvements + testing (#784)

Basically a port of https://github.com/firebase/firebase-js-sdk/commit/a1e346ff93c6cbcc0a1b3b33f0fbc3a7b66e7e12 and https://github.com/firebase/firebase-js-sdk/commit/fce4168309f42aa038125f39818fbf654b65b05f

* Introduces a DelayedCallback helper class in FSTDispatchQueue to encapsulate delayed callback logic.
* Adds cancellation support.
* Updates the idle timer in FSTStream to use new cancellation support.
* Adds a FSTTimerId enum for identifying delayed operations on the queue and uses it to identify our existing backoff and idle timers.
* Added containsDelayedCallback: and runDelayedCallbacksUntil: methods to FSTDispatchQueue which can be used from tests to check for the presence of a callback or to schedule them to run early.
* Removes FSTTestDispatchQueue and changes idle tests to use new test methods.

2 files changed, 248 insertions, 5 deletions

diff --git a/Firestore/Source/Util/FSTDispatchQueue.h b/Firestore/Source/Util/FSTDispatchQueue.h
index fe87887..9b28c9c 100644
--- a/Firestore/Source/Util/FSTDispatchQueue.h
+++ b/Firestore/Source/Util/FSTDispatchQueue.h
@@ -18,6 +18,34 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+/**
+ * Well-known "timer" IDs used when scheduling delayed callbacks on the FSTDispatchQueue. These IDs
+ * can then be used from tests to check for the presence of callbacks or to run them early.
+ */
+typedef NS_ENUM(NSInteger, FSTTimerID) {
+  FSTTimerIDAll,  // Sentinel value to be used with runDelayedCallbacksUntil: to run all blocks.
+  FSTTimerIDListenStreamIdle,
+  FSTTimerIDListenStreamConnection,
+  FSTTimerIDWriteStreamIdle,
+  FSTTimerIDWriteStreamConnection
+};
+
+/**
+ * Handle to a callback scheduled via [FSTDispatchQueue dispatchAfterDelay:]. Supports cancellation
+ * via the cancel method.
+ */
+@interface FSTDelayedCallback : NSObject
+
+/**
+ * Cancels the callback if it hasn't already been executed or canceled.
+ *
+ * As long as the callback has not yet been run, calling cancel() (from a callback already running
+ * on the dispatch queue) provides a guarantee that the operation will not be run.
+ */
+- (void)cancel;
+
+@end
+
 @interface FSTDispatchQueue : NSObject
 
 /** Creates and returns an FSTDispatchQueue wrapping the specified dispatch_queue_t. */
@@ -56,12 +84,32 @@ NS_ASSUME_NONNULL_BEGIN
  * Schedules a callback after the specified delay.
  *
  * Unlike dispatchAsync: this method does not require you to dispatch to a different queue than
- * the current one (thus it is equivalent to a raw dispatch_after()).
+ * the current one.
+ *
+ * The returned FSTDelayedCallback handle can be used to cancel the callback prior to its running.
  *
  * @param block The block to run.
  * @param delay The delay (in seconds) after which to run the block.
+ * @param timerID An FSTTimerID that can be used from tests to check for the presence of this
+ *   callback or to schedule it to run early.
+ * @return A FSTDelayedCallback instance that can be used for cancellation.
+ */
+- (FSTDelayedCallback *)dispatchAfterDelay:(NSTimeInterval)delay
+                                   timerID:(FSTTimerID)timerID
+                                     block:(void (^)(void))block;
+
+/**
+ * For Tests: Determine if a delayed callback with a particular FSTTimerID exists.
+ */
+- (BOOL)containsDelayedCallbackWithTimerID:(FSTTimerID)timerID;
+
+/**
+ * For Tests: Runs delayed callbacks early, blocking until completion.
+ *
+ * @param lastTimerID Only delayed callbacks up to and including one that was scheduled using this
+ *   FSTTimerID will be run. Method throws if no matching callback exists.
  */
-- (void)dispatchAfterDelay:(NSTimeInterval)delay block:(void (^)(void))block;
+- (void)runDelayedCallbacksUntil:(FSTTimerID)lastTimerID;
 
 /** The underlying wrapped dispatch_queue_t */
 @property(nonatomic, strong, readonly) dispatch_queue_t queue;
diff --git a/Firestore/Source/Util/FSTDispatchQueue.mm b/Firestore/Source/Util/FSTDispatchQueue.mm
index 6ce5d74..5bd7f27 100644
--- a/Firestore/Source/Util/FSTDispatchQueue.mm
+++ b/Firestore/Source/Util/FSTDispatchQueue.mm
@@ -21,8 +21,138 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+/**
+ * removeDelayedCallback is used by FSTDelayedCallback and so we pre-declare it before the rest of
+ * the FSTDispatchQueue private interface.
+ */
 @interface FSTDispatchQueue ()
+- (void)removeDelayedCallback:(FSTDelayedCallback *)callback;
+@end
+
+#pragma mark - FSTDelayedCallback
+
+/**
+ * Represents a callback scheduled to be run in the future on an FSTDispatchQueue.
+ *
+ * It is created via [FSTDelayedCallback createAndScheduleWithQueue].
+ *
+ * Supports cancellation (via cancel) and early execution (via skipDelay).
+ */
+@interface FSTDelayedCallback ()
+
+@property(nonatomic, strong, readonly) FSTDispatchQueue *queue;
+@property(nonatomic, assign, readonly) FSTTimerID timerID;
+@property(nonatomic, assign, readonly) NSTimeInterval targetTime;
+@property(nonatomic, copy) void (^callback)();
+/** YES if the callback has been run or canceled. */
+@property(nonatomic, getter=isDone) BOOL done;
+
+/**
+ * Creates and returns an FSTDelayedCallback that has been scheduled on the provided queue with the
+ * provided delay.
+ *
+ * @param queue The FSTDispatchQueue to run the callback on.
+ * @param timerID A FSTTimerID identifying the type of the delayed callback.
+ * @param delay The delay before the callback should be scheduled.
+ * @param callback The callback block to run.
+ * @return The created FSTDelayedCallback instance.
+ */
++ (instancetype)createAndScheduleWithQueue:(FSTDispatchQueue *)queue
+                                   timerID:(FSTTimerID)timerID
+                                     delay:(NSTimeInterval)delay
+                                  callback:(void (^)(void))callback;
+
+/**
+ * Queues the callback to run immediately (if it hasn't already been run or canceled).
+ */
+- (void)skipDelay;
+
+@end
+
+@implementation FSTDelayedCallback
+
+- (instancetype)initWithQueue:(FSTDispatchQueue *)queue
+                      timerID:(FSTTimerID)timerID
+                   targetTime:(NSTimeInterval)targetTime
+                     callback:(void (^)(void))callback {
+  if (self = [super init]) {
+    _queue = queue;
+    _timerID = timerID;
+    _targetTime = targetTime;
+    _callback = callback;
+    _done = NO;
+  }
+  return self;
+}
+
++ (instancetype)createAndScheduleWithQueue:(FSTDispatchQueue *)queue
+                                   timerID:(FSTTimerID)timerID
+                                     delay:(NSTimeInterval)delay
+                                  callback:(void (^)(void))callback {
+  NSTimeInterval targetTime = [[NSDate date] timeIntervalSince1970] + delay;
+  FSTDelayedCallback *delayedCallback = [[FSTDelayedCallback alloc] initWithQueue:queue
+                                                                          timerID:timerID
+                                                                       targetTime:targetTime
+                                                                         callback:callback];
+  [delayedCallback startWithDelay:delay];
+  return delayedCallback;
+}
+
+/**
+ * Starts the timer. This is called immediately after construction by createAndScheduleWithQueue.
+ */
+- (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];
+  });
+}
+
+- (void)skipDelay {
+  [self.queue dispatchAsyncAllowingSameQueue:^{
+    [self delayDidElapse];
+  }];
+}
+
+- (void)cancel {
+  [self.queue verifyIsCurrentQueue];
+  if (!self.isDone) {
+    // PORTING NOTE: There's no way to actually cancel the dispatched callback, but it'll be a no-op
+    // since we set done to YES.
+    [self markDone];
+  }
+}
+
+- (void)delayDidElapse {
+  [self.queue verifyIsCurrentQueue];
+  if (!self.isDone) {
+    [self markDone];
+    self.callback();
+  }
+}
+
+/**
+ * Marks this delayed callback as done, and notifies the FSTDispatchQueue that it should be removed.
+ */
+- (void)markDone {
+  self.done = YES;
+  [self.queue removeDelayedCallback:self];
+}
+
+@end
+
+#pragma mark - FSTDispatchQueue
+
+@interface FSTDispatchQueue ()
+
+/**
+ * Callbacks scheduled to be queued in the future. Callbacks are automatically removed after they
+ * are run or canceled.
+ */
+@property(nonatomic, strong, readonly) NSMutableArray<FSTDelayedCallback *> *delayedCallbacks;
+
 - (instancetype)initWithQueue:(dispatch_queue_t)queue NS_DESIGNATED_INITIALIZER;
+
 @end
 
 @implementation FSTDispatchQueue
@@ -34,6 +164,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (instancetype)initWithQueue:(dispatch_queue_t)queue {
   if (self = [super init]) {
     _queue = queue;
+    _delayedCallbacks = [NSMutableArray array];
   }
   return self;
 }
@@ -56,9 +187,73 @@ NS_ASSUME_NONNULL_BEGIN
   dispatch_async(self.queue, block);
 }
 
-- (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);
+- (FSTDelayedCallback *)dispatchAfterDelay:(NSTimeInterval)delay
+                                   timerID:(FSTTimerID)timerID
+                                     block:(void (^)(void))block {
+  // While not necessarily harmful, we currently don't expect to have multiple callbacks with the
+  // same timerID in the queue, so defensively reject them.
+  FSTAssert(![self containsDelayedCallbackWithTimerID:timerID],
+            @"Attempted to schedule multiple callbacks with id %ld", (unsigned long)timerID);
+  FSTDelayedCallback *delayedCallback = [FSTDelayedCallback createAndScheduleWithQueue:self
+                                                                               timerID:timerID
+                                                                                 delay:delay
+                                                                              callback:block];
+  [self.delayedCallbacks addObject:delayedCallback];
+  return delayedCallback;
+}
+
+- (BOOL)containsDelayedCallbackWithTimerID:(FSTTimerID)timerID {
+  NSUInteger matchIndex = [self.delayedCallbacks
+      indexOfObjectPassingTest:^BOOL(FSTDelayedCallback *obj, NSUInteger idx, BOOL *stop) {
+        return obj.timerID == timerID;
+      }];
+  return matchIndex != NSNotFound;
+}
+
+- (void)runDelayedCallbacksUntil:(FSTTimerID)lastTimerID {
+  dispatch_semaphore_t doneSemaphore = dispatch_semaphore_create(0);
+
+  [self dispatchAsync:^{
+    FSTAssert(lastTimerID == FSTTimerIDAll || [self containsDelayedCallbackWithTimerID:lastTimerID],
+              @"Attempted to run callbacks until missing timer ID: %ld",
+              (unsigned long)lastTimerID);
+
+    [self sortDelayedCallbacks];
+    for (FSTDelayedCallback *callback in self.delayedCallbacks) {
+      [callback skipDelay];
+      if (lastTimerID != FSTTimerIDAll && callback.timerID == lastTimerID) {
+        break;
+      }
+    }
+
+    // Now that the callbacks are queued, we want to enqueue an additional item to release the
+    // 'done' semaphore.
+    [self dispatchAsyncAllowingSameQueue:^{
+      dispatch_semaphore_signal(doneSemaphore);
+    }];
+  }];
+
+  dispatch_semaphore_wait(doneSemaphore, DISPATCH_TIME_FOREVER);
+}
+
+// NOTE: For performance we could store the callbacks sorted (e.g. using std::priority_queue),
+// but this sort only happens in tests (if runDelayedCallbacksUntil: is called), and the size
+// is guaranteed to be small since we don't allow duplicate TimerIds (of which there are only 4).
+- (void)sortDelayedCallbacks {
+  // We want to run callbacks in the same order they'd run if they ran naturally.
+  [self.delayedCallbacks
+      sortUsingComparator:^NSComparisonResult(FSTDelayedCallback *a, FSTDelayedCallback *b) {
+        return a.targetTime < b.targetTime
+                   ? NSOrderedAscending
+                   : a.targetTime > b.targetTime ? NSOrderedDescending : NSOrderedSame;
+      }];
+}
+
+/** Called by FSTDelayedCallback when a callback is run or canceled. */
+- (void)removeDelayedCallback:(FSTDelayedCallback *)callback {
+  NSUInteger index = [self.delayedCallbacks indexOfObject:callback];
+  FSTAssert(index != NSNotFound, @"Delayed callback not found.");
+  [self.delayedCallbacks removeObjectAtIndex:index];
 }
 
 #pragma mark - Private Methods