SERVER-46494 improve perf of '$function'

8 files changed, 133 insertions, 16 deletions

diff --git a/etc/perf.yml b/etc/perf.yml
index 88d8b95b94f..bbe3550c448 100644
--- a/etc/perf.yml
+++ b/etc/perf.yml
@@ -531,7 +531,7 @@ tasks:
       vars:
         includeFilter1: "aggregation"
         includeFilter2: "regression"
-        excludeFilter: "none"
+        excludeFilter: "js"
         threads: "1"
         readCmd: false
     - func: "analyze"
@@ -571,7 +571,23 @@ tasks:
     - func: "analyze"
       vars:
         reports_analysis: true
-
+- name: javascript
+  depends_on:
+  - variant: linux-wt-standalone
+    name: compile
+  commands:
+    - func: "git get project"
+    - func: "start server"
+    - func: "run perf tests"
+      vars:
+        includeFilter1: "js"
+        includeFilter2: "aggregation"
+        excludeFilter: "none"
+        threads: "1 2 4 8"
+        readCmd: true
+    - func: "analyze"
+      vars:
+        reports_analysis: true
 buildvariants:
 - name: linux-wt-standalone
   display_name: Standalone Linux inMemory
@@ -622,6 +638,7 @@ buildvariants:
   - name: misc
   - name: singleThreaded
   - name: pipeline-updates
+  - name: javascript
 
 - name: linux-wt-repl
   display_name: 1-Node ReplSet Linux inMemory
diff --git a/jstests/noPassthrough/expression_function_kill.js b/jstests/noPassthrough/expression_function_kill.js
new file mode 100644
index 00000000000..398f08ce4a5
--- /dev/null
+++ b/jstests/noPassthrough/expression_function_kill.js
@@ -0,0 +1,82 @@
+/**
+ * Tests where/function can be interrupted through maxTimeMS and query knob.
+ */
+(function() {
+"use strict";
+
+const mongodOptions = {};
+const conn = MongoRunner.runMongod(mongodOptions);
+
+let db = conn.getDB("test_where_function_interrupt");
+let coll = db.getCollection("foo");
+
+let expensiveFunction = function() {
+    sleep(1000);
+    return true;
+};
+assert.commandWorked(coll.insert(Array.from({length: 1000}, _ => ({}))));
+
+let checkInterrupt = function(cursor) {
+    let err = assert.throws(function() {
+        cursor.itcount();
+    }, [], "expected interrupt error due to maxTimeMS being exceeded");
+    assert.commandFailedWithCode(
+        err, [ErrorCodes.MaxTimeMSExpired, ErrorCodes.Interrupted, ErrorCodes.InternalError]);
+};
+
+let tests = [
+    {
+        // Test that $where can be interrupted with a maxTimeMS of 100 ms.
+        timeout: 100,
+        query: {$where: expensiveFunction},
+        err: checkInterrupt,
+    },
+    {
+        // Test that $function can be interrupted with a maxTimeMS of 100 ms.
+        timeout: 100,
+        query: {
+            $expr: {
+                $function: {
+                    body: expensiveFunction,
+                    args: [],
+                    lang: 'js',
+                }
+            }
+        },
+        err: checkInterrupt
+    },
+    {
+
+        // Test that $function can be interrupted by a query knob of 100 ms.
+        pre: function() {
+            assert.commandWorked(
+                db.adminCommand({setParameter: 1, internalQueryJavaScriptFnTimeoutMillis: 100}));
+        },
+        query: {
+            $expr: {
+                $function: {
+                    body: expensiveFunction,
+                    args: [],
+                    lang: 'js',
+                }
+            }
+        },
+        err: checkInterrupt
+    },
+];
+
+tests.forEach(function(testCase) {
+    if (testCase.pre) {
+        testCase.pre();
+    }
+
+    let cursor = coll.find(testCase.query);
+
+    if (testCase.timeout) {
+        cursor.maxTimeMS(testCase.timeout);
+    }
+    testCase.err(cursor);
+});
+
+MongoRunner.stopMongod(conn);
+})();
diff --git a/src/mongo/db/pipeline/expression_function.cpp b/src/mongo/db/pipeline/expression_function.cpp
index 3db3c19d239..ff5dcd4e29f 100644
--- a/src/mongo/db/pipeline/expression_function.cpp
+++ b/src/mongo/db/pipeline/expression_function.cpp
@@ -111,6 +111,8 @@ Value ExpressionFunction::evaluate(const Document& root, Variables* variables) c
     auto jsExec = getExpressionContext()->getJsExecWithScope(_assignFirstArgToThis);
     auto scope = jsExec->getScope();
 
