Code review fixes

3 files changed, 38 insertions, 23 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 b36921d37ac..7f595b4a541 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
@@ -467,6 +467,29 @@ public class NodeRepository extends AbstractComponent {
         }
     }
 
+    /*
+     * This method is used to enable a smooth rollout of dynamic docker flavor allocations. Once we have switch
+     * everything this can be simplified to only deleting the node.
+     *
+     * Should only be called by node-admin for docker containers
+     */
+    public List<Node> markNodeAvailableForNewAllocation(String hostname) {
+        Node node = getNode(hostname).orElseThrow(() -> new NotFoundException("No node with hostname \"" + hostname + '"'));
+        if (node.flavor().getType() != Flavor.Type.DOCKER_CONTAINER) {
+            throw new IllegalArgumentException(
+                    "Cannot make " + hostname + " available for new allocation, must be a docker container node");
+        } else if (node.state() != Node.State.dirty) {
+            throw new IllegalArgumentException(
+                    "Cannot make " + hostname + " available for new allocation, must be in state dirty, but was in " + node.state());
+        }
+
+        if (dynamicAllocationEnabled()) {
+            return removeRecursively(node, true);
+        } else {
+            return setReady(Collections.singletonList(node));
+        }
+    }
+
     /**
      * Removes all the nodes that are children of hostname before finally removing the hostname itself.
      *
@@ -474,15 +497,18 @@ public class NodeRepository extends AbstractComponent {
      */
     public List<Node> removeRecursively(String hostname) {
         Node node = getNode(hostname).orElseThrow(() -> new NotFoundException("No node with hostname \"" + hostname + '"'));
+        return removeRecursively(node, false);
+    }
 
+    private List<Node> removeRecursively(Node node, boolean force) {
         try (Mutex lock = lockUnallocated()) {
             List<Node> removed = node.type() != NodeType.host ?
                     new ArrayList<>() :
-                    getChildNodes(hostname).stream()
-                            .filter(child -> allowedToRemove(child, true))
+                    getChildNodes(node.hostname()).stream()
+                            .filter(child -> force || verifyRemovalIsAllowed(child, true))
                             .collect(Collectors.toList());
 
-            if (allowedToRemove(node, false)) removed.add(node);
+            if (force || verifyRemovalIsAllowed(node, false)) removed.add(node);
             db.removeNodes(removed);
 
             return removed;
@@ -498,14 +524,15 @@ public class NodeRepository extends AbstractComponent {
      *    If only removing the container node: node in state ready
      *    If also removing the parent node: child is in state provisioned|failed|parked|ready
      */
-    private boolean allowedToRemove(Node nodeToRemove, boolean deletingAsChild) {
-        if (nodeToRemove.flavor().getType() == Flavor.Type.DOCKER_CONTAINER && !deletingAsChild) {
+    private boolean verifyRemovalIsAllowed(Node nodeToRemove, boolean deletingAsChild) {
+        // TODO: Enable once controller no longer deletes child nodes manually
+        /*if (nodeToRemove.flavor().getType() == Flavor.Type.DOCKER_CONTAINER && !deletingAsChild) {
             if (nodeToRemove.state() != Node.State.ready) {
                 throw new IllegalArgumentException(
                         String.format("Docker container node %s can only be removed when in state ready", nodeToRemove.hostname()));
             }
 
-        } else if (nodeToRemove.flavor().getType() == Flavor.Type.DOCKER_CONTAINER) {
+        } else */ if (nodeToRemove.flavor().getType() == Flavor.Type.DOCKER_CONTAINER) {
             List<Node.State> legalStates = Arrays.asList(Node.State.provisioned, Node.State.failed, Node.State.parked, Node.State.ready);
 
             if (! legalStates.contains(nodeToRemove.state())) {
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 1cb997b48fb..b16ce5f818e 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
@@ -125,23 +125,10 @@ public class NodesApiHandler extends LoggingRequestHandler {
             return new MessageResponse("Moved " + lastElement(path) + " to active");
         }
         else if (path.startsWith("/nodes/v2/state/availablefornewallocations/")) {
-            /*
-             * This is a temporary "state" or rest call that we use to enable a smooth rollout of
-             * dynamic docker flavor allocations. Once we have switch everything we remove this
-             * and change the code in the nodeadmin to delete directly (remember to allow deletion of ready nodes then).
-             *
-             * Should only be called by node-admin for docker containers (the docker constraint is
-             * enforced in the remove method)
-             */
             String hostname = lastElement(path);
-            nodeRepository.setReady(hostname);
-
-            if (nodeRepository.dynamicAllocationEnabled()) {
-                List<Node> removedNodes = nodeRepository.removeRecursively(hostname);
-                return new MessageResponse("Removed " + removedNodes.stream().map(Node::hostname).collect(Collectors.joining(", ")));
-            } else {
-                return new MessageResponse("Moved " + hostname + " to ready");
-            }
+            List<Node> available = nodeRepository.markNodeAvailableForNewAllocation(hostname);
+            return new MessageResponse("Marked following nodes as available for new allocation: " +
+                    available.stream().map(Node::hostname).collect(Collectors.joining(", ")));
         }
 
         throw new NotFoundException("Cannot put to path '" + path + "'");
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 fe3f12b105e..8eec56a1c00 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
@@ -8,6 +8,7 @@ import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.TenantName;
 import com.yahoo.path.Path;
 import com.yahoo.vespa.hosted.provision.node.Agent;
+import org.junit.Ignore;
 import org.junit.Test;
 
 import java.nio.charset.StandardCharsets;
@@ -69,7 +70,7 @@ public class NodeRepositoryTest {
         assertTrue(tester.nodeRepository().dynamicAllocationEnabled());
     }
 
-    @Test
+    @Test @Ignore // TODO: Enable once controller no longer deletes child nodes manually
     public void only_allow_docker_containers_remove_in_ready() {
         NodeRepositoryTester tester = new NodeRepositoryTester();
         tester.addNode("id1", "host1", "docker", NodeType.tenant);