Support querying for one node from multiple states at the time

6 files changed, 37 insertions, 38 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 11892c69c1c..d7a5594659f 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
@@ -79,21 +79,24 @@ public class NodeRepository extends AbstractComponent {
 
     // ---------------- Query API ----------------------------------------------------------------
 
-    /** Finds and returns the node with the given hostname */
-    public Optional<Node> getNode(String hostname) {
-        for (Node.State state : Node.State.values()) {
-            Optional<Node> node = getNode(state, hostname);
-            if (node.isPresent())
-                return node;
-        }
-        return Optional.empty();
-    }
-
-    /** Finds and returns the node with the given state and hostname, or empty if not found */
-    public Optional<Node> getNode(Node.State state, String hostname) {
-        return zkClient.getNode(state, hostname);
+    /** 
+     * Finds and returns the node with the hostname in any of the given states, or empty if not found 
+     * 
+     * @param hostname the full host name of the node
+     * @param inState the states the node may be in. If no states are given, it will be returned from any state
+     * @return the node, or empty if it was not found in any of the given states
+     */
+    public Optional<Node> getNode(String hostname, Node.State ... inState) {
+        return zkClient.getNode(hostname, inState);
     }
 
+    /**
+     * Finds and returns the nodes of the given type in any of the given states.
+     *
+     * @param type the node type to return
+     * @param inState the states to return nodes from. If no states are given, all nodes of the given type is returned
+     * @return the node, or empty if it was not found in any of the given states
+     */
     public List<Node> getNodes(Node.Type type, Node.State ... inState) {
         return zkClient.getNodes(inState).stream().filter(node -> node.type().equals(type)).collect(Collectors.toList());
     }
@@ -185,9 +188,7 @@ public class NodeRepository extends AbstractComponent {
      * Use this to recycle failed nodes which have been repaired or put on hold. 
      */
     public Node deallocate(String hostname) {
-        Optional<Node> nodeToDeallocate = getNode(Node.State.failed, hostname);
-        if ( ! nodeToDeallocate.isPresent())
-            nodeToDeallocate = getNode(Node.State.parked, hostname);
+        Optional<Node> nodeToDeallocate = getNode(hostname, Node.State.failed, Node.State.parked);
         if ( ! nodeToDeallocate.isPresent())
             throw new IllegalArgumentException("Could not deallocate " + hostname + ": No such node in the failed or parked state");
         return deallocate(Collections.singletonList(nodeToDeallocate.get())).get(0);
@@ -238,9 +239,7 @@ public class NodeRepository extends AbstractComponent {
      * @return true if the node was removed, false if it was not found
      */
     public boolean remove(String hostname) {
-        Optional<Node> nodeToRemove = getNode(Node.State.failed, hostname);
-        if ( ! nodeToRemove.isPresent())
-            nodeToRemove = getNode(Node.State.parked, hostname);
+        Optional<Node> nodeToRemove = getNode(hostname, Node.State.failed, Node.State.parked);
         if ( ! nodeToRemove.isPresent()) 
             return false;
         try (Mutex lock = lock(nodeToRemove.get())) {
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeFailer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeFailer.java
index ba2586a95fd..9a116478882 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeFailer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeFailer.java
@@ -97,7 +97,7 @@ public class NodeFailer extends Maintainer {
         for (ApplicationInstance<ServiceMonitorStatus> application : serviceMonitor.queryStatusOfAllApplicationInstances().values()) {
             for (ServiceCluster<ServiceMonitorStatus> cluster : application.serviceClusters()) {
                 for (ServiceInstance<ServiceMonitorStatus> service : cluster.serviceInstances()) {
-                    Optional<Node> node = nodeRepository().getNode(Node.State.active, service.hostName().s());
+                    Optional<Node> node = nodeRepository().getNode(service.hostName().s(), Node.State.active);
                     if ( ! node.isPresent()) continue; // we also get status from infrastructure nodes, which are not in the repo
 
                     if (service.serviceStatus().equals(ServiceMonitorStatus.DOWN))
@@ -120,7 +120,7 @@ public class NodeFailer extends Maintainer {
         if (node.history().event(History.Event.Type.down).isPresent()) return node; // already down: Don't change down timestamp
 
         try (Mutex lock = nodeRepository().lock(node.allocation().get().owner())) {
-            node = nodeRepository().getNode(Node.State.active, node.hostname()).get(); // re-get inside lock
+            node = nodeRepository().getNode(node.hostname(), Node.State.active).get(); // re-get inside lock
             return nodeRepository().write(node.setDown(clock.instant()));
         }
     }
@@ -129,7 +129,7 @@ public class NodeFailer extends Maintainer {
         if ( ! node.history().event(History.Event.Type.down).isPresent()) return;
 
         try (Mutex lock = nodeRepository().lock(node.allocation().get().owner())) {
-            node = nodeRepository().getNode(Node.State.active, node.hostname()).get(); // re-get inside lock
+            node = nodeRepository().getNode(node.hostname(), Node.State.active).get(); // re-get inside lock
             nodeRepository().write(node.setUp());
         }
     }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
index 813941de1eb..751d93c8235 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
@@ -187,7 +187,7 @@ public class CuratorDatabaseClient {
             states = Node.State.values();
         for (Node.State state : states) {
             for (String hostname : curatorDatabase.getChildren(toPath(state))) {
-                final Optional<Node> node = getNode(state, hostname);
+                Optional<Node> node = getNode(hostname, state);
                 if (node.isPresent()) nodes.add(node.get()); // node might disappear between getChildren and getNode
             }
         }
@@ -202,8 +202,15 @@ public class CuratorDatabaseClient {
     }
 
     /** Returns a particular node, or empty if there is no such node in this state */
-    public Optional<Node> getNode(Node.State state, String hostname) {
-        return curatorDatabase.getData(toPath(state, hostname)).map((data) -> nodeSerializer.fromJson(state, data));
+    public Optional<Node> getNode(String hostname, Node.State ... states) {
+        if (states.length == 0)
+            states = Node.State.values();
+        for (Node.State state : states) {
+            Optional<byte[]> nodeData = curatorDatabase.getData(toPath(state, hostname));
+            if (nodeData.isPresent())
+                return nodeData.map((data) -> nodeSerializer.fromJson(state, data));
+        }
+        return Optional.empty();
     }
 
     private Path toPath(Node.State nodeState) { return root.append(toDir(nodeState)); }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/legacy/ProvisionResource.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/legacy/ProvisionResource.java
index 867d45b4fcd..884f97e2c2b 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/legacy/ProvisionResource.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/legacy/ProvisionResource.java
@@ -73,11 +73,9 @@ public class ProvisionResource {
     @PUT
     @Path("/node/ready")
     public void setReady(String hostName) {
-        if ( nodeRepository.getNode(Node.State.ready, hostName).isPresent()) return; // node already 'ready'
+        if ( nodeRepository.getNode(hostName, Node.State.ready).isPresent()) return; // node already 'ready'
 
-        Optional<Node> node = nodeRepository.getNode(Node.State.provisioned, hostName);
-        if ( ! node.isPresent())
-            node = nodeRepository.getNode(Node.State.dirty, hostName);
+        Optional<Node> node = nodeRepository.getNode(hostName, Node.State.provisioned, Node.State.dirty);
         if ( ! node.isPresent())
             throw new IllegalArgumentException("Could not set " + hostName + " ready: Not registered as provisioned or dirty");
 
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 b137d10b311..c8bed1b9b9a 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
@@ -211,16 +211,11 @@ public class NodesApiHandler extends LoggingRequestHandler {
     // TODO: Move most of this to node repo
     public String setNodeReady(String path) {
         String hostname = lastElement(path);
-        if ( nodeRepository.getNode(Node.State.ready, hostname).isPresent())
+        if ( nodeRepository.getNode(hostname, Node.State.ready).isPresent())
             return "Nothing done; " + hostname + " is already ready";
 
-        Optional<Node> node = nodeRepository.getNode(Node.State.provisioned, hostname);
-        if ( ! node.isPresent())
-            node = nodeRepository.getNode(Node.State.dirty, hostname);
-        if ( ! node.isPresent())
-            node = nodeRepository.getNode(Node.State.failed, hostname);
-        if ( ! node.isPresent())
-            node = nodeRepository.getNode(Node.State.parked, hostname);
+        Optional<Node> node = nodeRepository.getNode(hostname, Node.State.provisioned, Node.State.dirty, Node.State.failed, Node.State.parked);
+
         if ( ! node.isPresent())
             throw new IllegalArgumentException("Could not set " + hostname + " ready: Not registered as provisioned, dirty, failed or parked");
 
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
index 8eec59590c2..737c19cb3a7 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
@@ -164,7 +164,7 @@ public class ProvisioningTester implements AutoCloseable {
     }
 
     public void fail(HostSpec host) {
-        int beforeFailCount = nodeRepository.getNode(Node.State.active, host.hostname()).get().status().failCount();
+        int beforeFailCount = nodeRepository.getNode(host.hostname(), Node.State.active).get().status().failCount();
         Node failedNode = nodeRepository.fail(host.hostname());
         assertTrue(nodeRepository.getNodes(Node.Type.tenant, Node.State.failed).contains(failedNode));
         assertEquals(beforeFailCount + 1, failedNode.status().failCount());