SERVER-40567 benchRun support for pipeline updates

This reverts commit 1b8a9f5dc5c3314042b55e7415a2a25045b32a94.

21 files changed, 159 insertions, 54 deletions

diff --git a/buildscripts/resmokeconfig/suites/multi_shard_local_read_write_multi_stmt_txn_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/multi_shard_local_read_write_multi_stmt_txn_jscore_passthrough.yml
index 7c55c988dfc..2406054200f 100644
--- a/buildscripts/resmokeconfig/suites/multi_shard_local_read_write_multi_stmt_txn_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/multi_shard_local_read_write_multi_stmt_txn_jscore_passthrough.yml
@@ -327,6 +327,7 @@ selector:
   # op3 will still be using the snapshot from op1, and not see op2 at all.
   - jstests/core/cursora.js
   - jstests/core/bench_test1.js
+  - jstests/core/benchrun_pipeline_updates.js
 
   # The tests below use applyOps, SERVER-1439.
   - jstests/core/list_collections1.js
diff --git a/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_jscore_passthrough.yml
index 362cdfa1a83..5b5b716b0dd 100644
--- a/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_jscore_passthrough.yml
@@ -354,6 +354,7 @@ selector:
   # op3 will still be using the snapshot from op1, and not see op2 at all.
   - jstests/core/cursora.js
   - jstests/core/bench_test1.js
+  - jstests/core/benchrun_pipeline_updates.js
 
   # The tests below use applyOps, SERVER-1439.
   - jstests/core/list_collections1.js
diff --git a/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_kill_primary_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_kill_primary_jscore_passthrough.yml
index 46375674947..a22bbe5e58f 100644
--- a/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_kill_primary_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_kill_primary_jscore_passthrough.yml
@@ -339,6 +339,7 @@ selector:
   # op3 will still be using the snapshot from op1, and not see op2 at all.
   - jstests/core/cursora.js
   - jstests/core/bench_test1.js
+  - jstests/core/benchrun_pipeline_updates.js
 
   # auto_retry_on_network_error.js will timeout with assert.soon and give a different error from
   # what test expects. Excluding from suite since it doesn't really do any database operations, so
diff --git a/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_stepdown_primary_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_stepdown_primary_jscore_passthrough.yml
index a15339fdcd3..1512af22c06 100644
--- a/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_stepdown_primary_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/multi_shard_multi_stmt_txn_stepdown_primary_jscore_passthrough.yml
@@ -348,6 +348,7 @@ selector:
   # op3 will still be using the snapshot from op1, and not see op2 at all.
   - jstests/core/cursora.js
   - jstests/core/bench_test1.js
+  - jstests/core/benchrun_pipeline_updates.js
 
   # auto_retry_on_network_error.js will timeout with assert.soon and give a different error from
   # what test expects. Excluding from suite since it doesn't really do any database operations, so
diff --git a/buildscripts/resmokeconfig/suites/multi_stmt_txn_jscore_passthrough_with_migration.yml b/buildscripts/resmokeconfig/suites/multi_stmt_txn_jscore_passthrough_with_migration.yml
index 4e664d4e999..3a0ed31639c 100644
--- a/buildscripts/resmokeconfig/suites/multi_stmt_txn_jscore_passthrough_with_migration.yml
+++ b/buildscripts/resmokeconfig/suites/multi_stmt_txn_jscore_passthrough_with_migration.yml
@@ -187,6 +187,7 @@ selector:
   # transaction and can cause test to timeout because of the 3 hr transactions lifetime setting and
   # moveChunk getting blocked trying to acquire collection X lock while doing chunk refresh.
   - jstests/core/bench_test*.js
+  - jstests/core/benchrun_pipeline_updates.js # benchRun() used for writes
 
   ##
   ## Some aggregation stages don't support snapshot readconcern.
@@ -370,6 +371,7 @@ selector:
   # op3 will still be using the snapshot from op1, and not see op2 at all.
   - jstests/core/cursora.js
   - jstests/core/bench_test1.js
+  - jstests/core/benchrun_pipeline_updates.js
 
   # The tests below use applyOps, SERVER-1439.
   - jstests/core/list_collections1.js
