SERVER-45439 Add support for per operation hedging via an explicit read preference

6 files changed, 252 insertions, 17 deletions

diff --git a/jstests/sharding/read_pref_with_hedging_mode.js b/jstests/sharding/read_pref_with_hedging_mode.js
new file mode 100644
index 00000000000..c9afbc8c16f
--- /dev/null
+++ b/jstests/sharding/read_pref_with_hedging_mode.js
@@ -0,0 +1,39 @@
+/*
+ * Intergration test for read preference with hedging mode. The more comprehensive
+ * unit test can be found in dbtests/read_preference_test.cpp.
+ * @tags: [requires_fcv_44]
+ */
+(function() {
+
+const st = new ShardingTest({shards: 2});
+const dbName = "foo";
+const collName = "bar";
+const ns = dbName + "." + collName;
+const testDB = st.s.getDB(dbName);
+
+assert.commandWorked(st.s.adminCommand({enableSharding: dbName}));
+st.ensurePrimaryShard(dbName, st.shard0.shardName);
+assert.commandWorked(st.s.adminCommand({shardCollection: ns, key: {x: 1}}));
+
+assert.commandWorked(
+    testDB.runCommand({query: {find: collName}, $readPreference: {mode: "nearest", hedge: {}}}));
+
+assert.commandWorked(testDB.runCommand({
+    query: {distinct: collName, key: "x"},
+    $readPreference: {mode: "primaryPreferred", hedge: {enabled: true, delay: false}}
+}));
+
+assert.commandFailedWithCode(
+    testDB.runCommand({query: {count: collName}, $readPreference: {mode: "primary", hedge: {}}}),
+    ErrorCodes.InvalidOptions);
+
+assert.commandFailedWithCode(testDB.runCommand({
+    $query: {
+        explain: {aggregate: collName, pipeline: [], cursor: {}},
+    },
+    $readPreference: {mode: "secondary", hedge: {enabled: false, delay: true}}
+}),
+                             ErrorCodes.InvalidOptions);
+
+st.stop();
+})();
diff --git a/src/mongo/client/SConscript b/src/mongo/client/SConscript
index 760693b0842..84a6d2a5fb9 100644
--- a/src/mongo/client/SConscript
+++ b/src/mongo/client/SConscript
@@ -32,6 +32,7 @@ env.Library(
     ],
     source=[
         'read_preference.cpp',
+        env.Idlc('hedging_mode.idl')[0],
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/bson/util/bson_extract',
@@ -384,4 +385,3 @@ env.Library(
         'clientdriver_minimal'
     ],
 )
