Implement RebuildingOsUpgrader

8 files changed, 214 insertions, 44 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 56752bc8fd2..0c19cf99539 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
@@ -48,6 +48,11 @@ public class NodeList extends AbstractFilteringList<Node, NodeList> {
         return matching(node -> node.status().wantToRetire() && node.status().wantToDeprovision());
     }
 
+    /** Returns the subset of nodes that are being rebuilt */
+    public NodeList rebuilding() {
+        return matching(node -> node.status().wantToRetire() && node.status().wantToRebuild());
+    }
+
     /** Returns the subset of nodes which are removable */
     public NodeList removable() {
         return matching(node -> node.allocation().isPresent() && node.allocation().get().isRemovable());
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/Agent.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/Agent.java
index ed7af4b4f03..8a943400b37 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/Agent.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/Agent.java
@@ -23,6 +23,7 @@ public enum Agent {
     ReservationExpirer,
     DynamicProvisioningMaintainer,
     RetiringUpgrader,
+    RebuildingOsUpgrader,
     SpareCapacityMaintainer,
     SwitchRebalancer,
 
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/DelegatingOsUpgrader.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/DelegatingOsUpgrader.java
index 65215cecbdf..af17934a878 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/DelegatingOsUpgrader.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/DelegatingOsUpgrader.java
@@ -31,6 +31,7 @@ public class DelegatingOsUpgrader implements OsUpgrader {
     public DelegatingOsUpgrader(NodeRepository nodeRepository, int maxActiveUpgrades) {
         this.nodeRepository = Objects.requireNonNull(nodeRepository);
         this.maxActiveUpgrades = maxActiveUpgrades;
+        if (maxActiveUpgrades < 1) throw new IllegalArgumentException("maxActiveUpgrades must be positive");
     }
 
     @Override
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java
index 2ee7b324582..613738458c2 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java
@@ -4,7 +4,9 @@ package com.yahoo.vespa.hosted.provision.os;
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.curator.Lock;
+import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
+import com.yahoo.vespa.hosted.provision.node.Status;
 import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
 
 import java.time.Duration;
@@ -28,16 +30,20 @@ public class OsVersions {
 
     private static final Logger log = Logger.getLogger(OsVersions.class.getName());
 
+    private final NodeRepository nodeRepository;
     private final CuratorDatabaseClient db;
-    private final OsUpgrader upgrader;
+    private final boolean reprovisionToUpgradeOs;
+    private final int maxDelegatedUpgrades;
 
     public OsVersions(NodeRepository nodeRepository) {
-        this(nodeRepository, upgraderIn(nodeRepository));
+        this(nodeRepository, nodeRepository.zone().getCloud().reprovisionToUpgradeOs(), 30);
     }
 
-    OsVersions(NodeRepository nodeRepository, OsUpgrader upgrader) {
-        this.db = Objects.requireNonNull(nodeRepository).database();
-        this.upgrader = Objects.requireNonNull(upgrader);
+    OsVersions(NodeRepository nodeRepository, boolean reprovisionToUpgradeOs, int maxDelegatedUpgrades) {
+        this.nodeRepository = Objects.requireNonNull(nodeRepository);
+        this.db = nodeRepository.database();
+        this.reprovisionToUpgradeOs = reprovisionToUpgradeOs;
+        this.maxDelegatedUpgrades = maxDelegatedUpgrades;
 
         // Read and write all versions to make sure they are stored in the latest version of the serialized format
         try (var lock = db.lockOsVersionChange()) {
@@ -72,7 +78,10 @@ public class OsVersions {
     public void removeTarget(NodeType nodeType) {
         require(nodeType);
         writeChange((change) -> {
-            upgrader.disableUpgrade(nodeType);
+            Version target = Optional.ofNullable(change.targets().get(nodeType))
+                                     .map(OsVersionTarget::version)
+                                     .orElse(Version.emptyVersion);
+            chooseUpgrader(nodeType, target).disableUpgrade(nodeType);
             return change.withoutTarget(nodeType);
         });
     }
@@ -102,8 +111,9 @@ public class OsVersions {
     public void resumeUpgradeOf(NodeType nodeType, boolean resume) {
         require(nodeType);
         try (Lock lock = db.lockOsVersionChange()) {
-            var target = readChange().targets().get(nodeType);
+            OsVersionTarget target = readChange().targets().get(nodeType);
             if (target == null) return; // No target set for this type
+            OsUpgrader upgrader = chooseUpgrader(nodeType, target.version());
             if (resume) {
                 upgrader.upgradeTo(target);
             } else {
@@ -112,10 +122,21 @@ public class OsVersions {
         }
     }
 
-    private void requireUpgradeBudget(Optional<Duration> upgradeBudget) {
-        if (upgrader instanceof RetiringOsUpgrader && upgradeBudget.isEmpty()) {
-            throw new IllegalArgumentException("Zone requires a time budget for OS upgrades");
+    /** Returns the upgrader to use when upgrading given node type to target */
+    private OsUpgrader chooseUpgrader(NodeType nodeType, Version target) {
+        if (reprovisionToUpgradeOs) {
+            return new RetiringOsUpgrader(nodeRepository);
+        }
+        // Require rebuild if we have any nodes of this type on a major version lower than target
+        boolean rebuildRequired = nodeRepository.nodes().list(Node.State.active).nodeType(nodeType).stream()
+                                                .map(Node::status)
+                                                .map(Status::osVersion)
+                                                .anyMatch(osVersion -> osVersion.current().isPresent() &&
+                                                                       osVersion.current().get().getMajor() < target.getMajor());
+        if (rebuildRequired) {
+            return new RebuildingOsUpgrader(nodeRepository);
         }
+        return new DelegatingOsUpgrader(nodeRepository, maxDelegatedUpgrades);
     }
 
     private static void requireNonZero(Version version) {
@@ -130,11 +151,4 @@ public class OsVersions {
         }
     }
 
-    private static OsUpgrader upgraderIn(NodeRepository nodeRepository) {
-        if (nodeRepository.zone().getCloud().reprovisionToUpgradeOs()) {
-            return new RetiringOsUpgrader(nodeRepository);
-        }
-        return new DelegatingOsUpgrader(nodeRepository, 30);
-    }
-
 }
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
new file mode 100644
index 00000000000..0e10e9f44de
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RebuildingOsUpgrader.java
@@ -0,0 +1,48 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.os;
+
+import com.yahoo.component.Version;
+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.filter.NodeListFilter;
+
+import java.time.Instant;
+import java.util.Optional;
+import java.util.logging.Logger;
+
+/**
+ * An upgrader that retires and rebuilds hosts on stale OS versions. Retirement of each host is spread out in time,
+ * according to a time budget, to avoid potential service impact of retiring too many hosts close together.
+ *
+ * Used in cases where performing an OS upgrade requires rebuilding the host, e.g. when upgrading across major versions.
+ *
+ * @author mpolden
+ */
+public class RebuildingOsUpgrader extends RetiringOsUpgrader {
+
+    private static final Logger LOG = Logger.getLogger(RebuildingOsUpgrader.class.getName());
+
+    public RebuildingOsUpgrader(NodeRepository nodeRepository) {
+        super(nodeRepository);
+    }
+
+    protected void upgradeNodes(NodeList activeNodes, Version version, Instant instant) {
+        activeNodes.osVersionIsBefore(version)
+                   .not().rebuilding()
+                   .byIncreasingOsVersion()
+                   .first(1)
+                   .forEach(node -> rebuild(node, version, instant));
+    }
+
+    private void rebuild(Node host, Version target, Instant now) {
+        LOG.info("Retiring and rebuilding " + host + ": On stale OS version " +
+                 host.status().osVersion().current().map(Version::toFullString).orElse("<unset>") +
+                 ", want " + target);
+        nodeRepository.nodes().rebuild(host.hostname(), Agent.RebuildingOsUpgrader, now);
+        nodeRepository.nodes().upgradeOs(NodeListFilter.from(host), Optional.of(target));
+        nodeRepository.osVersions().writeChange((change) -> change.withRetirementAt(now, host.type()));
+    }
+
+}
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 23e96d65fc1..61d9c6b6b5d 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
@@ -26,7 +26,7 @@ public class RetiringOsUpgrader implements OsUpgrader {
 
     private static final Logger LOG = Logger.getLogger(RetiringOsUpgrader.class.getName());
 
-    private final NodeRepository nodeRepository;
+    protected final NodeRepository nodeRepository;
 
     public RetiringOsUpgrader(NodeRepository nodeRepository) {
         this.nodeRepository = nodeRepository;
@@ -44,11 +44,7 @@ public class RetiringOsUpgrader implements OsUpgrader {
         Instant retiredAt = target.lastRetiredAt().orElse(Instant.EPOCH);
         if (now.isBefore(retiredAt.plus(nodeBudget))) return; // Budget has not been spent yet
 
-        activeNodes.osVersionIsBefore(target.version())
-                   .not().deprovisioning()
-                   .byIncreasingOsVersion()
-                   .first(1)
-                   .forEach(node -> upgrade(node, target.version(), now));
+        upgradeNodes(activeNodes, target.version(), now);
     }
 
     @Override
@@ -56,8 +52,16 @@ public class RetiringOsUpgrader implements OsUpgrader {
         // No action needed in this implementation.
     }
 
+    protected void upgradeNodes(NodeList activeNodes, Version version, Instant instant) {
+        activeNodes.osVersionIsBefore(version)
+                   .not().deprovisioning()
+                   .byIncreasingOsVersion()
+                   .first(1)
+                   .forEach(node -> deprovision(node, version, instant));
+    }
+
     /** Upgrade given host by retiring and deprovisioning it */
-    private void upgrade(Node host, Version target, Instant now) {
+    private void deprovision(Node host, Version target, Instant now) {
         LOG.info("Retiring and deprovisioning " + host + ": On stale OS version " +
                  host.status().osVersion().current().map(Version::toFullString).orElse("<unset>") +
                  ", want " + target);
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java
index cc3fd75a22c..97b9393bdd4 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java
@@ -478,6 +478,7 @@ public class NodeSerializer {
             case "Rebalancer" : return Agent.Rebalancer;
             case "ReservationExpirer" : return Agent.ReservationExpirer;
             case "RetiringUpgrader" : return Agent.RetiringUpgrader;
+            case "RebuildingOsUpgrader" : return Agent.RebuildingOsUpgrader;
             case "SpareCapacityMaintainer": return Agent.SpareCapacityMaintainer;
             case "SwitchRebalancer": return Agent.SwitchRebalancer;
         }
@@ -498,6 +499,7 @@ public class NodeSerializer {
             case Rebalancer : return "Rebalancer";
             case ReservationExpirer : return "ReservationExpirer";
             case RetiringUpgrader: return "RetiringUpgrader";
+            case RebuildingOsUpgrader: return "RebuildingOsUpgrader";
             case SpareCapacityMaintainer: return "SpareCapacityMaintainer";
             case SwitchRebalancer: return "SwitchRebalancer";
         }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java
index 3b967a9245d..8c9e52e80f5 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java
@@ -38,7 +38,7 @@ public class OsVersionsTest {
 
     @Test
     public void upgrade() {
-        var versions = new OsVersions(tester.nodeRepository(), new DelegatingOsUpgrader(tester.nodeRepository(), Integer.MAX_VALUE));
+        var versions = new OsVersions(tester.nodeRepository());
         provisionInfraApplication(10);
         Supplier<NodeList> hostNodes = () -> tester.nodeRepository().nodes().list().nodeType(NodeType.host);
 
@@ -91,7 +91,7 @@ public class OsVersionsTest {
     public void max_active_upgrades() {
         int totalNodes = 20;
         int maxActiveUpgrades = 5;
-        var versions = new OsVersions(tester.nodeRepository(), new DelegatingOsUpgrader(tester.nodeRepository(), maxActiveUpgrades));
+        var versions = new OsVersions(tester.nodeRepository(), false, maxActiveUpgrades);
         provisionInfraApplication(totalNodes);
         Supplier<NodeList> hostNodes = () -> tester.nodeRepository().nodes().list().state(Node.State.active).hosts();
 
@@ -127,7 +127,7 @@ public class OsVersionsTest {
                                             .subList(0, maxActiveUpgrades);
             assertEquals("Nodes on lowest version are told to upgrade",
                          nodesUpgrading.asList(), nodesOnLowestVersion);
-            completeUpgradeOf(nodesUpgrading.asList());
+            completeReprovisionOf(nodesUpgrading.asList());
         }
 
         // Activating again after all nodes have upgraded does nothing
@@ -137,7 +137,7 @@ public class OsVersionsTest {
 
     @Test
     public void newer_upgrade_aborts_upgrade_to_stale_version() {
-        var versions = new OsVersions(tester.nodeRepository(), new DelegatingOsUpgrader(tester.nodeRepository(), Integer.MAX_VALUE));
+        var versions = new OsVersions(tester.nodeRepository());
         provisionInfraApplication(10);
         Supplier<NodeList> hostNodes = () -> tester.nodeRepository().nodes().list().hosts();
 
@@ -156,7 +156,7 @@ public class OsVersionsTest {
 
     @Test
     public void upgrade_by_retiring() {
-        var versions = new OsVersions(tester.nodeRepository(), new RetiringOsUpgrader(tester.nodeRepository()));
+        var versions = new OsVersions(tester.nodeRepository(), true, Integer.MAX_VALUE);
         var clock = (ManualClock) tester.nodeRepository().clock();
         int hostCount = 10;
         // Provision hosts and children
@@ -183,7 +183,7 @@ public class OsVersionsTest {
         versions.resumeUpgradeOf(NodeType.host, true);
         NodeList nodesDeprovisioning = hostNodes.get().deprovisioning();
         assertEquals(1, nodesDeprovisioning.size());
-        assertEquals(2, retiringChildrenOf(nodesDeprovisioning.asList().get(0)).size());
+        assertEquals(2, deprovisioningChildrenOf(nodesDeprovisioning.asList().get(0)).size());
 
         // Budget has been spent and another host is retired
         clock.advance(nodeBudget);
@@ -191,7 +191,7 @@ public class OsVersionsTest {
         assertEquals(2, hostNodes.get().deprovisioning().size());
 
         // Two nodes complete their upgrade by being reprovisioned
-        completeUpgradeOf(hostNodes.get().deprovisioning().asList());
+        completeReprovisionOf(hostNodes.get().deprovisioning().asList());
         assertEquals(2, hostNodes.get().onOsVersion(version1).size());
         // The remaining hosts complete their upgrade
         for (int i = 0; i < hostCount - 2; i++) {
@@ -199,13 +199,13 @@ public class OsVersionsTest {
             versions.resumeUpgradeOf(NodeType.host, true);
             nodesDeprovisioning = hostNodes.get().deprovisioning();
             assertEquals(1, nodesDeprovisioning.size());
-            assertEquals(2, retiringChildrenOf(nodesDeprovisioning.asList().get(0)).size());
-            completeUpgradeOf(nodesDeprovisioning.asList());
+            assertEquals(2, deprovisioningChildrenOf(nodesDeprovisioning.asList().get(0)).size());
+            completeReprovisionOf(nodesDeprovisioning.asList());
         }
 
         // All hosts upgraded and none are deprovisioning
         assertEquals(hostCount, hostNodes.get().onOsVersion(version1).not().deprovisioning().size());
-        assertEquals(hostCount, tester.nodeRepository().nodes().list().state(Node.State.deprovisioned).size());
+        assertEquals(hostCount, tester.nodeRepository().nodes().list(Node.State.deprovisioned).size());
         var lastRetiredAt = clock.instant().truncatedTo(ChronoUnit.MILLIS);
 
         // Resuming after everything has upgraded does nothing
@@ -221,7 +221,7 @@ public class OsVersionsTest {
 
     @Test
     public void upgrade_by_retiring_everything_at_once() {
-        var versions = new OsVersions(tester.nodeRepository(), new RetiringOsUpgrader(tester.nodeRepository()));
+        var versions = new OsVersions(tester.nodeRepository(), true, Integer.MAX_VALUE);
         int hostCount = 3;
         provisionInfraApplication(hostCount, NodeType.confighost);
         Supplier<NodeList> hostNodes = () -> tester.nodeRepository().nodes().list()
@@ -238,15 +238,94 @@ public class OsVersionsTest {
         // All hosts are deprovisioning
         assertEquals(hostCount, hostNodes.get().deprovisioning().size());
         // Nodes complete their upgrade by being reprovisioned
-        completeUpgradeOf(hostNodes.get().deprovisioning().asList(), NodeType.confighost);
+        completeReprovisionOf(hostNodes.get().deprovisioning().asList(), NodeType.confighost);
         assertEquals(hostCount, hostNodes.get().onOsVersion(version1).size());
     }
 
-    private NodeList retiringChildrenOf(Node parent) {
+    @Test
+    public void upgrade_by_rebuilding() {
+        var versions = new OsVersions(tester.nodeRepository(), false, Integer.MAX_VALUE);
+        var clock = tester.clock();
+        int hostCount = 10;
+        provisionInfraApplication(hostCount + 1);
+        Supplier<NodeList> hostNodes = () -> tester.nodeRepository().nodes().list().nodeType(NodeType.host);
+
+        // All hosts upgrade to first version. Upgrades are delegated
+        var version0 = Version.fromString("7.0");
+        versions.setTarget(NodeType.host, version0, Duration.ZERO, false);
+        setCurrentVersion(hostNodes.get().asList(), version0);
+
+        // One host is failed out
+        Node failedHost = tester.nodeRepository().nodes().fail(hostNodes.get().first().get().hostname(),
+                                                               Agent.system, getClass().getSimpleName());
+
+        // Target is set for new major version. Upgrade mechanism switches to rebuilding
+        var version1 = Version.fromString("8.0");
+        Duration totalBudget = Duration.ofHours(12);
+        Duration nodeBudget = totalBudget.dividedBy(hostCount);
+        versions.setTarget(NodeType.host, version1, totalBudget, false);
+        versions.resumeUpgradeOf(NodeType.host, true);
+
+        // One host starts rebuilding
+        assertEquals(1, hostNodes.get().rebuilding().size());
+
+        // Nothing happens on next resume as first host has not spent its budget
+        versions.resumeUpgradeOf(NodeType.host, true);
+        assertEquals(1, hostNodes.get().rebuilding().size());
+
+        // Budget has been spent and another host is rebuilt
+        clock.advance(nodeBudget);
+        versions.resumeUpgradeOf(NodeType.host, true);
+        NodeList hostsRebuilding = hostNodes.get().rebuilding();
+        assertEquals(2, hostsRebuilding.size());
+
+        // Hosts are rebuilt
+        completeRebuildOf(hostsRebuilding.asList(), NodeType.host);
+        assertEquals(2, hostNodes.get().onOsVersion(version1).size());
+
+        // The remaining hosts complete their upgrade
+        for (int i = 0; i < hostCount - 2; i++) {
+            clock.advance(nodeBudget);
+            versions.resumeUpgradeOf(NodeType.host, true);
+            hostsRebuilding = hostNodes.get().rebuilding();
+            assertEquals(1, hostsRebuilding.size());
+            completeRebuildOf(hostsRebuilding.asList(), NodeType.host);
+        }
+
+        // All hosts upgraded and none are rebuilding
+        assertEquals(hostCount, hostNodes.get().onOsVersion(version1).not().rebuilding().size());
+        assertEquals(hostCount, tester.nodeRepository().nodes().list(Node.State.active).size());
+
+        // Resuming after everything has upgraded has no effect
+        versions.resumeUpgradeOf(NodeType.host, true);
+        assertEquals(0, hostNodes.get().rebuilding().size());
+
+        // Next version is within same major. Upgrade mechanism switches to delegated
+        var version2 = Version.fromString("8.1");
+        versions.setTarget(NodeType.host, version2, totalBudget, false);
+        versions.resumeUpgradeOf(NodeType.host, true);
+        NodeList nonFailingHosts = hostNodes.get().except(failedHost);
+        assertTrue("Wanted version is set", nonFailingHosts.except(failedHost).stream()
+                                                    .allMatch(node -> node.status().osVersion().wanted().isPresent()));
+        setCurrentVersion(nonFailingHosts.asList(), version2);
+        assertEquals(10, hostNodes.get().except(failedHost).onOsVersion(version2).size());
+
+        // Failed host is reactivated
+        Node reactivatedHost = tester.nodeRepository().nodes().reactivate(failedHost.hostname(), Agent.system, getClass().getSimpleName());
+        assertEquals(version0, reactivatedHost.status().osVersion().current().get());
+
+        // Resuming upgrades reactivated host. Upgrade mechanism switches to rebuilding
+        clock.advance(nodeBudget);
+        versions.resumeUpgradeOf(NodeType.host, true);
+        hostsRebuilding = hostNodes.get().rebuilding();
+        assertEquals(List.of(reactivatedHost), hostsRebuilding.asList());
+        completeRebuildOf(hostsRebuilding.asList(), NodeType.host);
+    }
+
+    private NodeList deprovisioningChildrenOf(Node parent) {
         return tester.nodeRepository().nodes().list()
                      .childrenOf(parent)
-                     .matching(child -> child.status().wantToRetire() &&
-                                        child.status().wantToDeprovision());
+                     .deprovisioning();
     }
 
     private List<Node> provisionInfraApplication(int nodeCount) {
@@ -281,18 +360,17 @@ public class OsVersionsTest {
         tester.patchNodes(nodes, node -> node.with(node.status().withOsVersion(node.status().osVersion().withCurrent(Optional.of(currentVersion)))));
     }
 
-    private void completeUpgradeOf(List<Node> nodes) {
-        completeUpgradeOf(nodes, NodeType.host);
+    private void completeReprovisionOf(List<Node> nodes) {
+        completeReprovisionOf(nodes, NodeType.host);
     }
 
-    private void completeUpgradeOf(List<Node> nodes, NodeType nodeType) {
+    private void completeReprovisionOf(List<Node> nodes, NodeType nodeType) {
         // Complete upgrade by deprovisioning stale hosts and provisioning new ones
         tester.patchNodes(nodes, (node) -> {
             Optional<Version> wantedOsVersion = node.status().osVersion().wanted();
             if (node.status().wantToDeprovision()) {
-                // Complete upgrade by deprovisioning stale hosts and provisioning new ones
                 tester.nodeRepository().nodes().park(node.hostname(), false, Agent.system,
-                                                     OsVersionsTest.class.getSimpleName());
+                                                     getClass().getSimpleName());
                 tester.nodeRepository().nodes().removeRecursively(node.hostname());
                 node = provisionInfraApplication(1, nodeType).get(0);
             }
@@ -300,4 +378,21 @@ public class OsVersionsTest {
         });
     }
 
+    private void completeRebuildOf(List<Node> nodes, NodeType nodeType) {
+        // Complete upgrade by rebuilding stale hosts
+        tester.patchNodes(nodes, (node) -> {
+            Optional<Version> wantedOsVersion = node.status().osVersion().wanted();
+            if (node.status().wantToRebuild()) {
+                tester.nodeRepository().nodes().park(node.hostname(), false, Agent.system,
+                                                     getClass().getSimpleName());
+                tester.nodeRepository().nodes().removeRecursively(node.hostname());
+                node = tester.nodeRepository().nodes().restore(node.hostname(), Agent.system, getClass().getSimpleName());
+                node = tester.nodeRepository().nodes().setReady(node.hostname(), Agent.system, getClass().getSimpleName());
+                tester.prepareAndActivateInfraApplication(infraApplication, nodeType);
+                node = tester.nodeRepository().nodes().node(node.hostname()).get();
+            }
+            return node.with(node.status().withOsVersion(node.status().osVersion().withCurrent(wantedOsVersion)));
+        });
+    }
+
 }