SERVER-31508: Make Javascript sleep interruptable

5 files changed, 32 insertions, 15 deletions

diff --git a/src/mongo/scripting/mozjs/global.cpp b/src/mongo/scripting/mozjs/global.cpp
index 1cb3e463e35..15c183fe78c 100644
--- a/src/mongo/scripting/mozjs/global.cpp
+++ b/src/mongo/scripting/mozjs/global.cpp
@@ -40,12 +40,14 @@
 #include "mongo/scripting/mozjs/jsstringwrapper.h"
 #include "mongo/scripting/mozjs/objectwrapper.h"
 #include "mongo/scripting/mozjs/valuereader.h"
+#include "mongo/scripting/mozjs/valuewriter.h"
 #include "mongo/util/version.h"
 
 namespace mongo {
 namespace mozjs {
 
-const JSFunctionSpec GlobalInfo::freeFunctions[6] = {
+const JSFunctionSpec GlobalInfo::freeFunctions[7] = {
+    MONGO_ATTACH_JS_FUNCTION(sleep),
     MONGO_ATTACH_JS_FUNCTION(gc),
     MONGO_ATTACH_JS_FUNCTION(print),
     MONGO_ATTACH_JS_FUNCTION(version),
@@ -110,6 +112,18 @@ void GlobalInfo::Functions::gc::call(JSContext* cx, JS::CallArgs args) {
     args.rval().setUndefined();
 }
 
+void GlobalInfo::Functions::sleep::call(JSContext* cx, JS::CallArgs args) {
+    uassert(16259,
+            "sleep takes a single numeric argument -- sleep(milliseconds)",
+            args.length() == 1 && args.get(0).isNumber());
+
+    auto scope = getScope(cx);
+    int64_t duration = ValueWriter(cx, args.get(0)).toInt64();
+    scope->sleep(Milliseconds(duration));
+
+    args.rval().setUndefined();
+}
+
 MONGO_INITIALIZER(JavascriptPrintDomain)(InitializerContext*) {
     jsPrintLogDomain = logger::globalLogManager()->getNamedDomain("javascriptOutput");
     return Status::OK();
diff --git a/src/mongo/scripting/mozjs/global.h b/src/mongo/scripting/mozjs/global.h
index b77d3b9855f..b5c89975933 100644
--- a/src/mongo/scripting/mozjs/global.h
+++ b/src/mongo/scripting/mozjs/global.h
@@ -41,6 +41,7 @@ namespace mozjs {
  */
 struct GlobalInfo : public BaseInfo {
     struct Functions {
+        MONGO_DECLARE_JS_FUNCTION(sleep);
         MONGO_DECLARE_JS_FUNCTION(gc);
         MONGO_DECLARE_JS_FUNCTION(print);
         MONGO_DECLARE_JS_FUNCTION(version);
@@ -48,7 +49,7 @@ struct GlobalInfo : public BaseInfo {
         MONGO_DECLARE_JS_FUNCTION(getJSHeapLimitMB);
     };
 
-    static const JSFunctionSpec freeFunctions[6];
+    static const JSFunctionSpec freeFunctions[7];
 
     static const char* const className;
     static const unsigned classFlags = JSCLASS_GLOBAL_FLAGS;
diff --git a/src/mongo/scripting/mozjs/implscope.cpp b/src/mongo/scripting/mozjs/implscope.cpp
index 2c120610008..249dc489fb1 100644
--- a/src/mongo/scripting/mozjs/implscope.cpp
+++ b/src/mongo/scripting/mozjs/implscope.cpp
@@ -192,7 +192,11 @@ void MozJSImplScope::unregisterOperation() {
 }
 
 void MozJSImplScope::kill() {
-    _pendingKill.store(true);
+    {
+        std::unique_lock<std::mutex> lk(_sleepMutex);
+        _pendingKill.store(true);
+    }
+    _sleepCondition.notify_all();
     JS_RequestInterruptCallback(_runtime);
 }
 
@@ -794,6 +798,12 @@ void MozJSImplScope::gc() {
     JS_RequestInterruptCallback(_runtime);
 }
 
+void MozJSImplScope::sleep(Milliseconds ms) {
+    std::unique_lock<std::mutex> lk(_sleepMutex);
+
+    _sleepCondition.wait_for(lk, ms.toSystemDuration(), [this] { return _pendingKill.load(); });
+}
+
 void MozJSImplScope::localConnectForDbEval(OperationContext* opCtx, const char* dbName) {
 
     _runSafely([this, &opCtx, &dbName] {
diff --git a/src/mongo/scripting/mozjs/implscope.h b/src/mongo/scripting/mozjs/implscope.h
index cba83cba184..6017ca55b3c 100644
--- a/src/mongo/scripting/mozjs/implscope.h
+++ b/src/mongo/scripting/mozjs/implscope.h
@@ -113,6 +113,8 @@ public:
 
     void gc() override;
 
+    void sleep(Milliseconds ms);
+
     bool isJavaScriptProtectionEnabled() const;
 
     double getNumber(const char* field) override;
@@ -418,6 +420,8 @@ private:
     std::vector<JS::PersistentRootedValue> _funcs;
     InternedStringTable _internedStrings;
     std::atomic<bool> _pendingKill;
+    std::mutex _sleepMutex;
+    std::condition_variable _sleepCondition;
     std::string _error;
     unsigned int _opId;        // op id for this scope
     OperationContext* _opCtx;  // Op context for DbEval
diff --git a/src/mongo/scripting/utils.cpp b/src/mongo/scripting/utils.cpp
index e9c2e4aa477..68aa70826f7 100644
--- a/src/mongo/scripting/utils.cpp
+++ b/src/mongo/scripting/utils.cpp
@@ -50,17 +50,6 @@ static BSONObj native_hex_md5(const BSONObj& args, void* data) {
     return BSON("" << digestToString(d));
 }
 
-static BSONObj native_sleep(const mongo::BSONObj& args, void* data) {
-    uassert(16259,
-            "sleep takes a single numeric argument -- sleep(milliseconds)",
-            args.nFields() == 1 && args.firstElement().isNumber());
-    sleepmillis(static_cast<long long>(args.firstElement().number()));
-
-    BSONObjBuilder b;
-    b.appendUndefined("");
-    return b.obj();
-}
-
 static BSONObj native_tostrictjson(const mongo::BSONObj& args, void* data) {
     uassert(40275,
             "tostrictjson takes a single BSON object argument, and on optional boolean argument "
@@ -81,7 +70,6 @@ static BSONObj native_tostrictjson(const mongo::BSONObj& args, void* data) {
 
 void installGlobalUtils(Scope& scope) {
     scope.injectNative("hex_md5", native_hex_md5);
-    scope.injectNative("sleep", native_sleep);
     scope.injectNative("tostrictjson", native_tostrictjson);
 }