SERVER-65048 Support retryable writes in internal_transactions_test_command

(cherry picked from commit e2ae4c71e1c4e32e3c2d191c9bc07583926dcd5c)

3 files changed, 125 insertions, 16 deletions

diff --git a/jstests/sharding/internal_txns/transaction_api_test_command_basic.js b/jstests/sharding/internal_txns/transaction_api_test_command_basic.js
index 5b86cb2e961..16ac90ea5ba 100644
--- a/jstests/sharding/internal_txns/transaction_api_test_command_basic.js
+++ b/jstests/sharding/internal_txns/transaction_api_test_command_basic.js
@@ -23,27 +23,81 @@ let primary = rst.getPrimary();
 let db = primary.getDB(kDbName);
 let rstColl = db.getCollection(kCollName);
 
-function runTxn(connection, commandInfos, collection) {
-    const res = assert.commandWorked(
-        connection.adminCommand({testInternalTransactions: 1, commandInfos: commandInfos}));
-    jsTest.log(res);
+let stColl = st.s.getCollection(kNs);
 
+function verifyCompletedInsertCommandResult(commandInfos, response, collection) {
     let i = 0;
     commandInfos.forEach(commandInfo => {
-        assert.eq(1, res.responses[i].ok);
-        assert.eq(commandInfo.command.documents.length, res.responses[i].n);
+        // Verifies command response.
+        assert.eq(1, response.responses[i].ok);
+        assert.eq(commandInfo.command.documents.length, response.responses[i].n);
         commandInfo.command.documents.forEach(document => {
+            // Verifies documents were successfully inserted.
             assert.eq(document, collection.findOne(document));
         });
         ++i;
     });
 }
 
+function runTxn(connection, commandInfos, collection) {
+    const res = assert.commandWorked(
+        connection.adminCommand({testInternalTransactions: 1, commandInfos: commandInfos}));
+    verifyCompletedInsertCommandResult(commandInfos, res, collection);
+}
+
+function runRetryableWrite(connection, commandInfos, lsid) {
+    const txnNumber = NumberLong(0);
+    return assert.commandWorked(connection.adminCommand({
+        testInternalTransactions: 1,
+        commandInfos: commandInfos,
+        lsid: {id: lsid},
+        txnNumber: txnNumber
+    }));
+}
+
+function runRetryableInsert(connection, commandInfos, lsid, collection) {
+    const originalRes = runRetryableWrite(connection, commandInfos, lsid);
+    verifyCompletedInsertCommandResult(commandInfos, originalRes, collection);
+
+    // This retryable write is expected to return a response of length 1 as the command list should
+    // stop executing after the first insert.
+    const retryRes = runRetryableWrite(connection, commandInfos, lsid);
+    assert.eq(1, retryRes.responses.length);
+    assert.eq(1, retryRes.ok);
+    assert.eq(0, retryRes.responses[0].retriedStmtIds[0]);
+}
+
+function runRetryableFindAndModify(connection, commandInfos, lsid, collection) {
+    const originalRes = runRetryableWrite(connection, commandInfos, lsid);
+    assert.eq(commandInfos.length, originalRes.responses.length);
+
+    // Verify command response.
+    assert.eq(1, originalRes.responses[0].lastErrorObject.n);
+    assert.eq(commandInfos[0].command.query._id, originalRes.responses[0].lastErrorObject.upserted);
+
+    // Verify document from commandInfos[0] was upserted into database.
+    let document = commandInfos[0].command.query;
+    assert.eq(document, collection.findOne(document));
+
+    // Verify document from commandInfos[1] was inserted into database.
+    assert.eq(1, originalRes.responses[1].n);
+    document = commandInfos[1].command.documents[0];
+    assert.eq(document, collection.findOne(document));
+
+    // Retry transaction. This retryable write is expected to return a response of length 1 as the
+    // command list should stop executing after the first findAndModify command.
+    const retryRes = runRetryableWrite(connection, commandInfos, lsid);
+    assert.eq(1, retryRes.responses.length);
+    assert.eq(1, retryRes.ok);
+    assert.eq(0, retryRes.responses[0].retriedStmtId);
+}
+
 // Insert initial data.
-assert.commandWorked(st.s.getCollection(kNs).insert([{_id: 0}]));
+assert.commandWorked(stColl.insert([{_id: 0}]));
 assert.commandWorked(rstColl.insert([{_id: 0}]));
 
-const commandInfos0 = [{
+// Set of commandInfos that will be used in tests below.
+const commandInfosSingleInsert = [{
     dbName: kDbName,
     command: {
         insert: kCollName,
@@ -51,7 +105,7 @@ const commandInfos0 = [{
     }
 }];
 
-const commandInfos1 = [
+const commandInfosBatchInsert = [
     {
         dbName: kDbName,
         command: {
@@ -68,17 +122,65 @@ const commandInfos1 = [
     }
 ];
 
+const commandInfosRetryableBatchInsert = [
+    {
+        dbName: kDbName,
+        command: {
+            insert: kCollName,
+            documents: [{_id: 5}],
+            stmtId: NumberInt(0),
+        }
+    },
+    {
+        dbName: kDbName,
+        command: {
+            insert: kCollName,
+            documents: [{_id: 6}, {_id: 7}],
+            stmtId: NumberInt(1),
+        }
+    }
+];
+
+function commandInfosRetryableFindAndModify(collection) {
+    return [
+        {
+            dbName: kDbName,
+            command: {
+                findandmodify: collection.getName(),
+                query: {_id: 8},
+                update: {},
+                upsert: true,
+                stmtId: NumberInt(0),
+            }
+        },
+        {
+            dbName: kDbName,
+            command: {
+                insert: collection.getName(),
+                documents: [{_id: 9}],
+                stmtId: NumberInt(1),
+            }
+        }
+    ];
+}
+
 jsTest.log(
     "Insert documents without a session into a sharded cluster, using internal transactions test command.");
-runTxn(st.s, commandInfos0, st.s.getCollection(kNs));
-runTxn(st.s, commandInfos1, st.s.getCollection(kNs));
+runTxn(st.s, commandInfosSingleInsert, stColl);
+runTxn(st.s, commandInfosBatchInsert, stColl);
 
 jsTest.log(
     "Insert documents without a session into a replica set, using internal transactions test command.");
-runTxn(primary, commandInfos0, rstColl);
-runTxn(primary, commandInfos1, rstColl);
+runTxn(primary, commandInfosSingleInsert, rstColl);
+runTxn(primary, commandInfosBatchInsert, rstColl);
+
+jsTest.log("Testing retryable write targeting a mongos.");
+runRetryableInsert(st.s, commandInfosRetryableBatchInsert, UUID(), stColl);
+runRetryableFindAndModify(st.s, commandInfosRetryableFindAndModify(stColl), UUID(), stColl);
 
-// TODO SERVER-65048: Add testing for retryable writes and txns run in sessions.
+jsTest.log("Testing retryable write targeting a mongod.");
+runRetryableInsert(primary, commandInfosRetryableBatchInsert, UUID(), rstColl);
+runRetryableFindAndModify(primary, commandInfosRetryableFindAndModify(rstColl), UUID(), rstColl);
 
 rst.stopSet();
 st.stop();
diff --git a/src/mongo/db/transaction_validation.cpp b/src/mongo/db/transaction_validation.cpp
index cc136f0fa6f..6571711cc76 100644
--- a/src/mongo/db/transaction_validation.cpp
+++ b/src/mongo/db/transaction_validation.cpp
@@ -53,6 +53,7 @@ const StringMap<int> retryableWriteCommands = {{"clusterDelete", 1},
                                                {"findandmodify", 1},
                                                {"findAndModify", 1},
                                                {"insert", 1},
+                                               {"testInternalTransactions", 1},
                                                {"update", 1},
                                                {"_recvChunkStart", 1},
                                                {"_configsvrRemoveChunks", 1},
diff --git a/src/mongo/s/commands/internal_transactions_test_command.h b/src/mongo/s/commands/internal_transactions_test_command.h
index dc4b7fc00c8..416a2a2ca94 100644
--- a/src/mongo/s/commands/internal_transactions_test_command.h
+++ b/src/mongo/s/commands/internal_transactions_test_command.h
@@ -116,8 +116,14 @@ public:
                             // commands.
                             uassertStatusOK(getStatusFromWriteCommandReply(res));
                         }
-                        // TODO SERVER-65048: Check if result has retriedStmtId & retriedStmtIds
-                        // field, exit.
+
+                        // Exit if we are reexecuting commands in a retryable write, identified by a
+                        // populated retriedStmtId. eoo() is false if field is found.
+                        const auto isRetryStmt = !(res.getField("retriedStmtIds").eoo() &&
+                                                   res.getField("retriedStmtId").eoo());
+                        if (isRetryStmt) {
+                            break;
+                        }
                     }
                     return SemiFuture<void>::makeReady();
                 });