SERVER-51103 Allow ShardRegistry to distinguish between complete/partial connection string updates

6 files changed, 52 insertions, 8 deletions

diff --git a/src/mongo/client/connection_string.cpp b/src/mongo/client/connection_string.cpp
index 90adb5a5f44..f3dd2fd192c 100644
--- a/src/mongo/client/connection_string.cpp
+++ b/src/mongo/client/connection_string.cpp
@@ -27,6 +27,8 @@
  *    it in the license file.
  */
 
+#include <set>
+
 #include "mongo/platform/basic.h"
 
 #include "mongo/client/connection_string.h"
@@ -160,6 +162,15 @@ void ConnectionString::_finishInit() {
     _string = ss.str();
 }
 
+ConnectionString ConnectionString::makeUnionWith(const ConnectionString& other) {
+    invariant(type() == other.type());
+    invariant(getSetName() == other.getSetName());
+    std::set<HostAndPort> servers{_servers.begin(), _servers.end()};
+    servers.insert(other._servers.begin(), other._servers.end());
+    return ConnectionString(
+        type(), std::vector<HostAndPort>(servers.begin(), servers.end()), getSetName());
+}
+
 bool ConnectionString::operator==(const ConnectionString& other) const {
     if (_type != other._type) {
         return false;
diff --git a/src/mongo/client/connection_string.h b/src/mongo/client/connection_string.h
index 67f62116d16..98b66bef986 100644
--- a/src/mongo/client/connection_string.h
+++ b/src/mongo/client/connection_string.h
@@ -116,6 +116,15 @@ public:
     }
 
     /**
+     * Creates a new ConnectionString object which contains all the servers in either this
+     * ConnectionString or the given one.  Useful for "extending" a connection string with
+     * (potentially) new servers.
+     *
+     * The given ConnectionString must have the same type() and getSetName() as this one.
+     */
+    ConnectionString makeUnionWith(const ConnectionString& other);
+
+    /**
      * Returns true if two connection strings match in terms of their type and the exact order of
      * their hosts.
      */
diff --git a/src/mongo/db/s/sharding_initialization_mongod.cpp b/src/mongo/db/s/sharding_initialization_mongod.cpp
index 3c2bcd143e4..0cfebae2635 100644
--- a/src/mongo/db/s/sharding_initialization_mongod.cpp
+++ b/src/mongo/db/s/sharding_initialization_mongod.cpp
@@ -105,7 +105,10 @@ public:
             LOGV2(471691,
                   "Updating the shard registry with confirmed replica set",
                   "connectionString"_attr = connStr);
-            Grid::get(_serviceContext)->shardRegistry()->updateReplSetHosts(connStr);
+            Grid::get(_serviceContext)
+                ->shardRegistry()
+                ->updateReplSetHosts(connStr,
+                                     ShardRegistry::ConnectionStringUpdateType::kConfirmed);
         } catch (const ExceptionForCat<ErrorCategory::ShutdownError>& e) {
             LOGV2(471692, "Unable to update the shard registry", "error"_attr = e);
         }
