Improve selection criteria for srvMaxHosts (#1590)

Improve selection criteria for srvMaxHosts

This PR improves the selection criteria for srvMaxHosts. The net result is that it will keep existing nodes, discard deleted nodes, and fill up to srvMaxHosts from a random shuffle of newly discovered nodes from the SRV response. It also updates the prose test to match the spec.

JAVA-4400

Author: joykim1005

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

@@ -21,13 +21,15 @@
 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.List;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.stream.Collectors;
 
 import static com.mongodb.assertions.Assertions.assertNotNull;
 
@@ -38,7 +40,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 +58,33 @@ 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> activeHosts = getActivePriorHosts(latestSrvHosts);
+                int numNewHostsToAdd = srvMaxHosts - activeHosts.size();
+                activeHosts.addAll(addShuffledHosts(latestSrvHosts, activeHosts, numNewHostsToAdd));
+
+                return activeHosts;
+            }
+
+            private List<ServerAddress> getActivePriorHosts(final 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> addShuffledHosts(final Collection<ServerAddress> latestSrvHosts,
+                    final List<ServerAddress> activePriorHosts, final int numNewHostsToAdd) {
+                List<ServerAddress> addedHosts = new ArrayList<>(latestSrvHosts);
+                addedHosts.removeAll(activePriorHosts);
+                Collections.shuffle(addedHosts, ThreadLocalRandom.current());
+
+                return addedHosts.subList(0, numNewHostsToAdd);
             }
 
             @Override

Diff for: 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()));
     }