Read only relevant states when state filter is present

5 files changed, 92 insertions, 71 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/Nodes.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/Nodes.java
index 480fd72967e..7f57ec219ae 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/Nodes.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/Nodes.java
@@ -16,7 +16,7 @@ import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
 import com.yahoo.vespa.hosted.provision.NodeMutex;
 import com.yahoo.vespa.hosted.provision.maintenance.NodeFailer;
-import com.yahoo.vespa.hosted.provision.node.filter.StateFilter;
+import com.yahoo.vespa.hosted.provision.node.filter.NodeFilter;
 import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
 
 import java.time.Clock;
@@ -591,7 +591,7 @@ public class Nodes {
      * @return the nodes in their new state
      */
     public List<Node> restartActive(Predicate<Node> filter) {
-        return restart(StateFilter.from(Node.State.active).and(filter));
+        return restart(NodeFilter.in(Set.of(Node.State.active)).and(filter));
     }
 
     /**
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/filter/NodeFilter.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/filter/NodeFilter.java
new file mode 100644
index 00000000000..a65ec30264f
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/filter/NodeFilter.java
@@ -0,0 +1,63 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.node.filter;
+
+
+import com.yahoo.text.StringUtilities;
+import com.yahoo.vespa.hosted.provision.Node;
+
+import java.util.EnumSet;
+import java.util.Objects;
+import java.util.Set;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+/**
+ * A filter for nodes, matching by state. This should be the top-most filter so that the node-repository can determine
+ * which node states to read before testing additional filters.
+ *
+ * @author mpolden
+ */
+public class NodeFilter implements Predicate<Node> {
+
+    private final Set<Node.State> states;
+    private final Predicate<Node> filter;
+
+    private NodeFilter(Set<Node.State> states, Predicate<Node> filter) {
+        this.states = Objects.requireNonNull(states);
+        this.filter = Objects.requireNonNull(filter);
+    }
+
+    @Override
+    public boolean test(Node node) {
+        return states.contains(node.state()) && filter.test(node);
+    }
+
+    /** The node states to match */
+    public Set<Node.State> states() {
+        return states;
+    }
+
+    /** Returns a copy of this that matches with given filter */
+    public NodeFilter matching(Predicate<Node> filter) {
+        return new NodeFilter(states, filter);
+    }
+
+    /** Returns a node filter which matches a comma or space-separated list of states */
+    public static NodeFilter in(String states, boolean includeDeprovisioned) {
+        if (states == null) {
+            return NodeFilter.in(includeDeprovisioned
+                                         ? EnumSet.allOf(Node.State.class)
+                                         : EnumSet.complementOf(EnumSet.of(Node.State.deprovisioned)));
+        }
+        return NodeFilter.in(StringUtilities.split(states).stream()
+                                            .map(Node.State::valueOf)
+                                            .collect(Collectors.toSet()));
+    }
+
+    /** Returns a node filter matching given states */
+    public static NodeFilter in(Set<Node.State> states) {
+        return new NodeFilter(states, (ignored) -> true);
+    }
+
+
+}
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/filter/StateFilter.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/filter/StateFilter.java
deleted file mode 100644
index 9e3928ecbe5..00000000000
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/filter/StateFilter.java
+++ /dev/null
@@ -1,43 +0,0 @@
-// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.vespa.hosted.provision.node.filter;
-
-import com.yahoo.text.StringUtilities;
-import com.yahoo.vespa.hosted.provision.Node;
-
-import java.util.EnumSet;
-import java.util.Objects;
-import java.util.function.Predicate;
-import java.util.stream.Collectors;
-
-/**
- * A node filter which filters on node states.
- *
- * @author bratseth
- */
-public class StateFilter {
-
-    private StateFilter() {}
-
-    private static Predicate<Node> makePredicate(EnumSet<Node.State> states) {
-        Objects.requireNonNull(states, "state cannot be null, use an empty set");
-        return node -> states.contains(node.state());
-    }
-
-    /** Returns a copy of the given filter which only matches for the given state */
-    public static Predicate<Node> from(Node.State state) {
-        return makePredicate(EnumSet.of(state));
-    }
-
-    /** Returns a node filter which matches a comma or space-separated list of states */
-    public static Predicate<Node> from(String states, boolean includeDeprovisioned) {
-        if (states == null) {
-            return makePredicate(includeDeprovisioned ?
-                    EnumSet.allOf(Node.State.class) : EnumSet.complementOf(EnumSet.of(Node.State.deprovisioned)));
-        }
-
-        return makePredicate(StringUtilities.split(states).stream()
-                .map(Node.State::valueOf)
-                .collect(Collectors.toCollection(() -> EnumSet.noneOf(Node.State.class))));
-    }
-
-}
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) {