SERVER-43303 Re introduce mapReduce fsm tests in sharded suites

4 files changed, 121 insertions, 13 deletions

diff --git a/jstests/concurrency/fsm_workloads/map_reduce_inline.js b/jstests/concurrency/fsm_workloads/map_reduce_inline.js
index 94382b1fe70..268fc1460d7 100644
--- a/jstests/concurrency/fsm_workloads/map_reduce_inline.js
+++ b/jstests/concurrency/fsm_workloads/map_reduce_inline.js
@@ -61,25 +61,23 @@ var $config = (function() {
 
     var transitions = {init: {mapReduce: 1}, mapReduce: {mapReduce: 1}};
 
-    function makeDoc(keyLimit, valueLimit) {
-        return {
-            _id: new ObjectId(),
-            key: Random.randInt(keyLimit),
-            value: Random.randInt(valueLimit)
-        };
-    }
-
     function setup(db, collName, cluster) {
         var bulk = db[collName].initializeUnorderedBulkOp();
         for (var i = 0; i < this.numDocs; ++i) {
             // TODO: this actually does assume that there are no unique indexes
-            var doc = makeDoc(this.numDocs / 100, this.numDocs / 10);
-            bulk.insert(doc);
+            bulk.insert({
+                _id: i,
+                key: Random.randInt(this.numDocs / 100),
+                value: Random.randInt(this.numDocs / 10)
+            });
         }
 
         var res = bulk.execute();
         assertAlways.commandWorked(res);
         assertAlways.eq(this.numDocs, res.nInserted);
+
+        assert.commandWorked(
+            db.adminCommand({setParameter: 1, internalQueryUseAggMapReduce: true}));
     }
 
     return {
diff --git a/jstests/concurrency/fsm_workloads/map_reduce_reduce.js b/jstests/concurrency/fsm_workloads/map_reduce_reduce.js
index 86d3733df9c..99207cd5832 100644
--- a/jstests/concurrency/fsm_workloads/map_reduce_reduce.js
+++ b/jstests/concurrency/fsm_workloads/map_reduce_reduce.js
@@ -28,6 +28,12 @@ var $config = extendWorkload($config, function($config, $super) {
     }
 
     $config.states.init = function init(db, collName) {
+        // TODO SERVER-44150: Cannot run MR with output 'reduce' in agg since the 'whenMatched'
+        // pipeline will always run which can cause unexpected failures in the user-specified reduce
+        // function.
+        assert.commandWorked(
+            db.adminCommand({setParameter: 1, internalQueryUseAggMapReduce: false}));
+
         $super.states.init.apply(this, arguments);
 
         this.outCollName = uniqueCollectionName(prefix, this.tid);
@@ -40,7 +46,6 @@ var $config = extendWorkload($config, function($config, $super) {
                      "output collection '" + fullName + "' should exist");
 
         var options = {finalize: this.finalizer, out: {reduce: this.outCollName}};
-
         var res = db[collName].mapReduce(this.mapper, this.reducer, options);
         assertAlways.commandWorked(res);
     };
diff --git a/jstests/concurrency/fsm_workloads/map_reduce_replace_remove.js b/jstests/concurrency/fsm_workloads/map_reduce_replace_remove.js
index 1bd48dce230..30225fc7c0f 100644
--- a/jstests/concurrency/fsm_workloads/map_reduce_replace_remove.js
+++ b/jstests/concurrency/fsm_workloads/map_reduce_replace_remove.js
@@ -21,8 +21,7 @@ load('jstests/concurrency/fsm_workloads/map_reduce_replace.js');  // for $config
 var $config = extendWorkload($config, function($config, $super) {
     $config.states.remove = function remove(db, collName) {
         for (var i = 0; i < 20; ++i) {
-            var res = db[collName].remove({value: {$gte: Random.randInt(this.numDocs / 10)}},
-                                          {justOne: true});
+            var res = db[collName].remove({_id: Random.randInt(this.numDocs)}, {justOne: true});
             assertAlways.commandWorked(res);
             assertAlways.lte(0, res.nRemoved, tojson(res));
         }
diff --git a/jstests/concurrency/fsm_workloads/map_reduce_with_chunk_migrations.js b/jstests/concurrency/fsm_workloads/map_reduce_with_chunk_migrations.js
new file mode 100644
index 00000000000..df4cee9c472
--- /dev/null
+++ b/jstests/concurrency/fsm_workloads/map_reduce_with_chunk_migrations.js
@@ -0,0 +1,106 @@
+'use strict';
+
+/**
+ * map_reduce_with_chunk_migrations.js
+ *
+ * This tests exercises mapReduce on a collection during chunk migrations. If extending this
+ * workload, consider overriding the following:
+ *
+ * $config.data.collWithMigrations: collection to run chunk migrations against (default is the
+ * input collection of the mapReduce).
+ * $config.state.mapReduce: function to execute the mapReduce.
+ *
+ * @tags: [
+ *  requires_sharding, assumes_balancer_off,
+ *  assumes_autosplit_off,
+ *  requires_non_retryable_writes,
+ *  # mapReduce does not support afterClusterTime.
+ *  does_not_support_causal_consistency
+ * ]
+ */
+load('jstests/concurrency/fsm_libs/extend_workload.js');                     // for extendWorkload
+load('jstests/concurrency/fsm_workloads/sharded_moveChunk_partitioned.js');  // for $config
+
+var $config = extendWorkload($config, function($config, $super) {
+    // The base setup will insert 'partitionSize' number of documents per thread, evenly
+    // distributing across the chunks. Documents will only have the "_id" field.
+    $config.data.partitionSize = 50;
+    $config.threadCount = 2;
+    $config.iterations = 100;
+    $config.data.numDocs = $config.data.partitionSize * $config.threadCount;
+
+    // By default, the collection that will be sharded with concurrent chunk migrations will be the
+    // one that the aggregate is run against.
+    $config.data.collWithMigrations = $config.collName;
+    $config.data.resultsCollection = "map_reduce_with_chunk_migrations_out";
+
+    $config.transitions = {
+        init: {mapReduce: 1},
+        mapReduce: {
+            moveChunk: 0.2,
+            mapReduce: 0.8,
+        },
+        moveChunk: {mapReduce: 1},
+    };
+
+    /**
+     * Moves a random chunk in the target collection.
+     */
+    $config.states.moveChunk = function moveChunk(db, collName, connCache) {
+        $super.states.moveChunk.apply(this, [db, this.collWithMigrations, connCache]);
+    };
+
+    /**
+     * Executes a mapReduce with output mode "replace".
+     */
+    $config.states.mapReduce = function mapReduce(db, collName, connCache) {
+        const map = function() {
+            emit(this._id, 1);
+        };
+        const reduce = function(k, values) {
+            return Array.sum(values);
+        };
+
+        const res = db[collName].mapReduce(map, reduce, {out: {replace: this.resultsCollection}});
+        assertWhenOwnColl.commandWorked(res);
+
+        // TODO SERVER-43290 Support for cluster stats should be able to enable this check.
+        // assertWhenOwnColl.eq(
+        //     this.numDocs, res.counts.output, `Expected each _id to be output once:
+        //     ${tojson(res)}`);
+    };
+
+    /**
+     * Uses the base class init() to initialize this thread for both collections.
+     */
+    $config.states.init = function init(db, collName, connCache) {
+        $super.states.init.apply(this, [db, collName, connCache]);
+
+        // Init the target collection in a similar manner, if it is different than the default
+        // collection.
+        if (collName != this.collWithMigrations) {
+            $super.states.init.apply(this, [db, this.collWithMigrations, connCache]);
+        }
+    };
+
+    /**
+     * Initializes the aggregate collection and the target collection for chunk migrations as
+     * sharded with an even distribution across each thread ID.
+     */
+    $config.setup = function setup(db, collName, cluster) {
+        $super.setup.apply(this, [db, collName, cluster]);
+
+        if (collName != this.collWithMigrations) {
+            // Setup the target collection in a similar manner. Note that the FSM infrastructure
+            // will have already enabled sharded on collName, but we need to manually do it for the
+            // output collection.
+            cluster.shardCollection(db[this.collWithMigrations], this.shardKey, false);
+            $super.setup.apply(this, [db, this.collWithMigrations, cluster]);
+        }
+
+        assert.commandWorked(
+            db.adminCommand({setParameter: 1, internalQueryUseAggMapReduce: true}));
+    };
+
+    return $config;
+});