SERVER-57897 Add readPrefMode option to benchRun find/findOne ops

(cherry picked from commit b8b5d3de4df8e0f82491b5d5ccbc4b635f61aecf)

3 files changed, 143 insertions, 6 deletions

diff --git a/jstests/noPassthrough/benchrun_read_pref_mode.js b/jstests/noPassthrough/benchrun_read_pref_mode.js
new file mode 100644
index 00000000000..d9c6027f4eb
--- /dev/null
+++ b/jstests/noPassthrough/benchrun_read_pref_mode.js
@@ -0,0 +1,101 @@
+/**
+ * Verifies that readPrefMode param works for find/fineOne/query ops in benchRun().
+ *
+ * @tags: [requires_replication]
+ */
+
+(function() {
+"use strict";
+
+const rs = new ReplSetTest({nodes: 2});
+rs.startSet();
+rs.initiate();
+
+const primary = rs.getPrimary();
+const secondary = rs.getSecondary();
+const collName = primary.getDB(jsTestName()).getCollection("coll").getFullName();
+
+const verifyNoError = res => {
+    assert.eq(res.errCount, 0);
+    assert.gt(res.totalOps, 0);
+};
+
+const benchArgArray = [
+    {
+        ops: [{op: "find", readCmd: true, query: {}, ns: collName, readPrefMode: "primary"}],
+        parallel: 1,
+        host: primary.host
+    },
+    {
+        ops: [{
+            op: "findOne",
+            readCmd: true,
+            query: {},
+            ns: collName,
+            readPrefMode: "primaryPreferred"
+        }],
+        parallel: 1,
+        host: primary.host
+    },
+    {
+        ops: [{op: "find", readCmd: true, query: {}, ns: collName, readPrefMode: "secondary"}],
+        parallel: 1,
+        host: secondary.host
+    },
+    {
+        ops: [{
+            op: "findOne",
+            readCmd: true,
+            query: {},
+            ns: collName,
+            readPrefMode: "secondaryPreferred"
+        }],
+        parallel: 1,
+        host: secondary.host
+    },
+    {
+        ops: [{op: "query", readCmd: true, query: {}, ns: collName, readPrefMode: "nearest"}],
+        parallel: 1,
+        host: secondary.host
+    },
+];
+
+benchArgArray.forEach(benchArg => verifyNoError(benchRun(benchArg)));
+
+const invalidArgAndError = [
+    {
+        benchArg: {
+            ops: [{op: "find", readCmd: true, query: {}, ns: collName, readPrefMode: 1}],
+            parallel: 1,
+            host: primary.host
+        },
+        error: ErrorCodes.BadValue
+    },
+    {
+        benchArg: {
+            ops:
+                [{op: "find", readCmd: true, query: {}, ns: collName, readPrefMode: "invalidPref"}],
+            parallel: 1,
+            host: primary.host
+        },
+        error: ErrorCodes.BadValue
+    },
+    {
+        benchArg: {
+            ops: [
+                {op: "insert", writeCmd: true, doc: {a: 1}, ns: collName, readPrefMode: "primary"}
+            ],
+            parallel: 1,
+            host: primary.host
+        },
+        error: ErrorCodes.InvalidOptions
+    },
+];
+
+invalidArgAndError.forEach(argAndError => {
+    const res = assert.throws(() => benchRun(argAndError.benchArg));
+    assert.commandFailedWithCode(res, argAndError.error);
+});
+
+rs.stopSet();
+})();
diff --git a/src/mongo/shell/bench.cpp b/src/mongo/shell/bench.cpp
index 01d6c264f99..4c2bec187a4 100644
--- a/src/mongo/shell/bench.cpp
+++ b/src/mongo/shell/bench.cpp
@@ -34,6 +34,7 @@
 #include "mongo/shell/bench.h"
 
 #include <pcrecpp.h>
+#include <string>
 
 #include "mongo/base/shim.h"
 #include "mongo/client/dbclient_cursor.h"
