SERVER-74131 Inject emit function on every call to ExpressionInternalJsEmit::evaluate()

4 files changed, 61 insertions, 11 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index bb2261f573c..ba6c8ffa2ee 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -326,6 +326,8 @@ last-continuous:
     ticket: SERVER-73229
   - test_file: jstests/noPassthrough/ttl_operation_metrics_multi_dbs.js
     ticket: SERVER-75277
+  - test_file: jstests/core/map_reduce_subplanning.js
+    ticket: SERVER-74131
   suites: null
 last-lts:
   all:
@@ -723,4 +725,6 @@ last-lts:
     ticket: SERVER-73229
   - test_file: jstests/noPassthrough/ttl_operation_metrics_multi_dbs.js
     ticket: SERVER-75277
+  - test_file: jstests/core/map_reduce_subplanning.js
+    ticket: SERVER-74131
   suites: null
diff --git a/jstests/core/map_reduce_subplanning.js b/jstests/core/map_reduce_subplanning.js
new file mode 100644
index 00000000000..35fed994ba8
--- /dev/null
+++ b/jstests/core/map_reduce_subplanning.js
@@ -0,0 +1,47 @@
+// The test runs commands that are not allowed with security token: mapReduce.
+// @tags: [
+//   not_allowed_with_security_token,
+//   does_not_support_stepdowns,
+//   requires_fastcount,
+//   requires_getmore,
+//   requires_non_retryable_writes,
+//   # This test has statements that do not support non-local read concern.
+//   does_not_support_causal_consistency,
+//   # Uses mapReduce command.
+//   requires_scripting,
+// ]
+
+load("jstests/aggregation/extras/utils.js");  // For resultsEq
+(function() {
+"use strict";
+
+const coll = db.map_reduce_subplanning;
+coll.drop();
+db.getCollection("mrOutput").drop();
+
+coll.createIndex({a: 1, c: 1});
+coll.createIndex({b: 1, c: 1});
+coll.createIndex({a: 1});
+coll.createIndex({b: 1});
+
+assert.commandWorked(coll.insert({a: 2}));
+assert.commandWorked(coll.insert({b: 3}));
+assert.commandWorked(coll.insert({b: 3}));
+assert.commandWorked(coll.insert({a: 2, b: 3}));
+
+assert.commandWorked(coll.mapReduce(
+    function() {
+        if (!this.hasOwnProperty('a')) {
+            emit('a', 0);
+        } else {
+            emit('a', this.a);
+        }
+    },
+    function(key, vals) {
+        return vals.reduce((a, b) => a + b, 0);
+    },
+    {out: {merge: "mrOutput"}, query: {$or: [{a: 2}, {b: 3}]}}));
+
+assert(resultsEq([{"_id": "a", "value": 4}], db.getCollection("mrOutput").find().toArray()),
+       db.getCollection("mrOutput").find().toArray());
+})();
diff --git a/src/mongo/db/pipeline/expression_js_emit.cpp b/src/mongo/db/pipeline/expression_js_emit.cpp
index 2a2b9e3d462..8208c586ad8 100644
--- a/src/mongo/db/pipeline/expression_js_emit.cpp
+++ b/src/mongo/db/pipeline/expression_js_emit.cpp
@@ -141,11 +141,14 @@ Value ExpressionInternalJsEmit::evaluate(const Document& root, Variables* variab
     ExpressionContext* expCtx = getExpressionContext();
 
     auto jsExec = expCtx->getJsExecWithScope();
-    // Inject the native "emit" function to be called from the user-defined map function. This
-    // particular Expression/ExpressionContext may be reattached to a new OperationContext (and thus
-    // a new JS Scope) when used across getMore operations, so this method will handle that case for
-    // us by only injecting if we haven't already.
-    jsExec->injectEmitIfNecessary(emitFromJS, &_emitState);
+
+    // Inject the native "emit" function to be called from the user-defined map function.
+    //
+    // We reinject this function on every invocation of evaluate(), because there is a single
+    // JsExecution instance for the OperationContext, which may be shared by multiple aggregation
+    // pipelines and we need to ensure that the injected function still points to the valid
+    // contextual data ('_emitState').
+    jsExec->injectEmit(emitFromJS, &_emitState);
 
     // Although inefficient to "create" a new function every time we evaluate, this will usually end
     // up being a simple cache lookup. This is needed because the JS Scope may have been recreated
diff --git a/src/mongo/db/pipeline/javascript_execution.h b/src/mongo/db/pipeline/javascript_execution.h
index 924ae447216..d95e90f6eea 100644
--- a/src/mongo/db/pipeline/javascript_execution.h
+++ b/src/mongo/db/pipeline/javascript_execution.h
@@ -104,11 +104,8 @@ public:
      * Injects the given function 'emitFn' as a native JS function named 'emit', callable from
      * user-defined functions.
      */
-    void injectEmitIfNecessary(NativeFunction emitFn, void* data) {
-        if (!_emitCreated) {
-            _scope->injectNative("emit", emitFn, data);
-            _emitCreated = true;
-        }
+    void injectEmit(NativeFunction emitFn, void* data) {
+        _scope->injectNative("emit", emitFn, data);
     }
 
     Scope* getScope() {
@@ -118,7 +115,6 @@ public:
 private:
     BSONObj _scopeVars;
     std::unique_ptr<Scope> _scope;
-    bool _emitCreated = false;
     bool _storedProceduresLoaded = false;
     int _fnCallTimeoutMillis;