SERVER-36902 Abort transaction on shell exit, try 2

8 files changed, 155 insertions, 65 deletions

diff --git a/jstests/change_streams/apply_ops.js b/jstests/change_streams/apply_ops.js
index ee116d3495d..5212f16e14b 100644
--- a/jstests/change_streams/apply_ops.js
+++ b/jstests/change_streams/apply_ops.js
@@ -72,7 +72,7 @@
           ns: {db: db.getName(), coll: coll.getName()},
           operationType: "insert",
           lsid: session.getSessionId(),
-          txnNumber: NumberLong(session._txnNumber),
+          txnNumber: session.getTxnNumber_forTesting(),
         },
         {
           documentKey: {_id: 2},
@@ -80,7 +80,7 @@
           ns: {db: db.getName(), coll: coll.getName()},
           operationType: "insert",
           lsid: session.getSessionId(),
-          txnNumber: NumberLong(session._txnNumber),
+          txnNumber: session.getTxnNumber_forTesting(),
         },
         {
           documentKey: {_id: 1},
@@ -88,14 +88,14 @@
           operationType: "update",
           updateDescription: {removedFields: [], updatedFields: {a: 1}},
           lsid: session.getSessionId(),
-          txnNumber: NumberLong(session._txnNumber),
+          txnNumber: session.getTxnNumber_forTesting(),
         },
         {
           documentKey: {_id: kDeletedDocumentId},
           ns: {db: db.getName(), coll: coll.getName()},
           operationType: "delete",
           lsid: session.getSessionId(),
-          txnNumber: NumberLong(session._txnNumber),
+          txnNumber: session.getTxnNumber_forTesting(),
         },
         {
           operationType: "drop",
@@ -123,7 +123,7 @@
         ns: {db: db.getName(), coll: otherCollName},
         operationType: "insert",
         lsid: session.getSessionId(),
-        txnNumber: NumberLong(session._txnNumber),
+        txnNumber: session.getTxnNumber_forTesting(),
     });
 
     // Verify that a whole-db stream returns the expected sequence of changes, including the insert
@@ -141,7 +141,7 @@
         ns: {db: otherDbName, coll: otherDbCollName},
         operationType: "insert",
         lsid: session.getSessionId(),
-        txnNumber: NumberLong(session._txnNumber),
+        txnNumber: session.getTxnNumber_forTesting(),
     });
 
     // Verify that a whole-cluster stream returns the expected sequence of changes, including the
diff --git a/jstests/change_streams/apply_ops_resumability.js b/jstests/change_streams/apply_ops_resumability.js
index a8d25d775d4..619bd7152d6 100644
--- a/jstests/change_streams/apply_ops_resumability.js
+++ b/jstests/change_streams/apply_ops_resumability.js
@@ -62,7 +62,7 @@
           ns: {db: db.getName(), coll: coll.getName()},
           operationType: "insert",
           lsid: session.getSessionId(),
-          txnNumber: NumberLong(session._txnNumber),
+          txnNumber: session.getTxnNumber_forTesting(),
         },
         {
           documentKey: {_id: 2},
@@ -70,7 +70,7 @@
           ns: {db: db.getName(), coll: coll.getName()},
           operationType: "insert",
           lsid: session.getSessionId(),
-          txnNumber: NumberLong(session._txnNumber),
+          txnNumber: session.getTxnNumber_forTesting(),
         },
         {
           documentKey: {_id: 1},
@@ -78,7 +78,7 @@
           operationType: "update",
           updateDescription: {removedFields: [], updatedFields: {a: 1}},
           lsid: session.getSessionId(),
-          txnNumber: NumberLong(session._txnNumber),
+          txnNumber: session.getTxnNumber_forTesting(),
         },
         {
           documentKey: {_id: 3},
@@ -149,7 +149,7 @@
         ns: {db: db.getName(), coll: otherCollName},
         operationType: "insert",
         lsid: session.getSessionId(),
-        txnNumber: NumberLong(session._txnNumber),
+        txnNumber: session.getTxnNumber_forTesting(),
     });
 
     // Verify that a whole-db stream can be resumed from the middle of the transaction, and that it