@@ -232,17 +233,19 @@ int runQueryWithReadCommands(DBClientBase* conn,
                              boost::optional<TxnNumber> txnNumber,
                              std::unique_ptr<FindCommandRequest> findCommand,
                              Milliseconds delayBeforeGetMore,
+                             BSONObj readPrefObj,
                              BSONObj* objOut) {
     const auto dbName =
         findCommand->getNamespaceOrUUID().nss().value_or(NamespaceString()).db().toString();
 
     BSONObj findCommandResult;
+    BSONObj findCommandObj = findCommand->toBSON(readPrefObj);
     uassert(ErrorCodes::CommandFailed,
             str::stream() << "find command failed; reply was: " << findCommandResult,
             runCommandWithSession(
                 conn,
                 dbName,
-                findCommand->toBSON(BSONObj()),
+                findCommandObj,
                 // read command with txnNumber implies performing reads in a
                 // multi-statement transaction
                 txnNumber ? kStartTransactionOption | kMultiStatementTransactionOption : kNoOptions,
@@ -312,8 +315,13 @@ Timestamp getLatestClusterTime(DBClientBase* conn) {
         findCommand->getNamespaceOrUUID().nss().value_or(NamespaceString()).db().toString();
 
     BSONObj oplogResult;
-    int count = runQueryWithReadCommands(
-        conn, boost::none, boost::none, std::move(findCommand), Milliseconds(0), &oplogResult);
+    int count = runQueryWithReadCommands(conn,
+                                         boost::none,
+                                         boost::none,
+                                         std::move(findCommand),
+                                         Milliseconds(0),
+                                         BSONObj(),
+                                         &oplogResult);
     uassert(ErrorCodes::OperationFailed,
             str::stream() << "Find cmd on the oplog collection failed; reply was: " << oplogResult,
             count == 1);
@@ -641,6 +649,26 @@ BenchRunOp opFromBson(const BSONObj& op) {
                               << opType,
                 (opType == "find") || (opType == "query"));
             myOp.maxRandomMillisecondDelayBeforeGetMore = arg.numberInt();
+        } else if (name == "readPrefMode") {
+            uassert(
+                ErrorCodes::InvalidOptions,
+                str::stream() << "Field 'readPrefMode' is only valid for find op types. Type is "
+                              << opType,
+                (opType == "find") || (opType == "query") || (opType == "findOne"));
+            uassert(ErrorCodes::BadValue,
+                    str::stream() << "Field 'readPrefMode' should be a string, instead it's type: "
+                                  << typeName(arg.type()),
+                    arg.type() == BSONType::String);
+
+            ReadPreference mode;
+            try {
+                mode = ReadPreference_parse(IDLParserErrorContext("mode"), arg.str());
+            } catch (DBException& e) {
+                e.addContext("benchRun(): Could not parse readPrefMode argument");
+                throw;
+            }
+
+            myOp.readPrefObj = ReadPreferenceSetting(mode).toContainingBSON();
         } else {
             uassert(34394, str::stream() << "Benchrun op has unsupported field: " << name, false);
         }
@@ -1032,8 +1060,13 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                     txnNumberForOp = state->txnNumber;
                     state->inProgressMultiStatementTxn = true;
                 }
-                runQueryWithReadCommands(
-                    conn, lsid, txnNumberForOp, std::move(findCommand), Milliseconds(0), &result);
+                runQueryWithReadCommands(conn,
+                                         lsid,
+                                         txnNumberForOp,
+                                         std::move(findCommand),
+                                         Milliseconds(0),
+                                         readPrefObj,
+                                         &result);
             } else {
                 if (!this->sort.isEmpty()) {
                     fixedQuery = makeQueryLegacyCompatible(std::move(fixedQuery), this->sort);
@@ -1042,7 +1075,6 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                 result = conn->findOne(
                     this->ns, fixedQuery, nullptr, DBClientCursor::QueryOptionLocal_forceOpQuery);
             }
-
             if (!config.hideResults || this->showResult)
                 LOGV2_INFO(22796, "Result from benchRun thread [findOne]", "result"_attr = result);
         } break;
@@ -1149,6 +1181,7 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                                                  txnNumberForOp,
                                                  std::move(findCommand),
                                                  Milliseconds(delayBeforeGetMore),
+                                                 readPrefObj,
                                                  nullptr);
             } else {
                 if (!this->sort.isEmpty()) {
diff --git a/src/mongo/shell/bench.h b/src/mongo/shell/bench.h
index 64e3dc8e60a..8938383620d 100644
--- a/src/mongo/shell/bench.h
+++ b/src/mongo/shell/bench.h
@@ -139,6 +139,9 @@ struct BenchRunOp {
     // resources that a snapshot transaction would hold for a time.
     int maxRandomMillisecondDelayBeforeGetMore{0};
 
+    // Format: {$readPreference: {mode: modeStr}}.  Only mode field is allowed.
+    BSONObj readPrefObj;
+
     // This is an owned copy of the raw operation. All unowned members point into this.
     BSONObj myBsonOp;
 };