Port public shutdown (#3467)

* async_queue shutdown ready

* shutdown feature complete. need to fix shared_ptr for credprovider and stop executor leakage.

* stop executor leakage.

* Improved comments without shared_from_this

* shared_ptr everywhere!

* undo scheme change

* addressing comments #1

* addressing comments

* addressing comments #3

Author: Hui-Wu

Firestore/Example/Tests/API/FSTAPIHelpers.mm

@@ -65,7 +65,8 @@
                                                persistenceKey:"db123"
                                           credentialsProvider:nullptr
                                                   workerQueue:nullptr
-                                                  firebaseApp:nil];
+                                                  firebaseApp:nil
+                                             instanceRegistry:nil];
   });
 #pragma clang diagnostic pop
   return sharedInstance;

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

@@ -22,6 +22,10 @@
 #import "Firestore/Example/Tests/Util/FSTEventAccumulator.h"
 #import "Firestore/Example/Tests/Util/FSTIntegrationTestCase.h"
 #import "Firestore/Source/API/FIRFirestore+Internal.h"
+#import "Firestore/Source/Core/FSTFirestoreClient.h"
+#include "Firestore/core/test/firebase/firestore/testutil/app_testing.h"
+
+namespace testutil = firebase::firestore::testutil;
 
 using firebase::firestore::util::TimerId;
 
@@ -1315,4 +1319,66 @@ - (void)testClearPersistenceWhileRunningFails {
   [self awaitExpectations];
 }
 
+- (void)testRestartFirestoreLeadsToNewInstance {
+  FIRApp *app = testutil::AppForUnitTesting(util::MakeString([FSTIntegrationTestCase projectID]));
+  FIRFirestore *firestore = [FIRFirestore firestoreForApp:app];
+  FIRFirestore *sameInstance = [FIRFirestore firestoreForApp:app];
+  firestore.settings = [FSTIntegrationTestCase settings];
+
+  XCTAssertEqual(firestore, sameInstance);
+
+  NSDictionary<NSString *, id> *data =
+      @{@"owner" : @{@"name" : @"Jonny", @"email" : @"abc@xyz.com"}};
+  [self writeDocumentRef:[firestore documentWithPath:@"abc/123"] data:data];
+
+  [self shutdownFirestore:firestore];
+
+  // Create a new instance, check it's a different instance.
+  FIRFirestore *newInstance = [FIRFirestore firestoreForApp:app];
+  newInstance.settings = [FSTIntegrationTestCase settings];
+  XCTAssertNotEqual(firestore, newInstance);
+
+  // New instance still functions.
+  FIRDocumentSnapshot *snapshot =
+      [self readDocumentForRef:[newInstance documentWithPath:@"abc/123"]];
+  XCTAssertTrue([data isEqualToDictionary:[snapshot data]]);
+}
+
+- (void)testAppDeleteLeadsToFirestoreShutdown {
+  FIRApp *app = testutil::AppForUnitTesting(util::MakeString([FSTIntegrationTestCase projectID]));
+  FIRFirestore *firestore = [FIRFirestore firestoreForApp:app];
+  firestore.settings = [FSTIntegrationTestCase settings];
+  NSDictionary<NSString *, id> *data =
+      @{@"owner" : @{@"name" : @"Jonny", @"email" : @"abc@xyz.com"}};
+  [self writeDocumentRef:[firestore documentWithPath:@"abc/123"] data:data];
+
+  [self deleteApp:app];
+
+  FSTFirestoreClient *client = firestore.wrapped->client();
+  XCTAssertTrue([client isShutdown]);
+}
+
+- (void)testShutdownCanBeCalledMultipleTimes {
+  FIRApp *app = testutil::AppForUnitTesting(util::MakeString([FSTIntegrationTestCase projectID]));
+  FIRFirestore *firestore = [FIRFirestore firestoreForApp:app];
+
+  [firestore shutdownWithCompletion:[self completionForExpectationWithName:@"Shutdown1"]];
+  [self awaitExpectations];
+  XCTAssertThrowsSpecific(
+      {
+        [firestore disableNetworkWithCompletion:^(NSError *error){
+        }];
+      },
+      NSException, @"The client has already been shutdown.");
+
+  [firestore shutdownWithCompletion:[self completionForExpectationWithName:@"Shutdown2"]];
+  [self awaitExpectations];
+  XCTAssertThrowsSpecific(
+      {
+        [firestore enableNetworkWithCompletion:^(NSError *error){
+        }];
+      },
+      NSException, @"The client has already been shutdown.");
+}
+
 @end

