Use probe when getting permission to set storage node permanently down

7 files changed, 42 insertions, 21 deletions

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..94af5c036ad 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;
@@ -99,7 +99,8 @@ public class OrchestratorImpl implements Orchestrator {
     {
         this(new HostedVespaPolicy(new HostedVespaClusterPolicy(flagSource, zone),
                                    clusterControllerClientFactory,
-                                   applicationApiFactory),
+                                   applicationApiFactory,
+                                   flagSource),
              clusterControllerClientFactory,
              statusService,
              serviceMonitor,
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..df8b5029693 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 keepStorageNodeUp;
 
     public HostedVespaPolicy(HostedVespaClusterPolicy clusterPolicy,
                              ClusterControllerClientFactory clusterControllerClientFactory,
-                             ApplicationApiFactory applicationApiFactory) {
+                             ApplicationApiFactory applicationApiFactory,
+                             FlagSource flagSource) {
         this.clusterPolicy = clusterPolicy;
         this.clusterControllerClientFactory = clusterControllerClientFactory;
         this.applicationApiFactory = applicationApiFactory;
+        this.keepStorageNodeUp = Flags.KEEP_STORAGE_NODE_UP.bindTo(flagSource);
     }
 
     @Override
@@ -94,10 +101,13 @@ 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
+        // 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(context, ClusterControllerNodeState.DOWN);
+            storageNode.setNodeState(keepStorageNodeUp ? context.createSubcontextForSingleAppOp(true) : 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..506cd3a5334 100644
--- a/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java
+++ b/orchestrator/src/test/java/com/yahoo/vespa/orchestrator/OrchestratorImplTest.java
@@ -104,7 +104,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 +451,8 @@ public class OrchestratorImplTest {
 
         orchestrator = new OrchestratorImpl(new HostedVespaPolicy(new HostedVespaClusterPolicy(flagSource, zone),
                                                                   clusterControllerClientFactory,
-                                                                  applicationApiFactory),
+                                                                  applicationApiFactory,
+                                                                  flagSource),
                                             clusterControllerClientFactory,
                                             statusService,
                                             serviceMonitor,
@@ -509,7 +511,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/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..67f4cae76df 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"));
 
@@ -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).setNodeState(probeContext, ClusterControllerNodeState.DOWN);
+        order.verify(storageNode3).setNodeState(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).setNodeState(probeContext, ClusterControllerNodeState.DOWN);
+        order.verify(storageNode3).setNodeState(probeContext, ClusterControllerNodeState.DOWN);
 
         order.verify(applicationApi).getNodesInGroupWith(any());
         order.verify(applicationApi).setHostState(context, hostName1, HostStatus.PERMANENTLY_DOWN);