Move method to NodeList

2 files changed, 14 insertions, 16 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java
index 4e9468925b6..d74817fb8cd 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java
@@ -8,6 +8,7 @@ import com.yahoo.config.provision.ClusterResources;
 import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
+import com.yahoo.vespa.hosted.provision.node.ClusterId;
 
 import java.util.Comparator;
 import java.util.EnumSet;
@@ -225,6 +226,16 @@ public class NodeList extends AbstractFilteringList<Node, NodeList> {
         return stream().map(Node::hostname).collect(Collectors.toUnmodifiableSet());
     }
 
+    /** Returns the stateful clusters on nodes in this */
+    public Set<ClusterId> statefulClusters() {
+        return stream().filter(node -> node.allocation().isPresent() &&
+                                       node.allocation().get().membership().cluster().isStateful())
+                       .map(node -> new ClusterId(node.allocation().get().owner(),
+                                                  node.allocation().get().membership().cluster().id()))
+                       .collect(Collectors.toUnmodifiableSet());
+
+    }
+
     /**
      * Returns the cluster spec of the nodes in this, without any group designation
      *
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 20536b9dd9f..e373f5edbfb 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
@@ -9,7 +9,6 @@ import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.node.Agent;
-import com.yahoo.vespa.hosted.provision.node.Allocation;
 import com.yahoo.vespa.hosted.provision.node.ClusterId;
 import com.yahoo.vespa.hosted.provision.node.filter.NodeListFilter;
 
@@ -69,8 +68,8 @@ public class RebuildingOsUpgrader implements OsUpgrader {
 
         // Find stateful clusters with retiring nodes
         NodeList activeNodes = allNodes.state(Node.State.active);
-        Set<ClusterId> retiringClusters = statefulClustersOf(activeNodes.nodeType(target.nodeType().childNodeType())
-                                                                        .retiring());
+        Set<ClusterId> retiringClusters = new HashSet<>(activeNodes.nodeType(target.nodeType().childNodeType())
+                                                                   .retiring().statefulClusters());
 
         // Rebuild hosts not containing stateful clusters with retiring nodes, up to rebuild limit
         List<Node> hostsToRebuild = new ArrayList<>(rebuildLimit);
@@ -80,7 +79,7 @@ public class RebuildingOsUpgrader implements OsUpgrader {
                                                .byIncreasingOsVersion();
         for (Node host : candidates) {
             if (hostsToRebuild.size() == rebuildLimit) break;
-            Set<ClusterId> clustersOnHost = statefulClustersOf(activeNodes.childrenOf(host));
+            Set<ClusterId> clustersOnHost = activeNodes.childrenOf(host).statefulClusters();
             boolean canRebuild = Collections.disjoint(retiringClusters, clustersOnHost);
             if (canRebuild) {
                 hostsToRebuild.add(host);
@@ -98,18 +97,6 @@ public class RebuildingOsUpgrader implements OsUpgrader {
         nodeRepository.nodes().upgradeOs(NodeListFilter.from(host), Optional.of(target));
     }
 
-    private static Set<ClusterId> statefulClustersOf(NodeList nodes) {
-        Set<ClusterId> clusters = new HashSet<>();
-        for (Node node : nodes) {
-            if (node.type().isHost()) illegal("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 ClusterId(allocation.owner(), allocation.membership().cluster().id()));
-        }
-        return clusters;
-    }
-
     private static void illegal(String msg) {
         throw new IllegalArgumentException(msg);
     }