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.

14 files changed, 442 insertions, 170 deletions

diff --git a/Firestore/Example/Firestore.xcodeproj/project.pbxproj b/Firestore/Example/Firestore.xcodeproj/project.pbxproj
index 3b91c76..f230cb4 100644
--- a/Firestore/Example/Firestore.xcodeproj/project.pbxproj
+++ b/Firestore/Example/Firestore.xcodeproj/project.pbxproj
@@ -37,14 +37,12 @@
 		5492E03320213FFC00B64F25 /* FSTSyncEngineTestDriver.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E02E20213FFC00B64F25 /* FSTSyncEngineTestDriver.mm */; };
 		5492E03420213FFC00B64F25 /* FSTMemorySpecTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E02F20213FFC00B64F25 /* FSTMemorySpecTests.mm */; };
 		5492E03520213FFC00B64F25 /* FSTSpecTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E03020213FFC00B64F25 /* FSTSpecTests.mm */; };
-		5492E03B2021401F00B64F25 /* FSTTestDispatchQueue.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0362021401E00B64F25 /* FSTTestDispatchQueue.mm */; };
 		5492E03C2021401F00B64F25 /* XCTestCase+Await.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0372021401E00B64F25 /* XCTestCase+Await.mm */; };
 		5492E03D2021401F00B64F25 /* FSTAssertTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0382021401E00B64F25 /* FSTAssertTests.mm */; };
 		5492E03E2021401F00B64F25 /* FSTEventAccumulator.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0392021401F00B64F25 /* FSTEventAccumulator.mm */; };
 		5492E03F2021401F00B64F25 /* FSTHelpers.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E03A2021401F00B64F25 /* FSTHelpers.mm */; };
 		5492E041202143E700B64F25 /* FSTEventAccumulator.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0392021401F00B64F25 /* FSTEventAccumulator.mm */; };
 		5492E0422021440500B64F25 /* FSTHelpers.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E03A2021401F00B64F25 /* FSTHelpers.mm */; };
-		5492E0432021441E00B64F25 /* FSTTestDispatchQueue.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0362021401E00B64F25 /* FSTTestDispatchQueue.mm */; };
 		5492E0442021457E00B64F25 /* XCTestCase+Await.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E0372021401E00B64F25 /* XCTestCase+Await.mm */; };
 		5492E050202154AA00B64F25 /* FIRCollectionReferenceTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E045202154AA00B64F25 /* FIRCollectionReferenceTests.mm */; };
 		5492E051202154AA00B64F25 /* FIRQueryTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 5492E046202154AA00B64F25 /* FIRQueryTests.mm */; };
@@ -133,6 +131,7 @@
 		6003F5BA195388D20070C39A /* InfoPlist.strings in Resources */ = {isa = PBXBuildFile; fileRef = 6003F5B8195388D20070C39A /* InfoPlist.strings */; };
 		6ED54761B845349D43DB6B78 /* Pods_Firestore_Example.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 75A6FE51C1A02DF38F62FAAD /* Pods_Firestore_Example.framework */; };
 		71719F9F1E33DC2100824A3D /* LaunchScreen.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 71719F9D1E33DC2100824A3D /* LaunchScreen.storyboard */; };
+		7346E61D20325C6900FD6CEF /* FSTDispatchQueueTests.mm in Sources */ = {isa = PBXBuildFile; fileRef = 7346E61C20325C6900FD6CEF /* FSTDispatchQueueTests.mm */; };
 		873B8AEB1B1F5CCA007FD442 /* Main.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 873B8AEA1B1F5CCA007FD442 /* Main.storyboard */; };
 		AB356EF7200EA5EB0089B766 /* field_value_test.cc in Sources */ = {isa = PBXBuildFile; fileRef = AB356EF6200EA5EB0089B766 /* field_value_test.cc */; };
 		AB380CFB2019388600D97691 /* target_id_generator_test.cc in Sources */ = {isa = PBXBuildFile; fileRef = AB380CF82019382300D97691 /* target_id_generator_test.cc */; };
@@ -228,7 +227,6 @@
 		5492E02E20213FFC00B64F25 /* FSTSyncEngineTestDriver.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTSyncEngineTestDriver.mm; sourceTree = "<group>"; };
 		5492E02F20213FFC00B64F25 /* FSTMemorySpecTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTMemorySpecTests.mm; sourceTree = "<group>"; };
 		5492E03020213FFC00B64F25 /* FSTSpecTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTSpecTests.mm; sourceTree = "<group>"; };
-		5492E0362021401E00B64F25 /* FSTTestDispatchQueue.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTTestDispatchQueue.mm; sourceTree = "<group>"; };
 		5492E0372021401E00B64F25 /* XCTestCase+Await.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = "XCTestCase+Await.mm"; sourceTree = "<group>"; };
 		5492E0382021401E00B64F25 /* FSTAssertTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTAssertTests.mm; sourceTree = "<group>"; };
 		5492E0392021401F00B64F25 /* FSTEventAccumulator.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTEventAccumulator.mm; sourceTree = "<group>"; };
@@ -335,6 +333,7 @@
 		6003F5B9195388D20070C39A /* en */ = {isa = PBXFileReference; lastKnownFileType = text.plist.strings; name = en; path = en.lproj/InfoPlist.strings; sourceTree = "<group>"; };
 		69F6A10DBD6187489481CD76 /* Pods_Firestore_Tests.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = Pods_Firestore_Tests.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		71719F9E1E33DC2100824A3D /* Base */ = {isa = PBXFileReference; lastKnownFileType = file.storyboard; name = Base; path = Base.lproj/LaunchScreen.storyboard; sourceTree = "<group>"; };
