Make set node to dirty recursive

6 files changed, 90 insertions, 12 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
index 4bf7e70d06b..ade2c94a25c 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
@@ -40,6 +40,7 @@ import java.util.Set;
 import java.util.TreeSet;
 import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * The hosted Vespa production node repository, which stores its state in Zookeeper.
@@ -394,12 +395,31 @@ public class NodeRepository extends AbstractComponent {
         return db.writeTo(Node.State.dirty, node, agent, Optional.of(reason));
     }
 
-    public Node setDirty(String hostname, Agent agent, String reason) {
-        Node node = getNode(hostname, Node.State.provisioned, Node.State.failed, Node.State.parked).orElseThrow(() ->
-                new IllegalArgumentException("Could not deallocate " + hostname +
-                                             ": No such node in the provisioned, failed or parked state"));
+    public List<Node> dirtyRecursively(String hostname, Agent agent, String reason) {
+        Node nodeToDirty = getNode(hostname).orElseThrow(() ->
+                new IllegalArgumentException("Could not deallocate " + hostname + ": Node not found"));
 
-        return setDirty(node, agent, reason);
+        List<Node> nodesToDirty =
+                (nodeToDirty.type().isDockerHost() ?
+                        Stream.concat(getChildNodes(hostname).stream(), Stream.of(nodeToDirty)) :
+                        Stream.of(nodeToDirty))
+                .filter(node -> node.state() != Node.State.dirty)
+                .collect(Collectors.toList());
+
+        List<String> hostnamesNotAllowedToDirty = nodesToDirty.stream()
+                .filter(node -> node.state() != Node.State.provisioned)
+                .filter(node -> node.state() != Node.State.failed)
+                .filter(node -> node.state() != Node.State.parked)
+                .map(Node::hostname)
+                .collect(Collectors.toList());
+        if (!hostnamesNotAllowedToDirty.isEmpty()) {
+            throw new IllegalArgumentException("Could not deallocate " + hostname + ": " +
+                    String.join(", ", hostnamesNotAllowedToDirty) + " must be in either provisioned, failed or parked state");
+        }
+
+        return nodesToDirty.stream()
+                .map(node -> setDirty(node, agent, reason))
+                .collect(Collectors.toList());
     }
 
     /**
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesApiHandler.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesApiHandler.java
index 2f7d3120211..03777078251 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesApiHandler.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesApiHandler.java
@@ -34,7 +34,6 @@ import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.Executor;
 import java.util.logging.Level;
 import java.util.stream.Collectors;
 import javax.inject.Inject;
@@ -122,8 +121,9 @@ public class NodesApiHandler extends LoggingRequestHandler {
             return new MessageResponse("Moved " + parkedHostnames + " to parked");
         }
         else if (path.startsWith("/nodes/v2/state/dirty/")) {
-            nodeRepository.setDirty(lastElement(path), Agent.operator, "Dirtied through the nodes/v2 API");
-            return new MessageResponse("Moved " + lastElement(path) + " to dirty");
+            List<Node> dirtiedNodes = nodeRepository.dirtyRecursively(lastElement(path), Agent.operator, "Dirtied through the nodes/v2 API");
+            String dirtiedHostnames = dirtiedNodes.stream().map(Node::hostname).sorted().collect(Collectors.joining(", "));
+            return new MessageResponse("Moved " + dirtiedHostnames + " to dirty");
         }
         else if (path.startsWith("/nodes/v2/state/active/")) {
             nodeRepository.reactivate(lastElement(path), Agent.operator, "Reactivated through nodes/v2 API");
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
index a21d03ff603..931ddc5e051 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
@@ -108,7 +108,7 @@ public class MockNodeRepository extends NodeRepository {
         setReady(nodes, Agent.system, getClass().getSimpleName());
 
         fail("host5.yahoo.com", Agent.system, getClass().getSimpleName());
-        setDirty("host55.yahoo.com", Agent.system, getClass().getSimpleName());
+        dirtyRecursively("host55.yahoo.com", Agent.system, getClass().getSimpleName());
 
         ApplicationId zoneApp = ApplicationId.from(TenantName.from("zoneapp"), ApplicationName.from("zoneapp"), InstanceName.from("zoneapp"));
         ClusterSpec zoneCluster = ClusterSpec.request(ClusterSpec.Type.container,
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java
index 6e97f3d1269..5d217ce86e5 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java
@@ -11,7 +11,12 @@ import com.yahoo.vespa.hosted.provision.node.Agent;
 import org.junit.Test;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.HashSet;
 import java.util.Optional;
+import java.util.Set;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 import static junit.framework.TestCase.fail;
 import static org.junit.Assert.assertEquals;
@@ -72,7 +77,7 @@ public class NodeRepositoryTest {
             // Expected
         }
 
-        tester.nodeRepository().setDirty("host1", Agent.system, getClass().getSimpleName());
+        tester.nodeRepository().dirtyRecursively("host1", Agent.system, getClass().getSimpleName());
         tester.nodeRepository().setReady("host1", Agent.system, getClass().getSimpleName());
         tester.nodeRepository().removeRecursively("host1");
     }
@@ -89,7 +94,7 @@ public class NodeRepositoryTest {
         tester.addNode("node20", "node20", "host2", "docker", NodeType.tenant);
         assertEquals(6, tester.nodeRepository().getNodes().size());
 
-        tester.nodeRepository().setDirty("node11", Agent.system, getClass().getSimpleName());
+        tester.setNodeState("node11", Node.State.dirty);
 
         try {
             tester.nodeRepository().removeRecursively("host1");
@@ -111,4 +116,46 @@ public class NodeRepositoryTest {
         tester.nodeRepository().removeRecursively("host1");
         assertEquals(0, tester.nodeRepository().getNodes().size());
     }
+
+    @Test
+    public void dirty_host_only_if_we_can_dirty_children() {
+        NodeRepositoryTester tester = new NodeRepositoryTester();
+
+        tester.addNode("id1", "host1", "default", NodeType.host);
+        tester.addNode("id2", "host2", "default", NodeType.host);
+        tester.addNode("node10", "node10", "host1", "docker", NodeType.tenant);
+        tester.addNode("node11", "node11", "host1", "docker", NodeType.tenant);
+        tester.addNode("node12", "node12", "host1", "docker", NodeType.tenant);
+        tester.addNode("node20", "node20", "host2", "docker", NodeType.tenant);
+
+        tester.setNodeState("node11", Node.State.ready);
+        tester.setNodeState("node12", Node.State.active);
+        tester.setNodeState("node20", Node.State.failed);
+
+        assertEquals(6, tester.nodeRepository().getNodes().size());
+
+        // Should be OK to dirty host2 as it is in provisioned and its only child is in failed
+        tester.nodeRepository().dirtyRecursively("host2", Agent.system, NodeRepositoryTest.class.getSimpleName());
+        assertEquals(asSet("host2", "node20"), filterNodes(tester, node -> node.state() == Node.State.dirty));
+
+        // Cant dirty host1, node11 is ready and node12 is active
+        try {
+            tester.nodeRepository().dirtyRecursively("host1", Agent.system, NodeRepositoryTest.class.getSimpleName());
+            fail("Should not be able to dirty host1");
+        } catch (IllegalArgumentException ignored) { } // Expected;
+
+        assertEquals(asSet("host2", "node20"), filterNodes(tester, node -> node.state() == Node.State.dirty));
+    }
+
+    private static Set<String> asSet(String... elements) {
+        return new HashSet<>(Arrays.asList(elements));
+    }
+
+    private static Set<String> filterNodes(NodeRepositoryTester tester, Predicate<Node> filter) {
+        return tester.nodeRepository()
+                .getNodes().stream()
+                .filter(filter)
+                .map(Node::hostname)
+                .collect(Collectors.toSet());
+    }
 }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java
index 1a42a71863d..03ada1e7951 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java
@@ -9,6 +9,7 @@ import com.yahoo.config.provision.Zone;
 import com.yahoo.config.provisioning.FlavorsConfig;
 import com.yahoo.test.ManualClock;
 import com.yahoo.vespa.curator.mock.MockCurator;
+import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.provisioning.FlavorConfigBuilder;
 import com.yahoo.vespa.hosted.provision.testutils.MockNameResolver;
 
@@ -58,6 +59,16 @@ public class NodeRepositoryTester {
         return nodeRepository.addNodes(Collections.singletonList(node)).get(0);
     }
 
+    /**
+     * Moves a node directly to the given state without doing any validation, useful
+     * to create wanted test scenario without having to move every node through series
+     * of valid state transitions
+     */
+    public void setNodeState(String hostname, Node.State state) {
+        Node node = nodeRepository.getNode(hostname).orElseThrow(RuntimeException::new);
+        nodeRepository.database().writeTo(state, node, Agent.system, Optional.empty());
+    }
+
     private FlavorsConfig createConfig() {
         FlavorConfigBuilder b = new FlavorConfigBuilder();
         b.addFlavor("default", 2., 4., 100, Flavor.Type.BARE_METAL).cost(3);
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/monitoring/MetricsReporterTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/monitoring/MetricsReporterTest.java
index 622732460c0..5e7c6787973 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/monitoring/MetricsReporterTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/monitoring/MetricsReporterTest.java
@@ -125,7 +125,7 @@ public class MetricsReporterTest {
 
         Node dockerHost = Node.create("openStackId1", Collections.singleton("::1"), additionalIps, "dockerHost", Optional.empty(), nodeFlavors.getFlavorOrThrow("host"), NodeType.host);
         nodeRepository.addNodes(Collections.singletonList(dockerHost));
-        nodeRepository.setDirty("dockerHost", Agent.system, getClass().getSimpleName());
+        nodeRepository.dirtyRecursively("dockerHost", Agent.system, getClass().getSimpleName());
         nodeRepository.setReady("dockerHost", Agent.system, getClass().getSimpleName());
 
         Node container1 = Node.createDockerNode("openStackId1:1", Collections.singleton("::2"), Collections.emptySet(), "container1", Optional.of("dockerHost"), nodeFlavors.getFlavorOrThrow("docker"), NodeType.tenant);