@@ -130,7 +133,10 @@ public:
 
     void onPossibleSet(const State& state) noexcept final {
         try {
-            Grid::get(_serviceContext)->shardRegistry()->updateReplSetHosts(state.connStr);
+            Grid::get(_serviceContext)
+                ->shardRegistry()
+                ->updateReplSetHosts(state.connStr,
+                                     ShardRegistry::ConnectionStringUpdateType::kPossible);
         } catch (const DBException& ex) {
             LOGV2_DEBUG(22070,
                         2,
diff --git a/src/mongo/s/client/shard_registry.cpp b/src/mongo/s/client/shard_registry.cpp
index 842da344bec..46b7406744e 100644
--- a/src/mongo/s/client/shard_registry.cpp
+++ b/src/mongo/s/client/shard_registry.cpp
@@ -370,11 +370,17 @@ ShardRegistry::_getLatestConnStrings() const {
     return {{_latestConnStrings.begin(), _latestConnStrings.end()}, _rsmIncrement.load()};
 }
 
-void ShardRegistry::updateReplSetHosts(const ConnectionString& newConnString) {
-    invariant(newConnString.type() == ConnectionString::SET ||
-              newConnString.type() == ConnectionString::CUSTOM);  // For dbtests
+void ShardRegistry::updateReplSetHosts(const ConnectionString& givenConnString,
+                                       ConnectionStringUpdateType updateType) {
+    invariant(givenConnString.type() == ConnectionString::SET ||
+              givenConnString.type() == ConnectionString::CUSTOM);  // For dbtests
 
     stdx::lock_guard<Latch> lk(_mutex);
+    ConnectionString newConnString =
+        (updateType == ConnectionStringUpdateType::kPossible &&
+         _latestConnStrings.find(givenConnString.getSetName()) != _latestConnStrings.end())
+        ? _latestConnStrings[givenConnString.getSetName()].makeUnionWith(givenConnString)
+        : givenConnString;
     if (auto shard = _configShardData.findByRSName(newConnString.getSetName())) {
         auto newData = ShardRegistryData::createFromExisting(
             _configShardData, newConnString, _shardFactory.get());
diff --git a/src/mongo/s/client/shard_registry.h b/src/mongo/s/client/shard_registry.h
index 711c8e0d34b..e91b93fc07c 100644
--- a/src/mongo/s/client/shard_registry.h
+++ b/src/mongo/s/client/shard_registry.h
@@ -179,6 +179,11 @@ public:
     using ShardRemovalHook = std::function<void(const ShardId&)>;
 
     /**
+     * Used when informing updateReplSetHosts() of a new connection string for a shard.
+     */
+    enum class ConnectionStringUpdateType { kConfirmed, kPossible };
+
+    /**
      * Instantiates a new shard registry.
      *
      * @param shardFactory      Makes shards
@@ -256,7 +261,8 @@ public:
      * up the corresponding Shard object based on the replica set name, then updates the
      * ShardRegistry's notion of what hosts make up that shard.
      */
-    void updateReplSetHosts(const ConnectionString& newConnString);
+    void updateReplSetHosts(const ConnectionString& givenConnString,
+                            ConnectionStringUpdateType updateType);
 
     /**
      * Instantiates a new detached shard connection, which does not appear in the list of shards
diff --git a/src/mongo/s/mongos_main.cpp b/src/mongo/s/mongos_main.cpp
index e3440eb151e..4dfdc0414b1 100644
--- a/src/mongo/s/mongos_main.cpp
+++ b/src/mongo/s/mongos_main.cpp
@@ -521,7 +521,10 @@ public:
             LOGV2(471693,
                   "Updating the shard registry with confirmed replica set",
                   "connectionString"_attr = connStr);
-            Grid::get(_serviceContext)->shardRegistry()->updateReplSetHosts(connStr);
+            Grid::get(_serviceContext)
+                ->shardRegistry()
+                ->updateReplSetHosts(connStr,
+                                     ShardRegistry::ConnectionStringUpdateType::kConfirmed);
         } catch (const ExceptionForCat<ErrorCategory::ShutdownError>& e) {
             LOGV2(471694,
                   "Unable to update the shard registry with confirmed replica set",
@@ -547,7 +550,10 @@ public:
 
     void onPossibleSet(const State& state) noexcept final {
         try {
-            Grid::get(_serviceContext)->shardRegistry()->updateReplSetHosts(state.connStr);
+            Grid::get(_serviceContext)
+                ->shardRegistry()
+                ->updateReplSetHosts(state.connStr,
+                                     ShardRegistry::ConnectionStringUpdateType::kPossible);
         } catch (const DBException& ex) {
             LOGV2_DEBUG(22849,
                         2,