Only consider allocatable hosts

3 files changed, 11 insertions, 9 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 869a06ac530..4ff782eaf89 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
@@ -13,6 +13,7 @@ import java.util.EnumSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -107,7 +108,11 @@ public class NodeList implements Iterable<Node> {
 
     /** Returns the subset of nodes that are in any of the given state(s) */
     public NodeList state(Node.State first, Node.State... rest) {
-        EnumSet<Node.State> nodeStates = EnumSet.of(first, rest);
+        return state(EnumSet.of(first, rest));
+    }
+
+    /** Returns the subset of nodes that are in any of the given state(s) */
+    public NodeList state(Collection<Node.State> nodeStates) {
         return filter(node -> nodeStates.contains(node.state()));
     }
 
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/Rebalancer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/Rebalancer.java
index 4c3134a9382..e8d0839d60e 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/Rebalancer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/Rebalancer.java
@@ -10,11 +10,13 @@ import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.provisioning.DockerHostCapacity;
 import com.yahoo.vespa.hosted.provision.provisioning.HostResourcesCalculator;
+import com.yahoo.vespa.hosted.provision.provisioning.NodePrioritizer;
 
 import java.time.Clock;
 import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.Stream;
 
 public class Rebalancer extends Maintainer {
 
@@ -40,7 +42,7 @@ public class Rebalancer extends Maintainer {
     }
 
     private boolean zoneIsStable(NodeList allNodes) {
-        List<Node> active = allNodes.state(Node.State.active).asList();
+        NodeList active = allNodes.state(Node.State.active);
         if (active.stream().anyMatch(node -> node.allocation().get().membership().retired())) return false;
         if (active.stream().anyMatch(node -> node.status().wantToRetire())) return false;
         return true;
@@ -53,8 +55,8 @@ public class Rebalancer extends Maintainer {
     private Move findBestMove(NodeList allNodes) {
         DockerHostCapacity capacity = new DockerHostCapacity(allNodes, hostResourcesCalculator);
         Move bestMove = Move.none;
-        for (Node node : allNodes.state(Node.State.active).asList()) {
-            for (Node toHost : allNodes.nodeType(NodeType.host).asList()) {
+        for (Node node : allNodes.state(Node.State.active)) {
+            for (Node toHost : allNodes.state(NodePrioritizer.ALLOCATABLE_HOST_STATES).nodeType(NodeType.host)) {
                 if (node.parentHostname().isEmpty()) continue;
                 if (toHost.hostname().equals(node.parentHostname().get())) continue;
                 if ( ! capacity.freeCapacityOf(toHost).satisfies(node.flavor().resources())) continue;
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/PrioritizableNode.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/PrioritizableNode.java
index d6ac5eb2c57..72f8550d063 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/PrioritizableNode.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/PrioritizableNode.java
@@ -4,12 +4,7 @@ package com.yahoo.vespa.hosted.provision.provisioning;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.vespa.hosted.provision.Node;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
 import java.util.Optional;
-import java.util.stream.Collectors;
 
 import static com.yahoo.vespa.hosted.provision.provisioning.NodePrioritizer.ALLOCATABLE_HOST_STATES;