SERVER-35094 add new txn states and allow retrying commitTransaction cmd

3 files changed, 146 insertions, 39 deletions

diff --git a/jstests/core/txns/transaction_error_handling.js b/jstests/core/txns/transaction_error_handling.js
index d30e8c6003e..f2436f4cf64 100644
--- a/jstests/core/txns/transaction_error_handling.js
+++ b/jstests/core/txns/transaction_error_handling.js
@@ -16,6 +16,7 @@
     const sessionDb = session.getDatabase(dbName);
     const sessionColl = sessionDb.getCollection(collName);
 
+    jsTestLog("Test that we cannot abort or commit a nonexistant transaction.");
     // Cannot abort or commit a nonexistant transaction.
     try {
         session.commitTransaction();
@@ -29,7 +30,8 @@
         assert.eq(e.message, "There is no active transaction to abort on this session.");
     }
 
-    // Try to start a transaction with one already running.
+    // Try to start a transaction when the state is 'active'.
+    jsTestLog("Test that we cannot start a transaction with one already started or in progress.");
     session.startTransaction();
     try {
         session.startTransaction();
@@ -37,5 +39,80 @@
         assert.eq(e.message, "Transaction already in progress on this session.");
     }
 
+    // Try starting a transaction after inserting something.
+    assert.commandWorked(sessionColl.insert({_id: "insert-1"}));
+    // Try to start a transaction when the state is 'active'.
+    try {
+        session.startTransaction();
+    } catch (e) {
+        assert.eq(e.message, "Transaction already in progress on this session.");
+    }
+
+    // At this point, the transaction is still 'active'. We will commit this transaction and test
+    // that calling commitTransaction again should work while calling abortTransaction should not.
+    session.commitTransaction();
+
+    jsTestLog("Test that we can commit a transaction more than once.");
+    // The transaction state is 'committed'. We can call commitTransaction again in this state.
+    session.commitTransaction();
+
+    jsTestLog("Test that we cannot abort a transaction that has already been committed");
+    // We cannot call abortTransaction on a transaction that has already been committed.
+    try {
+        session.abortTransaction();
+    } catch (e) {
+        assert.eq(e.message, "Cannot call abortTransaction after calling commitTransaction.");
+    }
+
+    // Start a new transaction that will be aborted. Test that we cannot call commit or
+    // abortTransaction on a transaction that is in the 'aborted' state.
+    session.startTransaction();
+    assert.commandWorked(sessionColl.insert({_id: "insert-2"}));
+    session.abortTransaction();
+
+    jsTestLog("Test that we cannot commit a transaction that has already been aborted.");
+    // We cannot call commitTransaction on a transaction that has already been aborted.
+    try {
+        session.commitTransaction();
+    } catch (e) {
+        assert.eq(e.message, "Cannot call commitTransaction after calling abortTransaction.");
+    }
+
+    jsTestLog("Test that we cannot abort a transaction that has already been aborted.");
+    // We also cannot call abortTransaction on a transaction that has already been aborted.
+    try {
+        session.abortTransaction();
+    } catch (e) {
+        assert.eq(e.message, "Cannot call abortTransaction twice.");
+    }
+
+    jsTestLog(
+        "Test that a normal operation after committing a transaction changes the state to inactive.");
+    session.startTransaction();
+    assert.commandWorked(sessionColl.insert({_id: "insert-3"}));
+    // The transaction state should be changed to 'committed'.
+    session.commitTransaction();
+    // The transaction state should be changed to 'inactive'.
+    assert.commandWorked(sessionColl.insert({_id: "normal-insert"}));
+    try {
+        session.commitTransaction();
+    } catch (e) {
+        assert.eq(e.message, "There is no active transaction to commit on this session.");
+    }
+
+    jsTestLog(
+        "Test that a normal operation after aborting a transaction changes the state to inactive.");
+    session.startTransaction();
+    assert.commandWorked(sessionColl.insert({_id: "insert-4"}));
+    // The transaction state should be changed to 'aborted'.
+    session.abortTransaction();
+    // The transaction state should be changed to 'inactive'.
+    assert.commandWorked(sessionColl.insert({_id: "normal-insert-2"}));
+    try {
+        session.abortTransaction();
+    } catch (e) {
+        assert.eq(e.message, "There is no active transaction to abort on this session.");
+    }
+
     session.endSession();
 }());
