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.

4 files changed, 57 insertions, 39 deletions

diff --git a/Firestore/Source/Remote/FSTExponentialBackoff.h b/Firestore/Source/Remote/FSTExponentialBackoff.h
index 674b1ac..03558ee 100644
--- a/Firestore/Source/Remote/FSTExponentialBackoff.h
+++ b/Firestore/Source/Remote/FSTExponentialBackoff.h
@@ -16,7 +16,7 @@
 
 #import <Foundation/Foundation.h>
 
-@class FSTDispatchQueue;
+#import "Firestore/Source/Util/FSTDispatchQueue.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -29,7 +29,7 @@ NS_ASSUME_NONNULL_BEGIN
 @interface FSTExponentialBackoff : NSObject
 
 /**
- * Creates and returns a helper for running delayed tasks following an exponential backoff curve
+ * Initializes a helper for running delayed tasks following an exponential backoff curve
  * between attempts.
  *
  * Each delay is made up of a "base" delay which follows the exponential backoff curve, and a
@@ -37,6 +37,7 @@ NS_ASSUME_NONNULL_BEGIN
  * accidentally synchronizing their delays causing spikes of load to the backend.
  *
  * @param dispatchQueue The dispatch queue to run tasks on.
+ * @param timerID The ID to use when scheduling backoff operations on the FSTDispatchQueue.
  * @param initialDelay The initial delay (used as the base delay on the first retry attempt).
  *     Note that jitter will still be applied, so the actual delay could be as little as
  *     0.5*initialDelay.
@@ -45,13 +46,13 @@ NS_ASSUME_NONNULL_BEGIN
  * @param maxDelay The maximum base delay after which no further backoff is performed. Note that
  *     jitter will still be applied, so the actual delay could be as much as 1.5*maxDelay.
  */
-+ (instancetype)exponentialBackoffWithDispatchQueue:(FSTDispatchQueue *)dispatchQueue
-                                       initialDelay:(NSTimeInterval)initialDelay
-                                      backoffFactor:(double)backoffFactor
-                                           maxDelay:(NSTimeInterval)maxDelay;
+- (instancetype)initWithDispatchQueue:(FSTDispatchQueue *)dispatchQueue
+                              timerID:(FSTTimerID)timerID
+                         initialDelay:(NSTimeInterval)initialDelay
+                        backoffFactor:(double)backoffFactor
+                             maxDelay:(NSTimeInterval)maxDelay NS_DESIGNATED_INITIALIZER;
 
-- (instancetype)init
-    __attribute__((unavailable("Use exponentialBackoffWithDispatchQueue constructor method.")));
+- (instancetype)init NS_UNAVAILABLE;
 
 /**
  * Resets the backoff delay.
@@ -68,7 +69,8 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)resetToMax;
 
 /**
- * Waits for currentDelay seconds, increases the delay and runs the specified block.
+ * Waits for currentDelay seconds, increases the delay and runs the specified block. If there was
+ * a pending block waiting to be run already, it will be canceled.
  *
  * @param block The block to run.
  */
diff --git a/Firestore/Source/Remote/FSTExponentialBackoff.mm b/Firestore/Source/Remote/FSTExponentialBackoff.mm
index 8077357..dddf164 100644
--- a/Firestore/Source/Remote/FSTExponentialBackoff.mm
+++ b/Firestore/Source/Remote/FSTExponentialBackoff.mm
@@ -26,16 +26,14 @@
 using firebase::firestore::util::SecureRandom;
 
 @interface FSTExponentialBackoff ()
-- (instancetype)initWithDispatchQueue:(FSTDispatchQueue *)dispatchQueue
-                         initialDelay:(NSTimeInterval)initialDelay
-                        backoffFactor:(double)backoffFactor
-                             maxDelay:(NSTimeInterval)maxDelay NS_DESIGNATED_INITIALIZER;
 
 @property(nonatomic, strong) FSTDispatchQueue *dispatchQueue;
+@property(nonatomic, assign, readonly) FSTTimerID timerID;
 @property(nonatomic) double backoffFactor;
 @property(nonatomic) NSTimeInterval initialDelay;
 @property(nonatomic) NSTimeInterval maxDelay;
 @property(nonatomic) NSTimeInterval currentBase;