@@ -169,7 +169,7 @@
         ns: {db: otherDbName, coll: otherDbCollName},
         operationType: "insert",
         lsid: session.getSessionId(),
-        txnNumber: NumberLong(session._txnNumber),
+        txnNumber: session.getTxnNumber_forTesting(),
     });
 
     // Verify that a whole-cluster stream can be resumed from the middle of the transaction, and
diff --git a/jstests/replsets/transaction_table_multi_statement_txn.js b/jstests/replsets/transaction_table_multi_statement_txn.js
index 21d95b0eb7a..01ca8f89704 100644
--- a/jstests/replsets/transaction_table_multi_statement_txn.js
+++ b/jstests/replsets/transaction_table_multi_statement_txn.js
@@ -31,7 +31,7 @@
     assert.writeOK(coll.insert({_id: 1}));
     session.commitTransaction();
     const opTime = session.getOperationTime();
-    const txnNum = NumberLong(session._txnNumber);
+    const txnNum = session.getTxnNumber_forTesting();
     jsTestLog('Successfully committed transaction at operation time ' + tojson(opTime) +
               'with transaction number ' + txnNum);
 
diff --git a/src/mongo/scripting/mozjs/session.cpp b/src/mongo/scripting/mozjs/session.cpp
index 676e865b317..477ab4024dc 100644
--- a/src/mongo/scripting/mozjs/session.cpp
+++ b/src/mongo/scripting/mozjs/session.cpp
@@ -41,28 +41,70 @@
 #include "mongo/scripting/mozjs/valuereader.h"
 #include "mongo/scripting/mozjs/wrapconstrainedmethod.h"
 #include "mongo/util/log.h"
