Precompute host -> children mapping as it is frequently used and expensive to compute individually.

7 files changed, 111 insertions, 44 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodesAndHosts.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodesAndHosts.java
new file mode 100644
index 00000000000..dd9c7021b9e
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodesAndHosts.java
@@ -0,0 +1,58 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Wraps a NodeList and builds a host -> children mapping for faster access
+ * as that is done very frequently.
+ *
+ * @author baldersheim
+ */
+public class NodesAndHosts<NL extends NodeList> {
+    private final NL nodes;
+    private final Map<String, HostAndNodes> host2Nodes;
+
+    public static <L extends NodeList> NodesAndHosts<L> create(L nodes) {
+        return new NodesAndHosts<L>(nodes);
+    }
+
+    private NodesAndHosts(NL nodes) {
+        this.nodes = nodes;
+        host2Nodes = new HashMap<>();
+        nodes.forEach(node -> {
+            node.parentHostname().ifPresentOrElse(
+                    parent -> host2Nodes.computeIfAbsent(parent, key -> new HostAndNodes()).add(node),
+                    () -> host2Nodes.computeIfAbsent(node.hostname(), key -> new HostAndNodes()).setHost(node));
+        });
+    }
+
+    /// Return the NodeList used for construction
+    public NL nodes() { return nodes; }
+
+    public NodeList childrenOf(Node host) {
+        HostAndNodes hostAndNodes = host2Nodes.get(host.hostname());
+        return hostAndNodes != null ? NodeList.copyOf(hostAndNodes.children) : NodeList.of();
+    }
+
+    public Node parentOf(Node node) {
+        if (node.parentHostname().isEmpty()) return null;
+
+        HostAndNodes hostAndNodes = host2Nodes.get(node.parentHostname().get());
+        return hostAndNodes != null ? hostAndNodes.host : null;
+    }
+
+    private static class HostAndNodes {
+        private Node host;
+        private final List<Node> children;
+        HostAndNodes() {
+            this.host = null;
+            children = new ArrayList<>();
+        }
+        void setHost(Node host) { this.host = host; }
+        void add(Node child) { children.add(child); }
+    }
+}
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeMover.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeMover.java
index 3003e519fcd..fef6713b072 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeMover.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeMover.java
@@ -9,6 +9,7 @@ import com.yahoo.jdisc.Metric;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
+import com.yahoo.vespa.hosted.provision.NodesAndHosts;
 import com.yahoo.vespa.hosted.provision.provisioning.HostCapacity;
 
 import java.time.Duration;
@@ -39,7 +40,7 @@ public abstract class NodeMover<MOVE> extends NodeRepositoryMaintainer {
     }
 
     /** Returns a suggested move for given node */
