From f48f03f2f4779797d340df182f1beaf5d620565a Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Mon, 21 Jun 2021 16:14:51 -0400
Subject: [PATCH 1/3] Fix the ReferenceCountedFutureImpl bug where it can use
 FirestoreInternal after its deletion.

---
 .../integration_test_internal/CMakeLists.txt  |   1 +
 .../src/android/promise_android_test.cc       |  18 +++
 .../android/promise_factory_android_test.cc   | 132 ++++++++++++++++++
 .../src/firestore_test.cc                     |  33 +++--
 .../src/source_test.cc                        |  77 +++++-----
 firestore/src/android/firestore_android.h     |  80 +++++++++++
 firestore/src/android/promise_android.h       |  74 ++++++----
 .../src/android/promise_factory_android.h     |  44 +++---
 8 files changed, 361 insertions(+), 98 deletions(-)
 create mode 100644 firestore/integration_test_internal/src/android/promise_factory_android_test.cc

diff --git a/firestore/integration_test_internal/CMakeLists.txt b/firestore/integration_test_internal/CMakeLists.txt
index e3ad14aa9d..249b40222f 100644
--- a/firestore/integration_test_internal/CMakeLists.txt
+++ b/firestore/integration_test_internal/CMakeLists.txt
@@ -110,6 +110,7 @@ set(FIREBASE_INTEGRATION_TEST_ANDROID_TEST_SRCS
     src/android/geo_point_android_test.cc
     src/android/jni_runnable_android_test.cc
     src/android/promise_android_test.cc
+    src/android/promise_factory_android_test.cc
     src/android/settings_android_test.cc
     src/android/snapshot_metadata_android_test.cc
     src/android/timestamp_android_test.cc
diff --git a/firestore/integration_test_internal/src/android/promise_android_test.cc b/firestore/integration_test_internal/src/android/promise_android_test.cc
index c4d0e0a1cf..8c9d128039 100644
--- a/firestore/integration_test_internal/src/android/promise_android_test.cc
+++ b/firestore/integration_test_internal/src/android/promise_android_test.cc
@@ -403,6 +403,24 @@ TEST_F(PromiseTest, FutureNonVoidShouldCallCompletionWhenTaskCancels) {
   EXPECT_EQ(completion.result(), nullptr);
 }
 
+TEST_F(PromiseTest, RegisterForTaskShouldNotCrashIfFirestoreWasDeleted) {
+  jni::Env env = GetEnv();
+  auto promise = promises().MakePromise<void>();
+  DeleteFirestore(TestFirestore());
+
+  promise.RegisterForTask(env, AsyncFn::kFn, GetTask());
+}
+
+TEST_F(PromiseTest, GetFutureShouldNotCrashIfFirestoreWasDeleted) {
+  jni::Env env = GetEnv();
+  auto promise = promises().MakePromise<void>();
+  promise.RegisterForTask(env, AsyncFn::kFn, GetTask());
+  DeleteFirestore(TestFirestore());
+
+  auto future = promise.GetFuture();
+  EXPECT_EQ(future.status(), FutureStatus::kFutureStatusInvalid);
+}
+
 }  // namespace
 }  // namespace firestore
 }  // namespace firebase
diff --git a/firestore/integration_test_internal/src/android/promise_factory_android_test.cc b/firestore/integration_test_internal/src/android/promise_factory_android_test.cc
new file mode 100644
index 0000000000..a900610519
--- /dev/null
+++ b/firestore/integration_test_internal/src/android/promise_factory_android_test.cc
@@ -0,0 +1,132 @@
+// Copyright 2021 Google LLC
+
+#include "firestore/src/android/promise_factory_android.h"
+
+#include <utility>
+
+#include "android/firestore_integration_test_android.h"
+#include "android/task_completion_source.h"
+#include "firebase/future.h"
+#include "firestore/src/android/converter_android.h"
+#include "firestore/src/android/firestore_android.h"
+#include "firestore/src/jni/env.h"
+#include "firestore/src/jni/integer.h"
+#include "firestore/src/jni/ownership.h"
+#include "firestore/src/jni/task.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace firebase {
+namespace firestore {
+namespace {
+
+using jni::Env;
+using jni::Integer;
+using jni::Local;
+using jni::Task;
+
+// Since `PromiseFactory` acts as a "constructor" of `Promise` objects, its
+// ability to create `Promise` objects is thoroughly tested in the unit tests
+// for `Promise` and therefore the tests here only cover the other aspects of
+// `PromiseFactory`, such as move semantics.
+
+class PromiseFactoryTest : public FirestoreAndroidIntegrationTest {
+ public:
+  // An enum of asynchronous functions to use in tests, as required by
+  // `FutureManager`.
+  enum class AsyncFn {
+    kFn,
+    kCount,  // Must be the last enum value.
+  };
+
+ protected:
+  void AssertCreatesValidFutures(Env& env,
+                                 PromiseFactory<AsyncFn>& promise_factory) {
+    Local<TaskCompletionSource> tcs = TaskCompletionSource::Create(env);
+    Local<Task> task = tcs.GetTask(env);
+    Future<void> future =
+        promise_factory.NewFuture<void>(env, AsyncFn::kFn, task);
+    EXPECT_EQ(future.status(), FutureStatus::kFutureStatusPending);
+    tcs.SetResult(env, jni::Integer::Create(env, 42));
+    Await(future);
+    EXPECT_EQ(future.status(), FutureStatus::kFutureStatusComplete);
+  }
+
+  void AssertCreatesInvalidFutures(Env& env,
+                                   PromiseFactory<AsyncFn>& promise_factory) {
+    Local<TaskCompletionSource> tcs = TaskCompletionSource::Create(env);
+    Local<Task> task = tcs.GetTask(env);
+    Future<void> future =
+        promise_factory.NewFuture<void>(env, AsyncFn::kFn, task);
+    EXPECT_EQ(future.status(), FutureStatus::kFutureStatusInvalid);
+  }
+};
+
+TEST_F(PromiseFactoryTest, CopyConstructor) {
+  Firestore* firestore = TestFirestore();
+  PromiseFactory<AsyncFn> promise_factory1(GetFirestoreInternal(firestore));
+
+  PromiseFactory<AsyncFn> promise_factory2(promise_factory1);
+
+  Env env;
+  {
+    SCOPED_TRACE("promise_factory1");
+    AssertCreatesValidFutures(env, promise_factory1);
+  }
+  {
+    SCOPED_TRACE("promise_factory2");
+    AssertCreatesValidFutures(env, promise_factory2);
+  }
+}
+
+TEST_F(PromiseFactoryTest, CopyConstructorWithDeletedFirestore) {
+  Firestore* firestore = TestFirestore();
+  PromiseFactory<AsyncFn> promise_factory1(GetFirestoreInternal(firestore));
+  DeleteFirestore(firestore);
+
+  PromiseFactory<AsyncFn> promise_factory2(promise_factory1);
+
+  Env env;
+  {
+    SCOPED_TRACE("promise_factory1");
+    AssertCreatesInvalidFutures(env, promise_factory1);
+  }
+  {
+    SCOPED_TRACE("promise_factory2");
+    AssertCreatesInvalidFutures(env, promise_factory2);
+  }
+}
+
+TEST_F(PromiseFactoryTest, MoveConstructor) {
+  Firestore* firestore = TestFirestore();
+  PromiseFactory<AsyncFn> promise_factory1(GetFirestoreInternal(firestore));
+
+  PromiseFactory<AsyncFn> promise_factory2(std::move(promise_factory1));
+
+  Env env;
+  AssertCreatesValidFutures(env, promise_factory2);
+}
+
+TEST_F(PromiseFactoryTest, MoveConstructorWithDeletedFirestore) {
+  Firestore* firestore = TestFirestore();
+  PromiseFactory<AsyncFn> promise_factory1(GetFirestoreInternal(firestore));
+  DeleteFirestore(firestore);
+
+  PromiseFactory<AsyncFn> promise_factory2(std::move(promise_factory1));
+
+  Env env;
+  AssertCreatesInvalidFutures(env, promise_factory2);
+}
+
+TEST_F(PromiseFactoryTest, ShouldCreateInvalidPromisesIfFirestoreIsDeleted) {
+  Firestore* firestore = TestFirestore();
+  PromiseFactory<AsyncFn> promise_factory(GetFirestoreInternal(firestore));
+  DeleteFirestore(firestore);
+  Env env;
+
+  AssertCreatesInvalidFutures(env, promise_factory);
+}
+
+}  // namespace
+}  // namespace firestore
+}  // namespace firebase
diff --git a/firestore/integration_test_internal/src/firestore_test.cc b/firestore/integration_test_internal/src/firestore_test.cc
index a80145026d..158b2d0cdb 100644
--- a/firestore/integration_test_internal/src/firestore_test.cc
+++ b/firestore/integration_test_internal/src/firestore_test.cc
@@ -3,11 +3,8 @@
 #include "firebase/firestore.h"
 
 #include <algorithm>
