SERVER-46494 improve perf of '$function'

(cherry picked from commit 16889524ede36d890f863b63447ba0af9d75013e)

6 files changed, 32 insertions, 14 deletions

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 34729b87a33..ec4ddf3d125 100644
--- a/src/mongo/scripting/mozjs/implscope.cpp
+++ b/src/mongo/scripting/mozjs/implscope.cpp
@@ -629,9 +629,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 c774efd3bd3..cd0701ea403 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 {
@@ -412,6 +413,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");