1 files changed, 10 insertions, 14 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java
index fa64a2677f4..f4dcf7f6088 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java
@@ -3,6 +3,7 @@ package com.yahoo.vespa.hosted.controller.maintenance;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.CloudName;
+import com.yahoo.config.provision.zone.NodeSlice;
 import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.text.Text;
 import com.yahoo.vespa.hosted.controller.Controller;
@@ -54,15 +55,16 @@ public class OsUpgrader extends InfrastructureUpgrader<OsVersionTarget> {
     }
 
     @Override
-    protected boolean convergedOn(OsVersionTarget target, SystemApplication application, ZoneApi zone) {
-        return !currentVersion(zone, application, target.osVersion().version()).isBefore(target.osVersion().version());
+    protected boolean convergedOn(OsVersionTarget target, SystemApplication application, ZoneApi zone, NodeSlice nodeSlice) {
+        Version currentVersion = versionOf(nodeSlice, zone, application, Node::currentOsVersion).orElse(target.osVersion().version());
+        return !currentVersion.isBefore(target.osVersion().version());
     }
 
     @Override
     protected boolean expectUpgradeOf(Node node, SystemApplication application, ZoneApi zone) {
         return cloud.equals(zone.getCloudName()) && // Cloud is managed by this upgrader
                application.shouldUpgradeOs() &&     // Application should upgrade in this cloud
-               canUpgrade(node);                    // Node is in an upgradable state
+               canUpgrade(node, false);
     }
 
     @Override
@@ -83,29 +85,23 @@ public class OsUpgrader extends InfrastructureUpgrader<OsVersionTarget> {
                            .orElse(true);
     }
 
-    private Version currentVersion(ZoneApi zone, SystemApplication application, Version defaultVersion) {
-        return minVersion(zone, application, Node::currentOsVersion).orElse(defaultVersion);
-    }
-
     /** Returns the available upgrade budget for given zone */
     private Duration zoneBudgetOf(Duration totalBudget, ZoneApi zone) {
         if (!spendBudgetOn(zone)) return Duration.ZERO;
         long consecutiveZones = upgradePolicy.steps().stream()
-                                             .filter(parallelZones -> parallelZones.stream().anyMatch(this::spendBudgetOn))
+                                             .filter(step -> step.zones().stream().anyMatch(this::spendBudgetOn))
                                              .count();
         return totalBudget.dividedBy(consecutiveZones);
     }
 
     /** Returns whether to spend upgrade budget on given zone */
     private boolean spendBudgetOn(ZoneApi zone) {
-        if (!zone.getEnvironment().isProduction()) return false;
-        if (controller().zoneRegistry().systemZone().getVirtualId().equals(zone.getVirtualId())) return false; // Controller zone
-        return true;
+        return !controller().zoneRegistry().systemZone().getVirtualId().equals(zone.getVirtualId()); // Do not spend budget on controller zone
     }
 
-    /** Returns whether node is in a state where it can be upgraded */
-    public static boolean canUpgrade(Node node) {
-        return upgradableNodeStates.contains(node.state());
+    /** Returns whether node currently allows upgrades */
+    public static boolean canUpgrade(Node node, boolean includeDeferring) {
+        return (includeDeferring || !node.deferOsUpgrade()) && upgradableNodeStates.contains(node.state());
     }
 
     private static String name(CloudName cloud) {