-
diff --git a/src/mongo/client/hedging_mode.idl b/src/mongo/client/hedging_mode.idl
new file mode 100644
index 00000000000..9013fb8e168
--- /dev/null
+++ b/src/mongo/client/hedging_mode.idl
@@ -0,0 +1,48 @@
+# Copyright (C) 2020-present MongoDB, Inc.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the Server Side Public License, version 1,
+# as published by MongoDB, Inc.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# Server Side Public License for more details.
+#
+# You should have received a copy of the Server Side Public License
+# along with this program. If not, see
+# <http://www.mongodb.com/licensing/server-side-public-license>.
+#
+# As a special exception, the copyright holders give permission to link the
+# code of portions of this program with the OpenSSL library under certain
+# conditions as described in each individual source file and distribute
+# linked combinations including the program with the OpenSSL library. You
+# must comply with the Server Side Public License in all respects for
+# all of the code used other than as permitted herein. If you modify file(s)
+# with this exception, you may extend this exception to your version of the
+# file(s), but you are not obligated to do so. If you do not wish to do so,
+# delete this exception statement from your version. If you delete this
+# exception statement from all source files in the program, then also delete
+# it in the license file.
+#
+
+global:
+   cpp_namespace: "mongo"
+
+imports:
+    - "mongo/idl/basic_types.idl"
+
+structs:
+  HedgingMode:
+      description: The hedging mode for mongos.
+      fields:
+          enabled:
+              cpp_name: enabled
+              type: bool
+              description: Enable hedged reads.
+              default: true
+          delay:
+              cpp_name: delay
+              type: bool
+              description: Enable staggered reads.
+              default: true
diff --git a/src/mongo/client/read_preference.cpp b/src/mongo/client/read_preference.cpp
index b71825ad395..1676708d6f4 100644
--- a/src/mongo/client/read_preference.cpp
+++ b/src/mongo/client/read_preference.cpp
@@ -48,6 +48,7 @@ namespace {
 const char kModeFieldName[] = "mode";
 const char kTagsFieldName[] = "tags";
 const char kMaxStalenessSecondsFieldName[] = "maxStalenessSeconds";
+const char kHedgeFieldName[] = "hedge";
 
 const char kPrimaryOnly[] = "primary";
 const char kPrimaryPreferred[] = "primaryPreferred";
@@ -132,10 +133,12 @@ TagSet TagSet::primaryOnly() {
 
 ReadPreferenceSetting::ReadPreferenceSetting(ReadPreference pref,
                                              TagSet tags,
-                                             Seconds maxStalenessSeconds)
+                                             Seconds maxStalenessSeconds,
+                                             boost::optional<HedgingMode> hedgingMode)
     : pref(std::move(pref)),
       tags(std::move(tags)),
-      maxStalenessSeconds(std::move(maxStalenessSeconds)) {}
+      maxStalenessSeconds(std::move(maxStalenessSeconds)),
+      hedgingMode(std::move(hedgingMode)) {}
 
 ReadPreferenceSetting::ReadPreferenceSetting(ReadPreference pref, Seconds maxStalenessSeconds)
     : ReadPreferenceSetting(pref, defaultTagSetForMode(pref), maxStalenessSeconds) {}
@@ -160,6 +163,22 @@ StatusWith<ReadPreferenceSetting> ReadPreferenceSetting::fromInnerBSON(const BSO
     }
     mode = std::move(swReadPrefMode.getValue());
 
+    boost::optional<HedgingMode> hedgingMode;
+    if (auto hedgingModeEl = readPrefObj[kHedgeFieldName]) {
+        hedgingMode = HedgingMode::parse(IDLParserErrorContext(kHedgeFieldName),
+                                         hedgingModeEl.embeddedObject());
+        if (hedgingMode->getEnabled() && mode == ReadPreference::PrimaryOnly) {
+            return {
+                ErrorCodes::InvalidOptions,
+                str::stream() << "cannot enable hedging for $readPreference mode \"primaryOnly\""};
+        }
+
+        if (!hedgingMode->getEnabled() && hedgingMode->getDelay()) {
+            return {ErrorCodes::InvalidOptions,
+                    str::stream() << "cannot enable staggered reads without enabling hedging"};
+        }
+    }
+
     TagSet tags;
     BSONElement tagsElem;
     auto tagExtractStatus =
@@ -222,7 +241,7 @@ StatusWith<ReadPreferenceSetting> ReadPreferenceSetting::fromInnerBSON(const BSO
                                     << " can not be set for the primary mode");
     }
 
