Limit rebuild of infrastructure clusters

1 files changed, 16 insertions, 10 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RebuildingOsUpgrader.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RebuildingOsUpgrader.java
index 71f1af09930..fce9eb562c9 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RebuildingOsUpgrader.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RebuildingOsUpgrader.java
@@ -58,15 +58,21 @@ public class RebuildingOsUpgrader implements OsUpgrader {
         // No action needed in this implementation. Hosts that have started rebuilding cannot be halted
     }
 
+    /** Returns the number of hosts of given type that can be rebuilt concurrently */
+    private int upgradeLimit(NodeType hostType, NodeList hosts) {
+        int limit = hostType == NodeType.host ? maxRebuilds.value() : 1;
+        return Math.max(0, limit - hosts.rebuilding().size());
+    }
+
     private List<Node> rebuildableHosts(OsVersionTarget target, NodeList allNodes) {
         NodeList hostsOfTargetType = allNodes.nodeType(target.nodeType());
         NodeList activeHosts = hostsOfTargetType.state(Node.State.active);
-        int upgradeLimit = Math.max(0, maxRebuilds.value() - hostsOfTargetType.rebuilding().size());
+        int upgradeLimit = upgradeLimit(target.nodeType(), hostsOfTargetType);
 
         // Find stateful clusters with retiring nodes
         NodeList activeNodes = allNodes.state(Node.State.active);
-        Set<ClusterKey> retiringClusters = statefulClustersOf(activeNodes.nodeType(target.nodeType().childNodeType())
-                                                                         .retiring());
+        Set<ClusterId> retiringClusters = statefulClustersOf(activeNodes.nodeType(target.nodeType().childNodeType())
+                                                                        .retiring());
 
         // Upgrade hosts not running stateful clusters that are already retiring
         List<Node> hostsToUpgrade = new ArrayList<>(upgradeLimit);
@@ -75,7 +81,7 @@ public class RebuildingOsUpgrader implements OsUpgrader {
                                          .byIncreasingOsVersion();
         for (Node host : candidates) {
             if (hostsToUpgrade.size() == upgradeLimit) break;
-            Set<ClusterKey> clustersOnHost = statefulClustersOf(activeNodes.childrenOf(host));
+            Set<ClusterId> clustersOnHost = statefulClustersOf(activeNodes.childrenOf(host));
             boolean canUpgrade = Collections.disjoint(retiringClusters, clustersOnHost);
             if (canUpgrade) {
                 hostsToUpgrade.add(host);
@@ -93,24 +99,24 @@ public class RebuildingOsUpgrader implements OsUpgrader {
         nodeRepository.nodes().upgradeOs(NodeListFilter.from(host), Optional.of(target));
     }
 
-    private static Set<ClusterKey> statefulClustersOf(NodeList nodes) {
-        Set<ClusterKey> clusters = new HashSet<>();
+    private static Set<ClusterId> statefulClustersOf(NodeList nodes) {
+        Set<ClusterId> clusters = new HashSet<>();
         for (Node node : nodes) {
             if (node.type().isHost()) throw new IllegalArgumentException("All nodes must be children, got host " + node);
             if (node.allocation().isEmpty()) continue;
             Allocation allocation = node.allocation().get();
             if (!allocation.membership().cluster().isStateful()) continue;
-            clusters.add(new ClusterKey(allocation.owner(), allocation.membership().cluster().id()));
+            clusters.add(new ClusterId(allocation.owner(), allocation.membership().cluster().id()));
         }
         return clusters;
     }
 
-    private static class ClusterKey {
+    private static class ClusterId {
 
         private final ApplicationId application;
         private final ClusterSpec.Id cluster;
 
-        public ClusterKey(ApplicationId application, ClusterSpec.Id cluster) {
+        public ClusterId(ApplicationId application, ClusterSpec.Id cluster) {
             this.application = Objects.requireNonNull(application);
             this.cluster = Objects.requireNonNull(cluster);
         }
@@ -119,7 +125,7 @@ public class RebuildingOsUpgrader implements OsUpgrader {
         public boolean equals(Object o) {
             if (this == o) return true;
             if (o == null || getClass() != o.getClass()) return false;
-            ClusterKey that = (ClusterKey) o;
+            ClusterId that = (ClusterId) o;
             return application.equals(that.application) && cluster.equals(that.cluster);
         }