diff --git a/buildscripts/resmokeconfig/suites/read_concern_linearizable_passthrough.yml b/buildscripts/resmokeconfig/suites/read_concern_linearizable_passthrough.yml
index 5e6de583dbe..50b431c3456 100644
--- a/buildscripts/resmokeconfig/suites/read_concern_linearizable_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/read_concern_linearizable_passthrough.yml
@@ -13,6 +13,7 @@ selector:
   - jstests/core/capped_update.js
   # These tests use benchRun(), which isn't configured to use the overridden writeConcern.
   - jstests/core/bench_test*.js
+  - jstests/core/benchrun_pipeline_updates.js # benchRun() used for writes
   # fsync.js runs the {fsync: 1, lock: 1} command, which prevents a linearizable read from writing
   # an entry to the oplog.
   - jstests/core/fsync.js
diff --git a/buildscripts/resmokeconfig/suites/read_concern_majority_passthrough.yml b/buildscripts/resmokeconfig/suites/read_concern_majority_passthrough.yml
index 045020afb21..e341aef35f4 100644
--- a/buildscripts/resmokeconfig/suites/read_concern_majority_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/read_concern_majority_passthrough.yml
@@ -14,6 +14,7 @@ selector:
   - jstests/core/capped_update.js
   # These tests use benchRun(), which isn't configured to use the overridden writeConcern.
   - jstests/core/bench_test*.js
+  - jstests/core/benchrun_pipeline_updates.js # benchRun() used for writes
   exclude_with_any_tags:
   ##
   # The next three tags correspond to the special errors thrown by the
diff --git a/buildscripts/resmokeconfig/suites/replica_sets_kill_primary_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/replica_sets_kill_primary_jscore_passthrough.yml
index efd8ae3583a..dee4de899d6 100644
--- a/buildscripts/resmokeconfig/suites/replica_sets_kill_primary_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/replica_sets_kill_primary_jscore_passthrough.yml
@@ -62,6 +62,7 @@ selector:
   - jstests/core/ismaster.js
 
   - jstests/core/bench_test*.js # benchRun() used for writes
+  - jstests/core/benchrun_pipeline_updates.js # benchRun() used for writes
   - jstests/core/connection_string_validation.js # Does not expect a replica set connection string.
   - jstests/core/explain_large_bounds.js # Stepdown can timeout waiting for global lock.
   - jstests/core/insert2.js # Creates new mongo connection.
diff --git a/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_jscore_passthrough.yml
index dcc0a484991..7f63d067d1e 100644
--- a/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_jscore_passthrough.yml
@@ -267,6 +267,7 @@ selector:
   # op3 will still be using the snapshot from op1, and not see op2 at all.
   - jstests/core/cursora.js
   - jstests/core/bench_test1.js
+  - jstests/core/benchrun_pipeline_updates.js
 
   # Does not support tojson of command objects.
   - jstests/core/SERVER-23626.js
diff --git a/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_kill_primary_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_kill_primary_jscore_passthrough.yml
index db4e307a1f2..a595efb2a34 100644
--- a/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_kill_primary_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_kill_primary_jscore_passthrough.yml
@@ -283,6 +283,7 @@ selector:
   - jstests/core/ismaster.js
 
   - jstests/core/bench_test*.js # benchRun() used for writes
+  - jstests/core/benchrun_pipeline_updates.js # benchRun() used for writes
   - jstests/core/connection_string_validation.js # Does not expect a replica set connection string.
   - jstests/core/explain_large_bounds.js # Stepdown can timeout waiting for global lock.
   - jstests/core/insert2.js # Creates new mongo connection.
diff --git a/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_stepdown_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_stepdown_jscore_passthrough.yml
index 8756ea77d0c..5c5e1d1ada8 100644
--- a/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_stepdown_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_stepdown_jscore_passthrough.yml
@@ -283,6 +283,7 @@ selector:
   - jstests/core/ismaster.js
 
   - jstests/core/bench_test*.js # benchRun() used for writes