+@property(nonatomic, strong, nullable) FSTDelayedCallback *timerCallback;
 @end
 
 @implementation FSTExponentialBackoff {
@@ -43,11 +41,13 @@ using firebase::firestore::util::SecureRandom;
 }
 
 - (instancetype)initWithDispatchQueue:(FSTDispatchQueue *)dispatchQueue
+                              timerID:(FSTTimerID)timerID
                          initialDelay:(NSTimeInterval)initialDelay
                         backoffFactor:(double)backoffFactor
                              maxDelay:(NSTimeInterval)maxDelay {
   if (self = [super init]) {
     _dispatchQueue = dispatchQueue;
+    _timerID = timerID;
     _initialDelay = initialDelay;
     _backoffFactor = backoffFactor;
     _maxDelay = maxDelay;
@@ -57,16 +57,6 @@ using firebase::firestore::util::SecureRandom;
   return self;
 }
 
-+ (instancetype)exponentialBackoffWithDispatchQueue:(FSTDispatchQueue *)dispatchQueue
-                                       initialDelay:(NSTimeInterval)initialDelay
-                                      backoffFactor:(double)backoffFactor
-                                           maxDelay:(NSTimeInterval)maxDelay {
-  return [[FSTExponentialBackoff alloc] initWithDispatchQueue:dispatchQueue
-                                                 initialDelay:initialDelay
-                                                backoffFactor:backoffFactor
-                                                     maxDelay:maxDelay];
-}
-
 - (void)reset {
   _currentBase = 0;
 }
@@ -76,6 +66,9 @@ using firebase::firestore::util::SecureRandom;
 }
 
 - (void)backoffAndRunBlock:(void (^)(void))block {
+  if (self.timerCallback) {
+    [self.timerCallback cancel];
+  }
   // First schedule the block using the current base (which may be 0 and should be honored as such).
   NSTimeInterval delayWithJitter = _currentBase + [self jitterDelay];
   if (_currentBase > 0) {
@@ -83,7 +76,8 @@ using firebase::firestore::util::SecureRandom;
            _currentBase);
   }
 
-  [self.dispatchQueue dispatchAfterDelay:delayWithJitter block:block];
+  self.timerCallback =
+      [self.dispatchQueue dispatchAfterDelay:delayWithJitter timerID:self.timerID block:block];
 
   // Apply backoff factor to determine next delay and ensure it is within bounds.
   _currentBase *= _backoffFactor;
diff --git a/Firestore/Source/Remote/FSTStream.h b/Firestore/Source/Remote/FSTStream.h
index c390dbb..6630083 100644
--- a/Firestore/Source/Remote/FSTStream.h
+++ b/Firestore/Source/Remote/FSTStream.h
@@ -17,6 +17,7 @@
 #import <Foundation/Foundation.h>
 
 #import "Firestore/Source/Core/FSTTypes.h"
+#import "Firestore/Source/Util/FSTDispatchQueue.h"
 
 #include "Firestore/core/src/firebase/firestore/core/database_info.h"
 
