Update FirestoreCache key type

firestore/integration_test_internal/src/firestore_test.cc

@@ -87,6 +87,20 @@ TEST_F(FirestoreIntegrationTest, GetInstance) {
   delete auth;
 }
 
+TEST_F(FirestoreIntegrationTest, GetInstanceWithNamedDatabase) {
+  App* app = this->app();
+  EXPECT_NE(nullptr, app);
+
+  InitResult result;
+  Firestore* instance = Firestore::GetInstance(app, "foo", &result);
+  EXPECT_EQ(kInitResultSuccess, result);
+  EXPECT_NE(nullptr, instance);
+  EXPECT_EQ(app, instance->app());
+  // TODO(Mila): check the database name is "foo".
+
+  delete instance;
+}
+
 // Sanity test for stubs.
 TEST_F(FirestoreIntegrationTest, TestCanCreateCollectionAndDocumentReferences) {
   Firestore* db = TestFirestore();
@@ -1348,6 +1362,43 @@ TEST_F(FirestoreIntegrationTest, RestartFirestoreLeadsToNewInstance) {
   delete db1;
 }
 
+TEST_F(FirestoreIntegrationTest,
+       RestartFirestoreLeadsToNewNamedDatabaseInstance) {
+  // TODO(Mila): Run this test against emulator.
+  GTEST_SKIP();
+
+  App* app = App::GetInstance();
+  InitResult init_result;
+  Firestore* db1 = Firestore::GetInstance(app, "test-db", &init_result);
+  ASSERT_EQ(kInitResultSuccess, init_result);
+
+  // Create a document that we can use for verification later.
+  DocumentReference doc1 = db1->Collection("abc").Document();
+  const std::string doc_path = doc1.path();
+  EXPECT_THAT(doc1.Set({{"foo", FieldValue::String("bar")}}), FutureSucceeds());
+
+  // Terminate `db1` so that it will be removed from the instance cache.
+  EXPECT_THAT(db1->Terminate(), FutureSucceeds());
+
+  // Verify that GetInstance() returns a new instance since the old instance has
+  // been terminated.
+  Firestore* db2 = Firestore::GetInstance(app, "test-db", &init_result);
+  ASSERT_EQ(kInitResultSuccess, init_result);
+  EXPECT_NE(db1, db2);
+
+  // Verify that the new instance points to the same database by verifying that
+  // the document created with the old instance exists in the new instance.
+  DocumentReference doc2 = db2->Document(doc_path);
+  const DocumentSnapshot* snapshot2 = Await(doc2.Get(Source::kCache));
+  ASSERT_NE(snapshot2, nullptr);
+  EXPECT_TRUE(snapshot2->exists());
+  EXPECT_THAT(snapshot2->GetData(),
+              ContainerEq(MapFieldValue{{"foo", FieldValue::String("bar")}}));
+
+  delete db2;
+  delete db1;
+}
+
 TEST_F(FirestoreIntegrationTest, CanStopListeningAfterTerminate) {
   auto instance = TestFirestore();
   DocumentReference reference = instance->Document("abc/123");

firestore/integration_test_internal/src/integration_test.cc

@@ -236,6 +236,7 @@ void FirebaseFirestoreBasicTest::TearDown() {
   FirebaseTest::TearDown();
 }
 
+// TODO(Mila): initialize named DB
 void FirebaseFirestoreBasicTest::InitializeFirestore() {
   LogDebug("Initializing Firebase Firestore.");
 
@@ -261,6 +262,7 @@ void FirebaseFirestoreBasicTest::InitializeFirestore() {
   initialized_ = true;
 }
 
+// TODO(Mila): terminate named DB
 void FirebaseFirestoreBasicTest::TerminateFirestore() {
   if (!initialized_) return;
 

firestore/integration_test_internal/src/util/integration_test_util.cc

@@ -37,7 +37,8 @@ struct TestFriend {
   static FirestoreInternal* CreateTestFirestoreInternal(App* app) {
 #if !defined(__ANDROID__)
     return new FirestoreInternal(
-        app, /*database_id=*/"(default)",
+        app,
+        /*database_id=*/"(default)",  // TODO(Mila): use dynamic database ID
         absl::make_unique<credentials::EmptyAuthCredentialsProvider>(),
         absl::make_unique<credentials::EmptyAppCheckCredentialsProvider>());
 #else

firestore/integration_test_internal/src/validation_test.cc

@@ -472,6 +472,77 @@ TEST_F(ValidationTest,
   EXPECT_NE(Firestore::GetInstance("foo"), nullptr);
 }
 
+TEST_F(ValidationTest,
+       FirestoreGetInstanceCalledMultipleTimeReturnSameInstance) {
+  {
+    Firestore* instance1 = Firestore::GetInstance();
+    Firestore* instance2 = Firestore::GetInstance();
+    EXPECT_EQ(instance1, instance2);
+  }
+  {
+    Firestore* instance1 = Firestore::GetInstance(app());
+    Firestore* instance2 = Firestore::GetInstance(app());
+    EXPECT_EQ(instance1, instance2);
+  }
+  {
+    Firestore* instance1 = Firestore::GetInstance("foo");
+    Firestore* instance2 = Firestore::GetInstance("foo");
+    EXPECT_EQ(instance1, instance2);
+  }
+  {
+    Firestore* instance1 = Firestore::GetInstance(app(), "foo");
+    Firestore* instance2 = Firestore::GetInstance(app(), "foo");
+    EXPECT_EQ(instance1, instance2);
+  }
+}
+
+TEST_F(ValidationTest,
+       differentFirestoreGetInstanceCanGetSameDefaultFirestoreInstance) {
+  Firestore* instance1 = Firestore::GetInstance();
+  Firestore* instance2 = Firestore::GetInstance(app());
+  Firestore* instance3 = Firestore::GetInstance("(default)");
+  Firestore* instance4 = Firestore::GetInstance(app(), "(default)");
+
+  EXPECT_EQ(instance1, instance2);
+  EXPECT_EQ(instance1, instance3);
+  EXPECT_EQ(instance1, instance4);
+  EXPECT_EQ(instance2, instance3);
+  EXPECT_EQ(instance2, instance4);
+  EXPECT_EQ(instance3, instance4);
+}
+
+TEST_F(
+    ValidationTest,
+    differentFirestoreGetInstanceWithSameDatabaseNameShouldGetSameFirestoreInstance) {
+  Firestore* instance1 = Firestore::GetInstance("foo");
+  Firestore* instance2 = Firestore::GetInstance(app(), "foo");
+  EXPECT_EQ(instance1, instance2);
+}
+
+TEST_F(
+    ValidationTest,
+    differentFirestoreGetInstanceWithDifferentDatabaseNameShouldGetDifferentFirestoreInstance) {
+  {
+    Firestore* instance1 = Firestore::GetInstance();
+    Firestore* instance2 = Firestore::GetInstance("bar");
+    EXPECT_NE(instance1, instance2);
+  }
+  {
+    Firestore* instance1 = Firestore::GetInstance("foo");
+    Firestore* instance2 = Firestore::GetInstance("bar");
+    EXPECT_NE(instance1, instance2);
+  }
+  {
+    Firestore* instance1 = Firestore::GetInstance(app(), "foo");
+    Firestore* instance2 = Firestore::GetInstance(app(), "bar");
+    EXPECT_NE(instance1, instance2);
+  }
+  {
+    Firestore* instance1 = Firestore::GetInstance("foo");
+    Firestore* instance2 = Firestore::GetInstance(app(), "bar");
+    EXPECT_NE(instance1, instance2);
+  }
+}
 TEST_F(ValidationTest, CollectionPathsMustBeOddLength) {
   Firestore* db = TestFirestore();
   DocumentReference base_document = db->Document("foo/bar");

firestore/src/common/firestore.cc

@@ -63,25 +63,37 @@ const char* GetPlatform() {
 #endif
 }
 
+// Use the combination of the App address and database ID to identify a
+// firestore instance in cache.
+using FirestoreMap = std::map<std::pair<App*, std::string>, Firestore*>;
+
+FirestoreMap::key_type MakeKey(App* app, const char* database_id) {
+  return std::make_pair(app, std::string(database_id));
+}
+
 Mutex* g_firestores_lock = new Mutex();
-std::map<App*, Firestore*>* g_firestores = nullptr;
+FirestoreMap* g_firestores = nullptr;
 
 // Ensures that the cache is initialized.
 // Prerequisite: `g_firestores_lock` must be locked before calling this
 // function.
-std::map<App*, Firestore*>* FirestoreCache() {
+FirestoreMap* FirestoreCache() {
   if (!g_firestores) {
-    g_firestores = new std::map<App*, Firestore*>();
+    g_firestores = new FirestoreMap();
   }
   return g_firestores;
 }
 
 // Prerequisite: `g_firestores_lock` must be locked before calling this
 // function.
-Firestore* FindFirestoreInCache(App* app, InitResult* init_result_out) {
+Firestore* FindFirestoreInCache(App* app,
+                                const char* database_id,
+                                InitResult* init_result_out) {
   auto* cache = FirestoreCache();
 
-  auto found = cache->find(app);
+  FirestoreMap::key_type key = MakeKey(app, database_id);
+  FirestoreMap::iterator found = g_firestores->find(key);
+
   if (found != cache->end()) {
     if (init_result_out) *init_result_out = kInitResultSuccess;
     return found->second;
@@ -107,8 +119,8 @@ void ValidateApp(App* app) {
   }
 }
 
-void ValidateDatabase(const char* db_name) {
-  if (!db_name) {
+void ValidateDatabase(const char* database_id) {
+  if (!database_id) {
     SimpleThrowInvalidArgument(
         "Provided database ID must not be null. Use other "
         "firebase::App::GetInstance() if you'd like to use the default "
@@ -150,12 +162,13 @@ Firestore* Firestore::GetInstance(App* app,
   ValidateDatabase(db_name);
 
   MutexLock lock(*g_firestores_lock);
-  Firestore* from_cache = FindFirestoreInCache(app, init_result_out);
+  Firestore* from_cache = FindFirestoreInCache(app, db_name, init_result_out);
   if (from_cache) {
     return from_cache;
   }
 
-  return AddFirestoreToCache(new Firestore(app, db_name), init_result_out);
+  return AddFirestoreToCache(new Firestore(app, db_name), db_name,
+                             init_result_out);
 }
 
 Firestore* Firestore::CreateFirestore(App* app,
@@ -167,14 +180,15 @@ Firestore* Firestore::CreateFirestore(App* app,
 
   MutexLock lock(*g_firestores_lock);
 
-  Firestore* from_cache = FindFirestoreInCache(app, init_result_out);
+  Firestore* from_cache = FindFirestoreInCache(app, nullptr, init_result_out);
   SIMPLE_HARD_ASSERT(from_cache == nullptr,
                      "Firestore must not be created already");
 
-  return AddFirestoreToCache(new Firestore(internal), init_result_out);
+  return AddFirestoreToCache(new Firestore(internal), nullptr, init_result_out);
 }
 
 Firestore* Firestore::AddFirestoreToCache(Firestore* firestore,
+                                          const char* database_id,
                                           InitResult* init_result_out) {
   InitResult init_result = CheckInitialized(*firestore->internal_);
   if (init_result_out) {
@@ -185,7 +199,8 @@ Firestore* Firestore::AddFirestoreToCache(Firestore* firestore,
     return nullptr;
   }
 
-  FirestoreCache()->emplace(firestore->app(), firestore);
+  FirestoreMap::key_type key = MakeKey(firestore->app(), database_id);
+  FirestoreCache()->emplace(key, firestore);
   return firestore;
 }
 
@@ -226,6 +241,7 @@ void Firestore::DeleteInternal() {
   if (!internal_) return;
 
   App* my_app = app();
+  const std::string& database_id = internal_->database_id().database_id();
 
   // Only need to unregister if internal_ is initialized.
   if (internal_->initialized()) {
@@ -249,8 +265,11 @@ void Firestore::DeleteInternal() {
   delete internal_;
   internal_ = nullptr;
   // If a Firestore is explicitly deleted, remove it from our cache.
-  FirestoreCache()->erase(my_app);
-  // If it's the last one, delete the map.
+
+  FirestoreMap::key_type key = MakeKey(
+      my_app, database_id == "" ? kDefaultDatabase : database_id.c_str());
+
+  FirestoreCache()->erase(key);
   if (g_firestores->empty()) {
     delete g_firestores;
     g_firestores = nullptr;
@@ -362,7 +381,11 @@ Future<void> Firestore::EnableNetwork() {
 
 Future<void> Firestore::Terminate() {
   if (!internal_) return FailedFuture<void>();
-  FirestoreCache()->erase(app());
+  const std::string& database_id = internal_->database_id().database_id();
+  FirestoreMap::key_type key = MakeKey(
+      app(), database_id == "" ? kDefaultDatabase : database_id.c_str());
+
+  FirestoreCache()->erase(key);
   return internal_->Terminate();
 }
 

firestore/src/include/firebase/firestore.h

@@ -503,6 +503,7 @@ class Firestore {
                                     FirestoreInternal* internal,
                                     InitResult* init_result_out);
   static Firestore* AddFirestoreToCache(Firestore* firestore,
+                                        const char* database_id,
                                         InitResult* init_result_out);
 
   static void SetClientLanguage(const std::string& language_token);