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

3 files changed, 16 insertions, 0 deletions

diff --git a/src/mongo/shell/dbshell.cpp b/src/mongo/shell/dbshell.cpp
index 5df9fc23679..768c8558efd 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->setJSHeapMBLimit(shellGlobalParams.jsHeapMBLimit);
     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..324ad87fa3d 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(
+        "jsHeapMBLimit", "jsHeapMBLimit", 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("jsHeapMBLimit")) {
+        int jsHeapMBLimit = params["jsHeapMBLimit"].as<int>();
+        if (jsHeapMBLimit <= 0) {
+            StringBuilder sb;
+            sb << "ERROR: \"jsHeapMBLimit\" needs to be greater than 0";
+            return Status(ErrorCodes::BadValue, sb.str());
+        }
+        shellGlobalParams.jsHeapMBLimit = jsHeapMBLimit;
+    }
+
     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..296567025e3 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 jsHeapMBLimit = 0;
 };
 
 extern ShellGlobalParams shellGlobalParams;