Add support for upgrade budget in REST API

3 files changed, 29 insertions, 11 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
index 175091fa729..3e6bc4e96e5 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
@@ -49,6 +49,7 @@ import java.io.InputStream;
 import java.io.UncheckedIOException;
 import java.net.URI;
 import java.net.URISyntaxException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
@@ -360,12 +361,13 @@ public class NodesV2ApiHandler extends LoggingRequestHandler {
     private MessageResponse setTargetVersions(HttpRequest request) {
         NodeType nodeType = NodeType.valueOf(lastElement(request.getUri().getPath()).toLowerCase());
         Inspector inspector = toSlime(request.getData()).get();
-        List<String> messageParts = new ArrayList<>(3);
+        List<String> messageParts = new ArrayList<>(4);
 
         boolean force = inspector.field("force").asBool();
         Inspector versionField = inspector.field("version");
         Inspector osVersionField = inspector.field("osVersion");
         Inspector dockerImageField = inspector.field("dockerImage");
+        Inspector upgradeBudgetField = inspector.field("upgradeBudget");
 
         if (versionField.valid()) {
             Version version = Version.fromString(versionField.asString());
@@ -380,8 +382,19 @@ public class NodesV2ApiHandler extends LoggingRequestHandler {
                 messageParts.add("osVersion to null");
             } else {
                 Version osVersion = Version.fromString(v);
-                nodeRepository.osVersions().setTarget(nodeType, osVersion, force);
+                Optional<Duration> upgradeBudget = Optional.of(upgradeBudgetField)
+                                                           .filter(Inspector::valid)
+                                                           .map(Inspector::asString)
+                                                           .map(s -> {
+                                                               try {
+                                                                   return Duration.parse(s);
+                                                               } catch (Exception e) {
+                                                                   throw new IllegalArgumentException("Invalid duration '" + s + "'", e);
+                                                               }
+                                                           });
+                nodeRepository.osVersions().setTarget(nodeType, osVersion, upgradeBudget, force);
                 messageParts.add("osVersion to " + osVersion.toFullString());
+                upgradeBudget.ifPresent(d -> messageParts.add("upgradeBudget to " + d));
             }
         }
 
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/UpgradeResponse.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/UpgradeResponse.java
index be87a984941..1082e9cce60 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/UpgradeResponse.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/UpgradeResponse.java
@@ -39,7 +39,7 @@ public class UpgradeResponse extends HttpResponse {
         infrastructureVersions.getTargetVersions().forEach((nodeType, version) -> versionsObject.setString(nodeType.name(), version.toFullString()));
 
         Cursor osVersionsObject = root.setObject("osVersions");
-        osVersions.readChange().targets().forEach((nodeType, osVersion) -> osVersionsObject.setString(nodeType.name(), osVersion.toFullString()));
+        osVersions.readChange().targets().forEach((nodeType, target) -> osVersionsObject.setString(nodeType.name(), target.version().toFullString()));
 
 
         Cursor dockerImagesObject = root.setObject("dockerImages");
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java
index 414b12e77b9..fee2958b124 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java
@@ -1,36 +1,28 @@
 // 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.restapi;
 
-import com.yahoo.application.Networking;
-import com.yahoo.application.container.JDisc;
 import com.yahoo.application.container.handler.Request;
 import com.yahoo.application.container.handler.Response;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.TenantName;
-import com.yahoo.io.IOUtils;
 import com.yahoo.text.Utf8;
 import com.yahoo.vespa.applicationmodel.HostName;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.maintenance.OsUpgradeActivator;
-import com.yahoo.vespa.hosted.provision.testutils.ContainerConfig;
 import com.yahoo.vespa.hosted.provision.testutils.MockNodeRepository;
 import com.yahoo.vespa.hosted.provision.testutils.OrchestratorMock;
 import org.junit.After;
 import org.junit.Before;
-import org.junit.ComparisonFailure;
 import org.junit.Test;
 
-import java.io.File;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
-import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
-import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 
 /**
@@ -807,6 +799,19 @@ public class NodesV2ApiTest {
                        "{\"url\":\"http://localhost:8080/nodes/v2/node/dockerhost2.yahoo.com\"}," +
                        "{\"url\":\"http://localhost:8080/nodes/v2/node/dockerhost1.yahoo.com\"}" +
                        "]}");
+
+        // Schedule OS upgrade with budget
+        assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/host",
+                                   Utf8.toBytes("{\"osVersion\": \"7.42.1\", \"upgradeBudget\": \"PT24H\"}"),
+                                   Request.Method.PATCH),
+                       "{\"message\":\"Set osVersion to 7.42.1, upgradeBudget to PT24H for nodes of type host\"}");
+
+        // Invalid budget
+        tester.assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/host",
+                                          Utf8.toBytes("{\"osVersion\": \"7.42.1\", \"upgradeBudget\": \"foo\"}"),
+                                          Request.Method.PATCH),
+                              400,
+                              "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Invalid duration 'foo': Text cannot be parsed to a Duration\"}");
     }
 
     @Test