Send OS upgrade budget in applicable clouds

2 files changed, 11 insertions, 2 deletions

diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/configserver/NodeRepository.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/configserver/NodeRepository.java
index b5d49df7e9c..aebfab7cbff 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/configserver/NodeRepository.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/configserver/NodeRepository.java
@@ -14,6 +14,7 @@ import com.yahoo.vespa.hosted.controller.api.integration.noderepository.NodeMemb
 import com.yahoo.vespa.hosted.controller.api.integration.noderepository.NodeRepositoryNode;
 import com.yahoo.vespa.hosted.controller.api.integration.noderepository.NodeState;
 
+import java.time.Duration;
 import java.time.Instant;
 import java.util.Collection;
 import java.util.List;
@@ -76,7 +77,7 @@ public interface NodeRepository {
     void upgrade(ZoneId zone, NodeType type, Version version);
 
     /** Upgrade OS for all nodes of given type to a new version */
-    void upgradeOs(ZoneId zone, NodeType type, Version version);
+    void upgradeOs(ZoneId zone, NodeType type, Version version, Optional<Duration> upgradeBudget);
 
     /** Get target versions for upgrades in given zone */
     TargetVersions targetVersionsOf(ZoneId zone);
diff --git a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/noderepository/NodeUpgrade.java b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/noderepository/NodeUpgrade.java
index 21c1d23ba3a..b8e2a626c72 100644
--- a/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/noderepository/NodeUpgrade.java
+++ b/controller-api/src/main/java/com/yahoo/vespa/hosted/controller/api/integration/noderepository/NodeUpgrade.java
@@ -22,12 +22,16 @@ public class NodeUpgrade {
     @JsonProperty("force")
     private final boolean force;
 
+    @JsonProperty("upgradeBudget")
+    private final String upgradeBudget;
+
     @JsonCreator
     public NodeUpgrade(@JsonProperty("version") String version, @JsonProperty("osVersion") String osVersion,
-                       @JsonProperty("force") boolean force) {
+                       @JsonProperty("force") boolean force, @JsonProperty("upgradeBudget") String upgradeBudget) {
         this.version = version;
         this.osVersion = osVersion;
         this.force = force;
+        this.upgradeBudget = upgradeBudget;
     }
 
     public String getVersion() {
@@ -42,4 +46,8 @@ public class NodeUpgrade {
         return force;
     }
 
+    public String getUpgradeBudget() {
+        return upgradeBudget;
+    }
+
 }