Read only relevant states when state filter is present

2 files changed, 27 insertions, 26 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java
index d1909b7a8f7..922c8bc8e20 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java
@@ -16,6 +16,7 @@ import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.node.Address;
 import com.yahoo.vespa.hosted.provision.node.History;
 import com.yahoo.vespa.hosted.provision.node.TrustStoreItem;
+import com.yahoo.vespa.hosted.provision.node.filter.NodeFilter;
 import com.yahoo.vespa.orchestrator.Orchestrator;
 import com.yahoo.vespa.orchestrator.status.HostInfo;
 import com.yahoo.vespa.orchestrator.status.HostStatus;
@@ -26,7 +27,6 @@ import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
-import java.util.function.Predicate;
 
 /**
 * @author bratseth
@@ -42,7 +42,7 @@ class NodesResponse extends SlimeJsonResponse {
     /** The parent url of nodes */
     private final String nodeParentUrl;
 
-    private final Predicate<Node> filter;
+    private final NodeFilter filter;
     private final boolean recursive;
     private final Function<HostName, Optional<HostInfo>> orchestrator;
     private final NodeRepository nodeRepository;
@@ -58,9 +58,9 @@ class NodesResponse extends SlimeJsonResponse {
 
         Cursor root = slime.setObject();
         switch (responseType) {
-            case nodeList: nodesToSlime(root); break;
+            case nodeList: nodesToSlime(filter.states(), root); break;
             case stateList : statesToSlime(root); break;
-            case nodesInStateList: nodesToSlime(NodeSerializer.stateFrom(lastElement(parentUrl)), root); break;
+            case nodesInStateList: nodesToSlime(Set.of(NodeSerializer.stateFrom(lastElement(parentUrl))), root); break;
             case singleNode : nodeToSlime(lastElement(parentUrl), root); break;
             default: throw new IllegalArgumentException();
         }
@@ -88,19 +88,20 @@ class NodesResponse extends SlimeJsonResponse {
     private void toSlime(Node.State state, Cursor object) {
         object.setString("url", parentUrl + NodeSerializer.toString(state));
         if (recursive)
-            nodesToSlime(state, object);
+            nodesToSlime(Set.of(state), object);
     }
 
-    /** Outputs the nodes in the given state to a node array */
-    private void nodesToSlime(Node.State state, Cursor parentObject) {
+    /** Outputs the nodes in the given states to a node array */
+    private void nodesToSlime(Set<Node.State> statesToRead, Cursor parentObject) {
         Cursor nodeArray = parentObject.setArray("nodes");
-        toSlime(nodeRepository.nodes().list(state).sortedBy(Comparator.comparing(Node::type)), nodeArray);
-    }
-
-    /** Outputs all the nodes to a node array */
-    private void nodesToSlime(Cursor parentObject) {
-        Cursor nodeArray = parentObject.setArray("nodes");
-        toSlime(nodeRepository.nodes().list(), nodeArray);
+        boolean sortByNodeType = statesToRead.size() == 1;
+        statesToRead.stream().sorted().forEach(state -> {
+            NodeList nodes = nodeRepository.nodes().list(state);
+            if (sortByNodeType) {
+                nodes = nodes.sortedBy(Comparator.comparing(Node::type));
+            }
+            toSlime(nodes, nodeArray);
+        });
     }
 
     private void toSlime(NodeList nodes, Cursor array) {
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
index 67bb69b6191..15e1061f5e1 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
@@ -35,11 +35,11 @@ import com.yahoo.vespa.hosted.provision.node.Address;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.node.IP;
 import com.yahoo.vespa.hosted.provision.node.filter.ApplicationFilter;
+import com.yahoo.vespa.hosted.provision.node.filter.NodeFilter;
 import com.yahoo.vespa.hosted.provision.node.filter.NodeHostFilter;
 import com.yahoo.vespa.hosted.provision.node.filter.NodeOsVersionFilter;
 import com.yahoo.vespa.hosted.provision.node.filter.NodeTypeFilter;
 import com.yahoo.vespa.hosted.provision.node.filter.ParentHostFilter;
-import com.yahoo.vespa.hosted.provision.node.filter.StateFilter;
 import com.yahoo.vespa.hosted.provision.restapi.NodesResponse.ResponseType;
 import com.yahoo.vespa.orchestrator.Orchestrator;
 import com.yahoo.yolean.Exceptions;
@@ -56,7 +56,6 @@ import java.util.List;
 import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
-import java.util.function.Predicate;
 import java.util.logging.Level;
 import java.util.stream.Collectors;
 
@@ -322,16 +321,17 @@ public class NodesV2ApiHandler extends LoggingRequestHandler {
         return NodeSerializer.typeFrom(object.asString());
     }
 
-    public static Predicate<Node> toNodeFilter(HttpRequest request) {
-        return NodeHostFilter.from(HostFilter.from(request.getProperty("hostname"),
-                                                   request.getProperty("flavor"),
-                                                   request.getProperty("clusterType"),
-                                                   request.getProperty("clusterId")))
-                .and(ApplicationFilter.from(request.getProperty("application")))
-                .and(StateFilter.from(request.getProperty("state"), request.getBooleanProperty("includeDeprovisioned")))
-                .and(NodeTypeFilter.from(request.getProperty("type")))
-                .and(ParentHostFilter.from(request.getProperty("parentHost")))
-                .and(NodeOsVersionFilter.from(request.getProperty("osVersion")));
+    public static NodeFilter toNodeFilter(HttpRequest request) {
+        return NodeFilter.in(request.getProperty("state"),
+                             request.getBooleanProperty("includeDeprovisioned"))
+                         .matching(NodeHostFilter.from(HostFilter.from(request.getProperty("hostname"),
+                                                                       request.getProperty("flavor"),
+                                                                       request.getProperty("clusterType"),
+                                                                       request.getProperty("clusterId")))
+                                                 .and(ApplicationFilter.from(request.getProperty("application")))
+                                                 .and(NodeTypeFilter.from(request.getProperty("type")))
+                                                 .and(ParentHostFilter.from(request.getProperty("parentHost")))
+                                                 .and(NodeOsVersionFilter.from(request.getProperty("osVersion"))));
     }
 
     private static boolean isPatchOverride(HttpRequest request) {