Consider only non-transient states

2 files changed, 5 insertions, 7 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancer.java
index 9095250827d..112f1fa96f6 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancer.java
@@ -69,9 +69,11 @@ public class SwitchRebalancer extends NodeMover<Move> {
     private NodeList clusterOf(Node node, NodeList allNodes) {
         ApplicationId application = node.allocation().get().owner();
         ClusterSpec.Id cluster = node.allocation().get().membership().cluster().id();
-        // This considers all states to prevent unnecessary moves. E.g. we don't want to start moving nodes to a host
-        // which already contain a failed node in our cluster
-        return allNodes.owner(application).cluster(cluster);
+        // This considers some non-active states to prevent unnecessary moves. E.g. we don't want to start moving nodes
+        // to a host which already contain a failed node in our cluster
+        return allNodes.state(Node.State.reserved, Node.State.active, Node.State.failed, Node.State.parked)
+                       .owner(application)
+                       .cluster(cluster);
     }
 
     /** Returns whether allocatedNode is on an exclusive switch */
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancerTest.java
index 95b993edc6e..dd687c27a7f 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancerTest.java
@@ -92,8 +92,6 @@ public class SwitchRebalancerTest {
 
             // Retired node becomes inactive and makes zone stable
             deactivate(tester, retired);
-            nodesIn(cluster, tester).state(Node.State.dirty)
-                                    .forEach(node -> tester.nodeRepository().nodes().removeRecursively(node, true));
         }
         assertEquals("Rebalanced all clusters", clusters, rebalancedClusters);
 
@@ -187,8 +185,6 @@ public class SwitchRebalancerTest {
         tester.assertSwitches(Set.of(switch0, switch1, switch2, switch3), app, spec.id());
         retired = nodesIn(spec.id(), tester).state(Node.State.active).retired();
         deactivate(tester, retired);
-        nodesIn(spec.id(), tester).state(Node.State.dirty)
-                                  .forEach(node -> tester.nodeRepository().nodes().removeRecursively(node.hostname()));
 
         // Next iteration does nothing
         tester.clock().advance(SwitchRebalancer.waitTimeAfterPreviousDeployment);