diff --git a/src/mongo/shell/bulk_api.js b/src/mongo/shell/bulk_api.js
index 6caa50374e6..de377b5a0f5 100644
--- a/src/mongo/shell/bulk_api.js
+++ b/src/mongo/shell/bulk_api.js
@@ -881,7 +881,7 @@ var _bulk_api_module = (function() {
                 const session = collection.getDB().getSession();
                 if (serverSupportsRetryableWrites && session.getOptions().shouldRetryWrites() &&
                     session._serverSession.canRetryWrites(cmd) &&
-                    !session._serverSession.isInActiveTransaction()) {
+                    !session._serverSession.isTxnActive()) {
                     cmd = session._serverSession.assignTransactionNumber(cmd);
                 }
             }
diff --git a/src/mongo/shell/session.js b/src/mongo/shell/session.js
index dd7b09c0c97..b7e77d8e9f2 100644
--- a/src/mongo/shell/session.js
+++ b/src/mongo/shell/session.js
@@ -140,7 +140,7 @@ var {
         function canUseReadConcern(driverSession, cmdObj) {
             // Always attach the readConcern to the first statement of the transaction, whether it
             // is a read or a write.
-            if (driverSession._serverSession.isInActiveTransaction()) {
+            if (driverSession._serverSession.isTxnActive()) {
                 return driverSession._serverSession.isFirstStatement();
             }
 
@@ -269,17 +269,9 @@ var {
                 }
             }
 
-            // If startTransaction was called on the session, attach txn number and readConcern.
-            if (driverSession._serverSession.isInActiveTransaction()) {
-                // If we reconnect to a 3.6 server in the middle of a transaction, we
-                // catch it here.
-                if (!serverSupports(kWireVersionSupportingMultiDocumentTransactions)) {
-                    driverSession._serverSession._abortTransactionOnClientOnly();
-                    throw new Error(
-                        "Transactions are only supported on server versions 4.0 and greater.");
-                }
-                cmdObj = driverSession._serverSession.assignTxnInfo(cmdObj);
-            }
+            // All commands go through transaction code, which will determine if the command is a
+            // part of the current transaction and will assign transaction info accordingly.
+            cmdObj = driverSession._serverSession.assignTxnInfo(cmdObj);
 
             // Retryable writes code should execute only we are not in an active transaction.
             if (jsTest.options().alwaysInjectTransactionNumber &&
@@ -335,7 +327,7 @@ var {
             let numRetries =
                 (sessionOptions.shouldRetryWrites() && cmdObj.hasOwnProperty("txnNumber") &&
                  !jsTest.options().skipRetryOnNetworkError &&
-                 !driverSession._serverSession.isInActiveTransaction())
+                 !driverSession._serverSession.isTxnActive())
                 ? 1
                 : 0;
 
@@ -504,7 +496,12 @@ var {
         this.client = new SessionAwareClient(client);
         this.handle = client._startSession();
 
-        this.isInActiveTransaction = function isInActiveTransaction() {
+        function serverSupports(wireVersion) {
+            return client.getMinWireVersion() <= wireVersion &&
+                wireVersion <= client.getMaxWireVersion();
+        }
+
+        this.isTxnActive = function isTxnActive() {
             return _txnState === ServerSession.TransactionStates.kActive;
         };
 
@@ -654,6 +651,27 @@ var {
         };
 
         this.assignTxnInfo = function assignTxnInfo(cmdObj) {
+            // We will want to reset the transaction state to 'inactive' if a normal operation
+            // follows a committed or aborted transaction.
+            if ((_txnState === ServerSession.TransactionStates.kAborted) ||
+                (_txnState === ServerSession.TransactionStates.kCommitted &&
+                 Object.keys(cmdObj)[0] !== "commitTransaction")) {
+                _txnState = ServerSession.TransactionStates.kInactive;
+            }
+
+            // If we're not in an active transaction or performing a retry on commitTransaction,
+            // return early.
+            if (_txnState === ServerSession.TransactionStates.kInactive) {
+                return cmdObj;
+            }
+
+            // If we reconnect to a 3.6 server in the middle of a transaction, we
+            // catch it here.
+            if (!serverSupports(kWireVersionSupportingMultiDocumentTransactions)) {
+                _txnState = ServerSession.TransactionStates.kInactive;
+                throw new Error(
+                    "Transactions are only supported on server versions 4.0 and greater.");
+            }
             cmdObj = Object.assign({}, cmdObj);
 
             const cmdName = Object.keys(cmdObj)[0];
@@ -679,8 +697,8 @@ var {
             // Statement Id is required on all transaction operations.
             cmdObjUnwrapped.stmtId = new NumberInt(_nextStatementId);
 
-            // 'readConcern' and 'startTransaction' can only be specified on the first statement in
-            // a transaction.
+            // 'readConcern' and 'startTransaction' can only be specified on the first statement
+            // in a transaction.
             if (_nextStatementId == 0) {
                 cmdObjUnwrapped.startTransaction = true;
                 if (_txnOptions.getTxnReadConcern() !== undefined) {
@@ -710,13 +728,9 @@ var {
 
         this.startTransaction = function startTransaction(txnOptsObj) {
             // If the session is already in a transaction, raise an error.
-            if (this.isInActiveTransaction()) {
+            if (this.isTxnActive()) {
                 throw new Error("Transaction already in progress on this session.");
             }
-            function serverSupports(wireVersion) {
-                return client.getMinWireVersion() <= wireVersion &&
-                    wireVersion <= client.getMaxWireVersion();
-            }
             if (!serverSupports(kWireVersionSupportingMultiDocumentTransactions)) {
                 throw new Error(
                     "Transactions are only supported on server versions 4.0 and greater.");
@@ -728,8 +742,12 @@ var {
         };
 
         this.commitTransaction = function commitTransaction(driverSession) {
+            // If the transaction state is already 'aborted' we cannot try to commit it.
+            if (_txnState === ServerSession.TransactionStates.kAborted) {
+                throw new Error("Cannot call commitTransaction after calling abortTransaction.");
+            }
             // If the session has no active transaction, raise an error.
-            if (!this.isInActiveTransaction()) {
+            if (_txnState === ServerSession.TransactionStates.kInactive) {
                 throw new Error("There is no active transaction to commit on this session.");
             }
             // run commitTxn command
@@ -737,26 +755,32 @@ var {
         };
 
         this.abortTransaction = function abortTransaction(driverSession) {
+            // If the transaction state is already 'aborted' we cannot try to abort it again.
+            if (_txnState === ServerSession.TransactionStates.kAborted) {
+                throw new Error("Cannot call abortTransaction twice.");
+            }
+            // We cannot attempt to abort a transaction that has already been committed.
+            if (_txnState === ServerSession.TransactionStates.kCommitted) {
+                throw new Error("Cannot call abortTransaction after calling commitTransaction.");
+            }
             // If the session has no active transaction, raise an error.
-            if (!this.isInActiveTransaction()) {
+            if (_txnState === ServerSession.TransactionStates.kInactive) {
                 throw new Error("There is no active transaction to abort on this session.");
             }
             // run abortTxn command
             return endTransaction("abortTransaction", driverSession);
         };
 
-        // This is used to abort transactions if we've reconnected to a non-transaction-supporting
-        // server.
-        this._abortTransactionOnClientOnly = function _abortTransactionOnClientOnly() {
-            _txnState = ServerSession.TransactionStates.kInactive;
-        };
-
         const endTransaction = (commandName, driverSession) => {
             // If commitTransaction or abortTransaction is the first statement in a
             // transaction, it should not send a command to the server and should mark the
-            // transaction as inactive.
+            // transaction as 'committed' or 'aborted' accordingly.
             if (this.isFirstStatement()) {
-                _txnState = ServerSession.TransactionStates.kInactive;
+                if (commandName === "commitTransaction") {
+                    _txnState = ServerSession.TransactionStates.kCommitted;
+                } else {
+                    _txnState = ServerSession.TransactionStates.kAborted;
+                }
                 return {"ok": 1};
             }
 
@@ -771,25 +795,31 @@ var {
                 cmd.writeConcern = _txnOptions.getTxnWriteConcern();
             }
 
-            // If commit or abort raises an error, the transaction's state should still change
-            // to inactive.
+            // If commit or abort raises an error, the transaction's state should still change.
             let res;
             try {
                 // run command against the admin database.
                 res = this.client.runCommand(driverSession, "admin", cmd, 0);
             } finally {
-                _txnState = ServerSession.TransactionStates.kInactive;
+                if (commandName === "commitTransaction") {
+                    _txnState = ServerSession.TransactionStates.kCommitted;
+                } else {
+                    _txnState = ServerSession.TransactionStates.kAborted;
+                }
             }
             return res;
         };
     }
 
     // TransactionStates represents the state of the current transaction. The default state
-    // is `inactive` until startTransaction is called and changes the state to `active`.
-    // Calling a successful abort or commitTransaction will change the state to `inactive`.
+    // is 'inactive' until startTransaction is called and changes the state to 'active'.
+    // Calling abortTransaction or commitTransaction will change the state to 'aborted' or
+    // 'committed' respectively, even on error.
     ServerSession.TransactionStates = {
         kActive: 'active',
         kInactive: 'inactive',
+        kCommitted: 'committed',
+        kAborted: 'aborted',
     };
 
     function makeDriverSessionConstructor(implMethods, defaultOptions = {}) {
@@ -983,7 +1013,7 @@ var {
                           return cmdObj;
                       },
 
-                      isInActiveTransaction: function isInActiveTransaction() {
+                      isTxnActive: function isTxnActive() {
                           return false;
                       },