Adding goOnline/goOffline to the RemoteStore (#347)

7 files changed, 255 insertions, 53 deletions

diff --git a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
index 6a727ad..d22552a 100644
--- a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
@@ -18,6 +18,8 @@
 
 #import <XCTest/XCTest.h>
 
+#import "Core/FSTFirestoreClient.h"
+#import "FIRFirestore+Internal.h"
 #import "FSTIntegrationTestCase.h"
 
 @interface FIRDatabaseTests : FSTIntegrationTestCase
@@ -766,4 +768,81 @@
   XCTAssertEqualObjects([self.db documentWithPath:@"foo/bar/baz/qux"].documentID, @"qux");
 }
 
+- (void)testCanQueueWritesWhileOffline {
+  XCTestExpectation *writeEpectation = [self expectationWithDescription:@"successfull write"];
+  XCTestExpectation *networkExpectation = [self expectationWithDescription:@"enable network"];
+
+  FIRDocumentReference *doc = [self documentRef];
+  FIRFirestore *firestore = doc.firestore;
+  NSDictionary<NSString *, id> *data = @{@"a" : @"b"};
+
+  [firestore.client disableNetworkWithCompletion:^(NSError *error) {
+    XCTAssertNil(error);
+
+    [doc setData:data
+        completion:^(NSError *error) {
+          XCTAssertNil(error);
+          [writeEpectation fulfill];
+        }];
+
+    [firestore.client enableNetworkWithCompletion:^(NSError *error) {
+      XCTAssertNil(error);
+      [networkExpectation fulfill];
+    }];
+  }];
+
+  [self awaitExpectations];
+
+  XCTestExpectation *getExpectation = [self expectationWithDescription:@"successfull get"];
+  [doc getDocumentWithCompletion:^(FIRDocumentSnapshot *snapshot, NSError *error) {
+    XCTAssertNil(error);
+    XCTAssertEqualObjects(snapshot.data, data);
+    XCTAssertFalse(snapshot.metadata.isFromCache);
+
+    [getExpectation fulfill];
+  }];
+
+  [self awaitExpectations];
+}
+
+- (void)testCantGetDocumentsWhileOffline {
+  FIRDocumentReference *doc = [self documentRef];
+  FIRFirestore *firestore = doc.firestore;
+  NSDictionary<NSString *, id> *data = @{@"a" : @"b"};
+
+  XCTestExpectation *onlineExpectation = [self expectationWithDescription:@"online read"];
+  XCTestExpectation *networkExpectation = [self expectationWithDescription:@"network online"];
+
+  __weak FIRDocumentReference *weakDoc = doc;
+
+  [firestore.client disableNetworkWithCompletion:^(NSError *error) {
+    XCTAssertNil(error);
+    [doc setData:data
+        completion:^(NSError *_Nullable error) {
+          XCTAssertNil(error);
+
+          [weakDoc getDocumentWithCompletion:^(FIRDocumentSnapshot *snapshot, NSError *error) {
+            XCTAssertNil(error);
+
+            // Verify that we are not reading from cache.
+            XCTAssertFalse(snapshot.metadata.isFromCache);
+            [onlineExpectation fulfill];
+          }];
+        }];
+
+    [doc getDocumentWithCompletion:^(FIRDocumentSnapshot *snapshot, NSError *error) {
+      XCTAssertNil(error);
+
+      // Verify that we are reading from cache.
+      XCTAssertTrue(snapshot.metadata.fromCache);
+      XCTAssertEqualObjects(snapshot.data, data);
+      [firestore.client enableNetworkWithCompletion:^(NSError *error) {
+        [networkExpectation fulfill];
+      }];
+    }];
+  }];
+
+  [self awaitExpectations];
+}
+
 @end
diff --git a/Firestore/Example/Tests/Integration/API/FIRListenerRegistrationTests.m b/Firestore/Example/Tests/Integration/API/FIRListenerRegistrationTests.m
index 19771ff..5824314 100644
--- a/Firestore/Example/Tests/Integration/API/FIRListenerRegistrationTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRListenerRegistrationTests.m
@@ -18,6 +18,8 @@
 
 #import <XCTest/XCTest.h>
 
+#import "Core/FSTFirestoreClient.h"
+#import "FIRFirestore+Internal.h"
 #import "FSTIntegrationTestCase.h"
 
 @interface FIRListenerRegistrationTests : FSTIntegrationTestCase