-    return ReadPreferenceSetting(mode, tags, Seconds(maxStalenessSecondsValue));
+    return ReadPreferenceSetting(mode, tags, Seconds(maxStalenessSecondsValue), hedgingMode);
 }
 
 StatusWith<ReadPreferenceSetting> ReadPreferenceSetting::fromInnerBSON(const BSONElement& elem) {
@@ -250,6 +269,9 @@ void ReadPreferenceSetting::toInnerBSON(BSONObjBuilder* bob) const {
     if (maxStalenessSeconds.count() > 0) {
         bob->append(kMaxStalenessSecondsFieldName, maxStalenessSeconds.count());
     }
+    if (hedgingMode) {
+        bob->append(kHedgeFieldName, hedgingMode.get().toBSON());
+    }
 }
 
 std::string ReadPreferenceSetting::toString() const {
diff --git a/src/mongo/client/read_preference.h b/src/mongo/client/read_preference.h
index f4807110b14..5ab06c6f397 100644
--- a/src/mongo/client/read_preference.h
+++ b/src/mongo/client/read_preference.h
@@ -30,6 +30,7 @@
 #pragma once
 
 #include "mongo/bson/simple_bsonobj_comparator.h"
+#include "mongo/client/hedging_mode_gen.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/db/repl/optime.h"
@@ -137,15 +138,27 @@ struct ReadPreferenceSetting {
      *     object's copy of tag will have the iterator in the initial
      *     position).
      */
-    ReadPreferenceSetting(ReadPreference pref, TagSet tags, Seconds maxStalenessSeconds);
+    ReadPreferenceSetting(ReadPreference pref,
+                          TagSet tags,
+                          Seconds maxStalenessSeconds,
+                          boost::optional<HedgingMode> hedgingMode = boost::none);
     ReadPreferenceSetting(ReadPreference pref, Seconds maxStalenessSeconds);
     ReadPreferenceSetting(ReadPreference pref, TagSet tags);
     explicit ReadPreferenceSetting(ReadPreference pref);
     ReadPreferenceSetting() : ReadPreferenceSetting(ReadPreference::PrimaryOnly) {}
 
     inline bool equals(const ReadPreferenceSetting& other) const {
+        auto hedgingModeEquals = [](const boost::optional<HedgingMode>& hedgingModeA,
+                                    const boost::optional<HedgingMode>& hedgingModeB) -> bool {
+            if (hedgingModeA && hedgingModeB) {
+                return hedgingModeA->toBSON().woCompare(hedgingModeB->toBSON()) == 0;
+            }
+            return !hedgingModeA && !hedgingModeB;
+        };
+
         return (pref == other.pref) && (tags == other.tags) &&
-            (maxStalenessSeconds == other.maxStalenessSeconds) && (minOpTime == other.minOpTime);
+            (maxStalenessSeconds == other.maxStalenessSeconds) &&
+            hedgingModeEquals(hedgingMode, other.hedgingMode) && (minOpTime == other.minOpTime);
     }
 
     /**
@@ -179,11 +192,12 @@ struct ReadPreferenceSetting {
 
     /**
      * Parses a ReadPreferenceSetting from a BSON document of the form:
-     * { mode: <mode>, tags: <array of tags>, maxStalenessSeconds: Number }. The 'mode' element must
-     * be a string equal to either "primary", "primaryPreferred", "secondary", "secondaryPreferred",
-     * or "nearest". Although the tags array is intended to be an array of unique BSON documents, no
-     * further validation is performed on it other than checking that it is an array, and that it is
-     * empty if 'mode' is 'primary'.
+     * { mode: <mode>, tags: <array of tags>, maxStalenessSeconds: Number, hedge: <hedgingMode>}.
+     * The 'mode' element must be a string equal to either "primary", "primaryPreferred",
+     * "secondary", "secondaryPreferred", or "nearest". Although the tags array is intended to be an
+     * array of unique BSON documents, no further validation is performed on it other than checking
+     * that it is an array, and that it is empty if 'mode' is 'primary'. The 'hedge' element
+     * consists of the optional field "enabled" (default true) and "delay" (default true).
      */
     static StatusWith<ReadPreferenceSetting> fromInnerBSON(const BSONObj& readPrefSettingObj);
     static StatusWith<ReadPreferenceSetting> fromInnerBSON(const BSONElement& readPrefSettingObj);
@@ -208,6 +222,7 @@ struct ReadPreferenceSetting {
     ReadPreference pref;
     TagSet tags;
     Seconds maxStalenessSeconds{};
+    boost::optional<HedgingMode> hedgingMode;
     repl::OpTime minOpTime{};
 };
 
diff --git a/src/mongo/client/read_preference_test.cpp b/src/mongo/client/read_preference_test.cpp
index 92bc2516be0..d0a6a1e048c 100644
--- a/src/mongo/client/read_preference_test.cpp
+++ b/src/mongo/client/read_preference_test.cpp
@@ -39,10 +39,14 @@ using namespace mongo;
 
 static const Seconds kMinMaxStaleness = ReadPreferenceSetting::kMinimalMaxStalenessValue;
 
-void checkParse(const BSONObj& rpsObj, const ReadPreferenceSetting& expected) {
+ReadPreferenceSetting parse(const BSONObj& rpsObj) {
     const auto swRps = ReadPreferenceSetting::fromInnerBSON(rpsObj);
     ASSERT_OK(swRps.getStatus());
-    const auto rps = swRps.getValue();
+    return std::move(swRps.getValue());
+}
+
+void checkParse(const BSONObj& rpsObj, const ReadPreferenceSetting& expected) {
+    const auto rps = parse(rpsObj);
     ASSERT_TRUE(rps.equals(expected));
 }
 
@@ -86,6 +90,84 @@ TEST(ReadPreferenceSetting, ParseValid) {
                                      TagSet(BSON_ARRAY(BSON("dc"
                                                             << "ny"))),
                                      kMinMaxStaleness));
+
+    checkParse(BSON("mode"
+                    << "nearest"
+                    << "tags"
+                    << BSON_ARRAY(BSON("dc"
+                                       << "ny"))
+                    << "maxStalenessSeconds" << kMinMaxStaleness.count() << "hedge" << BSONObj()),
+               ReadPreferenceSetting(ReadPreference::Nearest,
+                                     TagSet(BSON_ARRAY(BSON("dc"
+                                                            << "ny"))),
+                                     kMinMaxStaleness,
+                                     HedgingMode()));
+}
+
+TEST(ReadPreferenceSetting, ParseHedgingMode) {
+    // No hedging mode.
+    auto rpsObj = BSON("mode"
+                       << "primary");
+    auto rps = parse(rpsObj);
+    ASSERT_TRUE(rps.pref == ReadPreference::PrimaryOnly);
+    ASSERT_FALSE(rps.hedgingMode.has_value());
+
+    // Default hedging mode.
+    rpsObj = BSON("mode"
+                  << "primaryPreferred"
+                  << "hedge" << BSONObj());
+    rps = parse(rpsObj);
+    ASSERT_TRUE(rps.pref == ReadPreference::PrimaryPreferred);
+    ASSERT_TRUE(rps.hedgingMode.has_value());
+    ASSERT_TRUE(rps.hedgingMode->getEnabled());
+    ASSERT_TRUE(rps.hedgingMode->getDelay());
+
+    rpsObj = BSON("mode"
+                  << "secondary"
+                  << "hedge" << BSON("enabled" << true));
+    rps = parse(rpsObj);
+    ASSERT_TRUE(rps.pref == ReadPreference::SecondaryOnly);
+    ASSERT_TRUE(rps.hedgingMode.has_value());
+    ASSERT_TRUE(rps.hedgingMode->getEnabled());
+    ASSERT_TRUE(rps.hedgingMode->getDelay());
+
+    rpsObj = BSON("mode"
+                  << "secondaryPreferred"
+                  << "hedge" << BSON("delay" << false));
+    rps = parse(rpsObj);
+    ASSERT_TRUE(rps.pref == ReadPreference::SecondaryPreferred);
+    ASSERT_TRUE(rps.hedgingMode.has_value());
+    ASSERT_TRUE(rps.hedgingMode->getEnabled());
+    ASSERT_FALSE(rps.hedgingMode->getDelay());
+
+    // Input hedging mode.
+    rpsObj = BSON("mode"
+                  << "nearest"
+                  << "hedge" << BSON("enabled" << true << "delay" << false));
+    rps = parse(rpsObj);
+    ASSERT_TRUE(rps.pref == ReadPreference::Nearest);
+    ASSERT_TRUE(rps.hedgingMode.has_value());
+    ASSERT_TRUE(rps.hedgingMode->getEnabled());
+    ASSERT_FALSE(rps.hedgingMode->getDelay());
+
+    rpsObj = BSON("mode"
+                  << "nearest"
+                  << "hedge" << BSON("enabled" << false << "delay" << false));
+    rps = parse(rpsObj);
+    ASSERT_TRUE(rps.pref == ReadPreference::Nearest);
+    ASSERT_TRUE(rps.hedgingMode.has_value());
+    ASSERT_FALSE(rps.hedgingMode->getEnabled());
+    ASSERT_FALSE(rps.hedgingMode->getDelay());
+
+    rpsObj = BSON("mode"
+                  << "primary"
+                  << "hedge" << BSON("enabled" << false << "delay" << false));
+    rps = parse(rpsObj);
+
+    ASSERT_TRUE(rps.pref == ReadPreference::PrimaryOnly);
+    ASSERT_TRUE(rps.hedgingMode.has_value());
+    ASSERT_FALSE(rps.hedgingMode->getEnabled());
+    ASSERT_FALSE(rps.hedgingMode->getDelay());
 }
 
 void checkParseFails(const BSONObj& rpsObj) {
@@ -110,15 +192,22 @@ TEST(ReadPreferenceSetting, NonEquality) {
                                          << "ca")
                                     << BSON("foo"
                                             << "bar")));
