Set distributor down while permanently down

10 files changed, 84 insertions, 56 deletions

diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorImpl.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorImpl.java
index 94af5c036ad..587875363d5 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorImpl.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorImpl.java
@@ -19,6 +19,7 @@ import com.yahoo.vespa.orchestrator.controller.ClusterControllerClientFactory;
 import com.yahoo.vespa.orchestrator.controller.ClusterControllerNodeState;
 import com.yahoo.vespa.orchestrator.model.ApplicationApi;
 import com.yahoo.vespa.orchestrator.model.ApplicationApiFactory;
+import com.yahoo.vespa.orchestrator.model.ContentService;
 import com.yahoo.vespa.orchestrator.model.NodeGroup;
 import com.yahoo.vespa.orchestrator.model.VespaModelUtil;
 import com.yahoo.vespa.orchestrator.policy.BatchHostStateChangeDeniedException;
@@ -426,7 +427,7 @@ public class OrchestratorImpl implements Orchestrator {
                 ClusterControllerClient client = clusterControllerClientFactory.createClient(clusterControllers, cluster.clusterId().s());
                 for (ServiceInstance service : cluster.serviceInstances()) {
                     try {
-                        if ( ! client.trySetNodeState(context, service.hostName(), VespaModelUtil.getStorageNodeIndex(service.configId()), MAINTENANCE))
+                        if ( ! client.trySetNodeState(context, service.hostName(), VespaModelUtil.getStorageNodeIndex(service.configId()), MAINTENANCE, ContentService.STORAGE_NODE, false))
                             return false;
                     }
                     catch (Exception e) {
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClient.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClient.java
index 98ca9f805b4..e563f36a488 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClient.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClient.java
@@ -5,6 +5,7 @@ import com.yahoo.vespa.applicationmodel.ApplicationInstanceId;
 import com.yahoo.vespa.applicationmodel.HostName;
 import com.yahoo.vespa.orchestrator.ApplicationStateChangeDeniedException;
 import com.yahoo.vespa.orchestrator.OrchestratorContext;
+import com.yahoo.vespa.orchestrator.model.ContentService;
 import com.yahoo.vespa.orchestrator.policy.HostStateChangeDeniedException;
 
 /**
@@ -19,7 +20,8 @@ public interface ClusterControllerClient {
      * @throws HostStateChangeDeniedException if operation fails, or is otherwise disallowed.
      */
     boolean trySetNodeState(OrchestratorContext context, HostName host, int storageNodeIndex,
-                            ClusterControllerNodeState wantedState) throws HostStateChangeDeniedException;
+                            ClusterControllerNodeState wantedState, ContentService contentService, boolean force)
+            throws HostStateChangeDeniedException;
 
     /**
      * Requests that a cluster controller sets the requested node to the requested state.
@@ -27,7 +29,8 @@ public interface ClusterControllerClient {
      * @throws HostStateChangeDeniedException if operation fails, or is disallowed.
      */
     void setNodeState(OrchestratorContext context, HostName host, int storageNodeIndex,
-                      ClusterControllerNodeState wantedState) throws HostStateChangeDeniedException;
+                      ClusterControllerNodeState wantedState, ContentService contentService, boolean force)
+            throws HostStateChangeDeniedException;
 
     /**
      * Requests that a cluster controller sets all nodes in the cluster to the requested state.
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientImpl.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientImpl.java
index 28ba259d2ae..25567cb00a1 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientImpl.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientImpl.java
@@ -18,6 +18,7 @@ import com.yahoo.vespa.applicationmodel.ApplicationInstanceId;
 import com.yahoo.vespa.applicationmodel.HostName;
 import com.yahoo.vespa.orchestrator.ApplicationStateChangeDeniedException;
 import com.yahoo.vespa.orchestrator.OrchestratorContext;
+import com.yahoo.vespa.orchestrator.model.ContentService;
 import com.yahoo.vespa.orchestrator.policy.HostStateChangeDeniedException;
 import com.yahoo.vespa.orchestrator.policy.HostedVespaPolicy;
 import com.yahoo.yolean.Exceptions;
@@ -53,14 +54,16 @@ public class ClusterControllerClientImpl implements ClusterControllerClient {
     }
 
     private boolean setNodeState(OrchestratorContext context, HostName host, int storageNodeIndex,
-                              ClusterControllerNodeState wantedState, boolean throwOnFailure) {
+                                 ClusterControllerNodeState wantedState, ContentService contentService,
+                                 Condition condition, boolean throwOnFailure) {
         try {
             ClusterControllerClientTimeouts timeouts = context.getClusterControllerTimeouts();
             Inspector response = client.send(strategy(hosts), Method.POST)
-                                       .at("cluster", "v2", clusterName, "storage", Integer.toString(storageNodeIndex))
+                                       .at("cluster", "v2", clusterName, contentService.nameInClusterController(),
+                                           Integer.toString(storageNodeIndex))
                                        .deadline(timeouts.readBudget())
                                        .parameters(() -> deadline(timeouts))
-                                       .body(stateChangeRequestBytes(wantedState, Condition.SAFE, context.isProbe()))
+                                       .body(stateChangeRequestBytes(wantedState, condition, context.isProbe()))
                                        .throwing(retryOnRedirect)
                                        .read(SlimeUtils::jsonToSlime).get();
             if ( ! response.field("wasModified").asBool()) {
@@ -99,13 +102,17 @@ public class ClusterControllerClientImpl implements ClusterControllerClient {
     }
 
     @Override
-    public boolean trySetNodeState(OrchestratorContext context, HostName host, int storageNodeIndex, ClusterControllerNodeState wantedState) throws HostStateChangeDeniedException {
-        return setNodeState(context, host, storageNodeIndex, wantedState, false);
+    public boolean trySetNodeState(OrchestratorContext context, HostName host, int storageNodeIndex,
+                                   ClusterControllerNodeState wantedState, ContentService contentService, boolean force)
+            throws HostStateChangeDeniedException {
+        return setNodeState(context, host, storageNodeIndex, wantedState, contentService, force ? Condition.FORCE : Condition.SAFE, false);
     }
 
     @Override
-    public void setNodeState(OrchestratorContext context, HostName host, int storageNodeIndex, ClusterControllerNodeState wantedState) throws HostStateChangeDeniedException {
-        setNodeState(context, host, storageNodeIndex, wantedState, true);
+    public void setNodeState(OrchestratorContext context, HostName host, int storageNodeIndex,
+                             ClusterControllerNodeState wantedState, ContentService contentService, boolean force)
+            throws HostStateChangeDeniedException {
+        setNodeState(context, host, storageNodeIndex, wantedState, contentService, force ? Condition.FORCE : Condition.SAFE, true);
     }
 
     @Override
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/StorageNode.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/StorageNode.java
index 863c817fe6f..671bd351b3b 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/StorageNode.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/StorageNode.java
@@ -8,5 +8,6 @@ import com.yahoo.vespa.orchestrator.policy.HostStateChangeDeniedException;
 
 public interface StorageNode extends Comparable<StorageNode> {
     HostName hostName();
-    void setNodeState(OrchestratorContext context, ClusterControllerNodeState wantedState) throws HostStateChangeDeniedException;
+    void setStorageNodeState(OrchestratorContext context, ClusterControllerNodeState wantedState) throws HostStateChangeDeniedException;
+    void forceDistributorState(OrchestratorContext context, ClusterControllerNodeState wantedState) throws HostStateChangeDeniedException;
 }
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/StorageNodeImpl.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/StorageNodeImpl.java
index f3e3fd0e674..29f7700ef54 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/StorageNodeImpl.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/StorageNodeImpl.java
@@ -41,7 +41,17 @@ public class StorageNodeImpl implements StorageNode {
     }
 
     @Override
-    public void setNodeState(OrchestratorContext context, ClusterControllerNodeState wantedNodeState)
+    public void setStorageNodeState(OrchestratorContext context, ClusterControllerNodeState wantedNodeState)
+            throws HostStateChangeDeniedException {
+        setNodeState(context, wantedNodeState, ContentService.STORAGE_NODE, false);
+    }
+
+    @Override
+    public void forceDistributorState(OrchestratorContext context, ClusterControllerNodeState wantedState) throws HostStateChangeDeniedException {
+        setNodeState(context, wantedState, ContentService.DISTRIBUTOR, true);
+    }
+
+    public void setNodeState(OrchestratorContext context, ClusterControllerNodeState wantedNodeState, ContentService contentService, boolean force)
             throws HostStateChangeDeniedException {
         // The "cluster name" used by the Cluster Controller IS the cluster ID.
         String clusterId = this.clusterId.s();
@@ -52,17 +62,18 @@ public class StorageNodeImpl implements StorageNode {
                 clusterControllers,
                 clusterId);
 
-        ConfigId configId = storageService.configId();
-        int nodeIndex = VespaModelUtil.getStorageNodeIndex(configId);
+        int nodeIndex = VespaModelUtil.getStorageNodeIndex(storageService.configId());
 
-        logger.log(Level.FINE, () -> "Setting cluster controller state for " +
-                "application " + applicationInstance.reference().asString() +
-                ", host " + hostName() +
-                ", cluster name " + clusterId +
-                ", node index " + nodeIndex +
-                ", node state " + wantedNodeState);
+        logger.log(Level.FINE, () -> (force ? "Force" : "Safe") +
+                                     " setting cluster controller state for " +
+                                     "application " + applicationInstance.reference().asString() +
+                                     ", host " + hostName() +
+                                     ", cluster name " + clusterId +
+                                     ", service " + contentService.nameInClusterController() +
+                                     ", node index " + nodeIndex +
+                                     ", node state " + wantedNodeState);
 
-        client.setNodeState(context, storageService.hostName(), nodeIndex, wantedNodeState);
+        client.setNodeState(context, storageService.hostName(), nodeIndex, wantedNodeState, contentService, force);
     }
 
     @Override
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicy.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicy.java
index df8b5029693..3f23be5e514 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicy.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicy.java
@@ -34,7 +34,7 @@ public class HostedVespaPolicy implements Policy {
     private final HostedVespaClusterPolicy clusterPolicy;
     private final ClusterControllerClientFactory clusterControllerClientFactory;
     private final ApplicationApiFactory applicationApiFactory;
-    private final BooleanFlag keepStorageNodeUp;
+    private final BooleanFlag keepStorageNodeUpFlag;
 
     public HostedVespaPolicy(HostedVespaClusterPolicy clusterPolicy,
                              ClusterControllerClientFactory clusterControllerClientFactory,
@@ -43,7 +43,7 @@ public class HostedVespaPolicy implements Policy {
         this.clusterPolicy = clusterPolicy;
         this.clusterControllerClientFactory = clusterControllerClientFactory;
         this.applicationApiFactory = applicationApiFactory;
-        this.keepStorageNodeUp = Flags.KEEP_STORAGE_NODE_UP.bindTo(flagSource);
+        this.keepStorageNodeUpFlag = Flags.KEEP_STORAGE_NODE_UP.bindTo(flagSource);
     }
 
     @Override
@@ -59,7 +59,7 @@ public class HostedVespaPolicy implements Policy {
         // Ask Cluster Controller to set storage nodes in maintenance, unless the node is already allowed
         // to be down (or permanently down) in case they are guaranteed to be in maintenance already.
         for (StorageNode storageNode : application.getNoRemarksStorageNodesInGroupInClusterOrder()) {
-            storageNode.setNodeState(context, ClusterControllerNodeState.MAINTENANCE);
+            storageNode.setStorageNodeState(context, ClusterControllerNodeState.MAINTENANCE);
         }
 
         // Ensure all nodes in the group are marked as allowed to be down
@@ -75,7 +75,7 @@ public class HostedVespaPolicy implements Policy {
             throws HostStateChangeDeniedException {
         // Always defer to Cluster Controller whether it's OK to resume storage node
         for (StorageNode storageNode : application.getSuspendedStorageNodesInGroupInReverseClusterOrder()) {
-            storageNode.setNodeState(context, ClusterControllerNodeState.UP);
+            storageNode.setStorageNodeState(context, ClusterControllerNodeState.UP);
         }
 
         // In particular, we're not modifying the state of PERMANENTLY_DOWN nodes.
@@ -101,13 +101,18 @@ public class HostedVespaPolicy implements Policy {
             clusterPolicy.verifyGroupGoingDownPermanentlyIsFine(cluster);
         }
 
+        boolean keepStorageNodeUp = keepStorageNodeUpFlag
+                .with(FetchVector.Dimension.APPLICATION_ID, applicationApi.applicationId().serializedForm())
+                .value();
+
         // Get permission from the Cluster Controller to remove the content nodes.
-        boolean keepStorageNodeUp = this.keepStorageNodeUp.with(FetchVector.Dimension.APPLICATION_ID,
-                                                                applicationApi.applicationId().serializedForm())
-                                                          .value();
         for (StorageNode storageNode : applicationApi.getStorageNodesInGroupInClusterOrder()) {
-            storageNode.setNodeState(keepStorageNodeUp ? context.createSubcontextForSingleAppOp(true) : context,
-                                     ClusterControllerNodeState.DOWN);
+            if (keepStorageNodeUp) {
+                storageNode.setStorageNodeState(context.createSubcontextForSingleAppOp(true), ClusterControllerNodeState.DOWN);
+                storageNode.forceDistributorState(context, ClusterControllerNodeState.DOWN);
+            } else {
+                storageNode.setStorageNodeState(context, ClusterControllerNodeState.DOWN);
+            }
         }
 
         // Ensure all nodes in the group are marked as permanently down
diff --git a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java
index 506cd3a5334..9fd7b9dade9 100644
--- a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java
@@ -24,6 +24,7 @@ import com.yahoo.vespa.orchestrator.controller.ClusterControllerClientFactory;
 import com.yahoo.vespa.orchestrator.controller.ClusterControllerClientFactoryMock;
 import com.yahoo.vespa.orchestrator.controller.ClusterControllerNodeState;
 import com.yahoo.vespa.orchestrator.model.ApplicationApiFactory;
+import com.yahoo.vespa.orchestrator.model.ContentService;
 import com.yahoo.vespa.orchestrator.model.NodeGroup;
 import com.yahoo.vespa.orchestrator.policy.BatchHostStateChangeDeniedException;
 import com.yahoo.vespa.orchestrator.policy.HostStateChangeDeniedException;
@@ -461,16 +462,16 @@ public class OrchestratorImplTest {
                                             applicationApiFactory,
                                             flagSource);
 
-        when(fooClient.trySetNodeState(any(), any(), eq(1), eq(ClusterControllerNodeState.MAINTENANCE))).thenReturn(true);
-        when(fooClient.trySetNodeState(any(), any(), eq(2), eq(ClusterControllerNodeState.MAINTENANCE))).thenReturn(true);
-        when(barClient.trySetNodeState(any(), any(), eq(0), eq(ClusterControllerNodeState.MAINTENANCE))).thenReturn(true);
-        when(barClient.trySetNodeState(any(), any(), eq(3), eq(ClusterControllerNodeState.MAINTENANCE))).thenReturn(true);
+        when(fooClient.trySetNodeState(any(), any(), eq(1), eq(ClusterControllerNodeState.MAINTENANCE), eq(ContentService.STORAGE_NODE), eq(false))).thenReturn(true);
+        when(fooClient.trySetNodeState(any(), any(), eq(2), eq(ClusterControllerNodeState.MAINTENANCE), eq(ContentService.STORAGE_NODE), eq(false))).thenReturn(true);
+        when(barClient.trySetNodeState(any(), any(), eq(0), eq(ClusterControllerNodeState.MAINTENANCE), eq(ContentService.STORAGE_NODE), eq(false))).thenReturn(true);
+        when(barClient.trySetNodeState(any(), any(), eq(3), eq(ClusterControllerNodeState.MAINTENANCE), eq(ContentService.STORAGE_NODE), eq(false))).thenReturn(true);
         assertTrue(orchestrator.isQuiescent(id));
 
-        when(fooClient.trySetNodeState(any(), any(), eq(2), eq(ClusterControllerNodeState.MAINTENANCE))).thenReturn(false);
+        when(fooClient.trySetNodeState(any(), any(), eq(2), eq(ClusterControllerNodeState.MAINTENANCE), eq(ContentService.STORAGE_NODE), eq(false))).thenReturn(false);
         assertFalse(orchestrator.isQuiescent(id));
 
-        when(fooClient.trySetNodeState(any(), any(), eq(2), eq(ClusterControllerNodeState.MAINTENANCE))).thenThrow(new RuntimeException());
+        when(fooClient.trySetNodeState(any(), any(), eq(2), eq(ClusterControllerNodeState.MAINTENANCE), eq(ContentService.STORAGE_NODE), eq(false))).thenThrow(new RuntimeException());
         assertFalse(orchestrator.isQuiescent(id));
     }
 
diff --git a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientFactoryMock.java b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientFactoryMock.java
index 597487ccfc5..a5348ad3d07 100644
--- a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientFactoryMock.java
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientFactoryMock.java
@@ -5,11 +5,10 @@ import com.yahoo.vespa.applicationmodel.ApplicationInstance;
 import com.yahoo.vespa.applicationmodel.ApplicationInstanceId;
 import com.yahoo.vespa.applicationmodel.ClusterId;
 import com.yahoo.vespa.applicationmodel.HostName;
-import com.yahoo.vespa.orchestrator.ApplicationStateChangeDeniedException;
 import com.yahoo.vespa.orchestrator.DummyServiceMonitor;
 import com.yahoo.vespa.orchestrator.OrchestratorContext;
+import com.yahoo.vespa.orchestrator.model.ContentService;
 import com.yahoo.vespa.orchestrator.model.VespaModelUtil;
-import com.yahoo.vespa.orchestrator.policy.HostStateChangeDeniedException;
 
 import java.util.HashMap;
 import java.util.List;
@@ -30,7 +29,7 @@ public class ClusterControllerClientFactoryMock implements ClusterControllerClie
         try {
             ClusterId clusterName = VespaModelUtil.getContentClusterName(appInstance, hostName);
             int storageNodeIndex = VespaModelUtil.getStorageNodeIndex(appInstance, hostName);
-            String globalMapKey = clusterName + "/" + storageNodeIndex;
+            String globalMapKey = clusterName + "/" + ContentService.STORAGE_NODE.nameInClusterController() + "/" + storageNodeIndex;
             return nodes.getOrDefault(globalMapKey, ClusterControllerNodeState.UP) == ClusterControllerNodeState.MAINTENANCE;
         } catch (Exception e) {
             //Catch all - meant to catch cases where the node is not part of a storage cluster
@@ -44,7 +43,7 @@ public class ClusterControllerClientFactoryMock implements ClusterControllerClie
             for (HostName host : hosts) {
                 ClusterId clusterName = VespaModelUtil.getContentClusterName(app, host);
                 int storageNodeIndex = VespaModelUtil.getStorageNodeIndex(app, host);
-                String globalMapKey = clusterName + "/" + storageNodeIndex;
+                String globalMapKey = clusterName + "/" + ContentService.STORAGE_NODE.nameInClusterController() + "/" + storageNodeIndex;
                 nodes.put(globalMapKey, ClusterControllerNodeState.UP);
             }
         }
@@ -53,12 +52,12 @@ public class ClusterControllerClientFactoryMock implements ClusterControllerClie
     @Override
     public ClusterControllerClient createClient(List<HostName> clusterControllers, String clusterName) {
         return new ClusterControllerClient() {
-            @Override public boolean trySetNodeState(OrchestratorContext context, HostName host, int storageNodeIndex, ClusterControllerNodeState wantedState) {
-                nodes.put(clusterName + "/" + storageNodeIndex, wantedState);
+            @Override public boolean trySetNodeState(OrchestratorContext context, HostName host, int storageNodeIndex, ClusterControllerNodeState wantedState, ContentService contentService, boolean force) {
+                nodes.put(clusterName + "/" + contentService.nameInClusterController() + "/" + storageNodeIndex, wantedState);
                 return true;
             }
-            @Override public void setNodeState(OrchestratorContext context, HostName host, int storageNodeIndex, ClusterControllerNodeState wantedState) {
-                trySetNodeState(context, host, storageNodeIndex, wantedState);
+            @Override public void setNodeState(OrchestratorContext context, HostName host, int storageNodeIndex, ClusterControllerNodeState wantedState, ContentService contentService, boolean force) {
+                trySetNodeState(context, host, storageNodeIndex, wantedState, contentService, false);
             }
             @Override public void setApplicationState(OrchestratorContext context, ApplicationInstanceId applicationId, ClusterControllerNodeState wantedState) {
                 nodes.replaceAll((key, state) -> key.startsWith(clusterName + "/") ? wantedState : state);
diff --git a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientImplTest.java b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientImplTest.java
index b8e078c60d4..81615f59be9 100644
--- a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientImplTest.java
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/controller/ClusterControllerClientImplTest.java
@@ -7,13 +7,13 @@ import com.yahoo.vespa.applicationmodel.ApplicationInstanceId;
 import com.yahoo.vespa.applicationmodel.HostName;
 import com.yahoo.vespa.orchestrator.ApplicationStateChangeDeniedException;
 import com.yahoo.vespa.orchestrator.OrchestratorContext;
+import com.yahoo.vespa.orchestrator.model.ContentService;
 import com.yahoo.vespa.orchestrator.policy.HostStateChangeDeniedException;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
 import java.io.IOException;
-import java.time.Clock;
 import java.time.Duration;
 import java.util.List;
 
@@ -66,7 +66,7 @@ public class ClusterControllerClientImplTest {
                         return "{ \"wasModified\": true }";
                     },
                     200);
-        client.setNodeState(context, host, 2, DOWN);
+        client.setNodeState(context, host, 2, DOWN, ContentService.STORAGE_NODE, false);
 
         clock.advance(Duration.ofSeconds(9));
         wire.expect((url, body) -> {
@@ -79,7 +79,7 @@ public class ClusterControllerClientImplTest {
                     200);
         assertEquals("Changing the state of node would violate controller-set-node-state: Failed to set state to DOWN in cluster controller: because",
                      assertThrows(HostStateChangeDeniedException.class,
-                                  () -> client.setNodeState(context, host, 1, DOWN))
+                                  () -> client.setNodeState(context, host, 1, DOWN, ContentService.STORAGE_NODE, false))
                              .getMessage());
     }
 
@@ -93,7 +93,7 @@ public class ClusterControllerClientImplTest {
                         return "{ \"wasModified\": false, \"reason\": \"no reason\" }";
                     },
                     200);
-        assertFalse(client.trySetNodeState(OrchestratorContext.createContextForBatchProbe(clock), host, 2, MAINTENANCE));
+        assertFalse(client.trySetNodeState(OrchestratorContext.createContextForBatchProbe(clock), host, 2, MAINTENANCE, ContentService.STORAGE_NODE, false));
     }
 
     @Test
@@ -134,7 +134,7 @@ public class ClusterControllerClientImplTest {
         assertEquals("Changing the state of node would violate deadline: Timeout while waiting for setNodeState(2, UP) " +
                      "against [host1, host2, host3]: Timed out after PT10S",
                      assertThrows(HostStateChangeDeniedException.class,
-                                  () -> client.setNodeState(context, host, 2, UP))
+                                  () -> client.setNodeState(context, host, 2, UP, ContentService.STORAGE_NODE, false))
                              .getMessage());
     }
 
@@ -154,7 +154,7 @@ public class ClusterControllerClientImplTest {
         assertEquals("Changing the state of node would violate controller-set-node-state: Failed setting node 2 in cluster cc to state UP: " +
                      "got status code 503 for POST http://host1:19050/cluster/v2/cc/storage/2?timeout=9.6",
                      assertThrows(HostStateChangeDeniedException.class,
-                                  () -> client.setNodeState(context, host, 2, UP))
+                                  () -> client.setNodeState(context, host, 2, UP, ContentService.STORAGE_NODE, false))
                              .getMessage());
     }
 
diff --git a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicyTest.java b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicyTest.java
index 67f4cae76df..a622142b873 100644
--- a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicyTest.java
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/policy/HostedVespaPolicyTest.java
@@ -87,8 +87,8 @@ public class HostedVespaPolicyTest {
         order.verify(clusterPolicy).verifyGroupGoingDownIsFine(clusterApi3);
 
         order.verify(applicationApi).getNoRemarksStorageNodesInGroupInClusterOrder();
-        order.verify(storageNode1).setNodeState(context, ClusterControllerNodeState.MAINTENANCE);
-        order.verify(storageNode3).setNodeState(context, ClusterControllerNodeState.MAINTENANCE);
+        order.verify(storageNode1).setStorageNodeState(context, ClusterControllerNodeState.MAINTENANCE);
+        order.verify(storageNode3).setStorageNodeState(context, ClusterControllerNodeState.MAINTENANCE);
 
         order.verify(applicationApi).getNodesInGroupWithStatus(HostStatus.NO_REMARKS);
         order.verify(applicationApi).setHostState(context, hostName1, HostStatus.ALLOWED_TO_BE_DOWN);
@@ -140,8 +140,8 @@ public class HostedVespaPolicyTest {
         order.verify(clusterPolicy).verifyGroupGoingDownPermanentlyIsFine(clusterApi3);
 
         order.verify(applicationApi).getStorageNodesInGroupInClusterOrder();
-        order.verify(storageNode1).setNodeState(probeContext, ClusterControllerNodeState.DOWN);
-        order.verify(storageNode3).setNodeState(probeContext, ClusterControllerNodeState.DOWN);
+        order.verify(storageNode1).setStorageNodeState(probeContext, ClusterControllerNodeState.DOWN);
+        order.verify(storageNode3).setStorageNodeState(probeContext, ClusterControllerNodeState.DOWN);
 
         order.verify(applicationApi).getNodesInGroupWith(any());
         order.verify(applicationApi).setHostState(context, hostName1, HostStatus.PERMANENTLY_DOWN);
@@ -193,8 +193,8 @@ public class HostedVespaPolicyTest {
         order.verify(clusterPolicy).verifyGroupGoingDownPermanentlyIsFine(clusterApi3);
 
         order.verify(applicationApi).getStorageNodesInGroupInClusterOrder();
-        order.verify(storageNode1).setNodeState(probeContext, ClusterControllerNodeState.DOWN);
-        order.verify(storageNode3).setNodeState(probeContext, ClusterControllerNodeState.DOWN);
+        order.verify(storageNode1).setStorageNodeState(probeContext, ClusterControllerNodeState.DOWN);
+        order.verify(storageNode3).setStorageNodeState(probeContext, ClusterControllerNodeState.DOWN);
 
         order.verify(applicationApi).getNodesInGroupWith(any());
         order.verify(applicationApi).setHostState(context, hostName1, HostStatus.PERMANENTLY_DOWN);