SERVER-22688 Provide a setParameter and shell option to control the js heap limit

9 files changed, 48 insertions, 13 deletions

diff --git a/src/mongo/scripting/engine.h b/src/mongo/scripting/engine.h
index df6760f9014..62c3e5da946 100644
--- a/src/mongo/scripting/engine.h
+++ b/src/mongo/scripting/engine.h
@@ -249,6 +249,9 @@ public:
     virtual void enableJavaScriptProtection(bool value) = 0;
     virtual bool isJavaScriptProtectionEnabled() const = 0;
 
+    virtual int getJSHeapLimitMB() const = 0;
+    virtual void setJSHeapLimitMB(int limit) = 0;
+
     static void setup();
     static void dropScopeCache();
 
diff --git a/src/mongo/scripting/mozjs/engine.cpp b/src/mongo/scripting/mozjs/engine.cpp
index 90038dbea35..7d1e59d7973 100644
--- a/src/mongo/scripting/mozjs/engine.cpp
+++ b/src/mongo/scripting/mozjs/engine.cpp
@@ -50,6 +50,7 @@ namespace {
 
 MONGO_EXPORT_SERVER_PARAMETER(disableJavaScriptJIT, bool, false);
 MONGO_EXPORT_SERVER_PARAMETER(javascriptProtection, bool, false);
+MONGO_EXPORT_SERVER_PARAMETER(jsHeapLimitMB, int, 1100);
 
 }  // namespace
 
@@ -137,6 +138,14 @@ bool MozJSScriptEngine::isJavaScriptProtectionEnabled() const {
     return javascriptProtection.load();
 }
 