-    auto rps = ReadPreferenceSetting(ReadPreference::Nearest, tagSet, kMinMaxStaleness);
+    auto hedgingMode = HedgingMode();
+    hedgingMode.setDelay(false);
+    auto rps =
+        ReadPreferenceSetting(ReadPreference::Nearest, tagSet, kMinMaxStaleness, hedgingMode);
 
-    auto unexpected1 =
-        ReadPreferenceSetting(ReadPreference::Nearest, TagSet::primaryOnly(), kMinMaxStaleness);
+    auto unexpected1 = ReadPreferenceSetting(
+        ReadPreference::Nearest, TagSet::primaryOnly(), kMinMaxStaleness, hedgingMode);
     ASSERT_FALSE(rps.equals(unexpected1));
 
     auto unexpected2 = ReadPreferenceSetting(
-        ReadPreference::Nearest, tagSet, Seconds(kMinMaxStaleness.count() + 1));
+        ReadPreference::Nearest, tagSet, Seconds(kMinMaxStaleness.count() + 1), hedgingMode);
     ASSERT_FALSE(rps.equals(unexpected2));
+
+    auto unexpected3 =
+        ReadPreferenceSetting(ReadPreference::Nearest, tagSet, kMinMaxStaleness, HedgingMode());
+    ASSERT_FALSE(rps.equals(unexpected3));
 }
 
 TEST(ReadPreferenceSetting, ParseInvalid) {
@@ -170,6 +259,18 @@ TEST(ReadPreferenceSetting, ParseInvalid) {
                          << "secondary"
                          << "maxStalenessSeconds" << Seconds::max().count()));
 
