SERVER-62016 support tagged write concerns

10 files changed, 186 insertions, 20 deletions

diff --git a/src/mongo/db/repl/repl_set_config.cpp b/src/mongo/db/repl/repl_set_config.cpp
index 25d57b0248a..e4fa3c80748 100644
--- a/src/mongo/db/repl/repl_set_config.cpp
+++ b/src/mongo/db/repl/repl_set_config.cpp
@@ -34,6 +34,7 @@
 #include "mongo/db/repl/repl_set_config.h"
 
 #include <algorithm>
+#include <fmt/format.h>
 #include <functional>
 
 #include "mongo/bson/util/bson_check.h"
@@ -527,6 +528,29 @@ StatusWith<ReplSetTagPattern> ReplSetConfig::findCustomWriteMode(StringData patt
     return StatusWith<ReplSetTagPattern>(iter->second);
 }
 
+StatusWith<ReplSetTagPattern> ReplSetConfig::makeCustomWriteMode(const BSONObj& wTags) const {
+    ReplSetTagPattern pattern = _tagConfig.makePattern();
+    for (auto e : wTags) {
+        const auto tagName = e.fieldNameStringData();
+        if (!e.isNumber()) {
+            return {
+                ErrorCodes::BadValue,
+                fmt::format(
+                    "Custom write mode only supports integer values, found: \"{}\" for tag: \"{}\"",
+                    e.toString(),
+                    tagName)};
+        }
+
+        const auto minNodesWithTag = e.safeNumberInt();
+        auto status = _tagConfig.addTagCountConstraintToPattern(&pattern, tagName, minNodesWithTag);
+        if (!status.isOK()) {
+            return status;
+        }
+    }
+
+    return pattern;
+}
+
 void ReplSetConfig::_calculateMajorities() {
     const int voters = std::count_if(
         begin(getMembers()), end(getMembers()), [](const auto& x) { return x.isVoter(); });
diff --git a/src/mongo/db/repl/repl_set_config.h b/src/mongo/db/repl/repl_set_config.h
index 0f2706bc830..cc9e50a7460 100644
--- a/src/mongo/db/repl/repl_set_config.h
+++ b/src/mongo/db/repl/repl_set_config.h
@@ -423,6 +423,15 @@ public:
     StatusWith<ReplSetTagPattern> findCustomWriteMode(StringData patternName) const;
 
     /**
+     * Returns a pattern constructed from a raw set of tags provided as the `w` value
+     * of a write concern.
+     *
+     * @returns `ErrorCodes::NoSuchKey` if a tag was provided which is not found in
+     * the local tag config.
+     */
+    StatusWith<ReplSetTagPattern> makeCustomWriteMode(const BSONObj& wTags) const;
+
+    /**
      * Returns the "tags configuration" for this replicaset.
      *
      * NOTE(schwerin): Not clear if this should be used other than for reporting/debugging.
diff --git a/src/mongo/db/repl/repl_set_config_test.cpp b/src/mongo/db/repl/repl_set_config_test.cpp
index 2ce4bb6fcc4..7a2c5fe33f3 100644
--- a/src/mongo/db/repl/repl_set_config_test.cpp
+++ b/src/mongo/db/repl/repl_set_config_test.cpp
@@ -1968,6 +1968,31 @@ TEST(ReplSetConfig, IsImplicitDefaultWriteConcernMajority) {
     ASSERT_OK(config.validate());
     ASSERT_FALSE(config.isImplicitDefaultWriteConcernMajority());
 }
+
+TEST(ReplSetConfig, MakeCustomWriteMode) {
+    auto config = ReplSetConfig::parse(BSON("_id"
+                                            << "rs0"
+                                            << "version" << 1 << "term" << 1.0 << "protocolVersion"
+                                            << 1 << "members"
+                                            << BSON_ARRAY(BSON("_id" << 0 << "host"
+                                                                     << "localhost:12345"
+                                                                     << "tags"
+                                                                     << BSON("NYC"
+                                                                             << "NY")))));
+
+    auto swPattern = config.makeCustomWriteMode(BSON("NYC"
+                                                     << "invalid value type"));
+    ASSERT_FALSE(swPattern.isOK());
+    ASSERT_EQ(swPattern.getStatus().code(), ErrorCodes::BadValue);
+
+    swPattern = config.makeCustomWriteMode(BSON("NonExistentTag" << 1));
+    ASSERT_FALSE(swPattern.isOK());
+    ASSERT_EQ(swPattern.getStatus().code(), ErrorCodes::NoSuchKey);
+
+    swPattern = config.makeCustomWriteMode(BSON("NYC" << 1));
+    ASSERT_TRUE(swPattern.isOK());
+}
+
 }  // namespace
 }  // namespace repl
 }  // namespace mongo
diff --git a/src/mongo/db/repl/replication_coordinator_impl.cpp b/src/mongo/db/repl/replication_coordinator_impl.cpp
index faa045cc782..f29887947b4 100644
--- a/src/mongo/db/repl/replication_coordinator_impl.cpp
+++ b/src/mongo/db/repl/replication_coordinator_impl.cpp
@@ -1900,6 +1900,11 @@ bool ReplicationCoordinatorImpl::_doneWaitingForReplication_inlock(
 
     const bool useDurableOpTime = writeConcern.syncMode == WriteConcernOptions::SyncMode::JOURNAL;
     if (writeConcern.wMode.empty()) {
+        if (writeConcern.wTags()) {
+            auto tagPattern = uassertStatusOK(_rsConfig.makeCustomWriteMode(*writeConcern.wTags()));
+            return _topCoord->haveTaggedNodesReachedOpTime(opTime, tagPattern, useDurableOpTime);
+        }
+
         return _topCoord->haveNumNodesReachedOpTime(
             opTime, writeConcern.wNumNodes, useDurableOpTime);
     }
diff --git a/src/mongo/db/repl/replication_coordinator_impl_test.cpp b/src/mongo/db/repl/replication_coordinator_impl_test.cpp
index 9a2e2b84c5c..9fcd9adf9b8 100644
--- a/src/mongo/db/repl/replication_coordinator_impl_test.cpp
+++ b/src/mongo/db/repl/replication_coordinator_impl_test.cpp
@@ -1374,6 +1374,56 @@ TEST_F(ReplCoordTest,
     awaiter.reset();
 }
 
+TEST_F(ReplCoordTest, SupportTaggedWriteConcern) {
+    assertStartSuccess(BSON("_id"
+                            << "mySet"
+                            << "version" << 2 << "members"
+                            << BSON_ARRAY(BSON("_id" << 0 << "host"
+                                                     << "node1:12345"
+                                                     << "tags"
+                                                     << BSON("donor"
+                                                             << "node"))
+                                          << BSON("_id" << 1 << "host"
+                                                        << "node2:12345"
+                                                        << "tags"
+                                                        << BSON("recipient"
+                                                                << "two"))
+                                          << BSON("_id" << 2 << "host"
+                                                        << "node3:12345"
+                                                        << "tags"
+                                                        << BSON("recipient"
+                                                                << "three")))),
+                       HostAndPort("node1", 12345));
+
+    ASSERT_OK(getReplCoord()->setFollowerMode(MemberState::RS_SECONDARY));
+    replCoordSetMyLastAppliedOpTime(OpTimeWithTermOne(100, 1), Date_t() + Seconds(100));
+    replCoordSetMyLastDurableOpTime(OpTimeWithTermOne(100, 1), Date_t() + Seconds(100));
+    simulateSuccessfulV1Election();
+
+    OpTimeWithTermOne time1(100, 1);
+    OpTimeWithTermOne time2(100, 2);
+
+    auto writeConcern =
+        uassertStatusOK(WriteConcernOptions::parse(BSON("w" << BSON("recipient" << 2))));
+
+    ReplicationAwaiter awaiter(getReplCoord(), getServiceContext());
+    awaiter.setOpTime(time2);
+    awaiter.setWriteConcern(writeConcern);
+    awaiter.start();
+
+    // start nodes in a lagged state
+    ASSERT_OK(getReplCoord()->setLastAppliedOptime_forTest(2, 1, time1));
+    ASSERT_OK(getReplCoord()->setLastAppliedOptime_forTest(2, 2, time1));
+
+    // catch them up to time2
+    ASSERT_OK(getReplCoord()->setLastAppliedOptime_forTest(2, 1, time2));
+    ASSERT_OK(getReplCoord()->setLastAppliedOptime_forTest(2, 2, time2));
+
+    ReplicationCoordinator::StatusAndDuration sad = awaiter.getResult();
+    ASSERT_OK(sad.status);
+    awaiter.reset();
+}
+
 TEST_F(ReplCoordTest, NodeReturnsNotPrimaryWhenSteppingDownBeforeSatisfyingAWriteConcern) {
     // Test that a thread blocked in awaitReplication will be woken up and return PrimarySteppedDown
     // (a NotPrimaryError) if the node steps down while it is waiting.
diff --git a/src/mongo/db/write_concern_options.cpp b/src/mongo/db/write_concern_options.cpp
index c43da914179..5557cee8298 100644
--- a/src/mongo/db/write_concern_options.cpp
+++ b/src/mongo/db/write_concern_options.cpp
@@ -126,6 +126,9 @@ StatusWith<WriteConcernOptions> WriteConcernOptions::parse(const BSONObj& obj) t
                                         << repl::ReplSetConfig::kMaxMembers << ", found: " << wNum);
             }
             writeConcern.wNumNodes = static_cast<decltype(writeConcern.wNumNodes)>(*wNum);
+        } else if (auto tags = stdx::get_if<BSONObj>(&wVal)) {
+            writeConcern.wNumNodes = 0;
+            writeConcern._tags = *tags;
         } else {
             auto wMode = stdx::get_if<std::string>(&wVal);
             invariant(wMode);
@@ -194,7 +197,9 @@ StatusWith<WriteConcernOptions> WriteConcernOptions::extractWCFromCommand(const
 BSONObj WriteConcernOptions::toBSON() const {
     BSONObjBuilder builder;
 
-    if (wMode.empty()) {
+    if (_tags) {
+        builder.append("w", *_tags);
+    } else if (wMode.empty()) {
         builder.append("w", wNumNodes);
     } else {
         builder.append("w", wMode);
@@ -216,12 +221,13 @@ BSONObj WriteConcernOptions::toBSON() const {
 }
 
 bool WriteConcernOptions::needToWaitForOtherNodes() const {
-    return !wMode.empty() || wNumNodes > 1;
+    return !wMode.empty() || wNumNodes > 1 || _tags;
 }
 
 bool WriteConcernOptions::operator==(const WriteConcernOptions& other) const {
-    return syncMode == other.syncMode && wMode == other.wMode && wNumNodes == other.wNumNodes &&
-        wDeadline == other.wDeadline && wTimeout == other.wTimeout &&
+    return (_tags ? other._tags && _tags->woCompare(*other._tags) == 0
+                  : wMode == other.wMode && wNumNodes == other.wNumNodes) &&
+        syncMode == other.syncMode && wDeadline == other.wDeadline && wTimeout == other.wTimeout &&
         _provenance == other._provenance;
 }
 
diff --git a/src/mongo/db/write_concern_options.h b/src/mongo/db/write_concern_options.h
index 5fbc73e840c..2eb0468b804 100644
--- a/src/mongo/db/write_concern_options.h
+++ b/src/mongo/db/write_concern_options.h
@@ -170,8 +170,13 @@ public:
 
     CheckCondition checkCondition = CheckCondition::OpTime;
 
+    boost::optional<BSONObj> wTags() const {
+        return _tags;
+    }
+
 private:
     ReadWriteConcernProvenance _provenance;
+    boost::optional<BSONObj> _tags;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/write_concern_options_test.cpp b/src/mongo/db/write_concern_options_test.cpp
index f15bd344625..4f2295feb7f 100644
--- a/src/mongo/db/write_concern_options_test.cpp
+++ b/src/mongo/db/write_concern_options_test.cpp
@@ -102,9 +102,10 @@ TEST(WriteConcernOptionsTest, ParseLeavesSyncModeAsUnsetIfFSyncIsFalse) {
     ASSERT_EQUALS(0, options.wTimeout);
 }
 
-TEST(WriteConcernOptionsTest, ParseReturnsFailedToParseIfWIsNotNumberOrString) {
-    auto status = WriteConcernOptions::parse(BSON("w" << BSONObj())).getStatus();
-    ASSERT_EQUALS(ErrorCodes::TypeMismatch, status);
+TEST(WriteConcernOptionsTest, ParseReturnsFailedToParseIfWIsNotNumberOrStringOrObject) {
+    auto status = WriteConcernOptions::parse(BSON("w" << true)).getStatus();
+    ASSERT_EQUALS(ErrorCodes::FailedToParse, status);
+    ASSERT_EQUALS("w has to be a number, string, or object", status.reason());
 }
 
 TEST(WriteConcernOptionsTest, ParseReturnsFailedToParseIfWIsNegativeOrExceedsMaxMembers) {
@@ -189,5 +190,37 @@ TEST(WriteConcernOptionsTest, ParseIgnoresSpecialFields) {
     _testIgnoreWriteConcernField("getLastError");
 }
 
+TEST(WriteConcernOptionsTest, ParseWithTags) {
+    auto status = WriteConcernOptions::parse(BSON("w" << BSON("abc"
+                                                              << "def")))
+                      .getStatus();
+    ASSERT_EQUALS(ErrorCodes::FailedToParse, status);
+    ASSERT_EQUALS("tags must be a single level document with only number values", status.reason());
+
+    auto sw = WriteConcernOptions::parse(BSON("w" << BSON("abc" << 1)));
+    ASSERT_OK(sw.getStatus());
+
+    WriteConcernOptions wc = sw.getValue();
+    ASSERT_EQ(wc.wNumNodes, 0);
+    ASSERT_TRUE(wc.wMode.empty());
+    ASSERT_EQ(wc.wTimeout, WriteConcernOptions::kNoTimeout);
+    ASSERT_TRUE(wc.needToWaitForOtherNodes());
+
+    auto tags = wc.wTags();
+    ASSERT(tags.has_value());
+    ASSERT_BSONOBJ_EQ(*tags, BSON("abc" << 1));
+    ASSERT_BSONOBJ_EQ(wc.toBSON(), BSON("w" << BSON("abc" << 1) << "wtimeout" << 0));
+
+    auto wc2 = uassertStatusOK(WriteConcernOptions::parse(BSON("w" << BSON("abc" << 1))));
+    ASSERT(wc == wc2);
+    auto wc3 = uassertStatusOK(WriteConcernOptions::parse(BSON("w" << BSON("def" << 1))));
+    ASSERT(wc != wc3);
+    auto wc4 = uassertStatusOK(WriteConcernOptions::parse(BSON("w"
+                                                               << "majority")));
+    ASSERT(wc != wc4);
+    auto wc5 = uassertStatusOK(WriteConcernOptions::parse(BSON("w" << 2)));
+    ASSERT(wc != wc5);
+}
+
 }  // namespace
 }  // namespace mongo
diff --git a/src/mongo/idl/basic_types.h b/src/mongo/idl/basic_types.h
index 42619522f3b..a725e9a8c87 100644
--- a/src/mongo/idl/basic_types.h
+++ b/src/mongo/idl/basic_types.h
@@ -29,6 +29,7 @@
 
 #pragma once
 
+#include "mongo/util/assert_util.h"
 #include <boost/optional.hpp>
 
 #include "mongo/base/string_data.h"
@@ -191,21 +192,32 @@ public:
             return WriteConcernW{wEl.safeNumberLong()};
         } else if (wEl.type() == BSONType::String) {
             return WriteConcernW{wEl.str()};
+        } else if (wEl.type() == BSONType::Object) {
+            auto tags = wEl.Obj().getOwned();
+            auto valid =
+                std::all_of(tags.begin(), tags.end(), [](BSONElement e) { return e.isNumber(); });
+            uassert(ErrorCodes::FailedToParse,
+                    "tags must be a single level document with only number values",
+                    valid);
+
+            return WriteConcernW{std::move(tags)};
         } else if (wEl.eoo() || wEl.type() == BSONType::jstNULL ||
                    wEl.type() == BSONType::Undefined) {
             return WriteConcernW{};
         }
-        uasserted(ErrorCodes::FailedToParse, "w has to be a number or string");
+        uasserted(ErrorCodes::FailedToParse, "w has to be a number, string, or object");
     }
 
     void serializeWriteConcernW(StringData fieldName, BSONObjBuilder* builder) const {
         if (auto stringVal = stdx::get_if<std::string>(&_w)) {
             builder->append(fieldName, *stringVal);
-            return;
+        } else if (auto objVal = stdx::get_if<BSONObj>(&_w)) {
+            builder->append(fieldName, *objVal);
+        } else {
+            auto intVal = stdx::get_if<std::int64_t>(&_w);
+            invariant(intVal);
+            builder->appendNumber(fieldName, static_cast<long long>(*intVal));
         }
-        auto intVal = stdx::get_if<std::int64_t>(&_w);
-        invariant(intVal);
-        builder->appendNumber(fieldName, static_cast<long long>(*intVal));
     }
 
     WriteConcernW() : _w{1}, _usedDefaultConstructedW1{true} {};
@@ -214,14 +226,16 @@ public:
         return _usedDefaultConstructedW1;
     }
 
-    stdx::variant<std::string, std::int64_t> getValue() const {
+    stdx::variant<std::string, std::int64_t, BSONObj> getValue() const {
         return _w;
     }
 
 private:
     WriteConcernW(std::int64_t w) : _w{w}, _usedDefaultConstructedW1{false} {};
     WriteConcernW(std::string&& w) : _w(std::move(w)), _usedDefaultConstructedW1{false} {};
-    stdx::variant<std::string, std::int64_t> _w;
+    WriteConcernW(BSONObj&& tags) : _w(std::move(tags)), _usedDefaultConstructedW1{false} {};
+
+    stdx::variant<std::string, std::int64_t, BSONObj> _w;
     bool _usedDefaultConstructedW1;
 };
 
diff --git a/src/mongo/idl/basic_types.idl b/src/mongo/idl/basic_types.idl
index ae3041063cd..03cf00ba413 100644
--- a/src/mongo/idl/basic_types.idl
+++ b/src/mongo/idl/basic_types.idl
@@ -242,12 +242,7 @@ types:
         deserializer: mongo::IDLAnyTypeOwned::parseFromBSON
 
     writeConcernW:
-        bson_serialization_type:
-                                 - string
-                                 - int
-                                 - decimal
-                                 - double
-                                 - long
+        bson_serialization_type: any
         description: >-
             A string or integer representing the 'w' option in a document specifying write concern.
             See https://docs.mongodb.com/manual/reference/write-concern/"