+  - jstests/core/benchrun_pipeline_updates.js # benchRun() used for writes
   - jstests/core/connection_string_validation.js # Does not expect a replica set connection string.
   - jstests/core/explain_large_bounds.js # Stepdown can timeout waiting for global lock.
   - jstests/core/insert2.js # Creates new mongo connection.
diff --git a/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_terminate_primary_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_terminate_primary_jscore_passthrough.yml
index 467fe7a03f9..a60b23c4d32 100644
--- a/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_terminate_primary_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/replica_sets_multi_stmt_txn_terminate_primary_jscore_passthrough.yml
@@ -280,6 +280,7 @@ selector:
   - jstests/core/ismaster.js
 
   - jstests/core/bench_test*.js # benchRun() used for writes
+  - jstests/core/benchrun_pipeline_updates.js # benchRun() used for writes
   - jstests/core/connection_string_validation.js # Does not expect a replica set connection string.
   - jstests/core/explain_large_bounds.js # Stepdown can timeout waiting for global lock.
   - jstests/core/insert2.js # Creates new mongo connection.
diff --git a/buildscripts/resmokeconfig/suites/replica_sets_terminate_primary_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/replica_sets_terminate_primary_jscore_passthrough.yml
index 406da904f3e..8bffc6057f3 100644
--- a/buildscripts/resmokeconfig/suites/replica_sets_terminate_primary_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/replica_sets_terminate_primary_jscore_passthrough.yml
@@ -62,6 +62,7 @@ selector:
   - jstests/core/ismaster.js
 
   - jstests/core/bench_test*.js # benchRun() used for writes
+  - jstests/core/benchrun_pipeline_updates.js # benchRun() used for writes
   - jstests/core/connection_string_validation.js # Does not expect a replica set connection string.
   - jstests/core/explain_large_bounds.js # Stepdown can timeout waiting for global lock.
   - jstests/core/insert2.js # Creates new mongo connection.
diff --git a/buildscripts/resmokeconfig/suites/retryable_writes_jscore_stepdown_passthrough.yml b/buildscripts/resmokeconfig/suites/retryable_writes_jscore_stepdown_passthrough.yml
index 454e456b85a..5f6261327ef 100644
--- a/buildscripts/resmokeconfig/suites/retryable_writes_jscore_stepdown_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/retryable_writes_jscore_stepdown_passthrough.yml
@@ -65,6 +65,7 @@ selector:
   - jstests/core/shell_connection_strings.js
 
   - jstests/core/bench_test*.js # benchRun() used for writes
+  - jstests/core/benchrun_pipeline_updates.js # benchRun() used for writes
   - jstests/core/connection_string_validation.js # Does not expect a replica set connection string.
   - jstests/core/explain_large_bounds.js # Stepdown can timeout waiting for global lock.
   - jstests/core/insert2.js # Creates new mongo connection.
diff --git a/buildscripts/resmokeconfig/suites/secondary_reads_passthrough.yml b/buildscripts/resmokeconfig/suites/secondary_reads_passthrough.yml
index 2e9df4f0a3c..6c8b8d151e2 100644
--- a/buildscripts/resmokeconfig/suites/secondary_reads_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/secondary_reads_passthrough.yml
@@ -44,6 +44,7 @@ selector:
   - jstests/core/fsync.js
   # benchRun cannot be overridden to be causally consistent
   - jstests/core/bench_test*.js
+  - jstests/core/benchrun_pipeline_updates.js
   # getMore is not causal consistent because afterClusterTime has been specified
   # by the preceding find command which establishes the cursor. Therefore,
   # getMore cannot be guaranteed to get executed after commands in between find
diff --git a/buildscripts/resmokeconfig/suites/sharded_multi_stmt_txn_jscore_passthrough.yml b/buildscripts/resmokeconfig/suites/sharded_multi_stmt_txn_jscore_passthrough.yml
index 3abbe965717..50962990b1b 100644
--- a/buildscripts/resmokeconfig/suites/sharded_multi_stmt_txn_jscore_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/sharded_multi_stmt_txn_jscore_passthrough.yml
@@ -300,6 +300,7 @@ selector:
   # op3 will still be using the snapshot from op1, and not see op2 at all.
   - jstests/core/cursora.js
   - jstests/core/bench_test1.js
