Improve selection criteria for srvMaxHosts

driver-core/src/main/com/mongodb/internal/connection/DnsMultiServerCluster.java

@@ -21,13 +21,16 @@
 import com.mongodb.connection.ClusterId;
 import com.mongodb.connection.ClusterSettings;
 import com.mongodb.connection.ClusterType;
+import com.mongodb.connection.ServerDescription;
 import com.mongodb.lang.Nullable;
 
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.stream.Collectors;
 
 import static com.mongodb.assertions.Assertions.assertNotNull;
 
@@ -38,7 +41,6 @@ public final class DnsMultiServerCluster extends AbstractMultiServerCluster {
     private final DnsSrvRecordMonitor dnsSrvRecordMonitor;
     private volatile MongoException srvResolutionException;
 
-
     public DnsMultiServerCluster(final ClusterId clusterId, final ClusterSettings settings, final ClusterableServerFactory serverFactory,
                                  final DnsSrvRecordMonitorFactory dnsSrvRecordMonitorFactory) {
         super(clusterId, settings, serverFactory);
@@ -57,17 +59,35 @@ public void initialize(final Collection<ServerAddress> hosts) {
                 }
             }
 
-            private Collection<ServerAddress> applySrvMaxHosts(final Collection<ServerAddress> hosts) {
-                Collection<ServerAddress> newHosts = hosts;
+            private Collection<ServerAddress> applySrvMaxHosts(final Collection<ServerAddress> latestSrvHosts) {
                 Integer srvMaxHosts = getSettings().getSrvMaxHosts();
-                if (srvMaxHosts != null && srvMaxHosts > 0) {
-                    if (srvMaxHosts < hosts.size()) {
-                        List<ServerAddress> newHostsList = new ArrayList<>(hosts);
-                        Collections.shuffle(newHostsList, ThreadLocalRandom.current());
-                        newHosts = newHostsList.subList(0, srvMaxHosts);
-                    }
+                if (srvMaxHosts == null || srvMaxHosts <= 0 || latestSrvHosts.size() <= srvMaxHosts) {
+                    return new ArrayList<>(latestSrvHosts);
                 }
-                return newHosts;
+                List<ServerAddress> activePriorHosts = getActivePriorHosts(latestSrvHosts);
+                int numNewHostsToAdd = srvMaxHosts - activePriorHosts.size();
+                List<ServerAddress> result = getShuffledLatestSrvHosts(latestSrvHosts, activePriorHosts, numNewHostsToAdd);
+
+                return result;
+            }
+
+            private List<ServerAddress> getActivePriorHosts(Collection<ServerAddress> latestSrvHosts) {
+                List<ServerAddress> priorHosts = DnsMultiServerCluster.this.getCurrentDescription().getServerDescriptions().stream()
+                        .map(ServerDescription::getAddress).collect(Collectors.toList());
+                priorHosts.removeIf(host -> !latestSrvHosts.contains(host));
+
+                return priorHosts;
+            }
+
+            private List<ServerAddress> getShuffledLatestSrvHosts(final Collection<ServerAddress> latestSrvHosts,
+                    List<ServerAddress> activePriorHosts, int numNewHostsToAdd) {
+                List<ServerAddress> addedHosts = new ArrayList<>(latestSrvHosts);
+                addedHosts.removeAll(activePriorHosts);
+                Collections.shuffle(addedHosts, ThreadLocalRandom.current());
+                // add the shuffled latestSrvHosts to the activePriorHosts
+                activePriorHosts.addAll(addedHosts.subList(0, numNewHostsToAdd));
+
+                return activePriorHosts;
             }
 
             @Override

driver-core/src/test/unit/com/mongodb/internal/connection/SrvPollingProseTests.java

@@ -160,9 +160,10 @@ public void shouldUseAllRecordsWhenSrvMaxHostsIsGreaterThanOrEqualToNumSrvRecord
     public void shouldUseSrvMaxHostsWhenSrvMaxHostsIsLessThanNumSrvRecords() {
         int srvMaxHosts = 2;
         List<String> updatedHosts = asList(firstHost, thirdHost, fourthHost);
-
         initCluster(updatedHosts, srvMaxHosts);
+
         assertEquals(srvMaxHosts, clusterHostsSet().size());
+        assertTrue(updatedHosts.contains(firstHost));
         assertTrue(updatedHosts.containsAll(clusterHostsSet()));
     }