+#include <future>
 #include <memory>
-
-#if !defined(__ANDROID__)
-#include <future>  // NOLINT(build/c++11)
-#endif
 #include <stdexcept>
 
 #if defined(__ANDROID__)
@@ -1489,12 +1486,11 @@ TEST_F(FirestoreIntegrationTest, AuthWorks) {
   WriteDocument(doc, MapFieldValue{{"foo", FieldValue::Integer(43)}});
 }
 
-#if !defined(__ANDROID__)
 // This test is to ensure b/172986326 doesn't regress.
 // Note: this test only exists in C++.
-TEST_F(FirestoreIntegrationTest, FirestoreCanBeDeletedFromTransaction) {
-  auto* app = App::Create(this->app()->options(), "foo");
-  auto* db = Firestore::GetInstance(app);
+TEST_F(FirestoreIntegrationTest, FirestoreCanBeDeletedFromTransactionAsync) {
+  Firestore* db = TestFirestore();
+  DisownFirestore(db);
 
   auto future = db->RunTransaction(
       [](Transaction&, std::string&) { return Error::kErrorOk; });
@@ -1511,7 +1507,26 @@ TEST_F(FirestoreIntegrationTest, FirestoreCanBeDeletedFromTransaction) {
   callback_done.wait();
   deletion.wait();
 }
-#endif  // #if !defined(__ANDROID__)
+
+// This test is to ensure b/172986326 doesn't regress.
+// Note: this test only exists in C++.
+TEST_F(FirestoreIntegrationTest, FirestoreCanBeDeletedFromTransaction) {
+  Firestore* db = TestFirestore();
+  DisownFirestore(db);
+
+  auto future = db->RunTransaction(
+      [](Transaction&, std::string&) { return Error::kErrorOk; });
+
+  std::promise<void> callback_done_promise;
+  auto callback_done = callback_done_promise.get_future();
+  future.AddOnCompletion([&](const Future<void>&) mutable {
+    delete db;
+    callback_done_promise.set_value();
+  });
+
+  Await(future);
+  callback_done.wait();
+}
 
 #if defined(__ANDROID__)
 TEST_F(FirestoreAndroidIntegrationTest,
diff --git a/firestore/integration_test_internal/src/source_test.cc b/firestore/integration_test_internal/src/source_test.cc
index 8dd741cc7e..17cf76e0b4 100644
--- a/firestore/integration_test_internal/src/source_test.cc
+++ b/firestore/integration_test_internal/src/source_test.cc
@@ -15,12 +15,9 @@
 namespace firebase {
 namespace firestore {
 
-// TODO(b/187448376): Temporarily disabling all tests as they seem to time out
-// on Android.
-
 using SourceTest = FirestoreIntegrationTest;
 
-TEST_F(SourceTest, DISABLED_GetDocumentWhileOnlineWithDefaultGetOptions) {
+TEST_F(SourceTest, GetDocumentWhileOnlineWithDefaultGetOptions) {
   MapFieldValue initial_data = {{"key", FieldValue::String("value")}};
   DocumentReference doc_ref = DocumentWithData(initial_data);
 
@@ -34,7 +31,7 @@ TEST_F(SourceTest, DISABLED_GetDocumentWhileOnlineWithDefaultGetOptions) {
   EXPECT_EQ(initial_data, snapshot.GetData());
 }
 
-TEST_F(SourceTest, DISABLED_GetCollectionWhileOnlineWithDefaultGetOptions) {
+TEST_F(SourceTest, GetCollectionWhileOnlineWithDefaultGetOptions) {
   std::map<std::string, MapFieldValue> initial_docs = {
       {"doc1", {{"key1", FieldValue::String("value1")}}},
       {"doc2", {{"key2", FieldValue::String("value2")}}},
@@ -51,7 +48,7 @@ TEST_F(SourceTest, DISABLED_GetCollectionWhileOnlineWithDefaultGetOptions) {
   EXPECT_EQ(initial_docs, QuerySnapshotToMap(snapshot));
 }
 
-TEST_F(SourceTest, DISABLED_GetDocumentWhileOfflineWithDefaultGetOptions) {
+TEST_F(SourceTest, GetDocumentWhileOfflineWithDefaultGetOptions) {
   MapFieldValue initial_data = {{"key", FieldValue::String("value")}};
   DocumentReference doc_ref = DocumentWithData(initial_data);
 
@@ -68,7 +65,7 @@ TEST_F(SourceTest, DISABLED_GetDocumentWhileOfflineWithDefaultGetOptions) {
   EXPECT_EQ(initial_data, snapshot.GetData());
 }
 
-TEST_F(SourceTest, DISABLED_GetCollectionWhileOfflineWithDefaultGetOptions) {
+TEST_F(SourceTest, GetCollectionWhileOfflineWithDefaultGetOptions) {
   std::map<std::string, MapFieldValue> initial_docs = {
       {"doc1", {{"key1", FieldValue::String("value1")}}},
       {"doc2", {{"key2", FieldValue::String("value2")}}},
@@ -102,7 +99,7 @@ TEST_F(SourceTest, DISABLED_GetCollectionWhileOfflineWithDefaultGetOptions) {
   EXPECT_EQ(new_data, QuerySnapshotToMap(snapshot));
 }
 
-TEST_F(SourceTest, DISABLED_GetDocumentWhileOnlineWithSourceEqualToCache) {
+TEST_F(SourceTest, GetDocumentWhileOnlineWithSourceEqualToCache) {
   MapFieldValue initial_data = {{"key", FieldValue::String("value")}};
   DocumentReference doc_ref = DocumentWithData(initial_data);
 
@@ -117,7 +114,7 @@ TEST_F(SourceTest, DISABLED_GetDocumentWhileOnlineWithSourceEqualToCache) {
   EXPECT_EQ(initial_data, snapshot.GetData());
 }
 
-TEST_F(SourceTest, DISABLED_GetCollectionWhileOnlineWithSourceEqualToCache) {
+TEST_F(SourceTest, GetCollectionWhileOnlineWithSourceEqualToCache) {
   std::map<std::string, MapFieldValue> initial_docs = {
       {"doc1", {{"key1", FieldValue::String("value1")}}},
       {"doc2", {{"key2", FieldValue::String("value2")}}},
@@ -135,7 +132,7 @@ TEST_F(SourceTest, DISABLED_GetCollectionWhileOnlineWithSourceEqualToCache) {
   EXPECT_EQ(initial_docs, QuerySnapshotToMap(snapshot));
 }
 
-TEST_F(SourceTest, DISABLED_GetDocumentWhileOfflineWithSourceEqualToCache) {
+TEST_F(SourceTest, GetDocumentWhileOfflineWithSourceEqualToCache) {
   MapFieldValue initial_data = {{"key", FieldValue::String("value")}};
   DocumentReference doc_ref = DocumentWithData(initial_data);
 
@@ -151,7 +148,7 @@ TEST_F(SourceTest, DISABLED_GetDocumentWhileOfflineWithSourceEqualToCache) {
   EXPECT_EQ(initial_data, snapshot.GetData());
 }
 
-TEST_F(SourceTest, DISABLED_GetCollectionWhileOfflineWithSourceEqualToCache) {
+TEST_F(SourceTest, GetCollectionWhileOfflineWithSourceEqualToCache) {
   std::map<std::string, MapFieldValue> initial_docs = {
       {"doc1", {{"key1", FieldValue::String("value1")}}},
       {"doc2", {{"key2", FieldValue::String("value2")}}},
@@ -185,7 +182,7 @@ TEST_F(SourceTest, DISABLED_GetCollectionWhileOfflineWithSourceEqualToCache) {
   EXPECT_EQ(new_data, QuerySnapshotToMap(snapshot));
 }
 
-TEST_F(SourceTest, DISABLED_GetDocumentWhileOnlineWithSourceEqualToServer) {
+TEST_F(SourceTest, GetDocumentWhileOnlineWithSourceEqualToServer) {
   MapFieldValue initial_data = {{"key", FieldValue::String("value")}};
   DocumentReference doc_ref = DocumentWithData(initial_data);
 
@@ -199,7 +196,7 @@ TEST_F(SourceTest, DISABLED_GetDocumentWhileOnlineWithSourceEqualToServer) {
   EXPECT_EQ(initial_data, snapshot.GetData());
 }
 
-TEST_F(SourceTest, DISABLED_GetCollectionWhileOnlineWithSourceEqualToServer) {
+TEST_F(SourceTest, GetCollectionWhileOnlineWithSourceEqualToServer) {
   std::map<std::string, MapFieldValue> initial_docs = {
       {"doc1", {{"key1", FieldValue::String("value1")}}},
       {"doc2", {{"key2", FieldValue::String("value2")}}},
@@ -216,7 +213,7 @@ TEST_F(SourceTest, DISABLED_GetCollectionWhileOnlineWithSourceEqualToServer) {
   EXPECT_EQ(initial_docs, QuerySnapshotToMap(snapshot));
 }
 
-TEST_F(SourceTest, DISABLED_GetDocumentWhileOfflineWithSourceEqualToServer) {
+TEST_F(SourceTest, GetDocumentWhileOfflineWithSourceEqualToServer) {
   MapFieldValue initial_data = {{"key", FieldValue::String("value")}};
   DocumentReference doc_ref = DocumentWithData(initial_data);
 
@@ -229,7 +226,7 @@ TEST_F(SourceTest, DISABLED_GetDocumentWhileOfflineWithSourceEqualToServer) {
   EXPECT_EQ(future.error(), Error::kErrorUnavailable);
 }
 
-TEST_F(SourceTest, DISABLED_GetCollectionWhileOfflineWithSourceEqualToServer) {
+TEST_F(SourceTest, GetCollectionWhileOfflineWithSourceEqualToServer) {
   std::map<std::string, MapFieldValue> initial_docs = {
       {"doc1", {{"key1", FieldValue::String("value1")}}},
       {"doc2", {{"key2", FieldValue::String("value2")}}},
@@ -245,7 +242,7 @@ TEST_F(SourceTest, DISABLED_GetCollectionWhileOfflineWithSourceEqualToServer) {
   EXPECT_EQ(future.error(), Error::kErrorUnavailable);
 }
 
-TEST_F(SourceTest, DISABLED_GetDocumentWhileOfflineWithDifferentGetOptions) {
+TEST_F(SourceTest, GetDocumentWhileOfflineWithDifferentGetOptions) {
   MapFieldValue initial_data = {{"key", FieldValue::String("value")}};
   DocumentReference doc_ref = DocumentWithData(initial_data);
 
@@ -290,7 +287,7 @@ TEST_F(SourceTest, DISABLED_GetDocumentWhileOfflineWithDifferentGetOptions) {
   EXPECT_EQ(future.error(), Error::kErrorUnavailable);
 }
 
-TEST_F(SourceTest, DISABLED_GetCollectionWhileOfflineWithDifferentGetOptions) {
+TEST_F(SourceTest, GetCollectionWhileOfflineWithDifferentGetOptions) {
   std::map<std::string, MapFieldValue> initial_docs = {
       {"doc1", {{"key1", FieldValue::String("value1")}}},
       {"doc2", {{"key2", FieldValue::String("value2")}}},
@@ -352,7 +349,7 @@ TEST_F(SourceTest, DISABLED_GetCollectionWhileOfflineWithDifferentGetOptions) {
   EXPECT_EQ(future.error(), Error::kErrorUnavailable);
 }
 
-TEST_F(SourceTest, DISABLED_GetNonExistingDocWhileOnlineWithDefaultGetOptions) {
+TEST_F(SourceTest, GetNonExistingDocWhileOnlineWithDefaultGetOptions) {
   DocumentReference doc_ref = Document();
 
   Future<DocumentSnapshot> future = doc_ref.Get();
@@ -364,8 +361,7 @@ TEST_F(SourceTest, DISABLED_GetNonExistingDocWhileOnlineWithDefaultGetOptions) {
   EXPECT_FALSE(snapshot.metadata().has_pending_writes());
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingCollectionWhileOnlineWithDefaultGetOptions) {
+TEST_F(SourceTest, GetNonExistingCollectionWhileOnlineWithDefaultGetOptions) {
   CollectionReference col_ref = Collection();
 
   Future<QuerySnapshot> future = col_ref.Get();
@@ -378,8 +374,7 @@ TEST_F(SourceTest,
   EXPECT_FALSE(snapshot.metadata().has_pending_writes());
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingDocWhileOfflineWithDefaultGetOptions) {
+TEST_F(SourceTest, GetNonExistingDocWhileOfflineWithDefaultGetOptions) {
   DocumentReference doc_ref = Document();
 
   DisableNetwork();
@@ -389,10 +384,11 @@ TEST_F(SourceTest,
   EXPECT_EQ(future.error(), Error::kErrorUnavailable);
 }
 
-// TODO(b/112267729): We should raise a fromCache=true event with a
-// nonexistent snapshot, but because the default source goes through a normal
-// listener, we do not.
-TEST_F(SourceTest, DISABLED_GetDeletedDocWhileOfflineWithDefaultGetOptions) {
+
+TEST_F(SourceTest, GetDeletedDocWhileOfflineWithDefaultGetOptions) {
+  GTEST_SKIP() << "b/112267729: We should raise a fromCache=true event with "
+               << "a nonexistent snapshot, but because the default source goes "
+               << "through a normal listener, we do not.";
   DocumentReference doc_ref = Document();
   Await(doc_ref.Delete());
 
@@ -407,8 +403,7 @@ TEST_F(SourceTest, DISABLED_GetDeletedDocWhileOfflineWithDefaultGetOptions) {
   EXPECT_FALSE(snapshot.metadata().has_pending_writes());
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingCollectionWhileOfflineWithDefaultGetOptions) {
+TEST_F(SourceTest, GetNonExistingCollectionWhileOfflineWithDefaultGetOptions) {
   CollectionReference col_ref = Collection();
 
   DisableNetwork();
@@ -422,8 +417,7 @@ TEST_F(SourceTest,
   EXPECT_FALSE(snapshot.metadata().has_pending_writes());
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingDocWhileOnlineWithSourceEqualToCache) {
+TEST_F(SourceTest, GetNonExistingDocWhileOnlineWithSourceEqualToCache) {
   DocumentReference doc_ref = Document();
 
   // Attempt to get doc. This will fail since there's nothing in cache.
@@ -433,8 +427,7 @@ TEST_F(SourceTest,
   EXPECT_EQ(future.error(), Error::kErrorUnavailable);
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingCollectionWhileOnlineWithSourceEqualToCache) {
+TEST_F(SourceTest, GetNonExistingCollectionWhileOnlineWithSourceEqualToCache) {
   CollectionReference col_ref = Collection();
 
   Future<QuerySnapshot> future = col_ref.Get(Source::kCache);
@@ -447,8 +440,7 @@ TEST_F(SourceTest,
   EXPECT_FALSE(snapshot.metadata().has_pending_writes());
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingDocWhileOfflineWithSourceEqualToCache) {
+TEST_F(SourceTest, GetNonExistingDocWhileOfflineWithSourceEqualToCache) {
   DocumentReference doc_ref = Document();
 
   DisableNetwork();
@@ -459,7 +451,7 @@ TEST_F(SourceTest,
   EXPECT_EQ(future.error(), Error::kErrorUnavailable);
 }
 
-TEST_F(SourceTest, DISABLED_GetDeletedDocWhileOfflineWithSourceEqualToCache) {
+TEST_F(SourceTest, GetDeletedDocWhileOfflineWithSourceEqualToCache) {
   DocumentReference doc_ref = Document();
   Await(doc_ref.Delete());
 
@@ -474,8 +466,7 @@ TEST_F(SourceTest, DISABLED_GetDeletedDocWhileOfflineWithSourceEqualToCache) {
   EXPECT_FALSE(snapshot.metadata().has_pending_writes());
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingCollectionWhileOfflineWithSourceEqualToCache) {
+TEST_F(SourceTest, GetNonExistingCollectionWhileOfflineWithSourceEqualToCache) {
   CollectionReference col_ref = Collection();
 
   DisableNetwork();
@@ -489,8 +480,7 @@ TEST_F(SourceTest,
   EXPECT_FALSE(snapshot.metadata().has_pending_writes());
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingDocWhileOnlineWithSourceEqualToServer) {
+TEST_F(SourceTest, GetNonExistingDocWhileOnlineWithSourceEqualToServer) {
   DocumentReference doc_ref = Document();
 
   Future<DocumentSnapshot> future = doc_ref.Get(Source::kServer);
@@ -502,8 +492,7 @@ TEST_F(SourceTest,
   EXPECT_FALSE(snapshot.metadata().has_pending_writes());
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingCollectionWhileOnlineWithSourceEqualToServer) {
+TEST_F(SourceTest, GetNonExistingCollectionWhileOnlineWithSourceEqualToServer) {
   CollectionReference col_ref = Collection();
 
   Future<QuerySnapshot> future = col_ref.Get(Source::kServer);
@@ -516,8 +505,7 @@ TEST_F(SourceTest,
   EXPECT_FALSE(snapshot.metadata().has_pending_writes());
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingDocWhileOfflineWithSourceEqualToServer) {
+TEST_F(SourceTest, GetNonExistingDocWhileOfflineWithSourceEqualToServer) {
   DocumentReference doc_ref = Document();
 
   DisableNetwork();
@@ -528,8 +516,7 @@ TEST_F(SourceTest,
   EXPECT_EQ(future.error(), Error::kErrorUnavailable);
 }
 
-TEST_F(SourceTest,
-       DISABLED_GetNonExistingCollectionWhileOfflineWithSourceEqualToServer) {
+TEST_F(SourceTest, GetNonExistingCollectionWhileOfflineWithSourceEqualToServer) {
   CollectionReference col_ref = Collection();
 
   DisableNetwork();
diff --git a/firestore/src/android/firestore_android.h b/firestore/src/android/firestore_android.h
index 34ed9bb414..518f4bcd9f 100644
--- a/firestore/src/android/firestore_android.h
+++ b/firestore/src/android/firestore_android.h
@@ -7,6 +7,7 @@
 #include <functional>
 #include <list>
 #include <memory>
+#include <mutex>
 #include <unordered_set>
 
 #include "app/src/cleanup_notifier.h"
@@ -191,6 +192,85 @@ class FirestoreInternal {
   CleanupNotifier cleanup_;
 };
 
+// Holds a "weak reference" to a `FirestoreInternal` object.
+//
+// When the given `FirestoreInternal` object is deleted, this object will null
+// out its pointer to it.
+//
+// To use the encapsulated `FirestoreInternal` object, invoke the `Run()` or
+// `RunIfValid()` method with a callback that will be invoked with a pointer to
+// the `FirestoreInternal` object.
+class FirestoreInternalWeakReference {
+ public:
+  explicit FirestoreInternalWeakReference(FirestoreInternal* firestore)
+      : firestore_(firestore) {
+    if (firestore_) {
+      firestore_->cleanup().RegisterObject(this, CleanUp);
+    }
+  }
+
+  FirestoreInternalWeakReference(const FirestoreInternalWeakReference& rhs) {
+    std::lock_guard<std::recursive_mutex> lock(rhs.mutex_);
+    firestore_ = rhs.firestore_;
+    if (firestore_) {
+      firestore_->cleanup().RegisterObject(this, CleanUp);
+    }
+  }
+
+  FirestoreInternalWeakReference& operator=(
+      const FirestoreInternalWeakReference& rhs) = delete;
+
+  ~FirestoreInternalWeakReference() {
+    std::lock_guard<std::recursive_mutex> lock(mutex_);
+    if (firestore_) {
+      firestore_->cleanup().UnregisterObject(this);
+    }
+  }
+
+ private:
+  static void CleanUp(void* obj) {
+    auto* instance = static_cast<FirestoreInternalWeakReference*>(obj);
+    std::lock_guard<std::recursive_mutex> lock(instance->mutex_);
+    instance->firestore_ = nullptr;
+  }
+
+  mutable std::recursive_mutex mutex_;
+  FirestoreInternal* firestore_ = nullptr;
+
+  // Declare `Run()` after declaring `firestore_` so that `firestore_` can be
+  // referenced in the `decltype` return type of `Run()`.
+ public:
+  /**
+   * Invokes the given function with the referent `FirestoreInternal` object,
+   * specifying `nullptr` if the `FirestoreInternal` object has been deleted,
+   * returning whatever the given callback returns.
+   *
+   * If a non-null `FirestoreInternal` pointer was specified to the method, then
+   * that `FirestoreInternal` object will remain alive (i.e. will not be
+   * deleted) for the duration of the callback's execution.
+   *
+   * If this method is invoked concurrently then calls will be serialized in
+   * an undefined order. Recursive calls of this method are allowed and will not
+   * block.
+   */
+  template <typename CallbackT>
+  auto Run(CallbackT callback) -> decltype(callback(this->firestore_)) {
+    std::lock_guard<std::recursive_mutex> lock(mutex_);
+    return callback(firestore_);
+  }
+
+  /**
+   * Same as `Run()` except that it only invokes the given callback if the
+   * `FirestoreInternal` pointer is non-null.
+   */
+  void RunIfValid(std::function<void(FirestoreInternal*)> callback) {
+    std::lock_guard<std::recursive_mutex> lock(mutex_);
+    if (firestore_) {
+      callback(firestore_);
+    }
+  }
+};
+
 }  // namespace firestore
 }  // namespace firebase
 #endif  // FIREBASE_FIRESTORE_SRC_ANDROID_FIRESTORE_ANDROID_H_
diff --git a/firestore/src/android/promise_android.h b/firestore/src/android/promise_android.h
index a80aa05037..382e6107b1 100644
--- a/firestore/src/android/promise_android.h
+++ b/firestore/src/android/promise_android.h
@@ -51,8 +51,6 @@ class Promise {
                               PublicType* result) = 0;
   };
 
-  ~Promise() {}
-
   Promise(const Promise&) = delete;
   Promise& operator=(const Promise&) = delete;
 
@@ -69,33 +67,39 @@ class Promise {
         env.get(), task.get(), ResultCallback, completer, kApiIdentifier);
   }
 
-  Future<PublicType> GetFuture() { return MakeFuture(impl_, handle_); }
+  Future<PublicType> GetFuture() {
+    return firestore_ref_.Run([this](FirestoreInternal* firestore) {
+      if (firestore == nullptr) {
+        return Future<PublicType>{};
+      }
+      return MakeFuture(impl_, handle_);
+    });
+  }
 
  private:
   // The constructor is intentionally private.
   // Create instances with `PromiseFactory`.
-  Promise(ReferenceCountedFutureImpl* impl,
-          FirestoreInternal* firestore,
-          Completion* completion)
-      : completer_(make_unique<Completer<PublicType, InternalType>>(
-            impl, firestore, completion)),
-        impl_(impl) {}
+  Promise(const FirestoreInternalWeakReference& firestore_ref, ReferenceCountedFutureImpl* impl, Completion* completion) : firestore_ref_(firestore_ref), completer_(make_unique<Completer<PublicType, InternalType>>(firestore_ref, impl, completion)), impl_(impl) {
+  }
 
   template <typename PublicT>
   class CompleterBase {
    public:
-    CompleterBase(ReferenceCountedFutureImpl* impl,
-                  FirestoreInternal* firestore,
+    CompleterBase(const FirestoreInternalWeakReference& firestore_ref,
+                  ReferenceCountedFutureImpl* impl,
                   Completion* completion)
-        : impl_{impl}, firestore_{firestore}, completion_(completion) {}
+        : firestore_ref_{firestore_ref}, impl_{impl}, completion_(completion) {}
 
     virtual ~CompleterBase() = default;
 
-    FirestoreInternal* firestore() { return firestore_; }
-
     SafeFutureHandle<PublicT> Alloc(int fn_index) {
-      handle_ = impl_->SafeAlloc<PublicT>(fn_index);
-      return handle_;
+      return firestore_ref_.Run([this, fn_index](FirestoreInternal* firestore) {
+        if (firestore == nullptr) {
+          return SafeFutureHandle<PublicT>{};
+        }
+        handle_ = impl_->SafeAlloc<PublicT>(fn_index);
+        return handle_;
+      });
     }
 
     virtual void CompleteWithResult(jobject raw_result,
@@ -107,8 +111,12 @@ class Promise {
       jni::Object result(raw_result);
 
       if (result_code == ::firebase::util::kFutureResultSuccess) {
-        // When succeeded, result is the resolved object of the Future.
-        SucceedWithResult(env, result);
+        firestore_ref_.RunIfValid([this, &env, &result](FirestoreInternal* firestore) {
+          SucceedWithResult(env, result, firestore);
+        });
+
+        delete this;
+
         return;
       }
 
@@ -126,20 +134,26 @@ class Promise {
                                   result_code);
           break;
       }
-      this->impl_->Complete(this->handle_, error_code, status_message);
-      if (this->completion_ != nullptr) {
-        this->completion_->CompleteWith(error_code, status_message, nullptr);
+      firestore_ref_.RunIfValid(
+          [this, error_code, status_message](FirestoreInternal* firestore) {
+            impl_->Complete(handle_, error_code, status_message);
+          });
+      if (completion_ != nullptr) {
+        completion_->CompleteWith(error_code, status_message, nullptr);
       }
       delete this;
     }
 
     virtual void SucceedWithResult(jni::Env& env,
-                                   const jni::Object& result) = 0;
+                                   const jni::Object& result,
+                                   FirestoreInternal* firestore) = 0;
+
+   private:
+    FirestoreInternalWeakReference firestore_ref_;
 
    protected:
     SafeFutureHandle<PublicT> handle_;
     ReferenceCountedFutureImpl* impl_;  // not owning
-    FirestoreInternal* firestore_;      // not owning
     Completion* completion_;            // not owning
   };
 
@@ -151,9 +165,11 @@ class Promise {
    public:
     using CompleterBase<PublicT>::CompleterBase;
 
-    void SucceedWithResult(jni::Env& env, const jni::Object& result) override {
+    void SucceedWithResult(jni::Env& env,
+                           const jni::Object& result,
+                           FirestoreInternal* firestore) override {
       auto future_result =
-          MakePublic<PublicT, InternalT>(env, this->firestore_, result);
+          MakePublic<PublicT, InternalT>(env, firestore, result);
 
       this->impl_->CompleteWithResult(this->handle_, Error::kErrorOk,
                                       /*error_msg=*/"", future_result);
@@ -161,7 +177,6 @@ class Promise {
         this->completion_->CompleteWith(Error::kErrorOk, /*error_message*/ "",
                                         &future_result);
       }
-      delete this;
     }
   };
 
@@ -170,13 +185,14 @@ class Promise {
    public:
     using CompleterBase<void>::CompleterBase;
 
-    void SucceedWithResult(jni::Env& env, const jni::Object& result) override {
+    void SucceedWithResult(jni::Env& env,
+                           const jni::Object& result,
+                           FirestoreInternal*) override {
       this->impl_->Complete(this->handle_, Error::kErrorOk, /*error_msg=*/"");
       if (this->completion_ != nullptr) {
         this->completion_->CompleteWith(Error::kErrorOk, /*error_message*/ "",
                                         nullptr);
       }
-      delete this;
     }
   };
 
@@ -192,6 +208,8 @@ class Promise {
     }
   }
 
+  FirestoreInternalWeakReference firestore_ref_;
+
   std::unique_ptr<Completer<PublicType, InternalType>> completer_;
 
   // Keep these values separate from the Completer in case completion happens
diff --git a/firestore/src/android/promise_factory_android.h b/firestore/src/android/promise_factory_android.h
index 3e9a98b957..72ddeab6fe 100644
--- a/firestore/src/android/promise_factory_android.h
+++ b/firestore/src/android/promise_factory_android.h
@@ -18,19 +18,31 @@ template <typename EnumT>
 class PromiseFactory {
  public:
   explicit PromiseFactory(FirestoreInternal* firestore)
-      : firestore_(firestore) {
-    firestore_->future_manager().AllocFutureApi(this, ApiCount());
+      : firestore_ref_(firestore) {
+    firestore_ref_.RunIfValid([this](FirestoreInternal* firestore) {
+      firestore->future_manager().AllocFutureApi(this, ApiCount());
+    });
   }
 
-  PromiseFactory(const PromiseFactory& rhs) : firestore_(rhs.firestore_) {
-    firestore_->future_manager().AllocFutureApi(this, ApiCount());
+  PromiseFactory(const PromiseFactory& rhs)
+      : firestore_ref_(rhs.firestore_ref_) {
+    firestore_ref_.RunIfValid([this](FirestoreInternal* firestore) {
+      firestore->future_manager().AllocFutureApi(this, ApiCount());
+    });
   }
 
-  PromiseFactory(PromiseFactory&& rhs) : firestore_(rhs.firestore_) {
-    firestore_->future_manager().MoveFutureApi(&rhs, this);
+  PromiseFactory(PromiseFactory&& rhs)
+      : firestore_ref_(Move(rhs.firestore_ref_)) {
+    firestore_ref_.RunIfValid([this, &rhs](FirestoreInternal* firestore) {
+      firestore->future_manager().MoveFutureApi(&rhs, this);
+    });
   }
 
-  ~PromiseFactory() { firestore_->future_manager().ReleaseFutureApi(this); }
+  ~PromiseFactory() {
+    firestore_ref_.RunIfValid([this](FirestoreInternal* firestore) {
+      firestore->future_manager().ReleaseFutureApi(this);
+    });
+  }
 
   PromiseFactory& operator=(const PromiseFactory&) = delete;
   PromiseFactory& operator=(PromiseFactory&&) = delete;
@@ -42,8 +54,14 @@ class PromiseFactory {
   Promise<PublicT, InternalT, EnumT> MakePromise(
       typename Promise<PublicT, InternalT, EnumT>::Completion* completion =
           nullptr) {
-    return Promise<PublicT, InternalT, EnumT>{future_api(), firestore_,
-                                              completion};
+    return firestore_ref_.Run([this, completion](FirestoreInternal* firestore) {
+      ReferenceCountedFutureImpl* future_api = nullptr;
+      if (firestore) {
+        future_api = firestore->future_manager().GetFutureApi(this);
+      }
+      return Promise<PublicT, InternalT, EnumT>{firestore_ref_, future_api,
+                                                completion};
+    });
   }
 
   template <typename PublicT, typename InternalT = InternalType<PublicT>>
@@ -61,15 +79,9 @@ class PromiseFactory {
   }
 
  private:
-  // Gets the reference-counted Future implementation of this instance, which
-  // can be used to create a Future.
-  ReferenceCountedFutureImpl* future_api() {
-    return firestore_->future_manager().GetFutureApi(this);
-  }
-
   constexpr int ApiCount() const { return static_cast<int>(EnumT::kCount); }
 
-  FirestoreInternal* firestore_ = nullptr;
+  FirestoreInternalWeakReference firestore_ref_;
 };
 
 }  // namespace firestore

From 5804c0b023d6de2f21bf467826ec7c773abcb771 Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Thu, 29 Jul 2021 12:22:23 -0400
Subject: [PATCH 2/3] Format code

---
 .../integration_test_internal/src/source_test.cc |  4 ++--
 firestore/src/android/promise_android.h          | 16 +++++++++++-----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/firestore/integration_test_internal/src/source_test.cc b/firestore/integration_test_internal/src/source_test.cc
index 17cf76e0b4..09cd52f576 100644
--- a/firestore/integration_test_internal/src/source_test.cc
+++ b/firestore/integration_test_internal/src/source_test.cc
@@ -384,7 +384,6 @@ TEST_F(SourceTest, GetNonExistingDocWhileOfflineWithDefaultGetOptions) {
   EXPECT_EQ(future.error(), Error::kErrorUnavailable);
 }
 
-
 TEST_F(SourceTest, GetDeletedDocWhileOfflineWithDefaultGetOptions) {
   GTEST_SKIP() << "b/112267729: We should raise a fromCache=true event with "
                << "a nonexistent snapshot, but because the default source goes "
@@ -516,7 +515,8 @@ TEST_F(SourceTest, GetNonExistingDocWhileOfflineWithSourceEqualToServer) {
   EXPECT_EQ(future.error(), Error::kErrorUnavailable);
 }
 
-TEST_F(SourceTest, GetNonExistingCollectionWhileOfflineWithSourceEqualToServer) {
+TEST_F(SourceTest,
+       GetNonExistingCollectionWhileOfflineWithSourceEqualToServer) {
   CollectionReference col_ref = Collection();
 
   DisableNetwork();
diff --git a/firestore/src/android/promise_android.h b/firestore/src/android/promise_android.h
index 382e6107b1..d4891996c1 100644
--- a/firestore/src/android/promise_android.h
+++ b/firestore/src/android/promise_android.h
@@ -79,8 +79,13 @@ class Promise {
  private:
   // The constructor is intentionally private.
   // Create instances with `PromiseFactory`.
-  Promise(const FirestoreInternalWeakReference& firestore_ref, ReferenceCountedFutureImpl* impl, Completion* completion) : firestore_ref_(firestore_ref), completer_(make_unique<Completer<PublicType, InternalType>>(firestore_ref, impl, completion)), impl_(impl) {
-  }
+  Promise(const FirestoreInternalWeakReference& firestore_ref,
+          ReferenceCountedFutureImpl* impl,
+          Completion* completion)
+      : firestore_ref_(firestore_ref),
+        completer_(make_unique<Completer<PublicType, InternalType>>(
+            firestore_ref, impl, completion)),
+        impl_(impl) {}
 
   template <typename PublicT>
   class CompleterBase {
@@ -111,9 +116,10 @@ class Promise {
       jni::Object result(raw_result);
 
       if (result_code == ::firebase::util::kFutureResultSuccess) {
-        firestore_ref_.RunIfValid([this, &env, &result](FirestoreInternal* firestore) {
-          SucceedWithResult(env, result, firestore);
-        });
+        firestore_ref_.RunIfValid(
+            [this, &env, &result](FirestoreInternal* firestore) {
+              SucceedWithResult(env, result, firestore);
+            });
 
         delete this;
 

From 3dcc7b6aa16b4dff7fae2dcf14b64bc9863f716d Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Thu, 29 Jul 2021 15:10:30 -0400
Subject: [PATCH 3/3] Apply some code review feedback to `firestore_android.h`.

* Change `std::function<void(FirestoreInternal*)>` to a const reference.
* Change the argument in the `std::function` from `FirestoreInternal*` to `FirestoreInternal&`.
* Improve docs: s/are allowed/are allowed on the same thread/.
---
 firestore/src/android/firestore_android.h       |  8 ++++----
 firestore/src/android/promise_android.h         | 12 ++++++------
 firestore/src/android/promise_factory_android.h | 16 ++++++++--------
 3 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/firestore/src/android/firestore_android.h b/firestore/src/android/firestore_android.h
index 518f4bcd9f..3adb92da4f 100644
--- a/firestore/src/android/firestore_android.h
+++ b/firestore/src/android/firestore_android.h
@@ -250,8 +250,8 @@ class FirestoreInternalWeakReference {
    * deleted) for the duration of the callback's execution.
    *
    * If this method is invoked concurrently then calls will be serialized in
-   * an undefined order. Recursive calls of this method are allowed and will not
-   * block.
+   * an undefined order. Recursive calls of this method are allowed on the same
+   * thread and will not block.
    */
   template <typename CallbackT>
   auto Run(CallbackT callback) -> decltype(callback(this->firestore_)) {
@@ -263,10 +263,10 @@ class FirestoreInternalWeakReference {
    * Same as `Run()` except that it only invokes the given callback if the
    * `FirestoreInternal` pointer is non-null.
    */
-  void RunIfValid(std::function<void(FirestoreInternal*)> callback) {
+  void RunIfValid(const std::function<void(FirestoreInternal&)>& callback) {
     std::lock_guard<std::recursive_mutex> lock(mutex_);
     if (firestore_) {
-      callback(firestore_);
+      callback(*firestore_);
     }
   }
 };
diff --git a/firestore/src/android/promise_android.h b/firestore/src/android/promise_android.h
index d4891996c1..e39f337f16 100644
--- a/firestore/src/android/promise_android.h
+++ b/firestore/src/android/promise_android.h
@@ -117,7 +117,7 @@ class Promise {
 
       if (result_code == ::firebase::util::kFutureResultSuccess) {
         firestore_ref_.RunIfValid(
-            [this, &env, &result](FirestoreInternal* firestore) {
+            [this, &env, &result](FirestoreInternal& firestore) {
               SucceedWithResult(env, result, firestore);
             });
 
@@ -141,7 +141,7 @@ class Promise {
           break;
       }
       firestore_ref_.RunIfValid(
-          [this, error_code, status_message](FirestoreInternal* firestore) {
+          [this, error_code, status_message](FirestoreInternal& firestore) {
             impl_->Complete(handle_, error_code, status_message);
           });
       if (completion_ != nullptr) {
@@ -152,7 +152,7 @@ class Promise {
 
     virtual void SucceedWithResult(jni::Env& env,
                                    const jni::Object& result,
-                                   FirestoreInternal* firestore) = 0;
+                                   FirestoreInternal& firestore) = 0;
 
    private:
     FirestoreInternalWeakReference firestore_ref_;
@@ -173,9 +173,9 @@ class Promise {
 
     void SucceedWithResult(jni::Env& env,
                            const jni::Object& result,
-                           FirestoreInternal* firestore) override {
+                           FirestoreInternal& firestore) override {
       auto future_result =
-          MakePublic<PublicT, InternalT>(env, firestore, result);
+          MakePublic<PublicT, InternalT>(env, &firestore, result);
 
       this->impl_->CompleteWithResult(this->handle_, Error::kErrorOk,
                                       /*error_msg=*/"", future_result);
@@ -193,7 +193,7 @@ class Promise {
 
     void SucceedWithResult(jni::Env& env,
                            const jni::Object& result,
-                           FirestoreInternal*) override {
+                           FirestoreInternal&) override {
       this->impl_->Complete(this->handle_, Error::kErrorOk, /*error_msg=*/"");
       if (this->completion_ != nullptr) {
         this->completion_->CompleteWith(Error::kErrorOk, /*error_message*/ "",
diff --git a/firestore/src/android/promise_factory_android.h b/firestore/src/android/promise_factory_android.h
index 72ddeab6fe..bfd798b75b 100644
--- a/firestore/src/android/promise_factory_android.h
+++ b/firestore/src/android/promise_factory_android.h
@@ -19,28 +19,28 @@ class PromiseFactory {
  public:
   explicit PromiseFactory(FirestoreInternal* firestore)
       : firestore_ref_(firestore) {
-    firestore_ref_.RunIfValid([this](FirestoreInternal* firestore) {
-      firestore->future_manager().AllocFutureApi(this, ApiCount());
+    firestore_ref_.RunIfValid([this](FirestoreInternal& firestore) {
+      firestore.future_manager().AllocFutureApi(this, ApiCount());
     });
   }
 
   PromiseFactory(const PromiseFactory& rhs)
       : firestore_ref_(rhs.firestore_ref_) {
-    firestore_ref_.RunIfValid([this](FirestoreInternal* firestore) {
-      firestore->future_manager().AllocFutureApi(this, ApiCount());
+    firestore_ref_.RunIfValid([this](FirestoreInternal& firestore) {
+      firestore.future_manager().AllocFutureApi(this, ApiCount());
     });
   }
 
   PromiseFactory(PromiseFactory&& rhs)
       : firestore_ref_(Move(rhs.firestore_ref_)) {
-    firestore_ref_.RunIfValid([this, &rhs](FirestoreInternal* firestore) {
-      firestore->future_manager().MoveFutureApi(&rhs, this);
+    firestore_ref_.RunIfValid([this, &rhs](FirestoreInternal& firestore) {
+      firestore.future_manager().MoveFutureApi(&rhs, this);
     });
   }
 
   ~PromiseFactory() {
-    firestore_ref_.RunIfValid([this](FirestoreInternal* firestore) {
-      firestore->future_manager().ReleaseFutureApi(this);
+    firestore_ref_.RunIfValid([this](FirestoreInternal& firestore) {
+      firestore.future_manager().ReleaseFutureApi(this);
     });
   }