+    // createFunction is memoized in MozJSImplScope, so it's ok to call this for each
+    // eval call.
     ScriptingFunction func = jsExec->getScope()->createFunction(_funcSource.c_str());
     uassert(31265, "The body function did not evaluate", func);
 
diff --git a/src/mongo/db/pipeline/javascript_execution.cpp b/src/mongo/db/pipeline/javascript_execution.cpp
index 6ae5ccd9ea5..f9ce429ba76 100644
--- a/src/mongo/db/pipeline/javascript_execution.cpp
+++ b/src/mongo/db/pipeline/javascript_execution.cpp
@@ -49,7 +49,7 @@ JsExecution* JsExecution::get(OperationContext* opCtx,
                               boost::optional<int> jsHeapLimitMB) {
     auto& exec = getExec(opCtx);
     if (!exec) {
-        exec = std::make_unique<JsExecution>(scope, jsHeapLimitMB);
+        exec = std::make_unique<JsExecution>(opCtx, scope, jsHeapLimitMB);
         exec->getScope()->setLocalDB(database);
         if (loadStoredProcedures) {
             exec->getScope()->loadStored(opCtx, true);
@@ -66,9 +66,6 @@ JsExecution* JsExecution::get(OperationContext* opCtx,
 Value JsExecution::callFunction(ScriptingFunction func,
                                 const BSONObj& params,
                                 const BSONObj& thisObj) {
-    _scope->registerOperation(Client::getCurrent()->getOperationContext());
-    const auto guard = makeGuard([&] { _scope->unregisterOperation(); });
-
     int err = _scope->invoke(func, &params, &thisObj, _fnCallTimeoutMillis, false);
     uassert(
         31439, str::stream() << "js function failed to execute: " << _scope->getError(), err == 0);
@@ -81,9 +78,6 @@ Value JsExecution::callFunction(ScriptingFunction func,
 void JsExecution::callFunctionWithoutReturn(ScriptingFunction func,
                                             const BSONObj& params,
                                             const BSONObj& thisObj) {
-    _scope->registerOperation(Client::getCurrent()->getOperationContext());
-    const auto guard = makeGuard([&] { _scope->unregisterOperation(); });
-
     int err = _scope->invoke(func, &params, &thisObj, _fnCallTimeoutMillis, true);
     uassert(
         31470, str::stream() << "js function failed to execute: " << _scope->getError(), err == 0);
diff --git a/src/mongo/db/pipeline/javascript_execution.h b/src/mongo/db/pipeline/javascript_execution.h
index 163d1fcc974..924ae447216 100644
--- a/src/mongo/db/pipeline/javascript_execution.h
+++ b/src/mongo/db/pipeline/javascript_execution.h
@@ -61,16 +61,23 @@ public:
     /**
      * Construct with a thread-local scope and initialize with the given scope variables.
      */
-    explicit JsExecution(const BSONObj& scopeVars, boost::optional<int> jsHeapLimitMB = boost::none)
+    JsExecution(OperationContext* opCtx,
+                const BSONObj& scopeVars,
+                boost::optional<int> jsHeapLimitMB = boost::none)
         : _scope(getGlobalScriptEngine()->newScopeForCurrentThread(jsHeapLimitMB)) {
         _scopeVars = scopeVars.getOwned();
         _scope->init(&_scopeVars);
         _fnCallTimeoutMillis = internalQueryJavaScriptFnTimeoutMillis.load();
+        _scope->registerOperation(opCtx);
     }
 
+    ~JsExecution() {
+        _scope->unregisterOperation();
+    };
+
     /**
-     * Registers and invokes the javascript function given by 'func' with the arguments 'params' and
-     * input object 'thisObj'.
+     * Invokes the javascript function given by 'func' with the arguments 'params' and input object
+     * 'thisObj'.
      *
      * This method assumes that the desired function to execute does not return a value.
      */
@@ -79,14 +86,21 @@ public:
                                    const BSONObj& thisObj);
 
     /**
-     * Registers and invokes the javascript function given by 'func' with the arguments 'params' and
-     * input object 'thisObj'.
+     * Invokes the javascript function given by 'func' with the arguments 'params' and input object
+     * 'thisObj'.
      *
      * Returns the value returned by the function.
      */
     Value callFunction(ScriptingFunction func, const BSONObj& params, const BSONObj& thisObj);
 
     /**
+     * Creates a function in the owned Scope* if it hasn't been created yet.
+     */
+    ScriptingFunction createFunction(std::string funcCode) {
+        return _scope->createFunction(funcCode.c_str());
+    };
+
+    /**
      * Injects the given function 'emitFn' as a native JS function named 'emit', callable from
      * user-defined functions.
      */
diff --git a/src/mongo/db/pipeline/make_js_function.cpp b/src/mongo/db/pipeline/make_js_function.cpp
index eacb6d82d36..531f0bbab74 100644
--- a/src/mongo/db/pipeline/make_js_function.cpp
+++ b/src/mongo/db/pipeline/make_js_function.cpp
@@ -37,7 +37,8 @@ namespace mongo {
 // it as a JS function.
 ScriptingFunction makeJsFunc(boost::intrusive_ptr<ExpressionContext> expCtx,
                              const std::string& func) {
-    auto jsExec = expCtx->getJsExecWithScope();
+    auto jsExec =
+        expCtx->getJsExecWithScope();  // default arg forceLoadOfStoredProcedures is false here.
     ScriptingFunction parsedFunc = jsExec->getScope()->createFunction(func.c_str());
     uassert(
         31247, "The user-defined function failed to parse in the javascript engine", parsedFunc);
diff --git a/src/mongo/scripting/mozjs/implscope.cpp b/src/mongo/scripting/mozjs/implscope.cpp
index f06a20bfc0a..32883343ab0 100644
--- a/src/mongo/scripting/mozjs/implscope.cpp
+++ b/src/mongo/scripting/mozjs/implscope.cpp
@@ -630,9 +630,14 @@ bool hasFunctionIdentifier(StringData code) {
 ScriptingFunction MozJSImplScope::_createFunction(const char* raw) {
     return _runSafely([&] {
         JS::RootedValue fun(_context);
+        auto it = _funcCodeToHandleMap.find(StringData(raw));
+        if (it != _funcCodeToHandleMap.end()) {
+            return it->second;
+        }
         _MozJSCreateFunction(raw, &fun);
         _funcs.emplace_back(_context, fun.get());
-        return _funcs.size();
+        _funcCodeToHandleMap.emplace(raw, _funcs.size());
+        return ScriptingFunction(_funcs.size());
     });
 }
 
diff --git a/src/mongo/scripting/mozjs/implscope.h b/src/mongo/scripting/mozjs/implscope.h
index b842ae3cbef..86aa6c5dd05 100644
--- a/src/mongo/scripting/mozjs/implscope.h
+++ b/src/mongo/scripting/mozjs/implscope.h
@@ -67,6 +67,7 @@
 #include "mongo/scripting/mozjs/timestamp.h"
 #include "mongo/scripting/mozjs/uri.h"
 #include "mongo/stdx/unordered_set.h"
+#include "mongo/util/string_map.h"
 
 namespace mongo {
 namespace mozjs {
@@ -411,6 +412,7 @@ private:
     WrapType<GlobalInfo> _globalProto;
     JS::HandleObject _global;
     std::vector<JS::PersistentRootedValue> _funcs;
+    StringMap<ScriptingFunction> _funcCodeToHandleMap;
     InternedStringTable _internedStrings;
     Status _killStatus;
     mutable Mutex _mutex = MONGO_MAKE_LATCH("MozJSImplScope::_mutex");