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.

7 files changed, 133 insertions, 118 deletions

diff --git a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.mm b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.mm
index 3b6a67e..751e7ff 100644
--- a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.mm
+++ b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.mm
@@ -21,6 +21,7 @@
 #import "Firestore/Example/Tests/Util/FSTIntegrationTestCase.h"
 #import "Firestore/Source/API/FIRFirestore+Internal.h"
 #import "Firestore/Source/Core/FSTFirestoreClient.h"
+#import "Firestore/Source/Util/FSTDispatchQueue.h"
 
 @interface FIRDatabaseTests : FSTIntegrationTestCase
 @end
@@ -926,7 +927,7 @@
   FIRFirestore *firestore = doc.firestore;
 
   [self writeDocumentRef:doc data:@{@"foo" : @"bar"}];
-  [self waitForIdleFirestore:firestore];
+  [[self queueForFirestore:firestore] runDelayedCallbacksUntil:FSTTimerIDWriteStreamIdle];
   [self writeDocumentRef:doc data:@{@"foo" : @"bar"}];
 }
 
@@ -935,7 +936,7 @@
   FIRFirestore *firestore = doc.firestore;
 
   [self readSnapshotForRef:[self documentRef] requireOnline:YES];
-  [self waitForIdleFirestore:firestore];
+  [[self queueForFirestore:firestore] runDelayedCallbacksUntil:FSTTimerIDListenStreamIdle];
   [self readSnapshotForRef:[self documentRef] requireOnline:YES];
 }
 
diff --git a/Firestore/Example/Tests/Integration/FSTStreamTests.mm b/Firestore/Example/Tests/Integration/FSTStreamTests.mm
index 6259aff..a36361a 100644
--- a/Firestore/Example/Tests/Integration/FSTStreamTests.mm
+++ b/Firestore/Example/Tests/Integration/FSTStreamTests.mm
@@ -20,7 +20,6 @@
 
 #import "Firestore/Example/Tests/Util/FSTHelpers.h"
 #import "Firestore/Example/Tests/Util/FSTIntegrationTestCase.h"
-#import "Firestore/Example/Tests/Util/FSTTestDispatchQueue.h"
 #import "Firestore/Source/Auth/FSTEmptyCredentialsProvider.h"
 #import "Firestore/Source/Remote/FSTDatastore.h"
 #import "Firestore/Source/Remote/FSTStream.h"
@@ -133,7 +132,7 @@ using firebase::firestore::model::DatabaseId;
 
 @implementation FSTStreamTests {
   dispatch_queue_t _testQueue;
-  FSTTestDispatchQueue *_workerDispatchQueue;
+  FSTDispatchQueue *_workerDispatchQueue;
   DatabaseInfo _databaseInfo;
   FSTEmptyCredentialsProvider *_credentials;
   FSTStreamStatusDelegate *_delegate;
@@ -150,7 +149,7 @@ using firebase::firestore::model::DatabaseId;
                          DatabaseId::kDefaultDatabaseId);
 
   _testQueue = dispatch_queue_create("FSTStreamTestWorkerQueue", DISPATCH_QUEUE_SERIAL);
-  _workerDispatchQueue = [[FSTTestDispatchQueue alloc] initWithQueue:_testQueue];
+  _workerDispatchQueue = [[FSTDispatchQueue alloc] initWithQueue:_testQueue];
 
   _databaseInfo = DatabaseInfo(database_id, "test-key", util::MakeStringView(settings.host),
                                settings.sslEnabled);
@@ -272,10 +271,14 @@ using firebase::firestore::model::DatabaseId;
     [writeStream writeHandshake];
   }];
 