@@ -91,6 +92,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (instancetype)initWithDatabase:(const firebase::firestore::core::DatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
+               connectionTimerID:(FSTTimerID)connectionTimerID
+                     idleTimerID:(FSTTimerID)idleTimerID
                      credentials:(id<FSTCredentialsProvider>)credentials
             responseMessageClass:(Class)responseMessageClass NS_DESIGNATED_INITIALIZER;
 
@@ -142,8 +145,13 @@ NS_ASSUME_NONNULL_BEGIN
 - (void)stop;
 
 /**
- * Initializes the idle timer. If no write takes place within one minute, the GRPC stream will be
- * closed.
+ * Marks this stream as idle. If no further actions are performed on the stream for one minute, the
+ * stream will automatically close itself and notify the stream's close handler. The stream will
+ * then be in a non-started state, requiring the caller to start the stream again before further
+ * use.
+ *
+ * Only streams that are in state 'Open' can be marked idle, as all other states imply pending
+ * network operations.
  */
 - (void)markIdle;
 
diff --git a/Firestore/Source/Remote/FSTStream.mm b/Firestore/Source/Remote/FSTStream.mm
index c1479c5..f859fbb 100644
--- a/Firestore/Source/Remote/FSTStream.mm
+++ b/Firestore/Source/Remote/FSTStream.mm
@@ -113,7 +113,8 @@ typedef NS_ENUM(NSInteger, FSTStreamState) {
 
 @interface FSTStream ()
 
-@property(nonatomic, getter=isIdle) BOOL idle;
+@property(nonatomic, assign, readonly) FSTTimerID idleTimerID;
+@property(nonatomic, strong, nullable) FSTDelayedCallback *idleTimerCallback;
 @property(nonatomic, weak, readwrite, nullable) id delegate;
 
 @end
@@ -203,18 +204,22 @@ static const NSTimeInterval kIdleTimeout = 60.0;
 
 - (instancetype)initWithDatabase:(const DatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
+               connectionTimerID:(FSTTimerID)connectionTimerID
+                     idleTimerID:(FSTTimerID)idleTimerID
                      credentials:(id<FSTCredentialsProvider>)credentials
             responseMessageClass:(Class)responseMessageClass {
   if (self = [super init]) {
     _databaseInfo = database;
     _workerDispatchQueue = workerDispatchQueue;
+    _idleTimerID = idleTimerID;
     _credentials = credentials;
     _responseMessageClass = responseMessageClass;
 
-    _backoff = [FSTExponentialBackoff exponentialBackoffWithDispatchQueue:workerDispatchQueue
-                                                             initialDelay:kBackoffInitialDelay
-                                                            backoffFactor:kBackoffFactor
-                                                                 maxDelay:kBackoffMaxDelay];
+    _backoff = [[FSTExponentialBackoff alloc] initWithDispatchQueue:workerDispatchQueue
+                                                            timerID:connectionTimerID
+                                                       initialDelay:kBackoffInitialDelay
+                                                      backoffFactor:kBackoffFactor
+                                                           maxDelay:kBackoffMaxDelay];
     _state = FSTStreamStateInitial;
   }
   return self;
@@ -418,7 +423,7 @@ static const NSTimeInterval kIdleTimeout = 60.0;
 /** Called by the idle timer when the stream should close due to inactivity. */
 - (void)handleIdleCloseTimer {
   [self.workerDispatchQueue verifyIsCurrentQueue];
-  if (self.state == FSTStreamStateOpen && [self isIdle]) {
+  if ([self isOpen]) {
     // When timing out an idle stream there's no reason to force the stream into backoff when
     // it restarts so set the stream state to Initial instead of Error.
     [self closeWithFinalState:FSTStreamStateInitial error:nil];
@@ -427,18 +432,23 @@ static const NSTimeInterval kIdleTimeout = 60.0;
 
 - (void)markIdle {
   [self.workerDispatchQueue verifyIsCurrentQueue];
-  if (self.state == FSTStreamStateOpen) {
-    self.idle = YES;
-    [self.workerDispatchQueue dispatchAfterDelay:kIdleTimeout
-                                           block:^() {
-                                             [self handleIdleCloseTimer];
-                                           }];
+  // Starts the idle timer if we are in state 'Open' and are not yet already running a timer (in
+  // which case the previous idle timeout still applies).
+  if ([self isOpen] && !self.idleTimerCallback) {
+    self.idleTimerCallback = [self.workerDispatchQueue dispatchAfterDelay:kIdleTimeout
+                                                                  timerID:self.idleTimerID
+                                                                    block:^() {
+                                                                      [self handleIdleCloseTimer];
+                                                                    }];
   }
 }
 
 - (void)cancelIdleCheck {
   [self.workerDispatchQueue verifyIsCurrentQueue];
-  self.idle = NO;
+  if (self.idleTimerCallback) {
+    [self.idleTimerCallback cancel];
+    self.idleTimerCallback = nil;
+  }
 }
 
 /**
@@ -606,6 +616,8 @@ static const NSTimeInterval kIdleTimeout = 60.0;
                       serializer:(FSTSerializerBeta *)serializer {
   self = [super initWithDatabase:database
              workerDispatchQueue:workerDispatchQueue
+               connectionTimerID:FSTTimerIDListenStreamConnection
+                     idleTimerID:FSTTimerIDListenStreamIdle
                      credentials:credentials
             responseMessageClass:[GCFSListenResponse class]];
   if (self) {
@@ -689,6 +701,8 @@ static const NSTimeInterval kIdleTimeout = 60.0;
                       serializer:(FSTSerializerBeta *)serializer {
   self = [super initWithDatabase:database
              workerDispatchQueue:workerDispatchQueue
+               connectionTimerID:FSTTimerIDWriteStreamConnection
+                     idleTimerID:FSTTimerIDWriteStreamIdle
                      credentials:credentials
             responseMessageClass:[GCFSWriteResponse class]];
   if (self) {