Read upgradeBudget parameter in REST API

4 files changed, 24 insertions, 67 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiHandler.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiHandler.java
index 3a70eb71178..ac63f0cbcaf 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiHandler.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiHandler.java
@@ -27,7 +27,9 @@ import com.yahoo.yolean.Exceptions;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.UncheckedIOException;
+import java.time.Duration;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.StringJoiner;
 import java.util.logging.Level;
@@ -133,6 +135,7 @@ public class OsApiHandler extends AuditLoggingRequestHandler {
         Inspector root = requestData.get();
         Inspector versionField = root.field("version");
         Inspector cloudField = root.field("cloud");
+        Inspector upgradeBudgetField = root.field("upgradeBudget");
         boolean force = root.field("force").asBool();
         if (!versionField.valid() || !cloudField.valid()) {
             throw new IllegalArgumentException("Fields 'version' and 'cloud' are required");
@@ -145,12 +148,22 @@ public class OsApiHandler extends AuditLoggingRequestHandler {
         } catch (IllegalArgumentException e) {
             throw new IllegalArgumentException("Invalid version '" + versionField.asString() + "'", e);
         }
-
-        controller.upgradeOsIn(cloud, target, force);
+        Optional<Duration> upgradeBudget = Optional.of(upgradeBudgetField)
+                                                   .filter(Inspector::valid)
+                                                   .map(Inspector::asString).map(s -> {
+                    try {
+                        return Duration.parse(s);
+                    } catch (Exception e2) {
+                        throw new IllegalArgumentException("Invalid duration '" + s + "'", e2);
+                    }
+                });
+
+        controller.upgradeOsIn(cloud, target, upgradeBudget, force);
         Slime response = new Slime();
         Cursor cursor = response.setObject();
         cursor.setString("message", "Set target OS version for cloud '" + cloud.value() + "' to " +
-                                    target.toFullString());
+                                    target.toFullString() + upgradeBudget.map(d -> " with upgrade budget " + d)
+                                                                         .orElse(""));
         return response;
     }
 
@@ -163,7 +176,9 @@ public class OsApiHandler extends AuditLoggingRequestHandler {
         controller.osVersionStatus().versions().forEach((osVersion, nodeVersions) -> {
             Cursor currentVersionObject = versions.addObject();
             currentVersionObject.setString("version", osVersion.version().toFullString());
-            currentVersionObject.setBool("targetVersion", targets.stream().anyMatch(target -> target.osVersion().equals(osVersion)));
+            Optional<OsVersionTarget> target = targets.stream().filter(t -> t.osVersion().equals(osVersion)).findFirst();
+            currentVersionObject.setBool("targetVersion", target.isPresent());
+            target.flatMap(OsVersionTarget::upgradeBudget).ifPresent(budget -> currentVersionObject.setString("upgradeBudget", budget.toString()));
             currentVersionObject.setString("cloud", osVersion.cloud().value());
             Cursor nodesArray = currentVersionObject.setArray("nodes");
             nodeVersions.asMap().values().forEach(nodeVersion -> {
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java
index 42515fd0d3b..5534b6937b2 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java
@@ -38,7 +38,7 @@ public class OsApiTest extends ControllerContainerTest {
     private static final String responses = "src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/";
     private static final AthenzIdentity operator = AthenzUser.fromUserId("operatorUser");
     private static final Cloud cloud1 = new Cloud(CloudName.from("cloud1"), false, true, false, false);
-    private static final Cloud cloud2 = new Cloud(CloudName.from("cloud2"), false, true, false, false);
+    private static final Cloud cloud2 = new Cloud(CloudName.from("cloud2"), true, false, true, true);
     private static final ZoneApi zone1 = ZoneApiMock.newBuilder().withId("prod.us-east-3").with(cloud1.name()).build();
     private static final ZoneApi zone2 = ZoneApiMock.newBuilder().withId("prod.us-west-1").with(cloud1.name()).build();
     private static final ZoneApi zone3 = ZoneApiMock.newBuilder().withId("prod.eu-west-1").with(cloud2.name()).build();
@@ -72,8 +72,8 @@ public class OsApiTest extends ControllerContainerTest {
         // Upgrade OS to a different version in each cloud
         assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"7.5.2\", \"cloud\": \"cloud1\"}", Request.Method.PATCH),
                        "{\"message\":\"Set target OS version for cloud 'cloud1' to 7.5.2\"}", 200);
-        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"8.2.1\", \"cloud\": \"cloud2\"}", Request.Method.PATCH),
-                       "{\"message\":\"Set target OS version for cloud 'cloud2' to 8.2.1\"}", 200);
+        assertResponse(new Request("http://localhost:8080/os/v1/", "{\"version\": \"8.2.1\", \"cloud\": \"cloud2\", \"upgradeBudget\": \"PT24H\"}", Request.Method.PATCH),
+                       "{\"message\":\"Set target OS version for cloud 'cloud2' to 8.2.1 with upgrade budget PT24H\"}", 200);
 
         // Status is updated after some zones are upgraded
         upgradeAndUpdateStatus();
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-all-upgraded.json b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-all-upgraded.json
index 01af1bd70dd..5834e4cef4a 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-all-upgraded.json
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-all-upgraded.json
@@ -100,39 +100,10 @@
     {
       "version": "8.2.1",
       "targetVersion": true,
+      "upgradeBudget": "PT24H",
       "cloud": "cloud2",
       "nodes": [
         {
-          "hostname": "node-1-configserver-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
-          "hostname": "node-2-configserver-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
-          "hostname": "node-3-configserver-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
-          "hostname": "node-1-proxy-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
-          "hostname": "node-3-proxy-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
-          "hostname": "node-2-proxy-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
           "hostname": "node-1-tenant-host-prod.eu-west-1",
           "environment": "prod",
           "region": "eu-west-1"
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-partially-upgraded.json b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-partially-upgraded.json
index 2b907c1156c..c8833fea100 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-partially-upgraded.json
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/versions-partially-upgraded.json
@@ -110,36 +110,6 @@
       "cloud": "cloud2",
       "nodes": [
         {
-          "hostname": "node-1-configserver-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
-          "hostname": "node-2-configserver-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
-          "hostname": "node-3-configserver-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
-          "hostname": "node-1-proxy-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
-          "hostname": "node-3-proxy-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
-          "hostname": "node-2-proxy-host-prod.eu-west-1",
-          "environment": "prod",
-          "region": "eu-west-1"
-        },
-        {
           "hostname": "node-1-tenant-host-prod.eu-west-1",
           "environment": "prod",
           "region": "eu-west-1"
@@ -159,6 +129,7 @@
     {
       "version": "8.2.1",
       "targetVersion": true,
+      "upgradeBudget": "PT24H",
       "cloud": "cloud2",
       "nodes": []
     }