-  [_delegate awaitNotificationFromBlock:^{
+  [_workerDispatchQueue dispatchAsync:^{
     [writeStream markIdle];
+    XCTAssertTrue(
+        [_workerDispatchQueue containsDelayedCallbackWithTimerID:FSTTimerIDWriteStreamIdle]);
   }];
 
+  [_workerDispatchQueue runDelayedCallbacksUntil:FSTTimerIDWriteStreamIdle];
+
   dispatch_sync(_testQueue, ^{
     XCTAssertFalse([writeStream isOpen]);
   });
@@ -299,7 +302,11 @@ using firebase::firestore::model::DatabaseId;
   // Mark the stream idle, but immediately cancel the idle timer by issuing another write.
   [_delegate awaitNotificationFromBlock:^{
     [writeStream markIdle];
+    XCTAssertTrue(
+        [_workerDispatchQueue containsDelayedCallbackWithTimerID:FSTTimerIDWriteStreamIdle]);
     [writeStream writeMutations:_mutations];
+    XCTAssertFalse(
+        [_workerDispatchQueue containsDelayedCallbackWithTimerID:FSTTimerIDWriteStreamIdle]);
   }];
 
   dispatch_sync(_testQueue, ^{
diff --git a/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm b/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm
new file mode 100644
index 0000000..9f5b52d
--- /dev/null
+++ b/Firestore/Example/Tests/Util/FSTDispatchQueueTests.mm
@@ -0,0 +1,111 @@
+/*
+ * Copyright 2018 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import "Firestore/Source/Util/FSTDispatchQueue.h"
+
+#import <XCTest/XCTest.h>
+
+#import "Firestore/Example/Tests/Util/XCTestCase+Await.h"
+
+// In these generic tests the specific TimerIDs don't matter.
+static const FSTTimerID timerID1 = FSTTimerIDListenStreamConnection;
+static const FSTTimerID timerID2 = FSTTimerIDListenStreamIdle;
+static const FSTTimerID timerID3 = FSTTimerIDWriteStreamConnection;
+
+@interface FSTDispatchQueueTests : XCTestCase
+@end
+
+@implementation FSTDispatchQueueTests {
+  FSTDispatchQueue *_queue;
+  NSMutableArray *_completedSteps;
+  NSArray *_expectedSteps;
+  XCTestExpectation *_expectation;
+}
+
+- (void)setUp {
+  [super setUp];
+  dispatch_queue_t dispatch_queue =
+      dispatch_queue_create("FSTDispatchQueueTests", DISPATCH_QUEUE_SERIAL);
+  _queue = [[FSTDispatchQueue alloc] initWithQueue:dispatch_queue];
+  _completedSteps = [NSMutableArray array];
+  _expectedSteps = nil;
+}
+
+/**
+ * Helper to return a block that adds @(n) to _completedSteps when run and fulfils _expectation if
+ * the _completedSteps match the _expectedSteps.
+ */
+- (void (^)())blockForStep:(int)n {
+  return ^void() {
+    [self->_completedSteps addObject:@(n)];
+    if (self->_expectedSteps && self->_completedSteps.count >= self->_expectedSteps.count) {
+      XCTAssertEqualObjects(self->_completedSteps, self->_expectedSteps);
+      [self->_expectation fulfill];
+    }
+  };
+}
+
+- (void)testCanScheduleCallbacksInTheFuture {
+  _expectation = [self expectationWithDescription:@"Expected steps"];
+  _expectedSteps = @[ @1, @2, @3, @4 ];
+  [_queue dispatchAsync:[self blockForStep:1]];
+  [_queue dispatchAfterDelay:0.005 timerID:timerID1 block:[self blockForStep:4]];
+  [_queue dispatchAfterDelay:0.001 timerID:timerID2 block:[self blockForStep:3]];
+  [_queue dispatchAsync:[self blockForStep:2]];
+
+  [self awaitExpectations];
+}
+
+- (void)testCanCancelDelayedCallbacks {
+  _expectation = [self expectationWithDescription:@"Expected steps"];
+  _expectedSteps = @[ @1, @3 ];
+  // Queue everything from the queue to ensure nothing completes before we cancel.
+  [_queue dispatchAsync:^{
+    [_queue dispatchAsyncAllowingSameQueue:[self blockForStep:1]];
+    FSTDelayedCallback *step2Timer =
+        [_queue dispatchAfterDelay:.001 timerID:timerID1 block:[self blockForStep:2]];
+    [_queue dispatchAfterDelay:.005 timerID:timerID2 block:[self blockForStep:3]];
+
+    XCTAssertTrue([_queue containsDelayedCallbackWithTimerID:timerID1]);
+    [step2Timer cancel];
+    XCTAssertFalse([_queue containsDelayedCallbackWithTimerID:timerID1]);
+  }];
+
+  [self awaitExpectations];
+}
+
+- (void)testCanManuallyDrainAllDelayedCallbacksForTesting {
+  [_queue dispatchAsync:[self blockForStep:1]];
+  [_queue dispatchAfterDelay:20 timerID:timerID1 block:[self blockForStep:4]];
+  [_queue dispatchAfterDelay:10 timerID:timerID2 block:[self blockForStep:3]];
+  [_queue dispatchAsync:[self blockForStep:2]];
+
+  [_queue runDelayedCallbacksUntil:FSTTimerIDAll];
+  XCTAssertEqualObjects(_completedSteps, (@[ @1, @2, @3, @4 ]));
+}
+
+- (void)testCanManuallyDrainSpecificDelayedCallbacksForTesting {
+  [_queue dispatchAsync:[self blockForStep:1]];
+  [_queue dispatchAfterDelay:20 timerID:timerID1 block:[self blockForStep:5]];
+  [_queue dispatchAfterDelay:10 timerID:timerID2 block:[self blockForStep:3]];
+  [_queue dispatchAfterDelay:15 timerID:timerID3 block:[self blockForStep:4]];
+  [_queue dispatchAsync:[self blockForStep:2]];
+
+  [_queue runDelayedCallbacksUntil:timerID3];
+  XCTAssertEqualObjects(_completedSteps, (@[ @1, @2, @3, @4 ]));
+}
+
+@end
diff --git a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h
index e1820e2..9c80799 100644
--- a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h
+++ b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h
@@ -28,6 +28,7 @@
 @class FIRFirestoreSettings;
 @class FIRQuery;
 @class FSTEventAccumulator;
+@class FSTDispatchQueue;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -61,8 +62,6 @@ extern "C" {
 - (FIRCollectionReference *)collectionRefWithDocuments:
     (NSDictionary<NSString *, NSDictionary<NSString *, id> *> *)documents;
 
-- (void)waitForIdleFirestore:(FIRFirestore *)firestore;
-
 - (void)writeAllDocuments:(NSDictionary<NSString *, NSDictionary<NSString *, id> *> *)documents
              toCollection:(FIRCollectionReference *)collection;
 
@@ -87,6 +86,8 @@ extern "C" {
 
 - (void)enableNetwork;
 
+- (FSTDispatchQueue *)queueForFirestore:(FIRFirestore *)firestore;
+
 /**
  * "Blocks" the current thread/run loop until the block returns YES.
  * Should only be called on the main thread.
diff --git a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm
index df591b0..e34b2a5 100644
--- a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm
+++ b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm
@@ -30,7 +30,6 @@
 #import "Firestore/Source/Util/FSTDispatchQueue.h"
 
 #import "Firestore/Example/Tests/Util/FSTEventAccumulator.h"
-#import "Firestore/Example/Tests/Util/FSTTestDispatchQueue.h"
 
 #include "Firestore/core/src/firebase/firestore/model/database_id.h"
 #include "Firestore/core/src/firebase/firestore/util/string_apple.h"
@@ -133,7 +132,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (FIRFirestore *)firestoreWithProjectID:(NSString *)projectID {
   NSString *persistenceKey = [NSString stringWithFormat:@"db%lu", (unsigned long)_firestores.count];
 
-  FSTTestDispatchQueue *workerDispatchQueue = [FSTTestDispatchQueue
+  FSTDispatchQueue *workerDispatchQueue = [FSTDispatchQueue
       queueWith:dispatch_queue_create("com.google.firebase.firestore", DISPATCH_QUEUE_SERIAL)];
 
   FSTEmptyCredentialsProvider *credentialsProvider = [[FSTEmptyCredentialsProvider alloc] init];
@@ -155,14 +154,6 @@ NS_ASSUME_NONNULL_BEGIN
   return firestore;
 }
 
-- (void)waitForIdleFirestore:(FIRFirestore *)firestore {
-  XCTestExpectation *expectation = [self expectationWithDescription:@"idle"];
-  // Note that we wait on any task that is scheduled with a delay of 60s. Currently, the idle
-  // timeout is the only task that uses this delay.
-  [((FSTTestDispatchQueue *)firestore.workerDispatchQueue) fulfillOnExecution:expectation];
-  [self awaitExpectations];
-}
-
 - (void)shutdownFirestore:(FIRFirestore *)firestore {
   [firestore shutdownWithCompletion:[self completionForExpectationWithName:@"shutdown"]];
   [self awaitExpectations];
@@ -289,6 +280,10 @@ NS_ASSUME_NONNULL_BEGIN
   [self awaitExpectations];
 }
 
+- (FSTDispatchQueue *)queueForFirestore:(FIRFirestore *)firestore {
+  return firestore.workerDispatchQueue;
+}
+
 - (void)waitUntil:(BOOL (^)())predicate {
   NSTimeInterval start = [NSDate timeIntervalSinceReferenceDate];
   double waitSeconds = [self defaultExpectationWaitSeconds];
diff --git a/Firestore/Example/Tests/Util/FSTTestDispatchQueue.h b/Firestore/Example/Tests/Util/FSTTestDispatchQueue.h
deleted file mode 100644
index 7ecbbaf..0000000
--- a/Firestore/Example/Tests/Util/FSTTestDispatchQueue.h
+++ /dev/null
@@ -1,39 +0,0 @@
-/*
- * Copyright 2017 Google
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#import "Firestore/Source/Util/FSTDispatchQueue.h"
-
-@class XCTestExpectation;
-
-NS_ASSUME_NONNULL_BEGIN
-
-/**
- * Dispatch queue used in the integration tests that caps delayed executions at 1.0 seconds.
- */
-@interface FSTTestDispatchQueue : FSTDispatchQueue
-
-/** Creates and returns an FSTTestDispatchQueue wrapping the specified dispatch_queue_t. */
-+ (instancetype)queueWith:(dispatch_queue_t)dispatchQueue;
-
-/**
- * Registers a test expectation that is fulfilled when the next delayed callback finished
- * executing.
- */
-- (void)fulfillOnExecution:(XCTestExpectation *)expectation;
-
-@end
-
-NS_ASSUME_NONNULL_END
diff --git a/Firestore/Example/Tests/Util/FSTTestDispatchQueue.mm b/Firestore/Example/Tests/Util/FSTTestDispatchQueue.mm
deleted file mode 100644
index 8124cf2..0000000
--- a/Firestore/Example/Tests/Util/FSTTestDispatchQueue.mm
+++ /dev/null
@@ -1,61 +0,0 @@
-/*
- * Copyright 2017 Google
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#import "Firestore/Example/Tests/Util/FSTTestDispatchQueue.h"
-
-#import <XCTest/XCTestExpectation.h>
-
-#import "Firestore/Source/Util/FSTAssert.h"
-
-@interface FSTTestDispatchQueue ()
-
-@property(nonatomic, weak) XCTestExpectation* expectation;
-
-@end
-
-@implementation FSTTestDispatchQueue
-
-/** The delay used by the idle timeout */
-static const NSTimeInterval kIdleDispatchDelay = 60.0;
-
-/** The maximum delay we use in a test run. */
-static const NSTimeInterval kTestDispatchDelay = 1.0;
-
-+ (instancetype)queueWith:(dispatch_queue_t)dispatchQueue {
-  return [[FSTTestDispatchQueue alloc] initWithQueue:dispatchQueue];
-}
-
-- (instancetype)initWithQueue:(dispatch_queue_t)dispatchQueue {
-  return (self = [super initWithQueue:dispatchQueue]);
-}
-
-- (void)dispatchAfterDelay:(NSTimeInterval)delay block:(void (^)(void))block {
-  [super dispatchAfterDelay:MIN(delay, kTestDispatchDelay)
-                      block:^() {
-                        block();
-                        if (delay == kIdleDispatchDelay) {
-                          [_expectation fulfill];
-                          _expectation = nil;
-                        }
-                      }];
-}
-
-- (void)fulfillOnExecution:(XCTestExpectation*)expectation {
-  FSTAssert(_expectation == nil, @"Previous expectation still active");
-  _expectation = expectation;
-}
-
-@end