SERVER-64105 Support stmtIds in the transaction API

13 files changed, 174 insertions, 48 deletions

diff --git a/src/mongo/db/transaction_api.cpp b/src/mongo/db/transaction_api.cpp
index 2aaeec61b42..2d0f0b8ca79 100644
--- a/src/mongo/db/transaction_api.cpp
+++ b/src/mongo/db/transaction_api.cpp
@@ -29,10 +29,10 @@
 
 #define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kTransaction
 
-#include "mongo/platform/basic.h"
-
 #include "mongo/db/transaction_api.h"
 
+#include <fmt/format.h>
+
 #include "mongo/db/auth/authorization_session.h"
 #include "mongo/db/client.h"
 #include "mongo/db/commands/txn_cmds_gen.h"
@@ -41,9 +41,11 @@
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/db/operation_time_tracker.h"
+#include "mongo/db/ops/write_ops.h"
 #include "mongo/db/query/cursor_response.h"
 #include "mongo/db/repl/read_concern_args.h"
 #include "mongo/db/session_catalog.h"
+#include "mongo/db/transaction_validation.h"
 #include "mongo/db/write_concern_options.h"
 #include "mongo/executor/task_executor.h"
 #include "mongo/logv2/log.h"
@@ -274,7 +276,18 @@ SemiFuture<BSONObj> SEPTransactionClient::runCommand(StringData dbName, BSONObj
 
 SemiFuture<BatchedCommandResponse> SEPTransactionClient::runCRUDOp(
     const BatchedCommandRequest& cmd, std::vector<StmtId> stmtIds) const {
-    return runCommand(cmd.getNS().db(), cmd.toBSON())
+    invariant(!stmtIds.size() || (cmd.sizeWriteOps() == stmtIds.size()),
+              fmt::format("If stmtIds are specified, they must match the number of write ops. "
+                          "Found {} stmtId(s) and {} write op(s).",
+                          stmtIds.size(),
+                          cmd.sizeWriteOps()));
+
+    BSONObjBuilder cmdBob(cmd.toBSON());
+    if (stmtIds.size()) {
+        cmdBob.append(write_ops::WriteCommandRequestBase::kStmtIdsFieldName, stmtIds);
+    }
+
+    return runCommand(cmd.getNS().db(), cmdBob.obj())
         .thenRunOn(_executor)
         .then([](BSONObj reply) {
             uassertStatusOK(getStatusFromCommandResult(reply));
@@ -445,6 +458,25 @@ Transaction::ErrorHandlingStep Transaction::handleError(
 }
 
 void Transaction::prepareRequest(BSONObjBuilder* cmdBuilder) {
+    if (isInternalSessionForRetryableWrite(*_sessionInfo.getSessionId())) {
+        // Statement ids are meaningful in a transaction spawned on behalf of a retryable write, so
+        // every write in the transaction should explicitly specify an id. Either a positive number,
+        // which indicates retry history should be saved for the command, or kUninitializedStmtId
+        // (aka -1), which indicates retry history should not be saved. If statement ids are not
+        // explicitly sent, implicit ids may be inferred, which could lead to bugs if different
+        // commands have the same ids inferred.
+        uassert(
+            6410500,
+            str::stream()
+                << "In a retryable write transaction every retryable write command should have an "
+                   "explicit statement id, command: "
+                << redact(cmdBuilder->asTempObj()),
+            !isRetryableWriteCommand(
+                cmdBuilder->asTempObj().firstElement().fieldNameStringData()) ||
+                (cmdBuilder->hasField(write_ops::WriteCommandRequestBase::kStmtIdsFieldName) ||
+                 cmdBuilder->hasField(write_ops::WriteCommandRequestBase::kStmtIdFieldName)));
+    }
+
     stdx::lock_guard<Latch> lg(_mutex);
 
     _sessionInfo.serialize(cmdBuilder);
diff --git a/src/mongo/db/transaction_api.h b/src/mongo/db/transaction_api.h
index 03e9d85ad14..fd9c9d5c77a 100644
--- a/src/mongo/db/transaction_api.h
+++ b/src/mongo/db/transaction_api.h
@@ -92,6 +92,23 @@ public:
     /**
      * Helper method to run commands representable as a BatchedCommandRequest in the transaction
      * client's transaction.
+     *
+     * The given stmtIds are included in the sent command. If the API's transaction was spawned on
+     * behalf of a retryable write, the statement ids must be unique for each write in the
+     * transaction as the underlying servers will save history for each id the same as for a
+     * retryable write. A write can opt out of this by sending a -1 statement id, which is ignored.
+     *
+     * If a sent statement id had already been seen for this transaction, the write with that id
+     * won't apply a second time and instead returns its response from its original execution. That
+     * write's id will be in the batch response's "retriedStmtIds" array field.
+     *
+     * Users of this API for transactions spawned on behalf of retryable writes likely should
+     * include a stmtId for each write that should not execute twice and should check the
+     * "retriedStmtIds" in the returned BatchedCommandResponse to detect when a write had already
+     * applied, and thus the retryable write that spawned this transaction has already committed.
+     * Note that only one "pre" or "post" image can be stored per transaction, so only one
+     * findAndModify per transaction may have a non -1 statement id.
+     *
      */
     virtual SemiFuture<BatchedCommandResponse> runCRUDOp(const BatchedCommandRequest& cmd,
                                                          std::vector<StmtId> stmtIds) const = 0;
diff --git a/src/mongo/db/transaction_api_test.cpp b/src/mongo/db/transaction_api_test.cpp
index 3bb76ce7ff0..97a364ad46a 100644
--- a/src/mongo/db/transaction_api_test.cpp
+++ b/src/mongo/db/transaction_api_test.cpp
@@ -38,6 +38,7 @@
 #include "mongo/rpc/get_status_from_command_result.h"
 #include "mongo/s/is_mongos.h"
 #include "mongo/stdx/future.h"
+#include "mongo/unittest/death_test.h"
 #include "mongo/unittest/unittest.h"
 #include "mongo/util/executor_test_util.h"
 #include "mongo/util/fail_point.h"
@@ -1126,7 +1127,11 @@ TEST_F(TxnAPITest, ClientRetryableWrite_UsesRetryableInternalSession) {
                                  .runCommand("user"_sd,
                                              BSON("insert"
                                                   << "foo"
-                                                  << "documents" << BSON_ARRAY(BSON("x" << 1))))
+                                                  << "documents"
+                                                  << BSON_ARRAY(BSON("x" << 1))
+                                                  // Retryable transactions must include stmtIds for
+                                                  // retryable write commands.
+                                                  << "stmtIds" << BSON_ARRAY(1)))
                                  .get();
             ASSERT_EQ(insertRes["n"].Int(), 1);  // Verify the mocked response was returned.
             assertTxnMetadata(mockClient()->getLastSentRequest(),
@@ -1137,6 +1142,26 @@ TEST_F(TxnAPITest, ClientRetryableWrite_UsesRetryableInternalSession) {
                                     opCtx()->getLogicalSessionId(),
                                     opCtx()->getTxnNumber());
 
+            // Verify a non-retryable write command does not need to include stmtIds.
+            mockClient()->setNextCommandResponse(kOKCommandResponse);
+            auto findRes = txnClient
+                               .runCommand("user"_sd,
+                                           BSON("find"
+                                                << "foo"))
+                               .get();
+            ASSERT(findRes["ok"]);  // Verify the mocked response was returned.
+
+            // Verify the alternate format for stmtIds is allowed.
+            mockClient()->setNextCommandResponse(kOKInsertResponse);
+            insertRes =
+                txnClient
+                    .runCommand("user"_sd,
+                                BSON("insert"
+                                     << "foo"
+                                     << "documents" << BSON_ARRAY(BSON("x" << 1)) << "stmtId" << 1))
+                    .get();
+            ASSERT_EQ(insertRes["n"].Int(), 1);  // Verify the mocked response was returned.
+
             if (attempt == 0) {
                 firstAttemptLsid = getLsid(mockClient()->getLastSentRequest());
                 // Trigger transient error retry to verify the same session is used by the retry.
@@ -1163,5 +1188,31 @@ TEST_F(TxnAPITest, ClientRetryableWrite_UsesRetryableInternalSession) {
     ASSERT_EQ(lastRequest.firstElementFieldNameStringData(), "commitTransaction"_sd);
 }
 
+TEST_F(TxnAPITest, ClientRetryableWrite_RetryableWriteWithoutStmtIdFails) {
+    // This case is only currently supported on mongos.
+    // TODO SERVER-63747: Remove this once this restriction is lifted.
+    bool savedMongos = isMongos();
+    ON_BLOCK_EXIT([&] { setMongos(savedMongos); });
+    setMongos(true);
+
+    opCtx()->setLogicalSessionId(makeLogicalSessionIdForTest());
+    opCtx()->setTxnNumber(5);
+    resetTxnWithRetries();
+
+    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();
+
+            return SemiFuture<void>::makeReady();
+        });
+    ASSERT_EQ(swResult.getStatus(), ErrorCodes::duplicateCodeForTest(6410500));
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/db/transaction_validation.cpp b/src/mongo/db/transaction_validation.cpp
index 648639d8599..5fa6501e687 100644
--- a/src/mongo/db/transaction_validation.cpp
+++ b/src/mongo/db/transaction_validation.cpp
@@ -69,12 +69,12 @@ const StringMap<int> transactionCommands = {{"commitTransaction", 1},
                                             {"abortTransaction", 1},
                                             {"prepareTransaction", 1}};
 
+}  // namespace
+
 bool isRetryableWriteCommand(StringData cmdName) {
     return retryableWriteCommands.find(cmdName) != retryableWriteCommands.cend();
 }
 
-}  // namespace
-
 bool isTransactionCommand(StringData cmdName) {
     return transactionCommands.find(cmdName) != transactionCommands.cend();
 }
diff --git a/src/mongo/db/transaction_validation.h b/src/mongo/db/transaction_validation.h
index f2447f6ab3a..b3372dfe5b4 100644
--- a/src/mongo/db/transaction_validation.h
+++ b/src/mongo/db/transaction_validation.h
@@ -36,6 +36,11 @@
 namespace mongo {
 
 /**
+ * Returns true if the given command name can run as a retryable write.
+ */
+bool isRetryableWriteCommand(StringData cmdName);
+
+/**
  * Returns true if the given cmd name is a transaction control command.  These are also the only
  * commands allowed to specify write concern in a transaction.
  */
diff --git a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
index 66a7b120b47..2ff2f2ffccb 100644
--- a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
+++ b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
@@ -549,6 +549,7 @@ private:
                 // recursively sent through the service entry point.
                 parsedRequest.setLegacyRuntimeConstants(boost::none);
                 if (isRetryableWrite) {
+                    parsedRequest.setStmtId(0);
                     handleWouldChangeOwningShardErrorRetryableWrite(
                         opCtx, shardId, nss, parsedRequest, result);
                 } else {
diff --git a/src/mongo/s/commands/cluster_write_cmd.cpp b/src/mongo/s/commands/cluster_write_cmd.cpp
index 5dbb436c0d2..52d99053660 100644
--- a/src/mongo/s/commands/cluster_write_cmd.cpp
+++ b/src/mongo/s/commands/cluster_write_cmd.cpp
@@ -172,18 +172,19 @@ void handleWouldChangeOwningShardErrorRetryableWrite(OperationContext* opCtx,
 
     // Shared state for the transaction API use below.
     struct SharedBlock {
-        SharedBlock(NamespaceString nss_) : nss(nss_) {}
+        SharedBlock(BSONObj cmdObj_, NamespaceString nss_) : cmdObj(cmdObj_), nss(nss_) {}
 
+        BSONObj cmdObj;
         NamespaceString nss;
         BSONObj response;
     };
-    auto sharedBlock = std::make_shared<SharedBlock>(request->getNS());
+    BSONObjBuilder cmdWithStmtId(request->toBSON());
+    cmdWithStmtId.append(write_ops::WriteCommandRequestBase::kStmtIdFieldName, 0);
+    auto sharedBlock = std::make_shared<SharedBlock>(cmdWithStmtId.obj(), request->getNS());
 
     auto swCommitResult = txn.runSyncNoThrow(
-        opCtx,
-        [cmdObj = request->toBSON(), sharedBlock](const txn_api::TransactionClient& txnClient,
-                                                  ExecutorPtr txnExec) {
-            return txnClient.runCommand(sharedBlock->nss.db(), cmdObj)
+        opCtx, [sharedBlock](const txn_api::TransactionClient& txnClient, ExecutorPtr txnExec) {
+            return txnClient.runCommand(sharedBlock->nss.db(), sharedBlock->cmdObj)
                 .thenRunOn(txnExec)
                 .then([sharedBlock](auto res) {
                     uassertStatusOK(getStatusFromWriteCommandReply(res));
diff --git a/src/mongo/s/commands/document_shard_key_update_test.cpp b/src/mongo/s/commands/document_shard_key_update_test.cpp
index 4c297ce7bc8..dcaf62de0db 100644
--- a/src/mongo/s/commands/document_shard_key_update_test.cpp
+++ b/src/mongo/s/commands/document_shard_key_update_test.cpp
@@ -53,7 +53,7 @@ TEST_F(DocumentShardKeyUpdateTest, constructShardKeyDeleteCmdObj) {
     NamespaceString nss("test.foo");
     BSONObj updatePreImage = BSON("x" << 4 << "y" << 3 << "_id" << 20);
 
-    auto deleteCmdObj = constructShardKeyDeleteCmdObj(nss, updatePreImage, boost::none);
+    auto deleteCmdObj = constructShardKeyDeleteCmdObj(nss, updatePreImage);
 
     auto deletesObj = deleteCmdObj["deletes"].Array();
     ASSERT_EQ(deletesObj.size(), 1U);
@@ -69,7 +69,7 @@ TEST_F(DocumentShardKeyUpdateTest, constructShardKeyInsertCmdObj) {
     NamespaceString nss("test.foo");
     BSONObj updatePostImage = BSON("x" << 4 << "y" << 3 << "_id" << 20);
 
-    auto insertCmdObj = constructShardKeyInsertCmdObj(nss, updatePostImage, boost::none);
+    auto insertCmdObj = constructShardKeyInsertCmdObj(nss, updatePostImage);
 
     auto insertsObj = insertCmdObj["documents"].Array();
     ASSERT_EQ(insertsObj.size(), 1U);
diff --git a/src/mongo/s/commands/document_shard_key_update_util.cpp b/src/mongo/s/commands/document_shard_key_update_util.cpp
index 4524dc78689..2877ee9ec48 100644
--- a/src/mongo/s/commands/document_shard_key_update_util.cpp
+++ b/src/mongo/s/commands/document_shard_key_update_util.cpp
@@ -102,8 +102,7 @@ bool executeOperationsAsPartOfShardKeyUpdate(OperationContext* opCtx,
  * original document _id retrieved from 'updatePreImage'.
  */
 write_ops::DeleteCommandRequest createShardKeyDeleteOp(const NamespaceString& nss,
-                                                       const BSONObj& updatePreImage,
-                                                       boost::optional<StmtId> stmtId) {
+                                                       const BSONObj& updatePreImage) {
     write_ops::DeleteCommandRequest deleteOp(nss);
     deleteOp.setDeletes({[&] {
         write_ops::DeleteOpEntry entry;
@@ -111,9 +110,6 @@ write_ops::DeleteCommandRequest createShardKeyDeleteOp(const NamespaceString& ns
         entry.setMulti(false);
         return entry;
     }()});
-    if (stmtId) {
-        deleteOp.getWriteCommandRequestBase().setStmtId(*stmtId);
-    }
 
     return deleteOp;
 }
@@ -122,13 +118,9 @@ write_ops::DeleteCommandRequest createShardKeyDeleteOp(const NamespaceString& ns
  * Creates the insert op that will be used to insert the new document with the post-update image.
  */
 write_ops::InsertCommandRequest createShardKeyInsertOp(const NamespaceString& nss,
-                                                       const BSONObj& updatePostImage,
-                                                       boost::optional<StmtId> stmtId) {
+                                                       const BSONObj& updatePostImage) {
     write_ops::InsertCommandRequest insertOp(nss);
     insertOp.setDocuments({updatePostImage});
-    if (stmtId) {
-        insertOp.getWriteCommandRequestBase().setStmtId(*stmtId);
-    }
     return insertOp;
 }
 
@@ -142,8 +134,8 @@ bool updateShardKeyForDocumentLegacy(OperationContext* opCtx,
     auto updatePreImage = documentKeyChangeInfo.getPreImage().getOwned();
     auto updatePostImage = documentKeyChangeInfo.getPostImage().getOwned();
 
-    auto deleteCmdObj = constructShardKeyDeleteCmdObj(nss, updatePreImage, boost::none);
-    auto insertCmdObj = constructShardKeyInsertCmdObj(nss, updatePostImage, boost::none);
+    auto deleteCmdObj = constructShardKeyDeleteCmdObj(nss, updatePreImage);
+    auto insertCmdObj = constructShardKeyInsertCmdObj(nss, updatePostImage);
 
     return executeOperationsAsPartOfShardKeyUpdate(
         opCtx, deleteCmdObj, insertCmdObj, nss.db(), documentKeyChangeInfo.getShouldUpsert());
@@ -166,17 +158,13 @@ BSONObj commitShardKeyUpdateTransaction(OperationContext* opCtx) {
     return txnRouter.commitTransaction(opCtx, boost::none);
 }
 
-BSONObj constructShardKeyDeleteCmdObj(const NamespaceString& nss,
-                                      const BSONObj& updatePreImage,
-                                      boost::optional<StmtId> stmtId) {
-    auto deleteOp = createShardKeyDeleteOp(nss, updatePreImage, stmtId);
+BSONObj constructShardKeyDeleteCmdObj(const NamespaceString& nss, const BSONObj& updatePreImage) {
+    auto deleteOp = createShardKeyDeleteOp(nss, updatePreImage);
     return deleteOp.toBSON({});
 }
 
-BSONObj constructShardKeyInsertCmdObj(const NamespaceString& nss,
-                                      const BSONObj& updatePostImage,
-                                      boost::optional<StmtId> stmtId) {
-    auto insertOp = createShardKeyInsertOp(nss, updatePostImage, stmtId);
+BSONObj constructShardKeyInsertCmdObj(const NamespaceString& nss, const BSONObj& updatePostImage) {
+    auto insertOp = createShardKeyInsertOp(nss, updatePostImage);
     return insertOp.toBSON({});
 }
 
@@ -184,15 +172,14 @@ SemiFuture<bool> updateShardKeyForDocument(const txn_api::TransactionClient& txn
                                            ExecutorPtr txnExec,
                                            const NamespaceString& nss,
                                            const WouldChangeOwningShardInfo& changeInfo) {
-    // Use stmtId=1 for this delete (and 2 for the subsequent insert) because the original
-    // update/findAndModify that threw the WouldChangeOwningShard error used stmtId=0 to store the
-    // WouldChangeOwningShard sentinel noop entry.
     auto deleteCmdObj = documentShardKeyUpdateUtil::constructShardKeyDeleteCmdObj(
-        nss, changeInfo.getPreImage().getOwned(), {1});
+        nss, changeInfo.getPreImage().getOwned());
     auto deleteOpMsg = OpMsgRequest::fromDBAndBody(nss.db(), std::move(deleteCmdObj));
     auto deleteRequest = BatchedCommandRequest::parseDelete(std::move(deleteOpMsg));
 
-    return txnClient.runCRUDOp(deleteRequest, {})
+    // Retry history for this delete isn't necessary, but it can be part of a retryable transaction,
+    // so send it with the uninitialized sentinel statement id to opt out of storing history.
+    return txnClient.runCRUDOp(deleteRequest, {kUninitializedStmtId})
         .thenRunOn(txnExec)
         .then([&txnClient, &nss, &changeInfo](
                   auto deleteResponse) -> SemiFuture<BatchedCommandResponse> {
@@ -219,11 +206,13 @@ SemiFuture<bool> updateShardKeyForDocument(const txn_api::TransactionClient& txn
             }
 
             auto insertCmdObj = documentShardKeyUpdateUtil::constructShardKeyInsertCmdObj(
-                nss, changeInfo.getPostImage().getOwned(), {2});
+                nss, changeInfo.getPostImage().getOwned());
             auto insertOpMsg = OpMsgRequest::fromDBAndBody(nss.db(), std::move(insertCmdObj));
             auto insertRequest = BatchedCommandRequest::parseInsert(std::move(insertOpMsg));
 
-            return txnClient.runCRUDOp(insertRequest, {});
+            // Same as for the insert, retry history isn't necessary so opt out with a sentinel
+            // stmtId.
+            return txnClient.runCRUDOp(insertRequest, {kUninitializedStmtId});
         })
         .thenRunOn(txnExec)
         .then([&nss](auto insertResponse) {
diff --git a/src/mongo/s/commands/document_shard_key_update_util.h b/src/mongo/s/commands/document_shard_key_update_util.h
index 31825eece60..c386b39c0a0 100644
--- a/src/mongo/s/commands/document_shard_key_update_util.h
+++ b/src/mongo/s/commands/document_shard_key_update_util.h
@@ -112,9 +112,7 @@ BSONObj commitShardKeyUpdateTransaction(OperationContext* opCtx);
  * This method should not be called outside of this class. It is only temporarily exposed for
  * intermediary test coverage.
  */
-BSONObj constructShardKeyDeleteCmdObj(const NamespaceString& nss,
-                                      const BSONObj& updatePreImage,
-                                      boost::optional<StmtId> stmtId);
+BSONObj constructShardKeyDeleteCmdObj(const NamespaceString& nss, const BSONObj& updatePreImage);
 
 /*
  * Creates the BSONObj that will be used to insert the new document with the post-update image.
@@ -123,8 +121,6 @@ BSONObj constructShardKeyDeleteCmdObj(const NamespaceString& nss,
  * This method should not be called outside of this class. It is only temporarily exposed for
  * intermediary test coverage.
  */
-BSONObj constructShardKeyInsertCmdObj(const NamespaceString& nss,
-                                      const BSONObj& updatePostImage,
-                                      boost::optional<StmtId> stmtId);
+BSONObj constructShardKeyInsertCmdObj(const NamespaceString& nss, const BSONObj& updatePostImage);
 }  // namespace documentShardKeyUpdateUtil
 }  // namespace mongo
diff --git a/src/mongo/s/write_ops/batched_command_response.cpp b/src/mongo/s/write_ops/batched_command_response.cpp
index fe010d2a235..256f2e0767d 100644
--- a/src/mongo/s/write_ops/batched_command_response.cpp
+++ b/src/mongo/s/write_ops/batched_command_response.cpp
@@ -54,6 +54,7 @@ const BSONField<std::vector<WriteErrorDetail*>> BatchedCommandResponse::writeErr
 const BSONField<WriteConcernErrorDetail*> BatchedCommandResponse::writeConcernError(
     "writeConcernError");
 const BSONField<std::vector<std::string>> BatchedCommandResponse::errorLabels("errorLabels");
+const BSONField<std::vector<StmtId>> BatchedCommandResponse::retriedStmtIds("retriedStmtIds");
 
 BatchedCommandResponse::BatchedCommandResponse() {
     clear();
@@ -139,6 +140,10 @@ BSONObj BatchedCommandResponse::toBSON() const {
         builder.append(writeConcernError(), _wcErrDetails->toBSON());
     }
 
+    if (areRetriedStmtIdsSet()) {
+        builder.append(retriedStmtIds(), _retriedStmtIds);
+    }
+
     return builder.obj();
 }
 
@@ -229,6 +234,12 @@ bool BatchedCommandResponse::parseBSON(const BSONObj& source, string* errMsg) {
         return false;
     _errorLabels = std::move(tempErrorLabels);
 
+    std::vector<StmtId> tempRetriedStmtIds;
+    fieldState = FieldParser::extract(source, retriedStmtIds, &tempRetriedStmtIds, errMsg);
+    if (fieldState == FieldParser::FIELD_INVALID)
+        return false;
+    _retriedStmtIds = std::move(tempRetriedStmtIds);
+
     return true;
 }
 
@@ -452,4 +463,12 @@ const std::vector<std::string>& BatchedCommandResponse::getErrorLabels() const {
     return _errorLabels;
 }
 
+bool BatchedCommandResponse::areRetriedStmtIdsSet() const {
+    return !_retriedStmtIds.empty();
+}
+
+const std::vector<StmtId>& BatchedCommandResponse::getRetriedStmtIds() const {
+    return _retriedStmtIds;
+}
+
 }  // namespace mongo
diff --git a/src/mongo/s/write_ops/batched_command_response.h b/src/mongo/s/write_ops/batched_command_response.h
index 0381cfd40c2..0883f69db38 100644
--- a/src/mongo/s/write_ops/batched_command_response.h
+++ b/src/mongo/s/write_ops/batched_command_response.h
@@ -31,6 +31,7 @@
 
 #include "mongo/base/string_data.h"
 #include "mongo/db/jsobj.h"
+#include "mongo/db/logical_session_id.h"
 #include "mongo/db/repl/optime.h"
 #include "mongo/rpc/write_concern_error_detail.h"
 #include "mongo/s/write_ops/batched_upsert_detail.h"
@@ -54,6 +55,7 @@ public:
     static const BSONField<std::vector<WriteErrorDetail*>> writeErrors;
     static const BSONField<WriteConcernErrorDetail*> writeConcernError;
     static const BSONField<std::vector<std::string>> errorLabels;
+    static const BSONField<std::vector<StmtId>> retriedStmtIds;
 
     BatchedCommandResponse();
     ~BatchedCommandResponse();
@@ -126,6 +128,9 @@ public:
     bool isErrorLabelsSet() const;
     const std::vector<std::string>& getErrorLabels() const;
 
+    bool areRetriedStmtIdsSet() const;
+    const std::vector<StmtId>& getRetriedStmtIds() const;
+
 private:
     // Convention: (M)andatory, (O)ptional
 
@@ -167,6 +172,9 @@ private:
 
     // (O)  array containing the error labels in string format.
     std::vector<std::string> _errorLabels;
+
+    // (O)  Array containing the retried statement ids from the response.
+    std::vector<StmtId> _retriedStmtIds;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/s/write_ops/batched_command_response_test.cpp b/src/mongo/s/write_ops/batched_command_response_test.cpp
index 4d742968a1b..509af17d0db 100644
--- a/src/mongo/s/write_ops/batched_command_response_test.cpp
+++ b/src/mongo/s/write_ops/batched_command_response_test.cpp
@@ -55,17 +55,24 @@ TEST(BatchedCommandResponse, Basic) {
                     << "norepl"
                     << "errInfo" << BSON("a" << 1)));
 
+    auto retriedStmtIds = BSON_ARRAY(1 << 3);
+
     BSONObj origResponseObj =
         BSON(BatchedCommandResponse::n(0)
              << "opTime" << mongo::Timestamp(1ULL) << BatchedCommandResponse::writeErrors()
              << writeErrorsArray << BatchedCommandResponse::writeConcernError() << writeConcernError
-             << "ok" << 1.0);
+             << BatchedCommandResponse::retriedStmtIds() << retriedStmtIds << "ok" << 1.0);
 
     std::string errMsg;
     BatchedCommandResponse response;
     bool ok = response.parseBSON(origResponseObj, &errMsg);
     ASSERT_TRUE(ok);
 
+    ASSERT(response.areRetriedStmtIdsSet());
+    ASSERT_EQ(response.getRetriedStmtIds().size(), 2);
+    ASSERT_EQ(response.getRetriedStmtIds()[0], 1);
+    ASSERT_EQ(response.getRetriedStmtIds()[1], 3);
+
     BSONObj genResponseObj = BSONObjBuilder(response.toBSON()).append("ok", 1.0).obj();
 
     ASSERT_EQUALS(0, genResponseObj.woCompare(origResponseObj))