+#include "mongo/util/mongoutils/str.h"
 
 namespace mongo {
 namespace mozjs {
 
-const JSFunctionSpec SessionInfo::methods[3] = {
+const JSFunctionSpec SessionInfo::methods[8] = {
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(end, SessionInfo),
     MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(getId, SessionInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(getTxnState, SessionInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(setTxnState, SessionInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(getTxnNumber, SessionInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(setTxnNumber, SessionInfo),
+    MONGO_ATTACH_JS_CONSTRAINED_METHOD_NO_PROTO(incrementTxnNumber, SessionInfo),
     JS_FS_END,
 };
 
 const char* const SessionInfo::className = "Session";
-
 struct SessionHolder {
+    enum class TransactionState { kActive, kInactive, kCommitted, kAborted };
+    // txnNumber starts at -1 because when we increment it, the first transaction
+    // and retryable write will both have a txnNumber of 0.
     SessionHolder(std::shared_ptr<DBClientBase> client, BSONObj lsid)
-        : client(std::move(client)), lsid(std::move(lsid)) {}
+        : client(std::move(client)),
+          lsid(std::move(lsid)),
+          txnState(TransactionState::kInactive),
+          txnNumber(-1) {}
 
     std::shared_ptr<DBClientBase> client;
     BSONObj lsid;
+    TransactionState txnState;
+    std::int64_t txnNumber;
 };
 
 namespace {
 
+StringData transactionStateName(SessionHolder::TransactionState state) {
+    switch (state) {
+        case SessionHolder::TransactionState::kActive:
+            return "active"_sd;
+        case SessionHolder::TransactionState::kInactive:
+            return "inactive"_sd;
+        case SessionHolder::TransactionState::kCommitted:
+            return "committed"_sd;
+        case SessionHolder::TransactionState::kAborted:
+            return "aborted"_sd;
+    }
+
+    MONGO_UNREACHABLE;
+}
+
+SessionHolder::TransactionState transactionStateEnum(StringData name) {
+    if (name == "active"_sd) {
+        return SessionHolder::TransactionState::kActive;
+    } else if (name == "inactive"_sd) {
+        return SessionHolder::TransactionState::kInactive;
+    } else if (name == "committed"_sd) {
+        return SessionHolder::TransactionState::kCommitted;
+    } else if (name == "aborted"_sd) {
+        return SessionHolder::TransactionState::kAborted;
+    } else {
+        uasserted(ErrorCodes::BadValue, str::stream() << "Invalid TransactionState name: " << name);
+    }
+}
+
 SessionHolder* getHolder(JSObject* thisv) {
     return static_cast<SessionHolder*>(JS_GetPrivate(thisv));
 }
@@ -76,11 +118,22 @@ void endSession(SessionHolder* holder) {
         return;
     }
 
+    BSONObj out;
+
+    if (holder->txnState == SessionHolder::TransactionState::kActive) {
+        holder->txnState = SessionHolder::TransactionState::kAborted;
+        BSONObj abortObj = BSON("abortTransaction" << 1 << "lsid" << holder->lsid << "txnNumber"
+                                                   << holder->txnNumber
+                                                   << "autocommit"
+                                                   << false);
+
+        holder->client->runCommand("admin", abortObj, out);
+    }
+
     EndSessions es;
 
     es.setEndSessions({holder->lsid});
 
-    BSONObj out;
     holder->client->runCommand("admin", es.toBSON(), out);
 
     holder->client.reset();
@@ -126,6 +179,51 @@ void SessionInfo::Functions::getId::call(JSContext* cx, JS::CallArgs args) {
     ValueReader(cx, args.rval()).fromBSON(holder->lsid, nullptr, 1);
 }
 
+void SessionInfo::Functions::getTxnState::call(JSContext* cx, JS::CallArgs args) {
+    auto holder = getHolder(args);
+    invariant(holder);
+    uassert(ErrorCodes::BadValue, "getTxnState takes no arguments", args.length() == 0);
+
+    ValueReader(cx, args.rval()).fromStringData(transactionStateName(holder->txnState));
+}
+
+void SessionInfo::Functions::setTxnState::call(JSContext* cx, JS::CallArgs args) {
+    auto holder = getHolder(args);
+    invariant(holder);
+    uassert(ErrorCodes::BadValue, "setTxnState takes 1 argument", args.length() == 1);
+
+    auto arg = args.get(0);
+    holder->txnState = transactionStateEnum(ValueWriter(cx, arg).toString().c_str());
+    args.rval().setUndefined();
+}
+
+void SessionInfo::Functions::getTxnNumber::call(JSContext* cx, JS::CallArgs args) {
+    auto holder = getHolder(args);
+    invariant(holder);
+    uassert(ErrorCodes::BadValue, "getTxnNumber takes no arguments", args.length() == 0);
+
+    ValueReader(cx, args.rval()).fromInt64(holder->txnNumber);
+}
+
+void SessionInfo::Functions::setTxnNumber::call(JSContext* cx, JS::CallArgs args) {
+    auto holder = getHolder(args);
+    invariant(holder);
+    uassert(ErrorCodes::BadValue, "setTxnNumber takes 1 argument", args.length() == 1);
+
+    auto arg = args.get(0);
+    holder->txnNumber = ValueWriter(cx, arg).toInt64();
+    args.rval().setUndefined();
+}
+
+void SessionInfo::Functions::incrementTxnNumber::call(JSContext* cx, JS::CallArgs args) {
+    auto holder = getHolder(args);
+    invariant(holder);
+    uassert(ErrorCodes::BadValue, "incrementTxnNumber takes no arguments", args.length() == 0);
+
+    ++holder->txnNumber;
+    args.rval().setUndefined();
+}
+
 void SessionInfo::make(JSContext* cx,
                        JS::MutableHandleObject obj,
                        std::shared_ptr<DBClientBase> client,
diff --git a/src/mongo/scripting/mozjs/session.h b/src/mongo/scripting/mozjs/session.h
index 82f071d4241..8b00d77aff4 100644
--- a/src/mongo/scripting/mozjs/session.h
+++ b/src/mongo/scripting/mozjs/session.h
@@ -48,9 +48,14 @@ struct SessionInfo : public BaseInfo {
     struct Functions {
         MONGO_DECLARE_JS_FUNCTION(end);
         MONGO_DECLARE_JS_FUNCTION(getId);
+        MONGO_DECLARE_JS_FUNCTION(getTxnState);
+        MONGO_DECLARE_JS_FUNCTION(setTxnState);
+        MONGO_DECLARE_JS_FUNCTION(getTxnNumber);
+        MONGO_DECLARE_JS_FUNCTION(setTxnNumber);
+        MONGO_DECLARE_JS_FUNCTION(incrementTxnNumber);
     };
 
-    static const JSFunctionSpec methods[3];
+    static const JSFunctionSpec methods[8];
 
     static const char* const className;
     static const unsigned classFlags = JSCLASS_HAS_PRIVATE;
diff --git a/src/mongo/scripting/mozjs/valuereader.cpp b/src/mongo/scripting/mozjs/valuereader.cpp
index 378342e33c0..7f6637e147c 100644
--- a/src/mongo/scripting/mozjs/valuereader.cpp
+++ b/src/mongo/scripting/mozjs/valuereader.cpp
@@ -304,5 +304,13 @@ void ValueReader::fromDouble(double d) {
     }
 }
 
+void ValueReader::fromInt64(int64_t i) {
+    auto scope = getScope(_context);
+    JS::RootedObject num(_context);
+    scope->getProto<NumberLongInfo>().newObject(&num);
+    JS_SetPrivate(num, scope->trackedNew<int64_t>(i));
+    _value.setObjectOrNull(num);
+}
+
 }  // namespace mozjs
 }  // namespace mongo
diff --git a/src/mongo/scripting/mozjs/valuereader.h b/src/mongo/scripting/mozjs/valuereader.h
index 4bf591f8999..6c255e218db 100644
--- a/src/mongo/scripting/mozjs/valuereader.h
+++ b/src/mongo/scripting/mozjs/valuereader.h
@@ -53,6 +53,7 @@ public:
     void fromBSON(const BSONObj& obj, const BSONObj* parent, bool readOnly);
     void fromBSONArray(const BSONObj& obj, const BSONObj* parent, bool readOnly);
     void fromDouble(double d);
+    void fromInt64(int64_t i);
     void fromStringData(StringData sd);
     void fromDecimal128(Decimal128 decimal);
 
diff --git a/src/mongo/shell/session.js b/src/mongo/shell/session.js
index 9c68c35a063..d1587440545 100644
--- a/src/mongo/shell/session.js
+++ b/src/mongo/shell/session.js
@@ -486,17 +486,10 @@ var {
     // The server session maintains the state of a transaction, a monotonically increasing txn
     // number, and a transaction's read/write concerns.
     function ServerSession(client) {
-        // The default txnState is `inactive` until we call startTransaction.
-        let _txnState = ServerSession.TransactionStates.kInactive;
-
         let _txnOptions;
 
         // Keep track of the next available statement id of a transaction.
         let _nextStatementId = 0;
-
-        // _txnNumber starts at -1 because when we increment it, the first transaction
-        // and retryable write will both have a txnNumber of 0.
-        let _txnNumber = -1;
         let _lastUsed = new Date();
 
         this.client = new SessionAwareClient(client);
@@ -511,8 +504,11 @@ var {
                 wireVersion <= client.getMaxWireVersion();
         }
 
+        const hasTxnState = ((name) => this.handle.getTxnState() === name);
+        const setTxnState = ((name) => this.handle.setTxnState(name));
+
         this.isTxnActive = function isTxnActive() {
-            return _txnState === ServerSession.TransactionStates.kActive;
+            return hasTxnState("active");
         };
 
         this.isFirstStatement = function isFirstStatement() {
@@ -524,11 +520,11 @@ var {
         };
 
         this.getTxnNumber = function getTxnNumber() {
-            return _txnNumber;
+            return this.handle.getTxnNumber();
         };
 
         this.setTxnNumber_forTesting = function setTxnNumber_forTesting(newTxnNumber) {
-            _txnNumber = newTxnNumber;
+            this.handle.setTxnNumber(newTxnNumber);
         };
 
         this.getTxnOptions = function getTxnOptions() {
@@ -577,11 +573,8 @@ var {
             }
 
             if (!cmdObjUnwrapped.hasOwnProperty("txnNumber")) {
-                // Since there's no native support for adding NumberLong instances and getting back
-                // another NumberLong instance, converting from a 64-bit floating-point value to a
-                // 64-bit integer value will overflow at 2**53.
-                _txnNumber++;
-                cmdObjUnwrapped.txnNumber = new NumberLong(_txnNumber);
+                this.handle.incrementTxnNumber();
+                cmdObjUnwrapped.txnNumber = this.handle.getTxnNumber();
             }
 
             return cmdObj;
@@ -675,22 +668,21 @@ 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 ((hasTxnState("aborted")) ||
+                (hasTxnState("committed") && Object.keys(cmdObj)[0] !== "commitTransaction")) {
+                setTxnState("inactive");
             }
 
             // If we're not in an active transaction or performing a retry on commitTransaction,
             // return early.
-            if (_txnState === ServerSession.TransactionStates.kInactive) {
+            if (hasTxnState("inactive")) {
                 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;
+                setTxnState("inactive");
                 throw new Error(
                     "Transactions are only supported on server versions 4.0 and greater.");
             }
@@ -707,10 +699,7 @@ var {
             }
 
             if (!cmdObjUnwrapped.hasOwnProperty("txnNumber")) {
-                // Since there's no native support for adding NumberLong instances and getting back
-                // another NumberLong instance, converting from a 64-bit floating-point value to a
-                // 64-bit integer value will overflow at 2**53.
-                cmdObjUnwrapped.txnNumber = new NumberLong(_txnNumber);
+                cmdObjUnwrapped.txnNumber = this.handle.getTxnNumber();
             }
 
             // All operations of a multi-statement transaction must specify autocommit=false.
@@ -760,18 +749,18 @@ var {
                     "Transactions are only supported on server versions 4.0 and greater.");
             }
             _txnOptions = new TransactionOptions(txnOptsObj);
-            _txnState = ServerSession.TransactionStates.kActive;
+            setTxnState("active");
             _nextStatementId = 0;
-            _txnNumber++;
+            this.handle.incrementTxnNumber();
         };
 
         this.commitTransaction = function commitTransaction(driverSession) {
             // If the transaction state is already 'aborted' we cannot try to commit it.
-            if (_txnState === ServerSession.TransactionStates.kAborted) {
+            if (hasTxnState("aborted")) {
                 throw new Error("Cannot call commitTransaction after calling abortTransaction.");
             }
             // If the session has no active transaction, raise an error.
-            if (_txnState === ServerSession.TransactionStates.kInactive) {
+            if (hasTxnState("inactive")) {
                 throw new Error("There is no active transaction to commit on this session.");
             }
             // run commitTxn command
@@ -780,15 +769,15 @@ 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) {
+            if (hasTxnState("aborted")) {
                 throw new Error("Cannot call abortTransaction twice.");
             }
             // We cannot attempt to abort a transaction that has already been committed.
-            if (_txnState === ServerSession.TransactionStates.kCommitted) {
+            if (hasTxnState("committed")) {
                 throw new Error("Cannot call abortTransaction after calling commitTransaction.");
             }
             // If the session has no active transaction, raise an error.
-            if (_txnState === ServerSession.TransactionStates.kInactive) {
+            if (hasTxnState("inactive")) {
                 throw new Error("There is no active transaction to abort on this session.");
             }
             // run abortTxn command
@@ -801,14 +790,14 @@ var {
             // transaction as 'committed' or 'aborted' accordingly.
             if (this.isFirstStatement()) {
                 if (commandName === "commitTransaction") {
-                    _txnState = ServerSession.TransactionStates.kCommitted;
+                    setTxnState("committed");
                 } else {
-                    _txnState = ServerSession.TransactionStates.kAborted;
+                    setTxnState("aborted");
                 }
                 return {"ok": 1};
             }
 
-            let cmd = {[commandName]: 1, txnNumber: NumberLong(_txnNumber)};
+            let cmd = {[commandName]: 1, txnNumber: this.handle.getTxnNumber()};
             // writeConcern should only be specified on commit or abort. If a writeConcern is
             // not specified from the default transaction options, it will be inherited from
             // the session.
@@ -826,26 +815,15 @@ var {
                 res = this.client.runCommand(driverSession, "admin", cmd, 0);
             } finally {
                 if (commandName === "commitTransaction") {
-                    _txnState = ServerSession.TransactionStates.kCommitted;
+                    setTxnState("committed");
                 } else {
-                    _txnState = ServerSession.TransactionStates.kAborted;
+                    setTxnState("aborted");
                 }
             }
             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 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 = {}) {
         var driverSessionConstructor = function(client, options = defaultOptions) {
             let _options = options;