Retry OS upgrade if retirement is canceled

3 files changed, 5 insertions, 5 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsUpgrader.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsUpgrader.java
index 120b7f00b38..32fe9ba9f7b 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsUpgrader.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsUpgrader.java
@@ -34,12 +34,12 @@ public abstract class OsUpgrader {
     abstract void disableUpgrade(NodeType type);
 
     /** Returns the number of upgrade slots available for given target */
-    final int upgradeSlots(OsVersionTarget target, NodeList activeNodes) {
-        if (!activeNodes.stream().allMatch(node -> node.type() == target.nodeType())) {
+    final int upgradeSlots(OsVersionTarget target, NodeList candidates) {
+        if (!candidates.stream().allMatch(node -> node.type() == target.nodeType())) {
             throw new IllegalArgumentException("All node types must type of OS version target " + target.nodeType());
         }
         int max = target.nodeType() == NodeType.host ? maxActiveUpgrades.value() : 1;
-        int upgrading = activeNodes.changingOsVersionTo(target.version()).size();
+        int upgrading = candidates.changingOsVersionTo(target.version()).size();
         return Math.max(0, max - upgrading);
     }
 
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 f6779d08fd7..212bf5ffb12 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
@@ -67,7 +67,7 @@ public class RebuildingOsUpgrader extends OsUpgrader {
 
         // Rebuild hosts not containing stateful clusters with retiring nodes, up to rebuild limit
         NodeList activeHosts = hostsOfTargetType.state(Node.State.active);
-        int rebuildLimit = upgradeSlots(target, activeHosts);
+        int rebuildLimit = upgradeSlots(target, activeHosts.rebuilding(softRebuild));
         List<Node> hostsToRebuild = new ArrayList<>(rebuildLimit);
         NodeList candidates = activeHosts.not().rebuilding(softRebuild)
                                          .osVersionIsBefore(target.version())
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringOsUpgrader.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringOsUpgrader.java
index 4d98885b72c..de4915d60aa 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringOsUpgrader.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringOsUpgrader.java
@@ -56,7 +56,7 @@ public class RetiringOsUpgrader extends OsUpgrader {
                           .osVersionIsBefore(target.version())
                           .matching(node -> canUpgradeAt(instant, node))
                           .byIncreasingOsVersion()
-                          .first(upgradeSlots(target, activeNodes));
+                          .first(upgradeSlots(target, activeNodes.deprovisioning()));
     }
 
     /** Upgrade given host by retiring and deprovisioning it */