-    protected abstract MOVE suggestedMove(Node node, Node fromHost, Node toHost, NodeList allNodes);
+    protected abstract MOVE suggestedMove(Node node, Node fromHost, Node toHost, NodesAndHosts<? extends NodeList> allNodes);
 
     private static class HostWithResources {
         private final Node node;
@@ -55,17 +56,17 @@ public abstract class NodeMover<MOVE> extends NodeRepositoryMaintainer {
     }
 
     /** Find the best possible move */
-    protected final MOVE findBestMove(NodeList allNodes) {
+    protected final MOVE findBestMove(NodesAndHosts<? extends NodeList> allNodes) {
         HostCapacity capacity = new HostCapacity(allNodes, nodeRepository().resourcesCalculator());
         MOVE bestMove = emptyMove;
         // Shuffle nodes so we did not get stuck if the chosen move is consistently discarded. Node moves happen through
         // a soft request to retire (preferToRetire), which node allocation can disregard
-        NodeList activeNodes = allNodes.nodeType(NodeType.tenant)
+        NodeList activeNodes = allNodes.nodes().nodeType(NodeType.tenant)
                                        .state(Node.State.active)
                                        .shuffle(random);
-        Set<Node> spares = capacity.findSpareHosts(allNodes.asList(), nodeRepository().spareCount());
+        Set<Node> spares = capacity.findSpareHosts(allNodes.nodes().asList(), nodeRepository().spareCount());
         List<HostWithResources> hostResources = new ArrayList<>();
-        allNodes.matching(nodeRepository().nodes()::canAllocateTenantNodeTo).forEach(host -> hostResources.add(new HostWithResources(host, capacity.availableCapacityOf(host))));
+        allNodes.nodes().matching(nodeRepository().nodes()::canAllocateTenantNodeTo).forEach(host -> hostResources.add(new HostWithResources(host, capacity.availableCapacityOf(host))));
         for (Node node : activeNodes) {
             if (node.parentHostname().isEmpty()) continue;
             ApplicationId applicationId = node.allocation().get().owner();
@@ -76,7 +77,7 @@ public abstract class NodeMover<MOVE> extends NodeRepositoryMaintainer {
                 if (spares.contains(toHost.node)) continue; // Do not offer spares as a valid move as they are reserved for replacement of failed nodes
                 if ( ! toHost.hasCapacity(node.resources())) continue;
 
-                MOVE suggestedMove = suggestedMove(node, allNodes.parentOf(node).get(), toHost.node, allNodes);
+                MOVE suggestedMove = suggestedMove(node, allNodes.parentOf(node), toHost.node, allNodes);
                 bestMove = bestMoveOf(bestMove, suggestedMove);
             }
         }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/Rebalancer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/Rebalancer.java
index 50bb213ad73..bbe2a8e3e1d 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/Rebalancer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/Rebalancer.java
@@ -8,6 +8,7 @@ import com.yahoo.jdisc.Metric;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
+import com.yahoo.vespa.hosted.provision.NodesAndHosts;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.provisioning.HostCapacity;
 
@@ -40,15 +41,15 @@ public class Rebalancer extends NodeMover<Rebalancer.Move> {
         if (nodeRepository().zone().environment().isTest()) return 1.0; // Short lived deployments; no need to rebalance
 
         // Work with an unlocked snapshot as this can take a long time and full consistency is not needed
-        NodeList allNodes = nodeRepository().nodes().list();
+        NodesAndHosts<NodeList> allNodes = NodesAndHosts.create(nodeRepository().nodes().list());
         updateSkewMetric(allNodes);
-        if ( ! zoneIsStable(allNodes)) return 1.0;
+        if ( ! zoneIsStable(allNodes.nodes())) return 1.0;
         findBestMove(allNodes).execute(true, Agent.Rebalancer, deployer, metric, nodeRepository());
         return 1.0;
    }
 
     @Override
-    protected Move suggestedMove(Node node, Node fromHost, Node toHost, NodeList allNodes) {
+    protected Move suggestedMove(Node node, Node fromHost, Node toHost, NodesAndHosts<? extends NodeList> allNodes) {
         HostCapacity capacity = new HostCapacity(allNodes, nodeRepository().resourcesCalculator());
         double skewReductionAtFromHost = skewReductionByRemoving(node, fromHost, capacity);
         double skewReductionAtToHost = skewReductionByAdding(node, toHost, capacity);
@@ -63,11 +64,11 @@ public class Rebalancer extends NodeMover<Rebalancer.Move> {
     }
 
     /** We do this here rather than in MetricsReporter because it is expensive and frequent updates are unnecessary */
-    private void updateSkewMetric(NodeList allNodes) {
+    private void updateSkewMetric(NodesAndHosts<? extends NodeList> allNodes) {
         HostCapacity capacity = new HostCapacity(allNodes, nodeRepository().resourcesCalculator());
         double totalSkew = 0;
         int hostCount = 0;
-        for (Node host : allNodes.nodeType(NodeType.host).state(Node.State.active)) {
+        for (Node host : allNodes.nodes().nodeType(NodeType.host).state(Node.State.active)) {
             hostCount++;
             totalSkew += Node.skew(host.flavor().resources(), capacity.unusedCapacityOf(host));
         }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SpareCapacityMaintainer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SpareCapacityMaintainer.java
index 2447503515a..1603b0f0012 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SpareCapacityMaintainer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SpareCapacityMaintainer.java
@@ -9,6 +9,7 @@ import com.yahoo.jdisc.Metric;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
+import com.yahoo.vespa.hosted.provision.NodesAndHosts;
 import com.yahoo.vespa.hosted.provision.maintenance.MaintenanceDeployment.Move;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.provisioning.HostCapacity;
@@ -116,13 +117,13 @@ public class SpareCapacityMaintainer extends NodeRepositoryMaintainer {
         if (nodeWhichCantMove.isEmpty()) return List.of();
 
         Node node = nodeWhichCantMove.get();
-        NodeList allNodes = nodeRepository().nodes().list();
+        NodesAndHosts<NodeList> allNodes = NodesAndHosts.create(nodeRepository().nodes().list());
         // Allocation will assign the spareCount most empty nodes as "spares", which will not be allocated on
         // unless needed for node failing. Our goal here is to make room on these spares for the given node
         HostCapacity hostCapacity = new HostCapacity(allNodes, nodeRepository().resourcesCalculator());
-        Set<Node> spareHosts = hostCapacity.findSpareHosts(allNodes.hosts().satisfies(node.resources()).asList(),
+        Set<Node> spareHosts = hostCapacity.findSpareHosts(allNodes.nodes().hosts().satisfies(node.resources()).asList(),
                                                            nodeRepository().spareCount());
-        List<Node> hosts = allNodes.hosts().except(spareHosts).asList();
+        List<Node> hosts = allNodes.nodes().hosts().except(spareHosts).asList();
 
         CapacitySolver capacitySolver = new CapacitySolver(hostCapacity, maxIterations);
         List<Move> shortestMitigation = null;
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancer.java
index 44890f2f5af..86a4424cb73 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/SwitchRebalancer.java
@@ -8,6 +8,7 @@ import com.yahoo.jdisc.Metric;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
+import com.yahoo.vespa.hosted.provision.NodesAndHosts;
 import com.yahoo.vespa.hosted.provision.maintenance.MaintenanceDeployment.Move;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 
@@ -36,17 +37,17 @@ public class SwitchRebalancer extends NodeMover<Move> {
     protected double maintain() {
         if (!nodeRepository().nodes().isWorking()) return 0.0;
         if (!nodeRepository().zone().environment().isProduction()) return 1.0;
-        NodeList allNodes = nodeRepository().nodes().list(); // Lockless as strong consistency is not needed
-        if (!zoneIsStable(allNodes)) return 1.0;
+        NodesAndHosts<NodeList> allNodes = NodesAndHosts.create(nodeRepository().nodes().list()); // Lockless as strong consistency is not needed
+        if (!zoneIsStable(allNodes.nodes())) return 1.0;
 
         findBestMove(allNodes).execute(false, Agent.SwitchRebalancer, deployer, metric, nodeRepository());
         return 1.0;
     }
 
     @Override
-    protected Move suggestedMove(Node node, Node fromHost, Node toHost, NodeList allNodes) {
-        NodeList clusterNodes = clusterOf(node, allNodes);
-        NodeList clusterHosts = allNodes.parentsOf(clusterNodes);
+    protected Move suggestedMove(Node node, Node fromHost, Node toHost, NodesAndHosts<? extends NodeList> allNodes) {
+        NodeList clusterNodes = clusterOf(node, allNodes.nodes());
+        NodeList clusterHosts = allNodes.nodes().parentsOf(clusterNodes);
         if (onExclusiveSwitch(node, clusterHosts)) return Move.empty();
         if (!increasesExclusiveSwitches(clusterNodes, clusterHosts, toHost)) return Move.empty();
         return new Move(node, fromHost, toHost);
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostCapacity.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostCapacity.java
index 462187a5add..bc2d606aa63 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostCapacity.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostCapacity.java
@@ -5,6 +5,7 @@ import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
+import com.yahoo.vespa.hosted.provision.NodesAndHosts;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -22,15 +23,18 @@ import java.util.stream.Collectors;
  */
 public class HostCapacity {
 
-    private final NodeList allNodes;
+    private final NodesAndHosts<? extends NodeList> allNodes;
     private final HostResourcesCalculator hostResourcesCalculator;
 
     public HostCapacity(NodeList allNodes, HostResourcesCalculator hostResourcesCalculator) {
+        this(NodesAndHosts.create(Objects.requireNonNull(allNodes, "allNodes must be non-null")), hostResourcesCalculator);
+    }
+    public HostCapacity(NodesAndHosts<? extends NodeList> allNodes, HostResourcesCalculator hostResourcesCalculator) {
         this.allNodes = Objects.requireNonNull(allNodes, "allNodes must be non-null");
         this.hostResourcesCalculator = Objects.requireNonNull(hostResourcesCalculator, "hostResourcesCalculator must be non-null");
     }
 
-    public NodeList allNodes() { return allNodes; }
+    public NodeList allNodes() { return allNodes.nodes(); }
 
     /**
      * Spare hosts are the hosts in the system with the most free capacity. A zone may reserve a minimum number of spare
@@ -93,9 +97,9 @@ public class HostCapacity {
     /** Returns the number of available IP addresses on given host */
     int freeIps(Node host) {
         if (host.type() == NodeType.host) {
-            return host.ipConfig().pool().eventuallyUnusedAddressCount(allNodes);
+            return host.ipConfig().pool().eventuallyUnusedAddressCount(allNodes.nodes());
         }
-        return host.ipConfig().pool().findUnusedIpAddresses(allNodes).size();
+        return host.ipConfig().pool().findUnusedIpAddresses(allNodes.nodes()).size();
     }
 
     /** Returns the capacity of given host that is both free and usable */
@@ -122,7 +126,7 @@ public class HostCapacity {
         if (   requireIps && freeIps(host) == 0) return NodeResources.zero();
 
         NodeResources hostResources = hostResourcesCalculator.advertisedResourcesOf(host.flavor());
-        return allNodes.childrenOf(host).asList().stream()
+        return allNodes.childrenOf(host).stream()
                        .filter(node -> !(excludeInactive && inactiveOrRetired(node)))
                        .map(node -> node.flavor().resources().justNumbers())
                        .reduce(hostResources.justNumbers(), NodeResources::subtract)
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodePrioritizer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodePrioritizer.java
index bfc848d4031..64d0fd9942f 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodePrioritizer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodePrioritizer.java
@@ -7,6 +7,7 @@ import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.hosted.provision.LockedNodeList;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
+import com.yahoo.vespa.hosted.provision.NodesAndHosts;
 import com.yahoo.vespa.hosted.provision.node.Nodes;
 import com.yahoo.vespa.hosted.provision.persistence.NameResolver;
 
@@ -30,7 +31,7 @@ import java.util.stream.Collectors;
 public class NodePrioritizer {
 
     private final List<NodeCandidate> nodes = new ArrayList<>();
-    private final LockedNodeList allNodes;
+    private final NodesAndHosts<LockedNodeList> allNodes;
     private final HostCapacity capacity;
     private final NodeSpec requestedNodes;
     private final ApplicationId application;
@@ -44,21 +45,21 @@ public class NodePrioritizer {
     private final int currentClusterSize;
     private final Set<Node> spareHosts;
 
-    public NodePrioritizer(LockedNodeList allNodes, ApplicationId application, ClusterSpec clusterSpec, NodeSpec nodeSpec,
+    public NodePrioritizer(LockedNodeList allLockedNodes, ApplicationId application, ClusterSpec clusterSpec, NodeSpec nodeSpec,
                            int wantedGroups, boolean dynamicProvisioning, NameResolver nameResolver,
                            HostResourcesCalculator hostResourcesCalculator, int spareCount) {
-        this.allNodes = allNodes;
+        this.allNodes = NodesAndHosts.create(allLockedNodes);
         this.capacity = new HostCapacity(allNodes, hostResourcesCalculator);
         this.requestedNodes = nodeSpec;
         this.clusterSpec = clusterSpec;
         this.application = application;
         this.dynamicProvisioning = dynamicProvisioning;
         this.spareHosts = dynamicProvisioning ?
-                capacity.findSpareHostsInDynamicallyProvisionedZones(allNodes.asList()) :
-                capacity.findSpareHosts(allNodes.asList(), spareCount);
+                capacity.findSpareHostsInDynamicallyProvisionedZones(allNodes.nodes().asList()) :
+                capacity.findSpareHosts(allNodes.nodes().asList(), spareCount);
         this.nameResolver = nameResolver;
 
-        NodeList nodesInCluster = allNodes.owner(application).type(clusterSpec.type()).cluster(clusterSpec.id());
+        NodeList nodesInCluster = allNodes.nodes().owner(application).type(clusterSpec.type()).cluster(clusterSpec.id());
         NodeList nonRetiredNodesInCluster = nodesInCluster.not().retired();
         long currentGroups = nonRetiredNodesInCluster.state(Node.State.active).stream()
                 .flatMap(node -> node.allocation()
@@ -134,7 +135,7 @@ public class NodePrioritizer {
     private void addCandidatesOnExistingHosts() {
         if ( !canAllocateNew) return;
 
-        for (Node host : allNodes) {
+        for (Node host : allNodes.nodes()) {
             if ( ! Nodes.canAllocateTenantNodeTo(host, dynamicProvisioning)) continue;
             if (host.reservedTo().isPresent() && !host.reservedTo().get().equals(application.tenant())) continue;
             if (host.reservedTo().isPresent() && application.instance().isTester()) continue;
@@ -147,7 +148,7 @@ public class NodePrioritizer {
                                                    capacity.availableCapacityOf(host),
                                                    host,
                                                    spareHosts.contains(host),
-                                                   allNodes,
+                                                   allNodes.nodes(),
                                                    nameResolver));
         }
     }
@@ -155,7 +156,7 @@ public class NodePrioritizer {
     /** Add existing nodes allocated to the application */
     private void addApplicationNodes() {
         EnumSet<Node.State> legalStates = EnumSet.of(Node.State.active, Node.State.inactive, Node.State.reserved);
-        allNodes.asList().stream()
+        allNodes.nodes().stream()
                 .filter(node -> node.type() == requestedNodes.type())
                 .filter(node -> legalStates.contains(node.state()))
                 .filter(node -> node.allocation().isPresent())
@@ -168,7 +169,7 @@ public class NodePrioritizer {
 
     /** Add nodes already provisioned, but not allocated to any application */
     private void addReadyNodes() {
-        allNodes.asList().stream()
+        allNodes.nodes().stream()
                 .filter(node -> node.type() == requestedNodes.type())
                 .filter(node -> node.state() == Node.State.ready)
                 .map(node -> candidateFrom(node, false))
@@ -178,17 +179,17 @@ public class NodePrioritizer {
 
     /** Create a candidate from given pre-existing node */
     private NodeCandidate candidateFrom(Node node, boolean isSurplus) {
-        Optional<Node> parent = allNodes.parentOf(node);
-        if (parent.isPresent()) {
+        Node parent = allNodes.parentOf(node);
+        if (parent != null) {
             return NodeCandidate.createChild(node,
-                                             capacity.availableCapacityOf(parent.get()),
-                                             parent.get(),
-                                             spareHosts.contains(parent.get()),
+                                             capacity.availableCapacityOf(parent),
+                                             parent,
+                                             spareHosts.contains(parent),
                                              isSurplus,
                                              false,
-                                             parent.get().exclusiveToApplicationId().isEmpty()
+                                             parent.exclusiveToApplicationId().isEmpty()
                                              && requestedNodes.canResize(node.resources(),
-                                                                         capacity.availableCapacityOf(parent.get()),
+                                                                         capacity.availableCapacityOf(parent),
                                                                          topologyChange,
                                                                          currentClusterSize));
         }
@@ -215,9 +216,9 @@ public class NodePrioritizer {
      */
     private boolean canStillAllocate(Node node) {
         if (node.type() != NodeType.tenant || node.parentHostname().isEmpty()) return true;
-        Optional<Node> parent = allNodes.parentOf(node);
-        if (parent.isEmpty()) return false;
-        return Nodes.canAllocateTenantNodeTo(parent.get(), dynamicProvisioning);
+        Node parent = allNodes.parentOf(node);
+        if (parent == null) return false;
+        return Nodes.canAllocateTenantNodeTo(parent, dynamicProvisioning);
     }
 
 }