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

4 files changed, 151 insertions, 4 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..77329cb260b
--- /dev/null
+++ b/jstests/noPassthrough/benchrun_read_pref_mode.js
@@ -0,0 +1,110 @@
+/**
+ * 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.FailedToParse
+        },
+        {
+          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/client/read_preference.cpp b/src/mongo/client/read_preference.cpp
index 3229d4735e8..b43ead599d9 100644
--- a/src/mongo/client/read_preference.cpp
+++ b/src/mongo/client/read_preference.cpp
@@ -115,6 +115,11 @@ TagSet defaultTagSetForMode(ReadPreference mode) {
 
 }  // namespace
 
+ReadPreference ReadPreferenceMode_parse(StringData prefStr) {
+    auto sw = parseReadPreferenceMode(prefStr);
+    uassertStatusOK(sw);
+    return sw.getValue();
+}
 
 /**
  * Replica set refresh period on the task executor.
diff --git a/src/mongo/shell/bench.cpp b/src/mongo/shell/bench.cpp
index 511dcb4293a..7d0187e81bf 100644
--- a/src/mongo/shell/bench.cpp
+++ b/src/mongo/shell/bench.cpp
@@ -35,6 +35,7 @@
 #include "mongo/shell/bench.h"
 
 #include <pcrecpp.h>
+#include <string>
 
 #include "mongo/client/dbclientcursor.h"
 #include "mongo/db/namespace_string.h"
@@ -230,16 +231,22 @@ int runQueryWithReadCommands(DBClientBase* conn,
                              const boost::optional<LogicalSessionIdToClient>& lsid,
                              boost::optional<TxnNumber> txnNumber,
                              std::unique_ptr<QueryRequest> qr,
+                             BSONObj readPrefObj,
                              BSONObj* objOut) {
     const auto dbName = qr->nss().db().toString();
 
     BSONObj findCommandResult;
+    BSONObjBuilder findCommandBuilder;
+    qr->asFindCommand(&findCommandBuilder);
+    if (!readPrefObj.isEmpty()) {
+        findCommandBuilder.append("$readPreference", readPrefObj);
+    }
     uassert(ErrorCodes::CommandFailed,
             str::stream() << "find command failed; reply was: " << findCommandResult,
             runCommandWithSession(
                 conn,
                 dbName,
-                qr->asFindCommand(),
+                findCommandBuilder.obj(),
                 // read command with txnNumber implies performing reads in a
                 // multi-statement transaction
                 txnNumber ? kStartTransactionOption | kMultiStatementTransactionOption : kNoOptions,
@@ -559,6 +566,27 @@ BenchRunOp opFromBson(const BSONObj& op) {
             BSONObjBuilder valBuilder;
             valBuilder.append(arg);
             myOp.value = valBuilder.obj();
+        } 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);
+
+            extern ReadPreference ReadPreferenceMode_parse(StringData);
+            ReadPreference mode;
+            try {
+                mode = ReadPreferenceMode_parse(arg.str());
+            } catch (DBException& e) {
+                e.addContext("benchRun(): Could not parse readPrefMode argument");
+                throw;
+            }
+
+            myOp.readPrefObj = ReadPreferenceSetting(mode).toInnerBSON();
         } else {
             uassert(34394, str::stream() << "Benchrun op has unsupported field: " << name, false);
         }
@@ -936,7 +964,8 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                     txnNumberForOp = state->txnNumber;
                     state->inProgressMultiStatementTxn = true;
                 }
-                runQueryWithReadCommands(conn, lsid, txnNumberForOp, std::move(qr), &result);
+                runQueryWithReadCommands(
+                    conn, lsid, txnNumberForOp, std::move(qr), readPrefObj, &result);
             } else {
                 if (!this->sort.isEmpty()) {
                     fixedQuery = makeQueryLegacyCompatible(std::move(fixedQuery), this->sort);
@@ -1032,8 +1061,8 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                     txnNumberForOp = state->txnNumber;
                     state->inProgressMultiStatementTxn = true;
                 }
-                count =
-                    runQueryWithReadCommands(conn, lsid, txnNumberForOp, std::move(qr), nullptr);
+                count = runQueryWithReadCommands(
+                    conn, lsid, txnNumberForOp, std::move(qr), readPrefObj, nullptr);
             } else {
                 if (!this->sort.isEmpty()) {
                     fixedQuery = makeQueryLegacyCompatible(std::move(fixedQuery), this->sort);
diff --git a/src/mongo/shell/bench.h b/src/mongo/shell/bench.h
index cee301cfaa0..0bbe740de08 100644
--- a/src/mongo/shell/bench.h
+++ b/src/mongo/shell/bench.h
@@ -123,6 +123,9 @@ struct BenchRunOp {
     BSONObj writeConcern;
     BSONObj value;
 
+    // Format: {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;
 };