@@ -126,4 +128,35 @@
   [two remove];
 }
 
+- (void)testWatchSurvivesNetworkDisconnect {
+  XCTestExpectation *testExpectiation =
+      [self expectationWithDescription:@"testWatchSurvivesNetworkDisconnect"];
+
+  FIRCollectionReference *collectionRef = [self collectionRef];
+  FIRDocumentReference *docRef = [collectionRef documentWithAutoID];
+
+  FIRFirestore *firestore = collectionRef.firestore;
+
+  FIRQueryListenOptions *options = [[[FIRQueryListenOptions options]
+      includeDocumentMetadataChanges:YES] includeQueryMetadataChanges:YES];
+
+  [collectionRef addSnapshotListenerWithOptions:options
+                                       listener:^(FIRQuerySnapshot *snapshot, NSError *error) {
+                                         XCTAssertNil(error);
+                                         if (!snapshot.empty && !snapshot.metadata.fromCache) {
+                                           [testExpectiation fulfill];
+                                         }
+                                       }];
+
+  [firestore.client disableNetworkWithCompletion:^(NSError *error) {
+    XCTAssertNil(error);
+    [docRef setData:@{@"foo" : @"bar"}];
+    [firestore.client enableNetworkWithCompletion:^(NSError *error) {
+      XCTAssertNil(error);
+    }];
+  }];
+
+  [self awaitExpectations];
+}
+
 @end
diff --git a/Firestore/Example/Tests/Integration/FSTDatastoreTests.m b/Firestore/Example/Tests/Integration/FSTDatastoreTests.m
index bab8f44..54ab66b 100644
--- a/Firestore/Example/Tests/Integration/FSTDatastoreTests.m
+++ b/Firestore/Example/Tests/Integration/FSTDatastoreTests.m
@@ -174,7 +174,10 @@ NS_ASSUME_NONNULL_BEGIN
                                        credentials:_credentials];
 
   _remoteStore = [FSTRemoteStore remoteStoreWithLocalStore:_localStore datastore:_datastore];
-  [_remoteStore start];
+
+  [_testWorkerQueue dispatchAsync:^() {
+    [_remoteStore start];
+  }];
 }
 
 - (void)tearDown {
diff --git a/Firestore/Source/Core/FSTFirestoreClient.h b/Firestore/Source/Core/FSTFirestoreClient.h
index 45f13cc..21d61d4 100644
--- a/Firestore/Source/Core/FSTFirestoreClient.h
+++ b/Firestore/Source/Core/FSTFirestoreClient.h
@@ -56,6 +56,12 @@ NS_ASSUME_NONNULL_BEGIN
 /** Shuts down this client, cancels all writes / listeners, and releases all resources. */
 - (void)shutdownWithCompletion:(nullable FSTVoidErrorBlock)completion;
 
+/** Disables the network connection. Pending operations will not complete. */
+- (void)disableNetworkWithCompletion:(nullable FSTVoidErrorBlock)completion;
+
+/** Enables the network connection and requeues all pending operations. */
+- (void)enableNetworkWithCompletion:(nullable FSTVoidErrorBlock)completion;
+
 /** Starts listening to a query. */
 - (FSTQueryListener *)listenToQuery:(FSTQuery *)query
                             options:(FSTListenOptions *)options
diff --git a/Firestore/Source/Core/FSTFirestoreClient.m b/Firestore/Source/Core/FSTFirestoreClient.m
index 2066ce9..1a53197 100644
--- a/Firestore/Source/Core/FSTFirestoreClient.m
+++ b/Firestore/Source/Core/FSTFirestoreClient.m
@@ -187,6 +187,28 @@ NS_ASSUME_NONNULL_BEGIN
   [self.syncEngine userDidChange:user];
 }
 
