SERVER-25056 implscope checks interruptability against opctx

(cherry picked from commit 9ca2c6dcc96cf826b0ef06a8cb1590ec2bcaa8a1)

2 files changed, 23 insertions, 8 deletions

diff --git a/src/mongo/scripting/mozjs/implscope.cpp b/src/mongo/scripting/mozjs/implscope.cpp
index 9c298b4e52c..c5bd51ef5cd 100644
--- a/src/mongo/scripting/mozjs/implscope.cpp
+++ b/src/mongo/scripting/mozjs/implscope.cpp
@@ -95,10 +95,18 @@ bool gFirstRuntimeCreated = false;
 MONGO_TRIVIALLY_CONSTRUCTIBLE_THREAD_LOCAL MozJSImplScope* kCurrentScope;
 
 struct MozJSImplScope::MozJSEntry {
-    MozJSEntry(MozJSImplScope* scope) : ar(scope->_context), ac(scope->_context, scope->_global) {}
+    MozJSEntry(MozJSImplScope* scope)
+        : ar(scope->_context), ac(scope->_context, scope->_global), _scope(scope) {
+        ++_scope->_inOp;
+    }
+
+    ~MozJSEntry() {
+        --_scope->_inOp;
+    }
 
     JSAutoRequest ar;
     JSAutoCompartment ac;
+    MozJSImplScope* _scope;
 };
 
 void MozJSImplScope::_reportError(JSContext* cx, const char* message, JSErrorReport* report) {
@@ -136,23 +144,23 @@ std::string MozJSImplScope::getError() {
 }
 
 void MozJSImplScope::registerOperation(OperationContext* txn) {
-    invariant(_opId == 0);
+    invariant(_opCtx == nullptr);
 
     // getPooledScope may call registerOperation with a nullptr, so we have to
     // check for that here.
     if (!txn)
         return;
 
+    _opCtx = txn;
     _opId = txn->getOpID();
 
     _engine->registerOperation(txn, this);
 }
 
 void MozJSImplScope::unregisterOperation() {
-    if (_opId != 0) {
+    if (_opCtx) {
         _engine->unregisterOperation(_opId);
-
-        _opId = 0;
+        _opCtx = nullptr;
     }
 }
 
@@ -191,6 +199,14 @@ bool MozJSImplScope::_interruptCallback(JSContext* cx) {
     } else if (scope->isKillPending()) {
         scope->_status = Status(ErrorCodes::Interrupted, "Interrupted by the host");
     }
+    // If we are on the right thread, in the middle of an operation, and we have a registered opCtx,
+    // then we should check the opCtx for interrupts.
+    if ((scope->_mr._thread == PR_GetCurrentThread()) && (scope->_inOp > 0) && scope->_opCtx) {
+        auto status = scope->_opCtx->checkForInterruptNoAssert();
+        if (!status.isOK()) {
+            scope->_status = status;
+        }
+    }
 
     if (!scope->_status.isOK()) {
         scope->_engine->getDeadlineMonitor().stopDeadline(scope);
@@ -296,6 +312,7 @@ MozJSImplScope::MozJSImplScope(MozJSScriptEngine* engine)
       _pendingKill(false),
       _opId(0),
       _opCtx(nullptr),
+      _inOp(0),
       _pendingGC(false),
       _connectState(ConnectState::Not),
       _status(Status::OK()),
@@ -677,9 +694,6 @@ void MozJSImplScope::gc() {
 void MozJSImplScope::localConnectForDbEval(OperationContext* txn, const char* dbName) {
     MozJSEntry entry(this);
 
-    invariant(_opCtx == NULL);
-    _opCtx = txn;
-
     if (_connectState == ConnectState::External)
         uasserted(12510, "externalSetup already called, can't call localConnect");
     if (_connectState == ConnectState::Local) {
diff --git a/src/mongo/scripting/mozjs/implscope.h b/src/mongo/scripting/mozjs/implscope.h
index be125f3e987..3fac5e5bdef 100644
--- a/src/mongo/scripting/mozjs/implscope.h
+++ b/src/mongo/scripting/mozjs/implscope.h
@@ -366,6 +366,7 @@ private:
     std::string _error;
     unsigned int _opId;        // op id for this scope
     OperationContext* _opCtx;  // Op context for DbEval
+    std::size_t _inOp;
     std::atomic<bool> _pendingGC;
     ConnectState _connectState;
     Status _status;