SERVER-36902 Abort transaction on shell exit, try 2

4 files changed, 117 insertions, 5 deletions

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);