+- (void)disableNetworkWithCompletion:(nullable FSTVoidErrorBlock)completion {
+  [self.workerDispatchQueue dispatchAsync:^{
+    [self.remoteStore disableNetwork];
+    if (completion) {
+      [self.userDispatchQueue dispatchAsync:^{
+        completion(nil);
+      }];
+    }
+  }];
+}
+
+- (void)enableNetworkWithCompletion:(nullable FSTVoidErrorBlock)completion {
+  [self.workerDispatchQueue dispatchAsync:^{
+    [self.remoteStore enableNetwork];
+    if (completion) {
+      [self.userDispatchQueue dispatchAsync:^{
+        completion(nil);
+      }];
+    }
+  }];
+}
+
 - (void)shutdownWithCompletion:(nullable FSTVoidErrorBlock)completion {
   [self.workerDispatchQueue dispatchAsync:^{
     self.credentialsProvider.userChangeListener = nil;
diff --git a/Firestore/Source/Remote/FSTRemoteStore.h b/Firestore/Source/Remote/FSTRemoteStore.h
index 94208e1..0948cfa 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.h
+++ b/Firestore/Source/Remote/FSTRemoteStore.h
@@ -110,6 +110,12 @@ NS_ASSUME_NONNULL_BEGIN
 /** Shuts down the remote store, tearing down connections and otherwise cleaning up. */
 - (void)shutdown;
 
+/** Temporarily disables the network. The network can be re-enabled using 'enableNetwork:'. */
+- (void)disableNetwork;
+
+/** Re-enables the network. Only to be called as the counterpart to 'disableNetwork:'. */
+- (void)enableNetwork;
+
 /**
  * Tells the FSTRemoteStore that the currently authenticated user has changed.
  *
diff --git a/Firestore/Source/Remote/FSTRemoteStore.m b/Firestore/Source/Remote/FSTRemoteStore.m
index cea2ce8..85608e7 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.m
+++ b/Firestore/Source/Remote/FSTRemoteStore.m
@@ -57,6 +57,8 @@ static const NSUInteger kMaxPendingWrites = 10;
 @property(nonatomic, strong, readonly) FSTDatastore *datastore;
 
 #pragma mark Watch Stream
+// The watchStream is null when the network is disabled. The non-null check is performed by
+// isNetworkEnabled.
 @property(nonatomic, strong, nullable) FSTWatchStream *watchStream;
 
 /**
@@ -97,6 +99,8 @@ static const NSUInteger kMaxPendingWrites = 10;
 @property(nonatomic, assign) FSTOnlineState watchStreamOnlineState;
 
 #pragma mark Write Stream
+// The writeStream is null when the network is disabled. The non-null check is performed by
+// isNetworkEnabled.
 @property(nonatomic, strong, nullable) FSTWriteStream *writeStream;
 
 /**
@@ -136,10 +140,8 @@ static const NSUInteger kMaxPendingWrites = 10;
 }
 
 - (void)start {
-  [self setupStreams];
-
-  // Resume any writes
-  [self fillWritePipeline];
+  // For now, all setup is handled by enableNetwork(). We might expand on this in the future.
+  [self enableNetwork];
 }
 
 - (void)updateAndNotifyAboutOnlineState:(FSTOnlineState)watchStreamOnlineState {
@@ -150,68 +152,94 @@ static const NSUInteger kMaxPendingWrites = 10;
   }
 }
 
-- (void)setupStreams {
+#pragma mark Online/Offline state
+
+- (BOOL)isNetworkEnabled {
+  FSTAssert((self.watchStream == nil) == (self.writeStream == nil),
+            @"WatchStream and WriteStream should both be null or non-null");
+  return self.watchStream != nil;
+}
+
+- (void)enableNetwork {
+  FSTAssert(self.watchStream == nil, @"enableNetwork: called with non-null watchStream.");
+  FSTAssert(self.writeStream == nil, @"enableNetwork: called with non-null writeStream.");
+
+  // Create new streams (but note they're not started yet).
   self.watchStream = [self.datastore createWatchStreamWithDelegate:self];
   self.writeStream = [self.datastore createWriteStreamWithDelegate:self];
 
   // Load any saved stream token from persistent storage
   self.writeStream.lastStreamToken = [self.localStore lastStreamToken];
-}
 
-#pragma mark Shutdown
+  if ([self shouldStartWatchStream]) {
+    [self startWatchStream];
+  }
 
-- (void)shutdown {
-  FSTLog(@"FSTRemoteStore %p shutting down", (__bridge void *)self);
+  [self fillWritePipeline];  // This may start the writeStream.
 
-  self.watchStreamOnlineState = FSTOnlineStateUnknown;
-  [self cleanupWatchStreamState];
-  [self.watchStream stop];
-  [self.writeStream stop];
+  // We move back to the unknown state because we might not want to re-open the stream
+  [self updateAndNotifyAboutOnlineState:FSTOnlineStateUnknown];
 }
 
-- (void)userDidChange:(FSTUser *)user {
-  FSTLog(@"FSTRemoteStore %p changing users: %@", (__bridge void *)self, user);
+- (void)disableNetwork {
+  [self updateAndNotifyAboutOnlineState:FSTOnlineStateFailed];
 
-  // Clear pending writes because those are per-user. Watched targets persist across users so
-  // don't clear those.
-  _lastBatchSeen = kFSTBatchIDUnknown;
-  [self.pendingWrites removeAllObjects];
-
-  // Stop the streams. They promise not to call us back.
+  // NOTE: We're guaranteed not to get any further events from these streams (not even a close
+  // event).
   [self.watchStream stop];
   [self.writeStream stop];
 
-  [self cleanupWatchStreamState];
+  [self cleanUpWatchStreamState];
+  [self cleanUpWriteStreamState];
 
-  // Create new streams (but note they're not started yet).
-  [self setupStreams];
+  self.writeStream = nil;
+  self.watchStream = nil;
+}
 
-  // If there are any watchedTargets properly handle the stream restart now that FSTRemoteStore
-  // is ready to handle them.
-  if ([self shouldStartWatchStream]) {
-    [self.watchStream start];
+#pragma mark Shutdown
+
+- (void)shutdown {
+  FSTLog(@"FSTRemoteStore %p shutting down", (__bridge void *)self);
+
+  // Don't fire initial listener callbacks on shutdown.
+  self.onlineStateDelegate = nil;
+
+  // For now, all shutdown logic is handled by disableNetwork(). We might expand on this in the
+  // future.
+  if ([self isNetworkEnabled]) {
+    [self disableNetwork];
   }
+}
 
-  // Resume any writes
-  [self fillWritePipeline];
+- (void)userDidChange:(FSTUser *)user {
+  FSTLog(@"FSTRemoteStore %p changing users: %@", (__bridge void *)self, user);
 
-  // User change moves us back to the unknown state because we might not
-  // want to re-open the stream
-  [self updateAndNotifyAboutOnlineState:FSTOnlineStateUnknown];
+  // Tear down and re-create our network streams. This will ensure we get a fresh auth token
+  // for the new user and re-fill the write pipeline with new mutations from the LocalStore
+  // (since mutations are per-user).
+  [self disableNetwork];
+  [self enableNetwork];
 }
 
 #pragma mark Watch Stream
 
+- (void)startWatchStream {
+  FSTAssert([self shouldStartWatchStream],
+            @"startWatchStream: called when shouldStartWatchStream: is false.");
+  [self.watchStream start];
+}
+
 - (void)listenToTargetWithQueryData:(FSTQueryData *)queryData {
   NSNumber *targetKey = @(queryData.targetID);
   FSTAssert(!self.listenTargets[targetKey], @"listenToQuery called with duplicate target id: %@",
             targetKey);
 
   self.listenTargets[targetKey] = queryData;
-  if ([self.watchStream isOpen]) {
-    [self sendWatchRequestWithQueryData:queryData];
-  } else if (![self.watchStream isStarted]) {
+
+  if ([self shouldStartWatchStream]) {
     [self.watchStream start];
+  } else if ([self isNetworkEnabled] && [self.watchStream isOpen]) {
+    [self sendWatchRequestWithQueryData:queryData];
   }
 }
 
@@ -226,7 +254,7 @@ static const NSUInteger kMaxPendingWrites = 10;
   FSTAssert(queryData, @"unlistenToTarget: target not currently watched: %@", targetKey);
 
   [self.listenTargets removeObjectForKey:targetKey];
-  if ([self.watchStream isOpen]) {
+  if ([self isNetworkEnabled] && [self.watchStream isOpen]) {
     [self sendUnwatchRequestForTargetID:targetKey];
   }
 }
@@ -243,14 +271,14 @@ static const NSUInteger kMaxPendingWrites = 10;
 }
 
 /**
- * Returns whether the watch stream should be started because there are active targets trying to
- * be listened to.
+ * Returns YES if the network is enabled, the watch stream has not yet been started and there are
+ * active watch targets.
  */
 - (BOOL)shouldStartWatchStream {
-  return self.listenTargets.count > 0;
+  return [self isNetworkEnabled] && ![self.watchStream isStarted] && self.listenTargets.count > 0;
 }
 
-- (void)cleanupWatchStreamState {
+- (void)cleanUpWatchStreamState {
   // If the connection is closed then we'll never get a snapshot version for the accumulated
   // changes and so we'll never be able to complete the batch. When we start up again the server
   // is going to resend these changes anyway, so just toss the accumulated state.
@@ -298,9 +326,13 @@ static const NSUInteger kMaxPendingWrites = 10;
 }
 
 - (void)watchStreamDidClose:(NSError *_Nullable)error {
-  [self cleanupWatchStreamState];
+  FSTAssert([self isNetworkEnabled],
+            @"watchStreamDidClose should only be called when the network is enabled");
+
+  [self cleanUpWatchStreamState];
 
-  // If there was an error, retry the connection.
+  // If the watch stream closed due to an error, retry the connection if there are any active
+  // watch targets.
   if ([self shouldStartWatchStream]) {
     // If the connection fails before the stream has become healthy, consider the online state
     // failed. Otherwise consider the online state unknown and the next connection attempt will
@@ -313,8 +345,8 @@ static const NSUInteger kMaxPendingWrites = 10;
     }
     [self.watchStream start];
   } else {
-    // No need to restart watch stream because there are no active targets. The online state is set
-    // to unknown because there is no active attempt at establishing a connection.
+    // We don't need to restart the watch stream because there are no active targets. The online
+    // state is set to unknown because there is no active attempt at establishing a connection.
     [self updateAndNotifyAboutOnlineState:FSTOnlineStateUnknown];
   }
 }
@@ -439,6 +471,26 @@ static const NSUInteger kMaxPendingWrites = 10;
 
 #pragma mark Write Stream
 
+/**
+ * Returns YES if the network is enabled, the write stream has not yet been started and there are
+ * pending writes.
+ */
+- (BOOL)shouldStartWriteStream {
+  return [self isNetworkEnabled] && ![self.writeStream isStarted] && self.pendingWrites.count > 0;
+}
+
+- (void)startWriteStream {
+  FSTAssert([self shouldStartWriteStream],
+            @"startWriteStream: called when shouldStartWriteStream: is false.");
+
+  [self.writeStream start];
+}
+
+- (void)cleanUpWriteStreamState {
+  self.lastBatchSeen = kFSTBatchIDUnknown;
+  [self.pendingWrites removeAllObjects];
+}
+
 - (void)fillWritePipeline {
   while ([self canWriteMutations]) {
     FSTMutationBatch *batch = [self.localStore nextMutationBatchAfterBatchID:self.lastBatchSeen];
@@ -460,7 +512,7 @@ static const NSUInteger kMaxPendingWrites = 10;
  * to accept more.
  */
 - (BOOL)canWriteMutations {
-  return self.pendingWrites.count < kMaxPendingWrites;
+  return [self isNetworkEnabled] && self.pendingWrites.count < kMaxPendingWrites;
 }
 
 /** Given mutations to commit, actually commits them to the backend. */
@@ -468,13 +520,11 @@ static const NSUInteger kMaxPendingWrites = 10;
   FSTAssert([self canWriteMutations], @"commitBatch called when mutations can't be written");
   self.lastBatchSeen = batch.batchID;
 
-  if (!self.writeStream.isStarted) {
-    [self.writeStream start];
-  }
-
   [self.pendingWrites addObject:batch];
 
-  if (self.writeStream.handshakeComplete) {
+  if ([self shouldStartWriteStream]) {
+    [self startWriteStream];
+  } else if ([self isNetworkEnabled] && self.writeStream.handshakeComplete) {
     [self.writeStream writeMutations:batch.mutations];
   }
 }
@@ -533,6 +583,9 @@ static const NSUInteger kMaxPendingWrites = 10;
  * has been terminated by the client or the server.
  */
 - (void)writeStreamDidClose:(NSError *_Nullable)error {
+  FSTAssert([self isNetworkEnabled],
+            @"writeStreamDidClose: should only be called when the network is enabled");
+
   NSMutableArray *pendingWrites = self.pendingWrites;
   // Ignore close if there are no pending writes.
   if (pendingWrites.count == 0) {
@@ -552,7 +605,7 @@ static const NSUInteger kMaxPendingWrites = 10;
   }
 
   // The write stream might have been started by refilling the write pipeline for failed writes
-  if (pendingWrites.count > 0 && !self.writeStream.isStarted) {
+  if ([self shouldStartWriteStream]) {
     [self.writeStream start];
   }
 }