Fix a Firestore transaction crash if getDocument was called after terminateWithCompletion (#8760)

Author: dconeybe

Firestore/CHANGELOG.md

@@ -1,3 +1,7 @@
+# Unreleased
+- [fixed] Fix a crash if `[FIRTransaction getDocument]` was called after
+  `[FIRFirestore terminateWithCompletion]` (#8760).
+
 # v8.6.0
 - [changed] Internal refactor to improve serialization performance.
 - [changed] `DocumentSnapshot` objects consider the document's key and data for

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

@@ -1316,6 +1316,47 @@ - (void)testClientCallsAfterTerminationFail {
                           NSException, @"The client has already been terminated.");
 }
 
+- (void)testTransactionGetAfterTerminationFail {
+  FIRDocumentReference *doc = [self documentRef];
+  FIRFirestore *firestore = doc.firestore;
+  XCTestExpectation *expectation1 = [self expectationWithDescription:@"TransactionBlockStart"];
+  XCTestExpectation *expectation2 = [self expectationWithDescription:@"FirestoreTerminated"];
+  XCTestExpectation *expectation3 = [self expectationWithDescription:@"TransactionBlockDone"];
+  __block NSError *capturedError = nil;
+
+  [firestore
+      runTransactionWithBlock:^id _Nullable(FIRTransaction *transaction, NSError **error) {
+        // Tell the test thread that it can now call `terminateWithCompletion`.
+        [expectation1 fulfill];
+        // Wait for `terminateWithCompletion` to complete.
+        [self awaitExpectation:expectation2];
+        // Call `[transaction getDocument]` to make sure that it correctly reports the error.
+        [transaction getDocument:doc error:error];
+        // Save the `NSError` into a variable on which the test thread will perform assertions.
+        capturedError = *error;
+        // Tell the test thread that it can now perform assertions on the captured `NSError` object.
+        [expectation3 fulfill];
+        return @"I should have failed since the transaction was terminated";
+      }
+                   completion:^(id, NSError *){
+                   }];
+
+  // Wait for the transaction callback to start.
+  [self awaitExpectation:expectation1];
+
+  // Terminate the Firestore instance.
+  [firestore terminateWithCompletion:[self completionForExpectation:expectation2]];
+
+  // Wait for the transaction callback to "publish" the `NSError` from `[transaction getDocument]`.
+  [self awaitExpectation:expectation3];
+
+  // Verify that `[transaction getDocument]` set its `NSError` argument to a non-nil error.
+  XCTAssertNotNil(capturedError);
+  XCTAssertEqual(capturedError.code, FIRFirestoreErrorCodeFailedPrecondition);
+  XCTAssertEqualObjects(capturedError.userInfo[NSLocalizedDescriptionKey],
+                        @"The client has already been terminated.");
+}
+
 - (void)testMaintainsPersistenceAfterRestarting {
   FIRDocumentReference *doc = [self documentRef];
   FIRFirestore *firestore = doc.firestore;

Firestore/core/src/core/transaction.cc

@@ -17,6 +17,7 @@
 #include "Firestore/core/src/core/transaction.h"
 
 #include <algorithm>
+#include <memory>
 #include <unordered_set>
 #include <utility>
 
@@ -47,8 +48,8 @@ namespace firebase {
 namespace firestore {
 namespace core {
 
-Transaction::Transaction(Datastore* datastore)
-    : datastore_{NOT_NULL(datastore)} {
+Transaction::Transaction(std::shared_ptr<Datastore> datastore)
+    : datastore_{datastore} {
 }
 
 Status Transaction::RecordVersion(const Document& doc) {
@@ -90,7 +91,14 @@ void Transaction::Lookup(const std::vector<DocumentKey>& keys,
     return;
   }
 
-  datastore_->LookupDocuments(
+  std::shared_ptr<Datastore> datastore = datastore_.lock();
+  if (!datastore) {
+    callback(Status(Error::kErrorFailedPrecondition,
+                    "The client has already been terminated."));
+    return;
+  }
+
+  datastore->LookupDocuments(
       keys,
       [this, callback](const StatusOr<std::vector<Document>>& maybe_documents) {
         if (!maybe_documents.ok()) {
@@ -206,7 +214,15 @@ void Transaction::Commit(util::StatusCallback&& callback) {
     mutations_.push_back(VerifyMutation(key, CreatePrecondition(key)));
   }
   committed_ = true;
-  datastore_->CommitMutations(mutations_, std::move(callback));
+
+  std::shared_ptr<Datastore> datastore = datastore_.lock();
+  if (!datastore) {
+    callback(Status(Error::kErrorFailedPrecondition,
+                    "The client has already been terminated."));
+    return;
+  }
+
+  datastore->CommitMutations(mutations_, std::move(callback));
 }
 
 void Transaction::MarkPermanentlyFailed() {

Firestore/core/src/core/transaction.h

@@ -54,7 +54,7 @@ class Transaction {
       std::function<void(const util::StatusOr<std::vector<model::Document>>&)>;
 
   Transaction() = default;
-  explicit Transaction(remote::Datastore* datastore);
+  explicit Transaction(std::shared_ptr<remote::Datastore> datastore);
 
   /**
    * Takes a set of keys and asynchronously attempts to fetch all the documents
@@ -129,7 +129,7 @@ class Transaction {
   absl::optional<model::SnapshotVersion> GetVersion(
       const model::DocumentKey& key) const;
 
-  remote::Datastore* datastore_ = nullptr;
+  std::weak_ptr<remote::Datastore> datastore_;
 
   std::vector<model::Mutation> mutations_;
   bool committed_ = false;

Firestore/core/src/remote/remote_store.cc

@@ -533,7 +533,7 @@ bool RemoteStore::CanUseNetwork() const {
 }
 
 std::shared_ptr<Transaction> RemoteStore::CreateTransaction() {
-  return std::make_shared<Transaction>(datastore_.get());
+  return std::make_shared<Transaction>(datastore_);
 }
 
 DocumentKeySet RemoteStore::GetRemoteKeysForTarget(TargetId target_id) const {