set-node-state timeout in CC

6 files changed, 45 insertions, 5 deletions

diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/FleetController.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/FleetController.java
index 6b40577b1b0..56fe679fc6a 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/FleetController.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/FleetController.java
@@ -882,11 +882,11 @@ public class FleetController implements NodeStateOrHostInfoChangeHandler, NodeAd
      */
     private void scheduleVersionDependentTasksForFutureCompletion(int completeAtVersion) {
         // TODO expose and use monotonic clock instead of system clock
-        final long deadlineTimePointMs = timer.getCurrentTimeInMillis() + options.getMaxDeferredTaskVersionWaitTime().toMillis();
+        final long maxDeadlineTimePointMs = timer.getCurrentTimeInMillis() + options.getMaxDeferredTaskVersionWaitTime().toMillis();
         for (RemoteClusterControllerTask task : tasksPendingStateRecompute) {
             log.finest(() -> String.format("Adding task of type '%s' to be completed at version %d",
                     task.getClass().getName(), completeAtVersion));
-            taskCompletionQueue.add(new VersionDependentTaskCompletion(completeAtVersion, task, deadlineTimePointMs));
+            taskCompletionQueue.add(new VersionDependentTaskCompletion(completeAtVersion, task, maxDeadlineTimePointMs));
         }
         tasksPendingStateRecompute.clear();
     }
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/RemoteClusterControllerTask.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/RemoteClusterControllerTask.java
index e96209c083a..8382e127e13 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/RemoteClusterControllerTask.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/RemoteClusterControllerTask.java
@@ -5,6 +5,9 @@ import com.yahoo.vdslib.state.ClusterState;
 import com.yahoo.vespa.clustercontroller.core.listeners.NodeAddedOrRemovedListener;
 import com.yahoo.vespa.clustercontroller.core.listeners.NodeStateOrHostInfoChangeHandler;
 
