Do not spend upgrade budget on controller zone

2 files changed, 24 insertions, 7 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 e71fcf12b23..203c8187c2c 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,7 +3,6 @@ package com.yahoo.vespa.hosted.controller.maintenance;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.CloudName;
-import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.vespa.hosted.controller.Controller;
 import com.yahoo.vespa.hosted.controller.api.integration.configserver.Node;
@@ -89,12 +88,18 @@ public class OsUpgrader extends InfrastructureUpgrader<OsVersionTarget> {
 
     /** Returns the available upgrade budget for given zone */
     private Duration zoneBudgetOf(Duration totalBudget, ZoneApi zone) {
-        if (!zone.getEnvironment().isProduction()) return Duration.ZERO;
-        long consecutiveProductionZones = upgradePolicy.asList().stream()
-                                                       .filter(parallelZones -> parallelZones.stream().map(ZoneApi::getEnvironment)
-                                                                                             .anyMatch(Environment::isProduction))
-                                                       .count();
-        return totalBudget.dividedBy(consecutiveProductionZones);
+        if (!spendBudget(zone)) return Duration.ZERO;
+        long consecutiveZones = upgradePolicy.asList().stream()
+                                             .filter(parallelZones -> parallelZones.stream().anyMatch(this::spendBudget))
+                                             .count();
+        return totalBudget.dividedBy(consecutiveZones);
+    }
+
+    /** Returns whether to spend upgrade budget on given zone */
+    private boolean spendBudget(ZoneApi zone) {
+        if (!zone.getEnvironment().isProduction()) return false;
+        if (controller().zoneRegistry().systemZone().getVirtualId().equals(zone.getVirtualId())) return false; // Controller zone
+        return true;
     }
 
     /** Returns whether node is in a state where it can be upgraded */
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java
index 3e2fd4ec0b9..664a1fdc83c 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java
@@ -119,11 +119,13 @@ public class OsUpgraderTest {
     @Test
     public void upgrade_os_with_budget() {
         CloudName cloud = CloudName.from("cloud");
+        ZoneApi zone0 = zone("prod.us-north-42", "prod.controller", cloud);
         ZoneApi zone1 = zone("dev.us-east-1", cloud);
         ZoneApi zone2 = zone("prod.us-west-1", cloud);
         ZoneApi zone3 = zone("prod.us-central-1", cloud);
         ZoneApi zone4 = zone("prod.eu-west-1", cloud);
         UpgradePolicy upgradePolicy = UpgradePolicy.create()
+                                                   .upgrade(zone0)
                                                    .upgrade(zone1)
                                                    .upgradeInParallel(zone2, zone3)
                                                    .upgrade(zone4);
@@ -133,6 +135,7 @@ public class OsUpgraderTest {
         List<SystemApplication> nodeTypes = List.of(SystemApplication.configServerHost, SystemApplication.tenantHost);
         tester.configServer().bootstrap(List.of(zone1.getId(), zone2.getId(), zone3.getId(), zone4.getId()),
                                         nodeTypes);
+        tester.configServer().addNodes(List.of(zone0.getVirtualId()), List.of(SystemApplication.controllerHost));
 
         // Upgrade with budget
         Version version = Version.fromString("7.1");
@@ -141,7 +144,16 @@ public class OsUpgraderTest {
         statusUpdater.maintain();
         osUpgrader.maintain();
 
+        // Controllers upgrade first
+        osUpgrader.maintain();
+        assertWanted(version, SystemApplication.controllerHost, zone0);
+        assertEquals("Controller zone gets a zero budget", Duration.ZERO, upgradeBudget(zone0, SystemApplication.controllerHost, version));
+        completeUpgrade(version, SystemApplication.controllerHost, zone0);
+        statusUpdater.maintain();
+        assertEquals(3, nodesOn(version).size());
+
         // First zone upgrades
+        osUpgrader.maintain();
         for (var nodeType : nodeTypes) {
             assertEquals("Dev zone gets a zero budget", Duration.ZERO, upgradeBudget(zone1, nodeType, version));
             completeUpgrade(version, nodeType, zone1);