Merge pull request #23439 from vespa-engine/hakonhall/keep-storage-node-up-while-permanently-down

Keep storage node up while permanently down [run-systemtest]

15 files changed, 135 insertions, 66 deletions

diff --git a/flags/src/main/java/com/yahoo/vespa/flags/Flags.java b/flags/src/main/java/com/yahoo/vespa/flags/Flags.java
index 04c7dcf0595..7c940c95feb 100644
--- a/flags/src/main/java/com/yahoo/vespa/flags/Flags.java
+++ b/flags/src/main/java/com/yahoo/vespa/flags/Flags.java
@@ -58,6 +58,12 @@ public class Flags {
             "Takes effect at redeployment (requires restart)",
             ZONE_ID, APPLICATION_ID);
 
+    public static final UnboundBooleanFlag KEEP_STORAGE_NODE_UP = defineFeatureFlag(
+            "keep-storage-node-up", true,
+            List.of("hakonhall"), "2022-07-07", "2022-08-07",
+            "Whether to leave the storage node (with wanted state) UP while the node is permanently down.",
+            "Takes effect immediately for nodes transitioning to permanently down.",
+            ZONE_ID, APPLICATION_ID);
 
     public static final UnboundIntFlag MAX_UNCOMMITTED_MEMORY = defineIntFlag(
             "max-uncommitted-memory", 130000,
diff --git a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorContext.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorContext.java
index aa218755792..f798cea3572 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorContext.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorContext.java
@@ -47,7 +47,7 @@ public class OrchestratorContext implements AutoCloseable {
     /** Create an OrchestratorContext for an operation on a single application. */
     public static OrchestratorContext createContextForSingleAppOp(Clock clock) {
         return new OrchestratorContext(null, clock, TimeBudget.fromNow(clock, DEFAULT_TIMEOUT_FOR_SINGLE_OP),
-                false, false);
+                                       false, false);
     }
 
     public static OrchestratorContext createContextForAdminOp(Clock clock) {
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 b3244c1ac74..587875363d5 100644
--- a/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorImpl.java
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/OrchestratorImpl.java
@@ -1,8 +1,8 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.orchestrator;
 
-import com.yahoo.component.annotation.Inject;
 import com.yahoo.cloud.config.ConfigserverConfig;
+import com.yahoo.component.annotation.Inject;
 import com.yahoo.concurrent.UncheckedTimeoutException;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.Zone;
@@ -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;
@@ -99,7 +100,8 @@ public class OrchestratorImpl implements Orchestrator {
     {
         this(new HostedVespaPolicy(new HostedVespaClusterPolicy(flagSource, zone),
                                    clusterControllerClientFactory,
-                                   applicationApiFactory),
+                                   applicationApiFactory,
+                                   flagSource),
              clusterControllerClientFactory,
              statusService,
              serviceMonitor,
@@ -425,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/ContentService.java b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/ContentService.java
new file mode 100644
index 00000000000..f611bada264
--- /dev/null
+++ b/orchestrator/src/main/java/com/yahoo/vespa/orchestrator/model/ContentService.java
@@ -0,0 +1,14 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.orchestrator.model;
+
+public enum ContentService {
+    DISTRIBUTOR("distributor"), STORAGE_NODE("storage");
+
+    private final String nameInClusterController;
+
+    ContentService(String nameInClusterController) {
+        this.nameInClusterController = nameInClusterController;
+    }
+
+    public String nameInClusterController() { return nameInClusterController; }
+}
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 2ce62081f51..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
@@ -3,6 +3,10 @@ package com.yahoo.vespa.orchestrator.policy;
 
 import com.yahoo.vespa.applicationmodel.ApplicationInstance;
 import com.yahoo.vespa.applicationmodel.HostName;
+import com.yahoo.vespa.flags.BooleanFlag;
+import com.yahoo.vespa.flags.FetchVector;
+import com.yahoo.vespa.flags.FlagSource;
+import com.yahoo.vespa.flags.Flags;
 import com.yahoo.vespa.orchestrator.OrchestratorContext;
 import com.yahoo.vespa.orchestrator.controller.ClusterControllerClientFactory;
 import com.yahoo.vespa.orchestrator.controller.ClusterControllerNodeState;
@@ -30,13 +34,16 @@ public class HostedVespaPolicy implements Policy {
     private final HostedVespaClusterPolicy clusterPolicy;
     private final ClusterControllerClientFactory clusterControllerClientFactory;
     private final ApplicationApiFactory applicationApiFactory;
+    private final BooleanFlag keepStorageNodeUpFlag;
 
     public HostedVespaPolicy(HostedVespaClusterPolicy clusterPolicy,
                              ClusterControllerClientFactory clusterControllerClientFactory,
-                             ApplicationApiFactory applicationApiFactory) {
+                             ApplicationApiFactory applicationApiFactory,
+                             FlagSource flagSource) {
         this.clusterPolicy = clusterPolicy;
         this.clusterControllerClientFactory = clusterControllerClientFactory;
         this.applicationApiFactory = applicationApiFactory;
+        this.keepStorageNodeUpFlag = Flags.KEEP_STORAGE_NODE_UP.bindTo(flagSource);
     }
 
     @Override
@@ -52,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
@@ -68,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.
@@ -94,10 +101,18 @@ public class HostedVespaPolicy implements Policy {
             clusterPolicy.verifyGroupGoingDownPermanentlyIsFine(cluster);
         }
 
-        // Ask Cluster Controller to set storage nodes to DOWN.
-        // These storage nodes are guaranteed to be NO_REMARKS
+        boolean keepStorageNodeUp = keepStorageNodeUpFlag
+                .with(FetchVector.Dimension.APPLICATION_ID, applicationApi.applicationId().serializedForm())
+                .value();
+
+        // Get permission from the Cluster Controller to remove the content nodes.
         for (StorageNode storageNode : applicationApi.getStorageNodesInGroupInClusterOrder()) {
-            storageNode.setNodeState(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 f5b7771d5f4..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;
@@ -104,7 +105,8 @@ public class OrchestratorImplTest {
         clustercontroller = new ClusterControllerClientFactoryMock();
         orchestrator = new OrchestratorImpl(new HostedVespaPolicy(new HostedVespaClusterPolicy(flagSource, zone),
                                                                   clustercontroller,
-                                                                  applicationApiFactory),
+                                                                  applicationApiFactory,
+                                                                  flagSource),
                                             clustercontroller,
                                             statusService,
                                             new DummyServiceMonitor(),
@@ -450,7 +452,8 @@ public class OrchestratorImplTest {
 
         orchestrator = new OrchestratorImpl(new HostedVespaPolicy(new HostedVespaClusterPolicy(flagSource, zone),
                                                                   clusterControllerClientFactory,
-                                                                  applicationApiFactory),
+                                                                  applicationApiFactory,
+                                                                  flagSource),
                                             clusterControllerClientFactory,
                                             statusService,
                                             serviceMonitor,
@@ -459,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));
     }
 
@@ -509,7 +512,8 @@ public class OrchestratorImplTest {
 
         orchestrator = new OrchestratorImpl(new HostedVespaPolicy(new HostedVespaClusterPolicy(flagSource, zone),
                                                                   clusterControllerClientFactory,
-                                                                  applicationApiFactory),
+                                                                  applicationApiFactory,
+                                                                  flagSource),
                                             clusterControllerClientFactory,
                                             statusService,
                                             serviceMonitor,
diff --git a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorTest.java b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorTest.java
index 29d6463a7d7..323ae678b0b 100644
--- a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorTest.java
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorTest.java
@@ -57,7 +57,8 @@ public class OrchestratorTest {
         var timer = new TestTimer();
         var clustercontroller = new ClusterControllerClientFactoryMock();
         var applicationApiFactory = new ApplicationApiFactory(3, 5, timer.toUtcClock());
-        var policy = new HostedVespaPolicy(new HostedVespaClusterPolicy(flagSource, zone), clustercontroller, applicationApiFactory);
+        var clusterPolicy = new HostedVespaClusterPolicy(flagSource, zone);
+        var policy = new HostedVespaPolicy(clusterPolicy, clustercontroller, applicationApiFactory, flagSource);
         var zone = new Zone(SystemName.cd, Environment.prod, RegionName.from("cd-us-east-1"));
         this.superModelManager = new MySuperModelProvider();
         var duperModel = new DuperModel();
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/model/ModelTestUtils.java b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/model/ModelTestUtils.java
index 291f96e3dc3..f2e2972ae9f 100644
--- a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/model/ModelTestUtils.java
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/model/ModelTestUtils.java
@@ -80,9 +80,9 @@ class ModelTestUtils {
             mock(Metric.class),
             new TestTimer(),
             new DummyAntiServiceMonitor());
-    private final Orchestrator orchestrator = new OrchestratorImpl(new HostedVespaPolicy(new HostedVespaClusterPolicy(flagSource, Zone.defaultZone()),
-                                                                                         clusterControllerClientFactory,
-                                                                                         applicationApiFactory()),
+    private final HostedVespaClusterPolicy clusterPolicy = new HostedVespaClusterPolicy(flagSource, Zone.defaultZone());
+    private final HostedVespaPolicy policy = new HostedVespaPolicy(clusterPolicy, clusterControllerClientFactory, applicationApiFactory(), flagSource);
+    private final Orchestrator orchestrator = new OrchestratorImpl(policy,
                                                                    clusterControllerClientFactory,
                                                                    statusService,
                                                                    serviceMonitor,
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 f01ce5a2227..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
@@ -5,6 +5,7 @@ package com.yahoo.vespa.orchestrator.policy;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.test.ManualClock;
 import com.yahoo.vespa.applicationmodel.HostName;
+import com.yahoo.vespa.flags.InMemoryFlagSource;
 import com.yahoo.vespa.orchestrator.OrchestrationException;
 import com.yahoo.vespa.orchestrator.OrchestratorContext;
 import com.yahoo.vespa.orchestrator.controller.ClusterControllerClient;
@@ -37,6 +38,7 @@ public class HostedVespaPolicyTest {
     private final ClusterControllerClient client = mock(ClusterControllerClient.class);
     private final ManualClock clock = new ManualClock();
     private final ApplicationApiFactory applicationApiFactory = new ApplicationApiFactory(3, 5, clock);
+    private final InMemoryFlagSource flagSource = new InMemoryFlagSource();
 
     @Before
     public void setUp() {
@@ -47,7 +49,7 @@ public class HostedVespaPolicyTest {
     public void testGrantSuspension() throws HostStateChangeDeniedException {
         final HostedVespaClusterPolicy clusterPolicy = mock(HostedVespaClusterPolicy.class);
         when(clusterPolicy.verifyGroupGoingDownIsFine(any())).thenReturn(SuspensionReasons.nothingNoteworthy());
-        final HostedVespaPolicy policy = new HostedVespaPolicy(clusterPolicy, clientFactory, applicationApiFactory);
+        final HostedVespaPolicy policy = new HostedVespaPolicy(clusterPolicy, clientFactory, applicationApiFactory, flagSource);
         final ApplicationApi applicationApi = mock(ApplicationApi.class);
         when(applicationApi.applicationId()).thenReturn(ApplicationId.fromSerializedForm("tenant:app:default"));
 
@@ -85,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);
@@ -99,7 +101,7 @@ public class HostedVespaPolicyTest {
     @Test
     public void testAcquirePermissionToRemove() throws OrchestrationException {
         final HostedVespaClusterPolicy clusterPolicy = mock(HostedVespaClusterPolicy.class);
-        final HostedVespaPolicy policy = new HostedVespaPolicy(clusterPolicy, clientFactory, applicationApiFactory);
+        final HostedVespaPolicy policy = new HostedVespaPolicy(clusterPolicy, clientFactory, applicationApiFactory, flagSource);
         final ApplicationApi applicationApi = mock(ApplicationApi.class);
         when(applicationApi.applicationId()).thenReturn(ApplicationId.fromSerializedForm("tenant:app:default"));
 
@@ -128,6 +130,8 @@ public class HostedVespaPolicyTest {
         InOrder order = inOrder(applicationApi, clusterPolicy, storageNode1, storageNode3);
 
         OrchestratorContext context = mock(OrchestratorContext.class);
+        OrchestratorContext probeContext = mock(OrchestratorContext.class);
+        when(context.createSubcontextForSingleAppOp(true)).thenReturn(probeContext);
         policy.acquirePermissionToRemove(context, applicationApi);
 
         order.verify(applicationApi).getClusters();
@@ -136,8 +140,8 @@ public class HostedVespaPolicyTest {
         order.verify(clusterPolicy).verifyGroupGoingDownPermanentlyIsFine(clusterApi3);
 
         order.verify(applicationApi).getStorageNodesInGroupInClusterOrder();
-        order.verify(storageNode1).setNodeState(context, ClusterControllerNodeState.DOWN);
-        order.verify(storageNode3).setNodeState(context, 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);
@@ -150,7 +154,7 @@ public class HostedVespaPolicyTest {
     @Test
     public void testAcquirePermissionToRemoveConfigServer() throws OrchestrationException {
         final HostedVespaClusterPolicy clusterPolicy = mock(HostedVespaClusterPolicy.class);
-        final HostedVespaPolicy policy = new HostedVespaPolicy(clusterPolicy, clientFactory, applicationApiFactory);
+        final HostedVespaPolicy policy = new HostedVespaPolicy(clusterPolicy, clientFactory, applicationApiFactory, flagSource);
         final ApplicationApi applicationApi = mock(ApplicationApi.class);
         when(applicationApi.applicationId()).thenReturn(ApplicationId.fromSerializedForm("tenant:app:default"));
 
@@ -179,6 +183,8 @@ public class HostedVespaPolicyTest {
         InOrder order = inOrder(applicationApi, clusterPolicy, storageNode1, storageNode3);
 
         OrchestratorContext context = mock(OrchestratorContext.class);
+        OrchestratorContext probeContext = mock(OrchestratorContext.class);
+        when(context.createSubcontextForSingleAppOp(true)).thenReturn(probeContext);
         policy.acquirePermissionToRemove(context, applicationApi);
 
         order.verify(applicationApi).getClusters();
@@ -187,8 +193,8 @@ public class HostedVespaPolicyTest {
         order.verify(clusterPolicy).verifyGroupGoingDownPermanentlyIsFine(clusterApi3);
 
         order.verify(applicationApi).getStorageNodesInGroupInClusterOrder();
-        order.verify(storageNode1).setNodeState(context, ClusterControllerNodeState.DOWN);
-        order.verify(storageNode3).setNodeState(context, 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);