SERVER-5722 Support 'sort' in benchrun

2 files changed, 31 insertions, 0 deletions

diff --git a/src/mongo/shell/bench.cpp b/src/mongo/shell/bench.cpp
index 67ef3f55996..1f5fd3e4775 100644
--- a/src/mongo/shell/bench.cpp
+++ b/src/mongo/shell/bench.cpp
@@ -131,6 +131,16 @@ BSONObj fixQuery(const BSONObj& obj, BsonTemplateEvaluator& btl) {
     return b.obj();
 }
 
+/**
+ * Adds a '$orderby' to the query document. Useful when running on the legacy reads path.
+ */
+BSONObj makeQueryLegacyCompatible(const BSONObj& query, const BSONObj& sortSpec) {
+    BSONObjBuilder bob;
+    bob.append("$query", query);
+    bob.append("$orderby", sortSpec);
+    return bob.obj();
+}
+
 bool runCommandWithSession(DBClientBase* conn,
                            const std::string& dbname,
                            const BSONObj& cmdObj,
@@ -541,6 +551,16 @@ BenchRunOp opFromBson(const BSONObj& op) {
                                   << opType,
                     (opType == "find") || (opType == "query"));
             myOp.skip = arg.numberInt();
+        } else if (name == "sort") {
+            uassert(ErrorCodes::BadValue,  // TODO
+                    str::stream()
+                        << "Field 'sort' is only valid for query, fineOne and find. Op type is "
+                        << opType,
+                    (opType == "findOne") || (opType == "query") || (opType == "find"));
+            uassert(ErrorCodes::BadValue,
+                    "Expected sort to be an object",
+                    arg.type() == BSONType::Object);
+            myOp.sort = arg.Obj();
         } else if (name == "showError") {
             myOp.showError = arg.trueValue();
         } else if (name == "showResult") {
@@ -987,6 +1007,7 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                 qr->setProj(this->projection);
                 qr->setLimit(1LL);
                 qr->setWantMore(false);
+                qr->setSort(this->sort);
                 if (config.useSnapshotReads) {
                     qr->setReadConcern(readConcernSnapshot);
                 }
@@ -1002,6 +1023,9 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                 runQueryWithReadCommands(
                     conn, lsid, txnNumberForOp, std::move(qr), Milliseconds(0), &result);
             } else {
+                if (!this->sort.isEmpty()) {
+                    fixedQuery = makeQueryLegacyCompatible(std::move(fixedQuery), this->sort);
+                }
                 BenchRunEventTrace _bret(&state->stats->findOneCounter);
                 result = conn->findOne(
                     this->ns, fixedQuery, nullptr, DBClientCursor::QueryOptionLocal_forceOpQuery);
@@ -1077,6 +1101,9 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                 if (this->batchSize) {
                     qr->setBatchSize(this->batchSize);
                 }
+                if (!this->sort.isEmpty()) {
+                    qr->setSort(this->sort);
+                }
                 BSONObjBuilder readConcernBuilder;
                 if (config.useSnapshotReads) {
                     readConcernBuilder.append("level", "snapshot");
@@ -1112,6 +1139,9 @@ void BenchRunOp::executeOnce(DBClientBase* conn,
                                                  Milliseconds(delayBeforeGetMore),
                                                  nullptr);
             } else {
+                if (!this->sort.isEmpty()) {
+                    fixedQuery = makeQueryLegacyCompatible(std::move(fixedQuery), this->sort);
+                }
                 // Use special query function for exhaust query option.
                 if (this->options & QueryOption_Exhaust) {
                     BenchRunEventTrace _bret(&state->stats->queryCounter);
diff --git a/src/mongo/shell/bench.h b/src/mongo/shell/bench.h
index 170023917ac..830e3c6924a 100644
--- a/src/mongo/shell/bench.h
+++ b/src/mongo/shell/bench.h
@@ -113,6 +113,7 @@ struct BenchRunOp {
     BSONObj query;
     bool safe = false;
     int skip = 0;
+    BSONObj sort;
     bool showError = false;
     bool showResult = false;
     std::string target;