+		7346E61C20325C6900FD6CEF /* FSTDispatchQueueTests.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = FSTDispatchQueueTests.mm; sourceTree = "<group>"; };
 		75A6FE51C1A02DF38F62FAAD /* Pods_Firestore_Example.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = Pods_Firestore_Example.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		873B8AEA1B1F5CCA007FD442 /* Main.storyboard */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = file.storyboard; name = Main.storyboard; path = Base.lproj/Main.storyboard; sourceTree = "<group>"; };
 		8E002F4AD5D9B6197C940847 /* Firestore.podspec */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text; name = Firestore.podspec; path = ../Firestore.podspec; sourceTree = "<group>"; };
@@ -364,7 +363,6 @@
 		B686F2B02024FFD70028D6BE /* resource_path_test.cc */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.cpp; path = resource_path_test.cc; sourceTree = "<group>"; };
 		CE00BABB5A3AAB44A4C209E2 /* Pods-Firestore_Tests.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-Firestore_Tests.debug.xcconfig"; path = "Pods/Target Support Files/Pods-Firestore_Tests/Pods-Firestore_Tests.debug.xcconfig"; sourceTree = "<group>"; };
 		D3CC3DC5338DCAF43A211155 /* README.md */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = net.daringfireball.markdown; name = README.md; path = ../README.md; sourceTree = "<group>"; };
-		D5B259DAA9149B80D6245B57 /* FSTTestDispatchQueue.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = FSTTestDispatchQueue.h; sourceTree = "<group>"; };
 		DB17FEDFB80770611A935A60 /* Pods-Firestore_IntegrationTests.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-Firestore_IntegrationTests.release.xcconfig"; path = "Pods/Target Support Files/Pods-Firestore_IntegrationTests/Pods-Firestore_IntegrationTests.release.xcconfig"; sourceTree = "<group>"; };
 		DE03B2E91F2149D600A30B9C /* Firestore_IntegrationTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = Firestore_IntegrationTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
 		DE03B3621F215E1600A30B9C /* CAcert.pem */ = {isa = PBXFileReference; lastKnownFileType = text; path = CAcert.pem; sourceTree = "<group>"; };
@@ -724,10 +722,9 @@
 				5492E03A2021401F00B64F25 /* FSTHelpers.mm */,
 				54E9281E1F33950B00C1953E /* FSTIntegrationTestCase.h */,
 				5491BC711FB44593008B3588 /* FSTIntegrationTestCase.mm */,
-				D5B259DAA9149B80D6245B57 /* FSTTestDispatchQueue.h */,
-				5492E0362021401E00B64F25 /* FSTTestDispatchQueue.mm */,
 				54E9282A1F339CAD00C1953E /* XCTestCase+Await.h */,
 				5492E0372021401E00B64F25 /* XCTestCase+Await.mm */,
+				7346E61C20325C6900FD6CEF /* FSTDispatchQueueTests.mm */,
 			);
 			path = Util;
 			sourceTree = "<group>";
@@ -1272,6 +1269,7 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				7346E61D20325C6900FD6CEF /* FSTDispatchQueueTests.mm in Sources */,
 				DE2EF0881F3D0B6E003D0CDC /* FSTTreeSortedDictionaryTests.m in Sources */,
 				ABE6637A201FA81900ED349A /* database_id_test.cc in Sources */,
 				5492E0AF2021552D00B64F25 /* FSTReferenceSetTests.mm in Sources */,
@@ -1360,7 +1358,6 @@
 				5492E052202154AB00B64F25 /* FIRGeoPointTests.mm in Sources */,
 				5492E0C72021557E00B64F25 /* FSTSerializerBetaTests.mm in Sources */,
 				5492E03520213FFC00B64F25 /* FSTSpecTests.mm in Sources */,
-				5492E03B2021401F00B64F25 /* FSTTestDispatchQueue.mm in Sources */,
 				5492E057202154AB00B64F25 /* FIRSnapshotMetadataTests.mm in Sources */,
 				54740A571FC914BA00713A1A /* secure_random_test.cc in Sources */,
 				5492E0BE2021555100B64F25 /* FSTMutationTests.mm in Sources */,
@@ -1388,7 +1385,6 @@
 				5492E07F202154EC00B64F25 /* FSTTransactionTests.mm in Sources */,
 				5492E075202154D600B64F25 /* FIRDatabaseTests.mm in Sources */,
 				5492E078202154D600B64F25 /* FIRWriteBatchTests.mm in Sources */,
-				5492E0432021441E00B64F25 /* FSTTestDispatchQueue.mm in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
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
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) {
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