Throw error in subsequent function calls after the client is shutdown (#2978)

Author: Brian Chen

Firestore/Example/Tests/Integration/API/FIRDatabaseTests.mm

@@ -1223,4 +1223,21 @@ - (void)testCanDisableNetwork {
   [self awaitExpectations];
 }
 
+- (void)testClientCallsAfterShutdownFail {
+  FIRDocumentReference *doc = [self documentRef];
+  FIRFirestore *firestore = doc.firestore;
+
+  [firestore enableNetworkWithCompletion:[self completionForExpectationWithName:@"Enable network"]];
+  [self awaitExpectations];
+  [firestore shutdownWithCompletion:[self completionForExpectationWithName:@"Shutdown"]];
+  [self awaitExpectations];
+
+  XCTAssertThrowsSpecific(
+      {
+        [firestore disableNetworkWithCompletion:^(NSError *error){
+        }];
+      },
+      NSException, @"The client has already been shutdown.");
+}
+
 @end

Firestore/Source/Core/FSTFirestoreClient.mm

@@ -62,6 +62,7 @@
 using firebase::firestore::api::DocumentSnapshot;
 using firebase::firestore::api::Settings;
 using firebase::firestore::api::SnapshotMetadata;
+using firebase::firestore::api::ThrowIllegalState;
 using firebase::firestore::auth::CredentialsProvider;
 using firebase::firestore::auth::User;
 using firebase::firestore::core::DatabaseInfo;
@@ -128,7 +129,8 @@ @implementation FSTFirestoreClient {
   std::unique_ptr<Executor> _userExecutor;
   std::chrono::milliseconds _initialGcDelay;
   std::chrono::milliseconds _regularGcDelay;
-  BOOL _gcHasRun;
+  bool _gcHasRun;
+  std::atomic<bool> _isShutdown;
   _Nullable id<FSTLRUDelegate> _lruDelegate;
   DelayedOperation _lruCallback;
 }
@@ -165,7 +167,8 @@ - (instancetype)initWithDatabaseInfo:(const DatabaseInfo &)databaseInfo
     _credentialsProvider = credentialsProvider;
     _userExecutor = std::move(userExecutor);
     _workerQueue = std::move(workerQueue);
-    _gcHasRun = NO;
+    _gcHasRun = false;
+    _isShutdown = false;
     _initialGcDelay = FSTLruGcInitialDelay;
     _regularGcDelay = FSTLruGcRegularDelay;
 
@@ -270,7 +273,7 @@ - (void)scheduleLruGarbageCollection {
   std::chrono::milliseconds delay = _gcHasRun ? _regularGcDelay : _initialGcDelay;
   _lruCallback = _workerQueue->EnqueueAfterDelay(delay, TimerId::GarbageCollectionDelay, [self]() {
     [self->_localStore collectGarbage:self->_lruDelegate.gc];
-    self->_gcHasRun = YES;
+    self->_gcHasRun = true;
     [self scheduleLruGarbageCollection];
   });
 }
@@ -283,6 +286,7 @@ - (void)credentialDidChangeWithUser:(const User &)user {
 }
 
 - (void)disableNetworkWithCallback:(util::StatusCallback)callback {
+  [self verifyNotShutdown];
   _workerQueue->Enqueue([self, callback] {
     _remoteStore->DisableNetwork();
     if (callback) {
@@ -292,6 +296,7 @@ - (void)disableNetworkWithCallback:(util::StatusCallback)callback {
 }
 
 - (void)enableNetworkWithCallback:(util::StatusCallback)callback {
+  [self verifyNotShutdown];
   _workerQueue->Enqueue([self, callback] {
     _remoteStore->EnableNetwork();
     if (callback) {
@@ -302,20 +307,29 @@ - (void)enableNetworkWithCallback:(util::StatusCallback)callback {
 
 - (void)shutdownWithCallback:(util::StatusCallback)callback {
   _workerQueue->Enqueue([self, callback] {
-    self->_credentialsProvider->SetCredentialChangeListener(nullptr);
+    if (!_isShutdown) {
+      self->_credentialsProvider->SetCredentialChangeListener(nullptr);
 
-    // If we've scheduled LRU garbage collection, cancel it.
-    if (self->_lruCallback) {
-      self->_lruCallback.Cancel();
+      // If we've scheduled LRU garbage collection, cancel it.
+      if (self->_lruCallback) {
+        self->_lruCallback.Cancel();
+      }
+      _remoteStore->Shutdown();
+      [self.persistence shutdown];
+      self->_isShutdown = true;
     }
-    _remoteStore->Shutdown();
-    [self.persistence shutdown];
     if (callback) {
       self->_userExecutor->Execute([=] { callback(Status::OK()); });
     }
   });
 }
 
+- (void)verifyNotShutdown {
+  if (_isShutdown) {
+    ThrowIllegalState("The client has already been shutdown.");
+  }
+}
+
 - (std::shared_ptr<QueryListener>)listenToQuery:(FSTQuery *)query
                                         options:(core::ListenOptions)options
                                        listener:(ViewSnapshot::SharedListener &&)listener {
@@ -327,11 +341,13 @@ - (void)shutdownWithCallback:(util::StatusCallback)callback {
 }
 
 - (void)removeListener:(const std::shared_ptr<QueryListener> &)listener {
+  [self verifyNotShutdown];
   _workerQueue->Enqueue([self, listener] { [self.eventManager removeListener:listener]; });
 }
 
 - (void)getDocumentFromLocalCache:(const DocumentReference &)doc
                          callback:(DocumentSnapshot::Listener &&)callback {
+  [self verifyNotShutdown];
   auto shared_callback = absl::ShareUniquePtr(std::move(callback));
   _workerQueue->Enqueue([self, doc, shared_callback] {
     FSTMaybeDocument *maybeDoc = [self.localStore readDocument:doc.key()];
@@ -362,6 +378,7 @@ - (void)getDocumentFromLocalCache:(const DocumentReference &)doc
 - (void)getDocumentsFromLocalCache:(FIRQuery *)query
                         completion:(void (^)(FIRQuerySnapshot *_Nullable query,
                                              NSError *_Nullable error))completion {
+  [self verifyNotShutdown];
   _workerQueue->Enqueue([self, query, completion] {
     DocumentMap docs = [self.localStore executeQuery:query.query];
 
@@ -391,6 +408,7 @@ - (void)writeMutations:(std::vector<FSTMutation *> &&)mutations
               callback:(util::StatusCallback)callback {
   // TODO(c++14): move `mutations` into lambda (C++14).
   _workerQueue->Enqueue([self, mutations, callback]() mutable {
+    [self verifyNotShutdown];
     if (mutations.empty()) {
       if (callback) {
         self->_userExecutor->Execute([=] { callback(Status::OK()); });
@@ -413,6 +431,7 @@ - (void)transactionWithRetries:(int)retries
                 resultCallback:(core::TransactionResultCallback)resultCallback {
   // Dispatch the result back onto the user dispatch queue.
   auto async_callback = [self, resultCallback](util::StatusOr<absl::any> maybe_value) {
+    [self verifyNotShutdown];
     if (resultCallback) {
       self->_userExecutor->Execute([=] { resultCallback(std::move(maybe_value)); });
     }

Firestore/core/src/firebase/firestore/api/firestore.mm

@@ -140,15 +140,10 @@
 }
 
 void Firestore::Shutdown(util::StatusCallback callback) {
-  if (!client_) {
-    if (callback) {
-      // We should be dispatching the callback on the user dispatch queue
-      // but if the client is nil here that queue was never created.
-      callback(Status::OK());
-    }
-  } else {
-    [client_ shutdownWithCallback:std::move(callback)];
-  }
+  // The client must be initialized to ensure that all subsequent API usage
+  // throws an exception.
+  EnsureClientConfigured();
+  [client_ shutdownWithCallback:std::move(callback)];
 }
 
 void Firestore::EnableNetwork(util::StatusCallback callback) {