Read node objects from unified path

8 files changed, 53 insertions, 89 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/MaintenanceDeployment.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/MaintenanceDeployment.java
index a94adc5aaae..bcc571355e3 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/MaintenanceDeployment.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/MaintenanceDeployment.java
@@ -190,7 +190,8 @@ class MaintenanceDeployment implements Closeable {
                     markPreferToRetire(node, false, agent, nodeRepository); // Necessary if this failed, no-op otherwise
 
                     // Immediately clean up if we reserved the node but could not activate or reserved a node on the wrong host
-                    expectedNewNode.flatMap(node -> nodeRepository.nodes().node(node.hostname(), Node.State.reserved))
+                    expectedNewNode.flatMap(node -> nodeRepository.nodes().node(node.hostname()))
+                                   .filter(node -> node.state() == Node.State.reserved)
                                    .ifPresent(node -> nodeRepository.nodes().deallocate(node, agent, "Expired by " + agent));
                 }
             }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeHealthTracker.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeHealthTracker.java
index b602d2ac7cd..94683d463af 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeHealthTracker.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeHealthTracker.java
@@ -96,7 +96,8 @@ public class NodeHealthTracker extends NodeRepositoryMaintainer {
 
     /** Get node by given hostname and application. The applicationLock must be held when calling this */
     private Optional<Node> getNode(String hostname, ApplicationId application, @SuppressWarnings("unused") Mutex applicationLock) {
-        return nodeRepository().nodes().node(hostname, Node.State.active)
+        return nodeRepository().nodes().node(hostname)
+                               .filter(node -> node.state() == Node.State.active)
                                .filter(node -> node.allocation().isPresent())
                                .filter(node -> node.allocation().get().owner().equals(application));
     }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java
index c0d0b220767..b2becc7ecfd 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java
@@ -131,8 +131,9 @@ public class IP {
          * @throws IllegalArgumentException if there are IP conflicts with existing nodes
          */
         public static List<Node> verify(List<Node> nodes, LockedNodeList allNodes) {
+            NodeList sortedNodes = allNodes.sortedBy(Comparator.comparing(Node::hostname));
             for (var node : nodes) {
-                for (var other : allNodes) {
+                for (var other : sortedNodes) {
                     if (node.equals(other)) continue;
                     if (canAssignIpOf(other, node)) continue;
 
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 b507d66e14f..5c2a6601ad8 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
@@ -76,11 +76,11 @@ public class Nodes {
     public void rewrite() {
         Instant start = clock.instant();
         int nodesWritten = 0;
-        for (Node.State state : Node.State.values()) {
-            List<Node> nodes = db.readNodes(state);
+        Map<Node.State, NodeList> nodes = list().groupingBy(Node::state);
+        for (var kv : nodes.entrySet()) {
             // TODO(mpolden): This should take the lock before writing
-            db.writeTo(state, nodes, Agent.system, Optional.empty());
-            nodesWritten += nodes.size();
+            db.writeTo(kv.getKey(), kv.getValue().asList(), Agent.system, Optional.empty());
+            nodesWritten += kv.getValue().size();
         }
         Instant end = clock.instant();
         log.log(Level.INFO, String.format("Rewrote %d nodes in %s", nodesWritten, Duration.between(start, end)));
@@ -88,24 +88,19 @@ public class Nodes {
 
     // ---------------- Query API ----------------------------------------------------------------
 
-    /**
-     * 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> node(String hostname, Node.State... inState) {
-        return db.readNode(hostname, inState);
+    /** Finds and returns the node with given hostname, or empty if not found */
+    public Optional<Node> node(String hostname) {
+        return db.readNode(hostname);
     }
 
     /**
-     * Returns a list of nodes in this repository in any of the given states
+     * Returns an unsorted list of all nodes in this repository, in any of the given states
      *
-     * @param inState the states to return nodes from. If no states are given, all nodes of the given type are returned
+     * @param inState the states to return nodes from. If no states are given, all nodes are returned
      */
     public NodeList list(Node.State... inState) {
-        return NodeList.copyOf(db.readNodes(inState));
+        NodeList nodes = NodeList.copyOf(db.readNodes());
+        return inState.length == 0 ? nodes : nodes.state(Set.of(inState));
     }
 
     /** Returns a locked list of all nodes in this repository */
@@ -768,13 +763,9 @@ public class Nodes {
         for (int i = 0; i < maxRetries; ++i) {
             Mutex lockToClose = lock(staleNode, timeout);
             try {
-                // As an optimization we first try finding the node in the same state
-                Optional<Node> freshNode = node(staleNode.hostname(), staleNode.state());
+                Optional<Node> freshNode = node(staleNode.hostname());
                 if (freshNode.isEmpty()) {
-                    freshNode = node(staleNode.hostname());
-                    if (freshNode.isEmpty()) {
-                        return Optional.empty();
-                    }
+                    return Optional.empty();
                 }
 
                 if (node.type() != NodeType.tenant ||
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CachingCurator.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CachingCurator.java
index 90d30a3a8bc..589468c48b8 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CachingCurator.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CachingCurator.java
@@ -208,7 +208,7 @@ public class CachingCurator {
         List<String> getChildren(Path path);
 
         /**
-         * Returns the a copy of the content of this child - which may be empty.
+         * Returns a copy of the content of this child - which may be empty.
          */
         Optional<byte[]> getData(Path path);
 
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDb.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDb.java
index cebc185360a..b9821b48fba 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDb.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDb.java
@@ -96,7 +96,7 @@ public class CuratorDb {
         db.create(nodesPath);
         // TODO(mpolden): Remove state paths after migration to nodesPath
         for (Node.State state : Node.State.values())
-            db.create(toPath(state));
+            db.create(toLegacyPath(state));
         db.create(applicationsPath);
         db.create(inactiveJobsPath);
         db.create(infrastructureVersionsPath);
@@ -117,7 +117,7 @@ public class CuratorDb {
             node = node.with(node.history().recordStateTransition(null, expectedState, agent, clock.instant()));
             // TODO(mpolden): Remove after migration to nodesPath
             byte[] serialized = nodeSerializer.toJson(node);
-            curatorTransaction.add(CuratorOperations.create(toPath(node).getAbsolute(), serialized));
+            curatorTransaction.add(CuratorOperations.create(toLegacyPath(node).getAbsolute(), serialized));
             curatorTransaction.add(CuratorOperations.create(nodePath(node).getAbsolute(), serialized));
         }
 
@@ -137,7 +137,7 @@ public class CuratorDb {
     /** Removes given nodes in transaction */
     public void removeNodes(List<Node> nodes, NestedTransaction transaction) {
         for (Node node : nodes) {
-            Path path = toPath(node.state(), node.hostname());
+            Path path = toLegacyPath(node.state(), node.hostname());
             CuratorTransaction curatorTransaction = db.newCuratorTransactionIn(transaction);
             // TODO(mpolden): Remove after migration to nodesPath
             curatorTransaction.add(CuratorOperations.delete(path.getAbsolute()));
@@ -237,8 +237,8 @@ public class CuratorDb {
     private void writeNode(Node.State toState, CuratorTransaction curatorTransaction, Node node, Node newNode) {
         byte[] nodeData = nodeSerializer.toJson(newNode);
         { // TODO(mpolden): Remove this after migration to nodesPath
-            String currentNodePath = toPath(node).getAbsolute();
-            String newNodePath = toPath(toState, newNode.hostname()).getAbsolute();
+            String currentNodePath = toLegacyPath(node).getAbsolute();
+            String newNodePath = toLegacyPath(toState, newNode.hostname()).getAbsolute();
             if (newNodePath.equals(currentNodePath)) {
                 curatorTransaction.add(CuratorOperations.setData(currentNodePath, nodeData));
             } else {
@@ -255,61 +255,39 @@ public class CuratorDb {
         return node.status();
     }
 
-    /**
-     * Returns all nodes which are in one of the given states.
-     * If no states are given this returns all nodes.
-     *
-     * @return the nodes in a mutable list owned by the caller
-     */
-    public List<Node> readNodes(Node.State ... states) {
-        List<Node> nodes = new ArrayList<>();
-        if (states.length == 0)
-            states = Node.State.values();
+    /** Returns all existing nodes */
+    public List<Node> readNodes() {
         CachingCurator.Session session = db.getSession();
-        for (Node.State state : states) {
-            for (String hostname : session.getChildren(toPath(state))) {
-                Optional<Node> node = readNode(session, hostname, state);
-                node.ifPresent(nodes::add); // node might disappear between getChildren and getNode
-            }
-        }
-        return nodes;
+        return session.getChildren(nodesPath).stream()
+                      .flatMap(hostname -> readNode(session, hostname).stream())
+                      .toList();
     }
 
-    /**
-     * Returns a particular node, or empty if this node is not in any of the given states.
-     * If no states are given this returns the node if it is present in any state.
-     */
-    public Optional<Node> readNode(CachingCurator.Session session, String hostname, Node.State ... states) {
-        if (states.length == 0)
-            states = Node.State.values();
-        for (Node.State state : states) {
-            Optional<byte[]> nodeData = session.getData(toPath(state, hostname));
-            if (nodeData.isPresent())
-                return nodeData.map((data) -> nodeSerializer.fromJson(state, data));
-        }
-        return Optional.empty();
+    private Optional<Node> readNode(CachingCurator.Session session, String hostname) {
+        return session.getData(nodePath(hostname)).map(nodeSerializer::fromJson);
     }
 
-    /** 
-     * Returns a particular node, or empty if this noe is not in any of the given states.
-     * If no states are given this returns the node if it is present in any state.
-     */
-    public Optional<Node> readNode(String hostname, Node.State ... states) {
-        return readNode(db.getSession(), hostname, states);
+    /** Read node with given hostname, if any such node exists */
+    public Optional<Node> readNode(String hostname) {
+        return readNode(db.getSession(), hostname);
     }
 
-    private Path toPath(Node.State nodeState) { return root.append(toDir(nodeState)); }
+    private Path toLegacyPath(Node.State nodeState) { return root.append(toDir(nodeState)); }
 
-    private Path toPath(Node node) {
+    private Path toLegacyPath(Node node) {
         return root.append(toDir(node.state())).append(node.hostname());
     }
 
-    private Path toPath(Node.State nodeState, String nodeName) {
+    private Path toLegacyPath(Node.State nodeState, String nodeName) {
         return root.append(toDir(nodeState)).append(nodeName);
     }
 
     private Path nodePath(Node node) {
-        return nodesPath.append(node.hostname());
+        return nodePath(node.hostname());
+    }
+
+    private Path nodePath(String hostname) {
+        return nodesPath.append(hostname);
     }
 
     /** Creates and returns the path to the lock for this application */
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java
index f448266b94b..39cccafb8ef 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java
@@ -40,7 +40,6 @@ import com.yahoo.vespa.hosted.provision.node.Status;
 import com.yahoo.vespa.hosted.provision.node.TrustStoreItem;
 
 import java.io.IOException;
-import java.nio.charset.StandardCharsets;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -266,19 +265,16 @@ public class NodeSerializer {
 
     // ---------------- Deserialization --------------------------------------------------
 
-    public Node fromJson(Node.State state, byte[] data) {
-        long key = Hashing.sipHash24().newHasher()
-                          .putString(state.name(), StandardCharsets.UTF_8)
-                          .putBytes(data).hash()
-                          .asLong();
+    public Node fromJson(byte[] data) {
+        long key = Hashing.sipHash24().newHasher().putBytes(data).hash().asLong();
         try {
-            return cache.get(key, () -> nodeFromSlime(state, SlimeUtils.jsonToSlime(data).get()));
+            return cache.get(key, () -> nodeFromSlime(SlimeUtils.jsonToSlime(data).get()));
         } catch (ExecutionException e) {
             throw new UncheckedExecutionException(e);
         }
     }
 
-    private Node nodeFromSlime(Node.State state, Inspector object) {
+    private Node nodeFromSlime(Inspector object) {
         Flavor flavor = flavorFromSlime(object);
         return new Node(object.field(idKey).asString(),
                         new IP.Config(ipAddressesFromSlime(object, ipAddressesKey),
@@ -288,7 +284,7 @@ public class NodeSerializer {
                         SlimeUtils.optionalString(object.field(parentHostnameKey)),
                         flavor,
                         statusFromSlime(object),
-                        state,
+                        nodeStateFromString(object.field(stateKey).asString()),
                         allocationFromSlime(flavor.resources(), object.field(instanceKey)),
                         historyFromSlime(object),
                         nodeTypeFromString(object.field(nodeTypeKey).asString()),
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 942f029bc6a..fd466108063 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
@@ -105,16 +105,12 @@ class NodesResponse extends SlimeJsonResponse {
     }
 
     /** Outputs the nodes in the given states to a node array */
-    private void nodesToSlime(Set<Node.State> statesToRead, Cursor parentObject) {
+    private void nodesToSlime(Set<Node.State> states, Cursor parentObject) {
         Cursor nodeArray = parentObject.setArray("nodes");
-        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);
-        });
+        NodeList nodes = nodeRepository.nodes().list()
+                                       .state(states)
+                                       .sortedBy(Comparator.comparing(Node::hostname));
+        toSlime(nodes, nodeArray);
     }
 
     private void toSlime(NodeList nodes, Cursor array) {