From 425aa6786b0dcda8bc68387d20437a83edaa6a03 Mon Sep 17 00:00:00 2001
From: Jason Zhang <jason.zhang@mongodb.com>
Date: Thu, 17 Mar 2022 00:39:19 +0000
Subject: SERVER-64328 Allow network errors to be retried in transaction api

---
 src/mongo/db/transaction_api.cpp      | 10 ++++--
 src/mongo/db/transaction_api_test.cpp | 58 +++++++++++++++++++++++++++++------
 2 files changed, 57 insertions(+), 11 deletions(-)

diff --git a/src/mongo/db/transaction_api.cpp b/src/mongo/db/transaction_api.cpp
index 5cc89639578..411c919943f 100644
--- a/src/mongo/db/transaction_api.cpp
+++ b/src/mongo/db/transaction_api.cpp
@@ -431,8 +431,14 @@ Transaction::ErrorHandlingStep Transaction::handleError(
 
     const auto& clientStatus = swResult.getStatus();
     if (!clientStatus.isOK()) {
-        // A network error before commit is a transient transaction error.
-        if (!hasStartedCommit && ErrorCodes::isNetworkError(clientStatus)) {
+        if (ErrorCodes::isNetworkError(clientStatus)) {
+            // A network error before commit is a transient transaction error, so we can retry the
+            // entire transaction. If there is a network error after a commit is sent, we can retry
+            // the commit command to either recommit if the operation failed or get the result of
+            // the successful commit.
+            if (hasStartedCommit) {
+                return ErrorHandlingStep::kRetryCommit;
+            }
             return ErrorHandlingStep::kRetryTransaction;
         }
         return ErrorHandlingStep::kAbortAndDoNotRetry;
diff --git a/src/mongo/db/transaction_api_test.cpp b/src/mongo/db/transaction_api_test.cpp
index 77f907e672a..05083f64e03 100644
--- a/src/mongo/db/transaction_api_test.cpp
+++ b/src/mongo/db/transaction_api_test.cpp
@@ -140,12 +140,16 @@ public:
         _lastSentRequest = cmdBob.obj();
 
         auto nextResponse = _nextResponse;
-        if (!_secondResponse.isEmpty()) {
-            _nextResponse = _secondResponse;
-            _secondResponse = BSONObj();
+        if (_secondResponse) {
+            _nextResponse = *_secondResponse;
+            _secondResponse.reset();
         }
-        _hooks->runReplyHook(nextResponse);
-        return SemiFuture<BSONObj>::makeReady(nextResponse);
+        if (!nextResponse.isOK()) {
+            return SemiFuture<BSONObj>::makeReady(nextResponse);
+        }
+        auto nextResponseRes = nextResponse.getValue();
+        _hooks->runReplyHook(nextResponseRes);
+        return SemiFuture<BSONObj>::makeReady(nextResponseRes);
     }
 
     virtual SemiFuture<BatchedCommandResponse> runCRUDOp(
@@ -166,18 +170,18 @@ public:
         return _lastSentRequest;
     }
 
-    void setNextCommandResponse(BSONObj res) {
+    void setNextCommandResponse(StatusWith<BSONObj> res) {
         _nextResponse = res;
     }
 
-    void setSecondCommandResponse(BSONObj res) {
+    void setSecondCommandResponse(StatusWith<BSONObj> res) {
         _secondResponse = res;
     }
 
 private:
     std::unique_ptr<TxnMetadataHooks> _hooks;
-    mutable BSONObj _nextResponse;
-    mutable BSONObj _secondResponse;
+    mutable StatusWith<BSONObj> _nextResponse{BSONObj()};
+    mutable boost::optional<StatusWith<BSONObj>> _secondResponse;
     mutable BSONObj _lastSentRequest;
 };
 
@@ -1362,5 +1366,41 @@ TEST_F(TxnAPITest, ClientTransaction_DoesNotRetryOnTransientErrors) {
     ASSERT_EQ(lastRequest.firstElementFieldNameStringData(), "insert"_sd);
 }
 
+TEST_F(TxnAPITest, OwnSession_HandleErrorRetryCommitOnNetworkError) {
+    auto swResult = txnWithRetries().runSyncNoThrow(
+        opCtx(), [&](const txn_api::TransactionClient& txnClient, ExecutorPtr txnExec) {
+            mockClient()->setNextCommandResponse(kOKInsertResponse);
+            auto insertRes = txnClient
+                                 .runCommand("user"_sd,
+                                             BSON("insert"
+                                                  << "foo"
+                                                  << "documents" << BSON_ARRAY(BSON("x" << 1))))
+                                 .get();
+
+            ASSERT_EQ(insertRes["n"].Int(), 1);  // Verify the mocked response was returned.
+            assertTxnMetadata(
+                mockClient()->getLastSentRequest(), 0 /* txnNumber */, true /* startTransaction */);
+            assertSessionIdMetadata(mockClient()->getLastSentRequest(), LsidAssertion::kStandalone);
+
+            // The commit response.
+            mockClient()->setNextCommandResponse(
+                Status(ErrorCodes::HostUnreachable, "Host Unreachable"));
+            mockClient()->setSecondCommandResponse(kOKCommandResponse);
+
+            return SemiFuture<void>::makeReady();
+        });
+    ASSERT(swResult.getStatus().isOK());
+    ASSERT(swResult.getValue().getEffectiveStatus().isOK());
+
+    auto lastRequest = mockClient()->getLastSentRequest();
+    assertTxnMetadata(lastRequest,
+                      0 /* txnNumber */,
+                      boost::none /* startTransaction */,
+                      boost::none /* readConcern */,
+                      WriteConcernOptions().toBSON() /* writeConcern */);
+    assertSessionIdMetadata(mockClient()->getLastSentRequest(), LsidAssertion::kStandalone);
+    ASSERT_EQ(lastRequest.firstElementFieldNameStringData(), "commitTransaction"_sd);
+}
+
 }  // namespace
 }  // namespace mongo
-- 
cgit v1.2.1