Firestore/Example/Tests/Integration/FSTDatastoreTests.mm

@@ -159,7 +159,6 @@ @interface FSTDatastoreTests : XCTestCase
 @implementation FSTDatastoreTests {
   std::shared_ptr<AsyncQueue> _testWorkerQueue;
   FSTLocalStore *_localStore;
-  EmptyCredentialsProvider _credentials;
 
   DatabaseInfo _databaseInfo;
   std::shared_ptr<Datastore> _datastore;
@@ -183,7 +182,8 @@ - (void)setUp {
   dispatch_queue_t queue = dispatch_queue_create(
       "com.google.firestore.FSTDatastoreTestsWorkerQueue", DISPATCH_QUEUE_SERIAL);
   _testWorkerQueue = std::make_shared<AsyncQueue>(absl::make_unique<ExecutorLibdispatch>(queue));
-  _datastore = std::make_shared<Datastore>(_databaseInfo, _testWorkerQueue, &_credentials);
+  _datastore = std::make_shared<Datastore>(_databaseInfo, _testWorkerQueue,
+                                           std::make_shared<EmptyCredentialsProvider>());
 
   _remoteStore =
       absl::make_unique<RemoteStore>(_localStore, _datastore, _testWorkerQueue, [](OnlineState) {});

Firestore/Example/Tests/SpecTests/FSTMockDatastore.h

@@ -38,7 +38,7 @@ class MockDatastore : public Datastore {
  public:
   MockDatastore(const core::DatabaseInfo& database_info,
                 const std::shared_ptr<util::AsyncQueue>& worker_queue,
-                auth::CredentialsProvider* credentials);
+                std::shared_ptr<auth::CredentialsProvider> credentials);
 
   std::shared_ptr<WatchStream> CreateWatchStream(WatchStreamCallback* callback) override;
   std::shared_ptr<WriteStream> CreateWriteStream(WriteStreamCallback* callback) override;
@@ -93,7 +93,7 @@ class MockDatastore : public Datastore {
   // reduces the number of test-only methods in `Datastore`.
   const core::DatabaseInfo* database_info_ = nullptr;
   std::shared_ptr<util::AsyncQueue> worker_queue_;
-  auth::CredentialsProvider* credentials_ = nullptr;
+  std::shared_ptr<auth::CredentialsProvider> credentials_;
 
   std::shared_ptr<MockWatchStream> watch_stream_;
   std::shared_ptr<MockWriteStream> write_stream_;

Firestore/Example/Tests/SpecTests/FSTMockDatastore.mm

@@ -65,7 +65,7 @@
 class MockWatchStream : public WatchStream {
  public:
   MockWatchStream(const std::shared_ptr<AsyncQueue>& worker_queue,
-                  CredentialsProvider* credentials_provider,
+                  std::shared_ptr<CredentialsProvider> credentials_provider,
                   FSTSerializerBeta* serializer,
                   GrpcConnection* grpc_connection,
                   WatchStreamCallback* callback,
@@ -156,7 +156,7 @@ void WriteWatchChange(const WatchChange& change, SnapshotVersion snap) {
 class MockWriteStream : public WriteStream {
  public:
   MockWriteStream(const std::shared_ptr<AsyncQueue>& worker_queue,
-                  CredentialsProvider* credentials_provider,
+                  std::shared_ptr<CredentialsProvider> credentials_provider,
                   FSTSerializerBeta* serializer,
                   GrpcConnection* grpc_connection,
                   WriteStreamCallback* callback,
@@ -239,7 +239,7 @@ int sent_mutations_count() const {
 
 MockDatastore::MockDatastore(const core::DatabaseInfo& database_info,
                              const std::shared_ptr<util::AsyncQueue>& worker_queue,
-                             auth::CredentialsProvider* credentials)
+                             std::shared_ptr<auth::CredentialsProvider> credentials)
     : Datastore{database_info, worker_queue, credentials, CreateNoOpConnectivityMonitor()},
       database_info_{&database_info},
       worker_queue_{worker_queue},

Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.mm

@@ -148,7 +148,6 @@ @implementation FSTSyncEngineTestDriver {
 
   DatabaseInfo _databaseInfo;
   User _currentUser;
-  EmptyCredentialsProvider _credentialProvider;
 
   std::shared_ptr<MockDatastore> _datastore;
 }
@@ -179,7 +178,8 @@ - (instancetype)initWithPersistence:(id<FSTPersistence>)persistence
     _persistence = persistence;
     _localStore = [[FSTLocalStore alloc] initWithPersistence:persistence initialUser:initialUser];
 
-    _datastore = std::make_shared<MockDatastore>(_databaseInfo, _workerQueue, &_credentialProvider);
+    _datastore = std::make_shared<MockDatastore>(_databaseInfo, _workerQueue,
+                                                 std::make_shared<EmptyCredentialsProvider>());
     _remoteStore = absl::make_unique<RemoteStore>(
         _localStore, _datastore, _workerQueue, [self](OnlineState onlineState) {
           [self.syncEngine applyChangedOnlineState:onlineState];

Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm

@@ -234,7 +234,8 @@ - (FIRFirestore *)firestoreWithApp:(FIRApp *)app {
                                 persistenceKey:util::MakeString(persistenceKey)
                            credentialsProvider:std::move(credentials_provider)
                                    workerQueue:std::move(workerQueue)
-                                   firebaseApp:app];
+                                   firebaseApp:app
+                              instanceRegistry:nil];
 
   firestore.settings = [FSTIntegrationTestCase settings];
   [_firestores addObject:firestore];

Firestore/Source/API/FIRFirestore+Internal.h

@@ -34,6 +34,14 @@ namespace util = firebase::firestore::util;
 
 NS_ASSUME_NONNULL_BEGIN
 
+/** Provides a registry management interface for FIRFirestore instances. */
+@protocol FSTFirestoreInstanceRegistry
+
+/** Removes the FIRFirestore instance with given database name from registry. */
+- (void)removeInstanceWithDatabase:(NSString *)database;
+
+@end
+
 @interface FIRFirestore (/* Init */)
 
 /**
@@ -42,9 +50,10 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (instancetype)initWithDatabaseID:(model::DatabaseId)databaseID
                     persistenceKey:(std::string)persistenceKey
-               credentialsProvider:(std::unique_ptr<auth::CredentialsProvider>)credentialsProvider
+               credentialsProvider:(std::shared_ptr<auth::CredentialsProvider>)credentialsProvider
                        workerQueue:(std::shared_ptr<util::AsyncQueue>)workerQueue
-                       firebaseApp:(FIRApp *)app;
+                       firebaseApp:(FIRApp *)app
+                  instanceRegistry:(nullable id<FSTFirestoreInstanceRegistry>)registry;
 @end
 
 /** Internal FIRFirestore API we don't want exposed in our public header files. */
@@ -56,14 +65,31 @@ NS_ASSUME_NONNULL_BEGIN
 + (FIRFirestore *)recoverFromFirestore:(std::shared_ptr<api::Firestore>)firestore;
 
 /**
- * Shutdown this `FIRFirestore`, releasing all resources (abandoning any outstanding writes,
- * removing all listens, closing all network connections, etc.).
+ * Shuts down this `FIRFirestore` instance.
+ *
+ * After shutdown only the `clearPersistence` method may be used. Any other method
+ * will throw an error.
+ *
+ * To restart after shutdown, simply create a new instance of FIRFirestore with
+ * `firestore` or `firestoreForApp` methods.
+ *
+ * Shutdown does not cancel any pending writes and any tasks that are awaiting a response from
+ * the server will not be resolved. The next time you start this instance, it will resume
+ * attempting to send these writes to the server.
+ *
+ * Note: Under normal circumstances, calling this method is not required. This
+ * method is useful only when you want to force this instance to release all of its resources or
+ * in combination with `clearPersistence` to ensure that all local state is destroyed
+ * between test runs.
  *
  * @param completion A block to execute once everything has shut down.
  */
+// TODO(b/135755126): Make this public.
 - (void)shutdownWithCompletion:(nullable void (^)(NSError *_Nullable error))completion
     NS_SWIFT_NAME(shutdown(completion:));
 
+- (void)shutdownInternalWithCompletion:(nullable void (^)(NSError *_Nullable error))completion;
+
 - (const std::shared_ptr<util::AsyncQueue> &)workerQueue;
 
 @property(nonatomic, assign, readonly) std::shared_ptr<api::Firestore> wrapped;

Firestore/Source/API/FIRFirestore.mm

@@ -72,6 +72,7 @@ @interface FIRFirestore ()
 @implementation FIRFirestore {
   std::shared_ptr<Firestore> _firestore;
   FIRFirestoreSettings *_settings;
+  __weak id<FSTFirestoreInstanceRegistry> _registry;
 }
 
 + (instancetype)firestore {
@@ -106,15 +107,17 @@ + (instancetype)firestoreForApp:(FIRApp *)app database:(NSString *)database {
 
 - (instancetype)initWithDatabaseID:(model::DatabaseId)databaseID
                     persistenceKey:(std::string)persistenceKey
-               credentialsProvider:(std::unique_ptr<CredentialsProvider>)credentialsProvider
+               credentialsProvider:(std::shared_ptr<CredentialsProvider>)credentialsProvider
                        workerQueue:(std::shared_ptr<AsyncQueue>)workerQueue
-                       firebaseApp:(FIRApp *)app {
+                       firebaseApp:(FIRApp *)app
+                  instanceRegistry:(nullable id<FSTFirestoreInstanceRegistry>)registry {
   if (self = [super init]) {
     _firestore = std::make_shared<Firestore>(std::move(databaseID), std::move(persistenceKey),
                                              std::move(credentialsProvider), std::move(workerQueue),
                                              (__bridge void *)self);
 
     _app = app;
+    _registry = registry;
 
     FSTPreConverterBlock block = ^id _Nullable(id _Nullable input) {
       if ([input isKindOfClass:[FIRDocumentReference class]]) {
@@ -298,10 +301,19 @@ + (FIRFirestore *)recoverFromFirestore:(std::shared_ptr<Firestore>)firestore {
   return (__bridge FIRFirestore *)firestore->extension();
 }
 
-- (void)shutdownWithCompletion:(nullable void (^)(NSError *_Nullable error))completion {
+- (void)shutdownInternalWithCompletion:(nullable void (^)(NSError *_Nullable error))completion {
   _firestore->Shutdown(util::MakeCallback(completion));
 }
 
+- (void)shutdownWithCompletion:(nullable void (^)(NSError *_Nullable error))completion {
+  id<FSTFirestoreInstanceRegistry> strongRegistry = _registry;
+  if (strongRegistry) {
+    [strongRegistry
+        removeInstanceWithDatabase:util::MakeNSString(_firestore->database_id().database_id())];
+  }
+  [self shutdownInternalWithCompletion:completion];
+}
+
 @end
 
 NS_ASSUME_NONNULL_END

Firestore/Source/API/FSTFirestoreComponent.h

@@ -16,6 +16,8 @@
 
 #import <Foundation/Foundation.h>
 
+#import "Firestore/Source/API/FIRFirestore+Internal.h"
+
 @class FIRApp;
 @class FIRFirestore;
 
@@ -35,7 +37,8 @@ NS_ASSUME_NONNULL_BEGIN
 
 /// A concrete implementation for FSTInstanceProvider to create Firestore instances and register
 /// with Core's component system.
-@interface FSTFirestoreComponent : NSObject <FSTFirestoreMultiDBProvider>
+@interface FSTFirestoreComponent
+    : NSObject <FSTFirestoreInstanceRegistry, FSTFirestoreMultiDBProvider>
 
 /// The FIRApp that instances will be set up with.
 @property(nonatomic, weak, readonly) FIRApp *app;
@@ -46,6 +49,8 @@ NS_ASSUME_NONNULL_BEGIN
 /// Default method for retrieving a Firestore instance, or creating one if it doesn't exist.
 - (FIRFirestore *)firestoreForDatabase:(NSString *)database;
 
+- (void)removeInstanceWithDatabase:(NSString *)database;
+
 /// Default initializer.
 - (instancetype)initWithApp:(FIRApp *)app NS_DESIGNATED_INITIALIZER;
 

Firestore/Source/API/FSTFirestoreComponent.mm

@@ -70,14 +70,18 @@ - (instancetype)initWithApp:(FIRApp *)app {
   return self;
 }
 
+- (NSString *)keyForDatabase:(NSString *)database {
+  return [NSString stringWithFormat:@"%@|%@", self.app.name, database];
+}
+
 #pragma mark - FSTInstanceProvider Conformance
 
 - (FIRFirestore *)firestoreForDatabase:(NSString *)database {
   if (!database) {
     ThrowInvalidArgument("Database identifier may not be nil.");
   }
 
-  NSString *key = [NSString stringWithFormat:@"%@|%@", self.app.name, database];
+  NSString *key = [self keyForDatabase:database];
 
   // Get the component from the container.
   @synchronized(self.instances) {
@@ -93,7 +97,7 @@ - (FIRFirestore *)firestoreForDatabase:(NSString *)database {
       auto workerQueue = absl::make_unique<AsyncQueue>(std::move(executor));
 
       id<FIRAuthInterop> auth = FIR_COMPONENT(FIRAuthInterop, self.app.container);
-      auto credentialsProvider = absl::make_unique<FirebaseCredentialsProvider>(self.app, auth);
+      auto credentialsProvider = std::make_shared<FirebaseCredentialsProvider>(self.app, auth);
 
       model::DatabaseId databaseID{util::MakeString(self.app.options.projectID),
                                    util::MakeString(database)};
@@ -102,19 +106,32 @@ - (FIRFirestore *)firestoreForDatabase:(NSString *)database {
                                             persistenceKey:std::move(persistenceKey)
                                        credentialsProvider:std::move(credentialsProvider)
                                                workerQueue:std::move(workerQueue)
-                                               firebaseApp:self.app];
+                                               firebaseApp:self.app
+                                          instanceRegistry:self];
       _instances[key] = firestore;
     }
-
     return firestore;
   }
 }
 
+- (void)removeInstanceWithDatabase:(NSString *)database {
+  @synchronized(_instances) {
+    NSString *key = [self keyForDatabase:database];
+    [_instances removeObjectForKey:key];
+  }
+}
+
 #pragma mark - FIRComponentLifecycleMaintainer
 
 - (void)appWillBeDeleted:(FIRApp *)app {
-  // Stop any actions and clean up resources since instances of Firestore associated with this app
-  // will be removed. Currently does not do anything.
+  NSDictionary<NSString *, FIRFirestore *> *instances;
+  @synchronized(_instances) {
+    instances = [_instances copy];
+    [_instances removeAllObjects];
+  }
+  for (NSString *key in instances) {
+    [instances[key] shutdownInternalWithCompletion:nil];
+  }
 }
 
 #pragma mark - Object Lifecycle

Firestore/Source/Core/FSTFirestoreClient.h

@@ -69,7 +69,7 @@ NS_ASSUME_NONNULL_BEGIN
 + (instancetype)clientWithDatabaseInfo:(const core::DatabaseInfo &)databaseInfo
                               settings:(const api::Settings &)settings
                    credentialsProvider:
-                       (auth::CredentialsProvider *)credentialsProvider  // no passing ownership
+                       (std::shared_ptr<auth::CredentialsProvider>)credentialsProvider
                           userExecutor:(std::shared_ptr<util::Executor>)userExecutor
                            workerQueue:(std::shared_ptr<util::AsyncQueue>)workerQueue;