+import java.time.Instant;
+import java.util.Optional;
+
 public abstract class RemoteClusterControllerTask {
 
     public static class Context {
@@ -65,6 +68,10 @@ public abstract class RemoteClusterControllerTask {
      */
     public void handleFailure(FailureCondition condition) {}
 
+    public Optional<Instant> getDeadline() {
+        return Optional.empty();
+    }
+
     public boolean isCompleted() {
         synchronized (monitor) {
             return completed;
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/VersionDependentTaskCompletion.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/VersionDependentTaskCompletion.java
index 5d6a4f66467..28df5a8e35a 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/VersionDependentTaskCompletion.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/VersionDependentTaskCompletion.java
@@ -16,10 +16,12 @@ class VersionDependentTaskCompletion {
     private final RemoteClusterControllerTask task;
     private final long deadlineTimePointMs;
 
-    VersionDependentTaskCompletion(long minimumVersion, RemoteClusterControllerTask task, long deadlineTimePointMs) {
+    VersionDependentTaskCompletion(long minimumVersion, RemoteClusterControllerTask task, long maxDeadlineTimePointMs) {
         this.minimumVersion = minimumVersion;
         this.task = task;
-        this.deadlineTimePointMs = deadlineTimePointMs;
+        this.deadlineTimePointMs = task.getDeadline().map(deadline ->
+                Math.max(0, Math.min(deadline.toEpochMilli(), maxDeadlineTimePointMs)))
+                .orElse(maxDeadlineTimePointMs);
     }
 
     long getMinimumVersion() {
@@ -30,7 +32,9 @@ class VersionDependentTaskCompletion {
         return task;
     }
 
-    long getDeadlineTimePointMs() { return deadlineTimePointMs; }
+    long getDeadlineTimePointMs() {
+        return deadlineTimePointMs;
+    }
 
     @Override
     public boolean equals(Object o) {
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequest.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequest.java
index 849d8cc6e7b..4d6738940a8 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequest.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequest.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.clustercontroller.core.restapiv2.requests;
 
 import com.yahoo.log.LogLevel;
+import com.yahoo.time.TimeBudget;
 import com.yahoo.vdslib.state.ClusterState;
 import com.yahoo.vdslib.state.Node;
 import com.yahoo.vdslib.state.NodeState;
@@ -21,8 +22,10 @@ import com.yahoo.vespa.clustercontroller.utils.staterestapi.requests.SetUnitStat
 import com.yahoo.vespa.clustercontroller.utils.staterestapi.response.SetResponse;
 import com.yahoo.vespa.clustercontroller.utils.staterestapi.response.UnitState;
 
+import java.time.Instant;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.logging.Logger;
 
 public class SetNodeStateRequest extends Request<SetResponse> {
@@ -33,6 +36,7 @@ public class SetNodeStateRequest extends Request<SetResponse> {
     private final SetUnitStateRequest.Condition condition;
     private final SetUnitStateRequest.ResponseWait responseWait;
     private final WantedStateSetter wantedState;
+    private final TimeBudget timeBudget;
 
     public SetNodeStateRequest(Id.Node id, SetUnitStateRequest setUnitStateRequest) {
         this(id, setUnitStateRequest, SetNodeStateRequest::setWantedState);
@@ -46,6 +50,7 @@ public class SetNodeStateRequest extends Request<SetResponse> {
         this.condition = setUnitStateRequest.getCondition();
         this.responseWait = setUnitStateRequest.getResponseWait();
         this.wantedState = wantedState;
+        this.timeBudget = setUnitStateRequest.timeBudget();
     }
 
     @Override
@@ -79,6 +84,11 @@ public class SetNodeStateRequest extends Request<SetResponse> {
     }
 
     @Override
+    public Optional<Instant> getDeadline() {
+        return timeBudget.deadline();
+    }
+
+    @Override
     public boolean isFailed() {
         // Failure to set a node state is propagated as a 200 with wasModified false.
         return super.isFailed() || (resultSet && !result.getWasModified());
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStatesForClusterRequest.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStatesForClusterRequest.java
index 1246ba62313..b4d189bcd55 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStatesForClusterRequest.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStatesForClusterRequest.java
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.clustercontroller.core.restapiv2.requests;
 
+import com.yahoo.time.TimeBudget;
 import com.yahoo.vdslib.distribution.ConfiguredNode;
 import com.yahoo.vdslib.state.Node;
 import com.yahoo.vdslib.state.NodeType;
@@ -14,7 +15,9 @@ import com.yahoo.vespa.clustercontroller.utils.staterestapi.requests.SetUnitStat
 import com.yahoo.vespa.clustercontroller.utils.staterestapi.response.SetResponse;
 import com.yahoo.vespa.clustercontroller.utils.staterestapi.response.UnitState;
 
+import java.time.Instant;
 import java.util.Map;
+import java.util.Optional;
 import java.util.logging.Logger;
 
 public class SetNodeStatesForClusterRequest extends Request<SetResponse> {
@@ -23,6 +26,7 @@ public class SetNodeStatesForClusterRequest extends Request<SetResponse> {
     private final Id.Cluster cluster;
     private final Map<String, UnitState> newStates;
     private final SetUnitStateRequest.Condition condition;
+    private final TimeBudget timeBudget;
 
 
     public SetNodeStatesForClusterRequest(Id.Cluster cluster, SetUnitStateRequest request) {
@@ -30,6 +34,7 @@ public class SetNodeStatesForClusterRequest extends Request<SetResponse> {
         this.cluster = cluster;
         this.newStates = request.getNewState();
         this.condition = request.getCondition();
+        this.timeBudget = request.timeBudget();
     }
 
     @Override
@@ -80,6 +85,11 @@ public class SetNodeStatesForClusterRequest extends Request<SetResponse> {
     }
 
     @Override
+    public Optional<Instant> getDeadline() {
+        return timeBudget.deadline();
+    }
+
+    @Override
     public boolean isFailed() {
         // Failure to set a node state is propagated as a 200 with wasModified false.
         return super.isFailed() || (resultSet && !result.getWasModified());
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/SetNodeStateTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/SetNodeStateTest.java
index 3f977273054..f3a4be5ac2f 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/SetNodeStateTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/SetNodeStateTest.java
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.clustercontroller.core.restapiv2;
 
+import com.yahoo.time.TimeBudget;
 import com.yahoo.vdslib.state.NodeType;
 import com.yahoo.vespa.clustercontroller.core.MasterInterface;
 import com.yahoo.vespa.clustercontroller.core.RemoteClusterControllerTask;
@@ -20,6 +21,8 @@ import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
+import java.time.Clock;
+import java.time.Duration;
 import java.util.LinkedHashMap;
 import java.util.Map;
 
@@ -42,6 +45,7 @@ public class SetNodeStateTest extends StateRestApiTest {
         private Map<String, UnitState> newStates = new LinkedHashMap<>();
         private Condition condition = Condition.FORCE;
         private ResponseWait responseWait = ResponseWait.WAIT_UNTIL_CLUSTER_ACKED;
+        private TimeBudget timeBudget = TimeBudget.fromNow(Clock.systemUTC(), Duration.ofSeconds(10));
 
         public SetUnitStateRequestImpl(String req) {
             super(req, 0);
@@ -89,6 +93,11 @@ public class SetNodeStateTest extends StateRestApiTest {
         public ResponseWait getResponseWait() {
             return responseWait;
         }
+
+        @Override
+        public TimeBudget timeBudget() {
+            return timeBudget;
+        }
     }
 
     private void verifyStateSet(String state, String reason) throws Exception {