+int MozJSScriptEngine::getJSHeapLimitMB() const {
+    return jsHeapLimitMB.load();
+}
+
+void MozJSScriptEngine::setJSHeapLimitMB(int limit) {
+    jsHeapLimitMB.store(limit);
+}
+
 void MozJSScriptEngine::registerOperation(OperationContext* txn, MozJSImplScope* scope) {
     stdx::lock_guard<stdx::mutex> giLock(_globalInterruptLock);
 
diff --git a/src/mongo/scripting/mozjs/engine.h b/src/mongo/scripting/mozjs/engine.h
index 8d6cf201e5f..84607597ccf 100644
--- a/src/mongo/scripting/mozjs/engine.h
+++ b/src/mongo/scripting/mozjs/engine.h
@@ -69,6 +69,9 @@ public:
     void enableJavaScriptProtection(bool value) override;
     bool isJavaScriptProtectionEnabled() const override;
 
+    int getJSHeapLimitMB() const override;
+    void setJSHeapLimitMB(int limit) override;
+
     void registerOperation(OperationContext* ctx, MozJSImplScope* scope);
     void unregisterOperation(unsigned int opId);
 
diff --git a/src/mongo/scripting/mozjs/global.cpp b/src/mongo/scripting/mozjs/global.cpp
index 105d9478dba..16e7f93ebab 100644
--- a/src/mongo/scripting/mozjs/global.cpp
+++ b/src/mongo/scripting/mozjs/global.cpp
@@ -35,6 +35,7 @@
 #include "mongo/base/init.h"
 #include "mongo/logger/logger.h"
 #include "mongo/logger/logstream_builder.h"
+#include "mongo/scripting/engine.h"
 #include "mongo/scripting/mozjs/implscope.h"
 #include "mongo/scripting/mozjs/jsstringwrapper.h"
 #include "mongo/scripting/mozjs/objectwrapper.h"
@@ -44,11 +45,12 @@
 namespace mongo {
 namespace mozjs {
 
-const JSFunctionSpec GlobalInfo::freeFunctions[5] = {
+const JSFunctionSpec GlobalInfo::freeFunctions[6] = {
     MONGO_ATTACH_JS_FUNCTION(gc),
     MONGO_ATTACH_JS_FUNCTION(print),
     MONGO_ATTACH_JS_FUNCTION(version),
     MONGO_ATTACH_JS_FUNCTION(buildInfo),
+    MONGO_ATTACH_JS_FUNCTION(getJSHeapLimitMB),
     JS_FS_END,
 };
 
@@ -96,6 +98,10 @@ void GlobalInfo::Functions::buildInfo::call(JSContext* cx, JS::CallArgs args) {
     ValueReader(cx, args.rval()).fromBSON(b.obj(), nullptr, false);
 }
 
+void GlobalInfo::Functions::getJSHeapLimitMB::call(JSContext* cx, JS::CallArgs args) {
+    ValueReader(cx, args.rval()).fromDouble(mongo::globalScriptEngine->getJSHeapLimitMB());
+}
+
 void GlobalInfo::Functions::gc::call(JSContext* cx, JS::CallArgs args) {
     auto scope = getScope(cx);
 
diff --git a/src/mongo/scripting/mozjs/global.h b/src/mongo/scripting/mozjs/global.h
index b90fa3c7682..b77d3b9855f 100644
--- a/src/mongo/scripting/mozjs/global.h
+++ b/src/mongo/scripting/mozjs/global.h
@@ -45,9 +45,10 @@ struct GlobalInfo : public BaseInfo {
         MONGO_DECLARE_JS_FUNCTION(print);
         MONGO_DECLARE_JS_FUNCTION(version);
         MONGO_DECLARE_JS_FUNCTION(buildInfo);
+        MONGO_DECLARE_JS_FUNCTION(getJSHeapLimitMB);
     };
 
-    static const JSFunctionSpec freeFunctions[5];
+    static const JSFunctionSpec freeFunctions[6];
 
     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 399a1e532f9..32562846c94 100644
--- a/src/mongo/scripting/mozjs/implscope.cpp
+++ b/src/mongo/scripting/mozjs/implscope.cpp
@@ -67,13 +67,6 @@ const char* const MozJSImplScope::kInvokeResult = "__returnValue";
 namespace {
 
 /**
- * The maximum amount of memory to be given out per thread to mozilla. We
- * manage this by trapping all calls to malloc, free, etc. and keeping track of
- * counts in some thread locals
- */
-const size_t kMallocMemoryLimit = 1024ul * 1024 * 1024 * 1.1;
-
-/**
  * The threshold (as a fraction of the max) after which garbage collection will be run during
  * interrupts.
  */
@@ -251,7 +244,13 @@ void MozJSImplScope::_gcCallback(JSRuntime* rt, JSGCStatus status, void* data) {
 }
 
 MozJSImplScope::MozRuntime::MozRuntime(const MozJSScriptEngine* engine) {
-    mongo::sm::reset(kMallocMemoryLimit);
+    /**
+     * The maximum amount of memory to be given out per thread to mozilla. We
+     * manage this by trapping all calls to malloc, free, etc. and keeping track of
+     * counts in some thread locals
+     */
+    size_t mallocMemoryLimit = 1024ul * 1024 * engine->getJSHeapLimitMB();
+    mongo::sm::reset(mallocMemoryLimit);
 
     // If this runtime isn't running on an NSPR thread, then it is
     // running on a mongo thread. In that case, we need to insert a
@@ -318,7 +317,7 @@ MozJSImplScope::MozRuntime::MozRuntime(const MozJSScriptEngine* engine) {
         }
 
         // The memory limit is in megabytes
-        JS_SetGCParametersBasedOnAvailableMemory(_runtime, kMallocMemoryLimit / (1024 * 1024));
+        JS_SetGCParametersBasedOnAvailableMemory(_runtime, engine->getJSHeapLimitMB());
     }
 
     _context = JS_NewContext(_runtime, kStackChunkSize);
@@ -780,8 +779,6 @@ void MozJSImplScope::externalSetup() {
     if (_connectState == ConnectState::Local)
         uasserted(12512, "localConnect already called, can't call externalSetup");
 
-    mongo::sm::reset(0);
-
     // install db access functions in the global object
     installDBAccess();
 
diff --git a/src/mongo/shell/dbshell.cpp b/src/mongo/shell/dbshell.cpp
index 5df9fc23679..06dac12f49c 100644
--- a/src/mongo/shell/dbshell.cpp
+++ b/src/mongo/shell/dbshell.cpp
@@ -708,6 +708,7 @@ int _main(int argc, char* argv[], char** envp) {
 
     mongo::ScriptEngine::setConnectCallback(mongo::shell_utils::onConnect);
     mongo::ScriptEngine::setup();
+    mongo::globalScriptEngine->setJSHeapLimitMB(shellGlobalParams.jsHeapLimitMB);
     mongo::globalScriptEngine->setScopeInitCallback(mongo::shell_utils::initScope);
     mongo::globalScriptEngine->enableJIT(!shellGlobalParams.nojit);
     mongo::globalScriptEngine->enableJavaScriptProtection(shellGlobalParams.javascriptProtection);
diff --git a/src/mongo/shell/shell_options.cpp b/src/mongo/shell/shell_options.cpp
index 1b378836c35..3a2b55da126 100644
--- a/src/mongo/shell/shell_options.cpp
+++ b/src/mongo/shell/shell_options.cpp
@@ -200,6 +200,9 @@ Status addMongoShellOptions(moe::OptionSection* options) {
     if (!ret.isOK())
         return ret;
 
+    options->addOptionChaining(
+        "jsHeapLimitMB", "jsHeapLimitMB", moe::Int, "set the js scope's heap size limit");
+
     return Status::OK();
 }
 
@@ -370,6 +373,16 @@ Status storeMongoShellOptions(const moe::Environment& params,
         }
     }
 
+    if (params.count("jsHeapLimitMB")) {
+        int jsHeapLimitMB = params["jsHeapLimitMB"].as<int>();
+        if (jsHeapLimitMB <= 0) {
+            StringBuilder sb;
+            sb << "ERROR: \"jsHeapLimitMB\" needs to be greater than 0";
+            return Status(ErrorCodes::BadValue, sb.str());
+        }
+        shellGlobalParams.jsHeapLimitMB = jsHeapLimitMB;
+    }
+
     if (shellGlobalParams.url == "*") {
         StringBuilder sb;
         sb << "ERROR: "
diff --git a/src/mongo/shell/shell_options.h b/src/mongo/shell/shell_options.h
index 62191bfb0c9..52b02ce1d34 100644
--- a/src/mongo/shell/shell_options.h
+++ b/src/mongo/shell/shell_options.h
@@ -73,6 +73,8 @@ struct ShellGlobalParams {
     std::string readMode = "compatibility";
 
     boost::optional<rpc::ProtocolSet> rpcProtocols = boost::none;
+
+    int jsHeapLimitMB = 0;
 };
 
 extern ShellGlobalParams shellGlobalParams;