+  - jstests/core/benchrun_pipeline_updates.js
 
   # The tests below use applyOps, SERVER-1439.
   - jstests/core/list_collections1.js
diff --git a/jstests/core/bench_test1.js b/jstests/core/bench_test1.js
index 683b5be4713..8e316c8b25e 100644
--- a/jstests/core/bench_test1.js
+++ b/jstests/core/bench_test1.js
@@ -4,43 +4,41 @@
 //   assumes_no_implicit_index_creation,
 //   uses_multiple_connections,
 // ]
-
-t = db.bench_test1;
-t.drop();
-
-t.insert({_id: 1, x: 1});
-t.insert({_id: 2, x: 1});
-
-ops = [
-    {op: "findOne", ns: t.getFullName(), query: {_id: 1}},
-    {op: "update", ns: t.getFullName(), query: {_id: 1}, update: {$inc: {x: 1}}}
-];
-
-seconds = 10;
-
-benchArgs = {
-    ops: ops,
-    parallel: 2,
-    seconds: seconds,
-    host: db.getMongo().host
-};
-
-if (jsTest.options().auth) {
-    benchArgs['db'] = 'admin';
-    benchArgs['username'] = jsTest.options().authUser;
-    benchArgs['password'] = jsTest.options().authPassword;
-}
-res = benchRun(benchArgs);
-
-assert.lte(seconds * res.update, t.findOne({_id: 1}).x * 1.5, "A1");
-
-assert.eq(1, t.getIndexes().length, "B1");
-benchArgs['ops'] = [{op: "createIndex", ns: t.getFullName(), key: {x: 1}}];
-benchArgs['parallel'] = 1;
-benchRun(benchArgs);
-assert.eq(2, t.getIndexes().length, "B2");
-benchArgs['ops'] = [{op: "dropIndex", ns: t.getFullName(), key: {x: 1}}];
-benchRun(benchArgs);
-assert.soon(function() {
-    return t.getIndexes().length == 1;
-});
+(function() {
+    "use strict";
+
+    const t = db.bench_test1;
+    t.drop();
+
+    t.insert({_id: 1, x: 1});
+    t.insert({_id: 2, x: 1});
+
+    const ops = [
+        {op: "findOne", ns: t.getFullName(), query: {_id: 1}},
+        {op: "update", ns: t.getFullName(), query: {_id: 1}, update: {$inc: {x: 1}}}
+    ];
+
+    const seconds = 10;
+
+    const benchArgs = {ops: ops, parallel: 2, seconds: seconds, host: db.getMongo().host};
+
+    if (jsTest.options().auth) {
+        benchArgs['db'] = 'admin';
+        benchArgs['username'] = jsTest.options().authUser;
+        benchArgs['password'] = jsTest.options().authPassword;
+    }
+    const res = benchRun(benchArgs);
+
+    assert.lte(seconds * res.update, t.findOne({_id: 1}).x * 1.5, "A1");
+
+    assert.eq(1, t.getIndexes().length, "B1");
+    benchArgs['ops'] = [{op: "createIndex", ns: t.getFullName(), key: {x: 1}}];
+    benchArgs['parallel'] = 1;
+    benchRun(benchArgs);
+    assert.eq(2, t.getIndexes().length, "B2");
+    benchArgs['ops'] = [{op: "dropIndex", ns: t.getFullName(), key: {x: 1}}];
+    benchRun(benchArgs);
+    assert.soon(function() {
+        return t.getIndexes().length == 1;
+    });
+}());
diff --git a/jstests/core/benchrun_pipeline_updates.js b/jstests/core/benchrun_pipeline_updates.js
new file mode 100644
index 00000000000..ef1ef719a0c
--- /dev/null
+++ b/jstests/core/benchrun_pipeline_updates.js
@@ -0,0 +1,54 @@
+/**
+ * Tests that benchRun can understand pipeline-style updates and findAndModifys.
+ *
+ * @tags: [uses_multiple_connections]
+ */
+(function() {
+    "use strict";
+    const coll = db.benchrun_pipeline_updates;
+    coll.drop();
+
+    assert.commandWorked(coll.insert({_id: 0, x: 0}));
+
+    // Test that a basic pipeline can be used by an update op.
+    let benchArgs = {
+        ops: [
+            {
+              op: "update",
+              ns: coll.getFullName(),
+              query: {_id: 0},
+              writeCmd: true,
+              update: [{$addFields: {x: {$add: ["$x", 1]}}}]
+            },
+        ],
+        parallel: 2,
+        seconds: 1,
+        host: db.getMongo().host,
+    };
+    if (jsTest.options().auth) {
+        benchArgs['db'] = 'admin';
+        benchArgs['username'] = jsTest.options().authUser;
+        benchArgs['password'] = jsTest.options().authPassword;
+    }
+    let res = benchRun(benchArgs);
+    assert.eq(res.errCount, 0);
+    assert.gt(
+        coll.findOne({_id: 0}).x, 2, "Expected at least one update to succeed and increment 'x'");
+
+    // Now test that the pipeline is still subject to benchRun's keyword replacement.
+
+    // Initialize x to something outside the range we'll expect it to be in below if the updates
+    // succeed.
+    assert.commandWorked(coll.updateOne({_id: 0}, {$set: {x: 100}}));
+    benchArgs.ops = [{
+        op: "update",
+        ns: coll.getFullName(),
+        query: {_id: 0},
+        writeCmd: true,
+        update: [{$project: {x: {$literal: {"#RAND_INT_PLUS_THREAD": [0, 2]}}}}]
+    }];
+    res = benchRun(benchArgs);
+    assert.eq(res.errCount, 0);
+    assert.lte(
+        coll.findOne({_id: 0}).x, 3, "Expected 'x' to be no more than 3 after randInt replacement");
+}());
diff --git a/src/mongo/shell/bench.cpp b/src/mongo/shell/bench.cpp
index 2b09ed443a5..78769755580 100644
--- a/src/mongo/shell/bench.cpp
+++ b/src/mongo/shell/bench.cpp
@@ -558,7 +558,7 @@ BenchRunOp opFromBson(const BSONObj& op) {
                     str::stream() << "Field 'update' is only valid for update op type. Op type is "
                                   << opType,
                     (opType == "update"));