+    // Hedging is not supported for read preference mode "primary".
+    checkParseFailsWithError(BSON("mode"
+                                  << "primary"
+                                  << "hedge" << BSONObj()),
+                             ErrorCodes::InvalidOptions);
+
+    // Cannot enable staggering without enabling hedging.
+    checkParseFailsWithError(BSON("mode"
+                                  << "primaryPreferred"
+                                  << "hedge" << BSON("enabled" << false << "delay" << true)),
+                             ErrorCodes::InvalidOptions);
+
     checkParseContainerFailsWithError(BSON("$query" << BSON("pang"
                                                             << "pong")
                                                     << "$readPreference" << 2),
@@ -204,6 +305,16 @@ TEST(ReadPreferenceSetting, Roundtrip) {
                                                            << BSON("foo"
                                                                    << "bar"))),
                                          kMinMaxStaleness));
+
+    auto hedgingMode = HedgingMode();
+    hedgingMode.setDelay(false);
+    checkRoundtrip(ReadPreferenceSetting(ReadPreference::Nearest,
+                                         TagSet(BSON_ARRAY(BSON("dc"
+                                                                << "ca")
+                                                           << BSON("foo"
+                                                                   << "bar"))),
+                                         kMinMaxStaleness,
+                                         hedgingMode));
 }
 
 }  // namespace