SERVER-46472 add maxTimeMSForHedgedReads

4 files changed, 52 insertions, 5 deletions

diff --git a/src/mongo/executor/network_interface_tl.cpp b/src/mongo/executor/network_interface_tl.cpp
index 76ff8e8cc7d..506aa01f517 100644
--- a/src/mongo/executor/network_interface_tl.cpp
+++ b/src/mongo/executor/network_interface_tl.cpp
@@ -45,6 +45,10 @@
 namespace mongo {
 namespace executor {
 
+namespace {
+static inline const std::string kMaxTimeMSOptionName = "maxTimeMS";
+}  // unnamed namespace
+
 /**
  * SynchronizedCounters is synchronized bucket of event counts for commands
  */
@@ -604,12 +608,36 @@ void NetworkInterfaceTL::RequestManager::trySend(
     LOGV2_DEBUG(4646300,
                 2,
                 "Sending request {request_id} with index {idx}",
-                "request_id"_attr = cmdState.lock()->requestOnAny.id,
+                "request_id"_attr = cmdStatePtr->requestOnAny.id,
                 "idx"_attr = idx);
 
     auto req = getNextRequest();
     if (req) {
-        req->send(std::move(swConn), {cmdStatePtr->requestOnAny, idx});
+        RemoteCommandRequest remoteReq({cmdStatePtr->requestOnAny, idx});
+        if (sentIdx.load() > 1) {  // this is a hedged read
+            invariant(remoteReq.hedgeOptions);
+            auto maxTimeMS = remoteReq.hedgeOptions->maxTimeMSForHedgedReads;
+            if (remoteReq.timeout == remoteReq.kNoTimeout ||
+                remoteReq.timeout > Milliseconds(maxTimeMS)) {
+                BSONObjBuilder updatedCmdBuilder;
+                for (const auto& elem : remoteReq.cmdObj) {
+                    if (elem.fieldNameStringData() != kMaxTimeMSOptionName) {
+                        updatedCmdBuilder.append(elem);
+                    }
+                }
+                updatedCmdBuilder.append(kMaxTimeMSOptionName, maxTimeMS);
+
+                remoteReq.cmdObj = updatedCmdBuilder.obj();
+                LOGV2_DEBUG(
+                    4647200,
+                    2,
+                    "Set  MaxTimeMS to {maxTimeMS} for request {request_id} with index {idx}",
+                    "maxTimeMS"_attr = maxTimeMS,
+                    "request_id"_attr = cmdStatePtr->requestOnAny.id,
+                    "idx"_attr = idx);
+            }
+        }
+        req->send(std::move(swConn), remoteReq);
     }
 }
 
diff --git a/src/mongo/executor/remote_command_request.h b/src/mongo/executor/remote_command_request.h
index dfa9e45a11b..78dc9be9c02 100644
--- a/src/mongo/executor/remote_command_request.h
+++ b/src/mongo/executor/remote_command_request.h
@@ -45,6 +45,7 @@ namespace executor {
 struct RemoteCommandRequestBase {
     struct HedgeOptions {
         size_t count;
+        int maxTimeMSForHedgedReads;
     };
 
     enum FireAndForgetMode { kOn, kOff };
diff --git a/src/mongo/s/hedge_options_util.cpp b/src/mongo/s/hedge_options_util.cpp
index 3781781a879..e1a77e4abf5 100644
--- a/src/mongo/s/hedge_options_util.cpp
+++ b/src/mongo/s/hedge_options_util.cpp
@@ -37,7 +37,8 @@ boost::optional<executor::RemoteCommandRequestOnAny::HedgeOptions> extractHedgeO
     const ReadPreferenceSetting& readPref) {
     if (gReadHedgingMode.load() == ReadHedgingMode::kOn && readPref.hedgingMode &&
         readPref.hedgingMode->getEnabled()) {
-        return executor::RemoteCommandRequestOnAny::HedgeOptions{1};
+        return executor::RemoteCommandRequestOnAny::HedgeOptions{1,
+                                                                 gMaxTimeMSForHedgedReads.load()};
     }
 
     return boost::none;
diff --git a/src/mongo/s/hedge_options_util_test.cpp b/src/mongo/s/hedge_options_util_test.cpp
index 6983fb6327d..e3e0a7abdb5 100644
--- a/src/mongo/s/hedge_options_util_test.cpp
+++ b/src/mongo/s/hedge_options_util_test.cpp
@@ -83,7 +83,8 @@ protected:
      */
     void checkHedgeOptions(const BSONObj& serverParameters,
                            const BSONObj& rspObj,
-                           const bool hedge) {
+                           const bool hedge,
+                           const int maxTimeMSForHedgedReads = kMaxTimeMSForHedgedReadsDefault) {
         setParameters(serverParameters);
 
         auto readPref = uassertStatusOK(ReadPreferenceSetting::fromInnerBSON(rspObj));
@@ -91,6 +92,7 @@ protected:
 
         if (hedge) {
             ASSERT_TRUE(hedgeOptions.has_value());
+            ASSERT_EQ(hedgeOptions->maxTimeMSForHedgedReads, maxTimeMSForHedgedReads);
         } else {
             ASSERT_FALSE(hedgeOptions.has_value());
         }
@@ -101,8 +103,12 @@ protected:
     static inline const std::string kCollName = "testColl";
 
     static inline const std::string kReadHedgingModeFieldName = "readHedgingMode";
+    static inline const std::string kMaxTimeMSForHedgedReadsFieldName = "maxTimeMSForHedgedReads";
+    static inline const int kMaxTimeMSForHedgedReadsDefault = 10;
 
-    static inline const BSONObj kDefaultParameters = BSON(kReadHedgingModeFieldName << "on");
+    static inline const BSONObj kDefaultParameters =
+        BSON(kReadHedgingModeFieldName << "on" << kMaxTimeMSForHedgedReadsFieldName
+                                       << kMaxTimeMSForHedgedReadsDefault);
 
 private:
     ServiceContext::UniqueServiceContext _serviceCtx = ServiceContext::make();
@@ -144,5 +150,16 @@ TEST_F(HedgeOptionsUtilTestFixture, ReadHedgingModeOff) {
     checkHedgeOptions(parameters, rspObj, false);
 }
 
+TEST_F(HedgeOptionsUtilTestFixture, MaxTimeMSForHedgedReads) {
+    const auto parameters =
+        BSON(kReadHedgingModeFieldName << "on" << kMaxTimeMSForHedgedReadsFieldName << 100);
+
+    const auto rspObj = BSON("mode"
+                             << "nearest"
+                             << "hedge" << BSONObj());
+
+    checkHedgeOptions(parameters, rspObj, true, 100);
+}
+
 }  // namespace
 }  // namespace mongo