-            myOp.update = arg.Obj();
+            myOp.update = write_ops::UpdateModification::parseFromBSON(arg);
         } else if (name == "upsert") {
             uassert(34392,
                     str::stream() << "Field 'upsert' is only valid for update op type. Op type is "
@@ -931,7 +931,7 @@ void BenchRunWorker::generateLoadOnConnection(DBClientBase* conn) {
                         return;
                 }
                 if (!_config->handleErrors && !op.handleError)
-                    return;
+                    throw;
 
                 sleepFor(_config->delayMillisOnFailedOperation);
 
@@ -1151,16 +1151,36 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
             {
                 BenchRunEventTrace _bret(&state->stats->updateCounter);
                 BSONObj query = fixQuery(this->query, *state->bsonTemplateEvaluator);
-                BSONObj update = fixQuery(this->update, *state->bsonTemplateEvaluator);
 
                 if (this->useWriteCmd) {
                     BSONObjBuilder builder;
                     builder.append("update", nsToCollectionSubstring(this->ns));
-                    BSONArrayBuilder docBuilder(builder.subarrayStart("updates"));
-                    docBuilder.append(BSON("q" << query << "u" << update << "multi" << this->multi
-                                               << "upsert"
-                                               << this->upsert));
-                    docBuilder.done();
+                    BSONArrayBuilder updateArray(builder.subarrayStart("updates"));
+                    {
+                        BSONObjBuilder singleUpdate;
+                        singleUpdate.append("q", query);
+                        switch (this->update.type()) {
+                            case write_ops::UpdateModification::Type::kClassic: {
+                                singleUpdate.append("u",
+                                                    fixQuery(this->update.getUpdateClassic(),
+                                                             *state->bsonTemplateEvaluator));
+                                break;
+                            }
+                            case write_ops::UpdateModification::Type::kPipeline: {
+                                BSONArrayBuilder pipelineBuilder(singleUpdate.subarrayStart("u"));
+                                for (auto&& stage : this->update.getUpdatePipeline()) {
+                                    pipelineBuilder.append(
+                                        fixQuery(stage, *state->bsonTemplateEvaluator));
+                                }
+                                pipelineBuilder.doneFast();
+                                break;
+                            }
+                        }
+                        singleUpdate.append("multi", this->multi);
+                        singleUpdate.append("upsert", this->upsert);
+                        updateArray.append(singleUpdate.done());
+                    }
+                    updateArray.doneFast();
                     builder.append("writeConcern", this->writeConcern);
 
                     boost::optional<TxnNumber> txnNumberForOp;
@@ -1176,11 +1196,16 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                                           txnNumberForOp,
                                           &result);
                 } else {
-                    auto toSend = makeUpdateMessage(this->ns,
-                                                    query,
-                                                    update,
-                                                    (this->upsert ? UpdateOption_Upsert : 0) |
-                                                        (this->multi ? UpdateOption_Multi : 0));
+                    uassert(
+                        30015,
+                        "cannot use legacy write protocol for anything but classic style updates",
+                        this->update.type() == write_ops::UpdateModification::Type::kClassic);
+                    auto toSend = makeUpdateMessage(
+                        this->ns,
+                        query,
+                        fixQuery(this->update.getUpdateClassic(), *state->bsonTemplateEvaluator),
+                        (this->upsert ? UpdateOption_Upsert : 0) |
+                            (this->multi ? UpdateOption_Multi : 0));
                     conn->say(toSend);
                     if (this->safe)
                         result = conn->getLastErrorDetailed();
diff --git a/src/mongo/shell/bench.h b/src/mongo/shell/bench.h
index 7683ad4b089..527beecae49 100644
--- a/src/mongo/shell/bench.h
+++ b/src/mongo/shell/bench.h
@@ -36,6 +36,7 @@
 #include "mongo/client/dbclient_base.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/db/logical_session_id.h"
+#include "mongo/db/ops/write_ops_parsers.h"
 #include "mongo/platform/atomic_word.h"
 #include "mongo/stdx/condition_variable.h"
 #include "mongo/stdx/mutex.h"
@@ -116,7 +117,7 @@ struct BenchRunOp {
     bool showResult = false;
     std::string target;
     bool throwGLE = false;
-    BSONObj update;
+    write_ops::UpdateModification update;
     bool upsert = false;
     bool useCheck = false;
     bool useReadCmd = false;
diff --git a/src/mongo/shell/dbshell.cpp b/src/mongo/shell/dbshell.cpp
index 633a99d8fc5..f7ff703b665 100644
--- a/src/mongo/shell/dbshell.cpp
+++ b/src/mongo/shell/dbshell.cpp
@@ -46,6 +46,7 @@
 #include "mongo/client/mongo_uri.h"
 #include "mongo/db/auth/sasl_command_constants.h"
 #include "mongo/db/client.h"
+#include "mongo/db/commands/test_commands_enabled.h"
 #include "mongo/db/log_process_details.h"
 #include "mongo/db/server_options.h"
 #include "mongo/logger/console_appender.h"
@@ -105,6 +106,15 @@ MONGO_INITIALIZER_WITH_PREREQUISITES(SetFeatureCompatibilityVersion42, ("EndStar
         ServerGlobalParams::FeatureCompatibility::Version::kFullyUpgradedTo42);
     return Status::OK();
 }
+
+// Initialize the testCommandsEnabled server parameter to true since the mongo shell does not have
+// any test-only commands that could cause harm to the server, and it may be necessary to enable
+// this to test certain features, for example through benchRun (see SERVER-40419).
+MONGO_INITIALIZER_WITH_PREREQUISITES(EnableShellTestCommands, ("EndStartupOptionSetup"))
+(InitializerContext* context) {
+    setTestCommandsEnabled(true);
+    return Status::OK();
+}
 const auto kAuthParam = "authSource"s;
 
 /**