Revert "Revert "Bratseth/select 1 host by index""

3 files changed, 47 insertions, 12 deletions

diff --git a/config-model/src/main/java/com/yahoo/vespa/model/HostResource.java b/config-model/src/main/java/com/yahoo/vespa/model/HostResource.java
index 5e74a2ebc8a..e6ed91165ca 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/HostResource.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/HostResource.java
@@ -7,7 +7,9 @@ import com.yahoo.config.provision.Flavor;
 
 import javax.annotation.Nullable;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -264,4 +266,30 @@ public class HostResource implements Comparable<HostResource> {
         return this.host.compareTo(other.host);
     }
 
+    /**
+     * Picks hosts by some mixture of host name and index 
+     * (where the mix of one or the other is decided by the last parameter).
+     */
+    public static List<HostResource> pickHosts(Collection<HostResource> hosts, int count, int targetHostsSelectedByIndex) {
+        targetHostsSelectedByIndex = Math.min(Math.min(targetHostsSelectedByIndex, count), hosts.size());
+
+        List<HostResource> hostsSortedByName = new ArrayList<>(hosts);
+        Collections.sort(hostsSortedByName);
+
+        List<HostResource> hostsSortedByIndex = new ArrayList<>(hosts);
+        hostsSortedByIndex.sort(Comparator.comparingInt(host -> host.primaryClusterMembership().get().index()));
+        return pickHosts(hostsSortedByName, hostsSortedByIndex, count, targetHostsSelectedByIndex);
+    }
+    public static List<HostResource> pickHosts(List<HostResource> hostsSelectedByName, List<HostResource> hostsSelectedByIndex, 
+                                               int count, int targetHostsSelectedByIndex) {
+        hostsSelectedByName = hostsSelectedByName.subList(0, Math.min(count - targetHostsSelectedByIndex, hostsSelectedByName.size()));
+        hostsSelectedByIndex.removeAll(hostsSelectedByName);
+        hostsSelectedByIndex = hostsSelectedByIndex.subList(0, Math.min(targetHostsSelectedByIndex, hostsSelectedByIndex.size()));
+
+        List<HostResource> finalHosts = new ArrayList<>();
+        finalHosts.addAll(hostsSelectedByName);
+        finalHosts.addAll(hostsSelectedByIndex);
+        return finalHosts;
+    }
+
 }
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/builder/xml/dom/DomAdminV4Builder.java b/config-model/src/main/java/com/yahoo/vespa/model/builder/xml/dom/DomAdminV4Builder.java
index cb8ec205395..f33c86134cb 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/builder/xml/dom/DomAdminV4Builder.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/builder/xml/dom/DomAdminV4Builder.java
@@ -19,6 +19,7 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 /**
  * Builds the admin model from a version 4 XML tag, or as a default when an admin 3 tag or no admin tag is used.
@@ -117,12 +118,11 @@ public class DomAdminV4Builder extends DomAdminBuilderBase {
 
     /** Returns the count first containers in the current model having isRetired set to the given value */
     private List<HostResource> sortedContainerHostsFrom(ContainerModel model, int count, boolean retired) {
-        List<HostResource> hosts = new ArrayList<>();
-        for (Container container : model.getCluster().getContainers())
-            if (retired == container.isRetired())
-                hosts.add(container.getHostResource());
-        Collections.sort(hosts);
-        return hosts.subList(0, Math.min(count, hosts.size()));
+        List<HostResource> hosts = model.getCluster().getContainers().stream()
+                                                                     .filter(container -> retired == container.isRetired())
+                                                                     .map(Container::getHostResource)
+                                                                     .collect(Collectors.toList());
+        return HostResource.pickHosts(hosts, count, 1);
     }
 
     private void createLogserver(Admin admin, Collection<HostResource> hosts) {
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/ContentCluster.java b/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/ContentCluster.java
index 7889b857fff..16b0674fc14 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/ContentCluster.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/ContentCluster.java
@@ -10,6 +10,7 @@ import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.Zone;
+import com.yahoo.lang.MutableInteger;
 import com.yahoo.vespa.config.content.MessagetyperouteselectorpolicyConfig;
 import com.yahoo.vespa.config.content.FleetcontrollerConfig;
 import com.yahoo.vespa.config.content.StorDistributionConfig;
@@ -328,9 +329,11 @@ public class ContentCluster extends AbstractConfigProducer implements StorDistri
         }
 
         private List<HostResource> drawControllerHosts(int count, StorageGroup rootGroup, Collection<ContainerModel> containers) {
-            List<HostResource> hosts = drawContentHostsRecursively(count, rootGroup);
+            List<HostResource> hostsByName = drawContentHostsRecursively(count, false, rootGroup);
+            List<HostResource> hostsByIndex = drawContentHostsRecursively(count, true, rootGroup);
             // if (hosts.size() < count) // supply with containers TODO: Currently disabled due to leading to topology change problems
             //     hosts.addAll(drawContainerHosts(count - hosts.size(), containers, new HashSet<>(hosts)));
+            List<HostResource> hosts = HostResource.pickHosts(hostsByName, hostsByIndex, count, 1);
             if (hosts.size() % 2 == 0) // ZK clusters of even sizes are less available (even in the size=2 case)
                 hosts = hosts.subList(0, hosts.size()-1);
             return hosts;
@@ -403,20 +406,24 @@ public class ContentCluster extends AbstractConfigProducer implements StorDistri
          */
         // Note: This method cannot be changed to draw different nodes without ensuring that it will draw nodes
         //       which overlaps with previously drawn nodes as this will prevent rolling upgrade
-        private List<HostResource> drawContentHostsRecursively(int count, StorageGroup group) {
+        private List<HostResource> drawContentHostsRecursively(int count, boolean byIndex, StorageGroup group) {
             Set<HostResource> hosts = new HashSet<>();
             if (group.getNodes().isEmpty()) {
                 int hostsPerSubgroup = (int)Math.ceil((double)count / group.getSubgroups().size());
                 for (StorageGroup subgroup : group.getSubgroups())
-                    hosts.addAll(drawContentHostsRecursively(hostsPerSubgroup, subgroup));
+                    hosts.addAll(drawContentHostsRecursively(hostsPerSubgroup, byIndex, subgroup));
             }
             else {
                 hosts.addAll(group.getNodes().stream()
-                    .filter(node -> ! node.isRetired()) // Avoid retired controllers to avoid surprises on expiry
-                    .map(StorageNode::getHostResource).collect(Collectors.toList()));
+                     .filter(node -> ! node.isRetired()) // Avoid retired controllers to avoid surprises on expiry
+                     .map(StorageNode::getHostResource).collect(Collectors.toList()));
             }
+
             List<HostResource> sortedHosts = new ArrayList<>(hosts);
-            Collections.sort(sortedHosts);
+            if (byIndex)
+                sortedHosts.sort(Comparator.comparingInt(host -> host.primaryClusterMembership().get().index()));
+            else // by name
+                Collections.sort(sortedHosts);
             sortedHosts = sortedHosts.subList(0, Math.min(count, hosts.size()));
             return sortedHosts;
         }