Refactor and read nodes from node repo further out.

5 files changed, 40 insertions, 33 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
index cdf117c11ce..70e5cf14b54 100644
--- 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
@@ -44,7 +44,10 @@ public class NodesAndHosts<NL extends NodeList> {
     public NL nodes() { return nodes; }
 
     public NodeList childrenOf(Node host) {
-        HostAndNodes hostAndNodes = host2Nodes.get(host.hostname());
+        return childrenOf(host.hostname());
+    }
+    public NodeList childrenOf(String hostname) {
+        HostAndNodes hostAndNodes = host2Nodes.get(hostname);
         return hostAndNodes != null ? NodeList.copyOf(hostAndNodes.children) : NodeList.of();
     }
 
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/DynamicProvisioningMaintainer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/DynamicProvisioningMaintainer.java
index 78686d8aeed..9c5b556b309 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/DynamicProvisioningMaintainer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/DynamicProvisioningMaintainer.java
@@ -22,6 +22,7 @@ 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.NodeRepository;
+import com.yahoo.vespa.hosted.provision.NodesAndHosts;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.node.History;
 import com.yahoo.vespa.hosted.provision.node.IP;
@@ -269,8 +270,8 @@ public class DynamicProvisioningMaintainer extends NodeRepositoryMaintainer {
                                                                    ArrayList<Node> mutableNodes) {
         for (int clusterIndex = 0; clusterIndex < preprovisionCapacity.size(); ++clusterIndex) {
             ClusterCapacity clusterCapacity = preprovisionCapacity.get(clusterIndex);
-            LockedNodeList nodeList = new LockedNodeList(mutableNodes, () -> {});
-            List<Node> candidates = findCandidates(clusterCapacity, clusterIndex, nodeList);
+            NodesAndHosts<LockedNodeList> nodesAndHosts = NodesAndHosts.create(new LockedNodeList(mutableNodes, () -> {}));
+            List<Node> candidates = findCandidates(clusterCapacity, clusterIndex, nodesAndHosts);
             int deficit = Math.max(0, clusterCapacity.count() - candidates.size());
             if (deficit > 0) {
                 return Optional.of(clusterCapacity.withCount(deficit));
@@ -283,7 +284,7 @@ public class DynamicProvisioningMaintainer extends NodeRepositoryMaintainer {
         return Optional.empty();
     }
 
-    private List<Node> findCandidates(ClusterCapacity clusterCapacity, int clusterIndex, LockedNodeList nodeList) {
+    private List<Node> findCandidates(ClusterCapacity clusterCapacity, int clusterIndex, NodesAndHosts<LockedNodeList> nodesAndHosts) {
         NodeResources nodeResources = toNodeResources(clusterCapacity);
 
         // We'll allocate each ClusterCapacity as a unique cluster in a dummy application
@@ -296,7 +297,7 @@ public class DynamicProvisioningMaintainer extends NodeRepositoryMaintainer {
         NodeSpec nodeSpec = NodeSpec.from(clusterCapacity.count(), nodeResources, false, true);
         int wantedGroups = 1;
 
-        NodePrioritizer prioritizer = new NodePrioritizer(nodeList, applicationId, clusterSpec, nodeSpec, wantedGroups,
+        NodePrioritizer prioritizer = new NodePrioritizer(nodesAndHosts, applicationId, clusterSpec, nodeSpec, wantedGroups,
                 true, nodeRepository().nameResolver(), nodeRepository().resourcesCalculator(),
                 nodeRepository().spareCount());
         List<NodeCandidate> nodeCandidates = prioritizer.collect(List.of());
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/GroupPreparer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/GroupPreparer.java
index aad32f64fa8..1aa6c2603ad 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/GroupPreparer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/GroupPreparer.java
@@ -10,6 +10,7 @@ import com.yahoo.transaction.Mutex;
 import com.yahoo.vespa.hosted.provision.LockedNodeList;
 import com.yahoo.vespa.hosted.provision.Node;
 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.HostProvisioner.HostSharing;
 
@@ -55,8 +56,9 @@ public class GroupPreparer {
         // Try preparing in memory without global unallocated lock. Most of the time there should be no changes and we
         // can return nodes previously allocated.
         {
+            NodesAndHosts<LockedNodeList> allNodesAndHosts = NodesAndHosts.create(nodeRepository.nodes().list(PROBE_LOCK));
             NodeAllocation probeAllocation = prepareAllocation(application, cluster, requestedNodes, surplusActiveNodes,
-                                                               indices::probeNext, wantedGroups, PROBE_LOCK);
+                                                               indices::probeNext, wantedGroups, allNodesAndHosts);
             if (probeAllocation.fulfilledAndNoChanges()) {
                 List<Node> acceptedNodes = probeAllocation.finalNodes();
                 surplusActiveNodes.removeAll(acceptedNodes);
@@ -69,9 +71,9 @@ public class GroupPreparer {
         // There were some changes, so re-do the allocation with locks
         try (Mutex lock = nodeRepository.nodes().lock(application);
              Mutex allocationLock = nodeRepository.nodes().lockUnallocated()) {
-
+            NodesAndHosts<LockedNodeList> allNodesAndHosts = NodesAndHosts.create(nodeRepository.nodes().list(allocationLock));
             NodeAllocation allocation = prepareAllocation(application, cluster, requestedNodes, surplusActiveNodes,
-                                                          indices::next, wantedGroups, allocationLock);
+                                                          indices::next, wantedGroups, allNodesAndHosts);
             NodeType hostType = allocation.nodeType().hostType();
             if (canProvisionDynamically(hostType)) {
                 HostSharing sharing = hostSharing(requestedNodes, hostType);
@@ -115,10 +117,10 @@ public class GroupPreparer {
 
     private NodeAllocation prepareAllocation(ApplicationId application, ClusterSpec cluster, NodeSpec requestedNodes,
                                              List<Node> surplusActiveNodes, Supplier<Integer> nextIndex, int wantedGroups,
-                                             Mutex allocationLock) {
-        LockedNodeList allNodes = nodeRepository.nodes().list(allocationLock);
-        NodeAllocation allocation = new NodeAllocation(allNodes, application, cluster, requestedNodes, nextIndex, nodeRepository);
-        NodePrioritizer prioritizer = new NodePrioritizer(allNodes,
+                                             NodesAndHosts<LockedNodeList> allNodesAndHosts) {
+
+        NodeAllocation allocation = new NodeAllocation(allNodesAndHosts, application, cluster, requestedNodes, nextIndex, nodeRepository);
+        NodePrioritizer prioritizer = new NodePrioritizer(allNodesAndHosts,
                                                           application,
                                                           cluster,
                                                           requestedNodes,
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java
index bdb7cb4b64d..d96613201c3 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java
@@ -11,6 +11,7 @@ import com.yahoo.config.provision.SystemName;
 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.node.Allocation;
 
@@ -40,7 +41,7 @@ class NodeAllocation {
     private static final Logger LOG = Logger.getLogger(NodeAllocation.class.getName());
 
     /** List of all nodes in node-repository */
-    private final NodeList allNodes;
+    private final NodesAndHosts<? extends NodeList> allNodesAndHosts;
 
     /** The application this list is for */
     private final ApplicationId application;
@@ -80,9 +81,9 @@ class NodeAllocation {
     private final NodeRepository nodeRepository;
     private final NodeResourceLimits nodeResourceLimits;
 
-    NodeAllocation(NodeList allNodes, ApplicationId application, ClusterSpec cluster, NodeSpec requestedNodes,
+    NodeAllocation(NodesAndHosts<? extends NodeList> allNodesAndHosts, ApplicationId application, ClusterSpec cluster, NodeSpec requestedNodes,
                    Supplier<Integer> nextIndex, NodeRepository nodeRepository) {
-        this.allNodes = allNodes;
+        this.allNodesAndHosts = allNodesAndHosts;
         this.application = application;
         this.cluster = cluster;
         this.requestedNodes = requestedNodes;
@@ -199,7 +200,7 @@ class NodeAllocation {
 
         // In non-dynamic provisioned zones we require that if either of the nodes on the host requires exclusivity,
         // then all the nodes on the host must have the same owner
-        for (Node nodeOnHost : allNodes.childrenOf(candidate.parentHostname().get())) {
+        for (Node nodeOnHost : allNodesAndHosts.childrenOf(candidate.parentHostname().get())) {
             if (nodeOnHost.allocation().isEmpty()) continue;
             if (requestedNodes.isExclusive() || nodeOnHost.allocation().get().membership().cluster().isExclusive()) {
                 if ( ! nodeOnHost.allocation().get().owner().equals(application)) return true;
@@ -273,7 +274,7 @@ class NodeAllocation {
     }
 
     private Node resize(Node node) {
-        NodeResources hostResources = allNodes.parentOf(node).get().flavor().resources();
+        NodeResources hostResources = allNodesAndHosts.parentOf(node).get().flavor().resources();
         return node.with(new Flavor(requestedNodes.resources().get()
                                                   .with(hostResources.diskSpeed())
                                                   .with(hostResources.storageType())),
@@ -324,7 +325,7 @@ class NodeAllocation {
         if (hostType == NodeType.host) return nodeRepository.database().readProvisionIndices(count);
 
         // Infrastructure hosts have fixed indices, starting at 1
-        Set<Integer> currentIndices = allNodes.nodeType(hostType)
+        Set<Integer> currentIndices = allNodesAndHosts.nodes().nodeType(hostType)
                                               .hostnames()
                                               .stream()
                                               // TODO(mpolden): Use cluster index instead of parsing hostname, once all
@@ -424,7 +425,7 @@ class NodeAllocation {
         if (nodeType() == NodeType.tenant) return accepted;
         // Infrastructure nodes are always allocated by type. Count all nodes as accepted so that we never exceed
         // the wanted number of nodes for the type.
-        return allNodes.nodeType(nodeType()).size();
+        return allNodesAndHosts.nodes().nodeType(nodeType()).size();
     }
 
     /** Prefer to retire nodes we want the least */
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 1f618a18d24..c7a2d27ca00 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
@@ -31,7 +31,7 @@ import java.util.stream.Collectors;
 public class NodePrioritizer {
 
     private final List<NodeCandidate> nodes = new ArrayList<>();
-    private final NodesAndHosts<LockedNodeList> allNodes;
+    private final NodesAndHosts<LockedNodeList> allNodesAndHosts;
     private final HostCapacity capacity;
     private final NodeSpec requestedNodes;
     private final ApplicationId application;
@@ -45,21 +45,21 @@ public class NodePrioritizer {
     private final int currentClusterSize;
     private final Set<Node> spareHosts;
 
-    public NodePrioritizer(LockedNodeList allLockedNodes, ApplicationId application, ClusterSpec clusterSpec, NodeSpec nodeSpec,
+    public NodePrioritizer(NodesAndHosts<LockedNodeList> allNodesAndHosts, ApplicationId application, ClusterSpec clusterSpec, NodeSpec nodeSpec,
                            int wantedGroups, boolean dynamicProvisioning, NameResolver nameResolver,
                            HostResourcesCalculator hostResourcesCalculator, int spareCount) {
-        this.allNodes = NodesAndHosts.create(allLockedNodes);
-        this.capacity = new HostCapacity(allNodes, hostResourcesCalculator);
+        this.allNodesAndHosts = allNodesAndHosts;
+        this.capacity = new HostCapacity(this.allNodesAndHosts, hostResourcesCalculator);
         this.requestedNodes = nodeSpec;
         this.clusterSpec = clusterSpec;
         this.application = application;
         this.dynamicProvisioning = dynamicProvisioning;
         this.spareHosts = dynamicProvisioning ?
-                capacity.findSpareHostsInDynamicallyProvisionedZones(allNodes.nodes().asList()) :
-                capacity.findSpareHosts(allNodes.nodes().asList(), spareCount);
+                capacity.findSpareHostsInDynamicallyProvisionedZones(this.allNodesAndHosts.nodes().asList()) :
+                capacity.findSpareHosts(this.allNodesAndHosts.nodes().asList(), spareCount);
         this.nameResolver = nameResolver;
 
-        NodeList nodesInCluster = allNodes.nodes().owner(application).type(clusterSpec.type()).cluster(clusterSpec.id());
+        NodeList nodesInCluster = this.allNodesAndHosts.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()
@@ -135,7 +135,7 @@ public class NodePrioritizer {
     private void addCandidatesOnExistingHosts() {
         if ( !canAllocateNew) return;
 
-        for (Node host : allNodes.nodes()) {
+        for (Node host : allNodesAndHosts.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;
@@ -143,12 +143,12 @@ public class NodePrioritizer {
             if ( ! host.exclusiveToClusterType().map(clusterSpec.type()::equals).orElse(true)) continue;
             if (spareHosts.contains(host) && !canAllocateToSpareHosts) continue;
             if ( ! capacity.hasCapacity(host, requestedNodes.resources().get())) continue;
-            if ( ! allNodes.childrenOf(host).owner(application).cluster(clusterSpec.id()).isEmpty()) continue;
+            if ( ! allNodesAndHosts.childrenOf(host).owner(application).cluster(clusterSpec.id()).isEmpty()) continue;
             nodes.add(NodeCandidate.createNewChild(requestedNodes.resources().get(),
                                                    capacity.availableCapacityOf(host),
                                                    host,
                                                    spareHosts.contains(host),
-                                                   allNodes.nodes(),
+                                                   allNodesAndHosts.nodes(),
                                                    nameResolver));
         }
     }
@@ -156,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.nodes().stream()
+        allNodesAndHosts.nodes().stream()
                 .filter(node -> node.type() == requestedNodes.type())
                 .filter(node -> legalStates.contains(node.state()))
                 .filter(node -> node.allocation().isPresent())
@@ -169,7 +169,7 @@ public class NodePrioritizer {
 
     /** Add nodes already provisioned, but not allocated to any application */
     private void addReadyNodes() {
-        allNodes.nodes().stream()
+        allNodesAndHosts.nodes().stream()
                 .filter(node -> node.type() == requestedNodes.type())
                 .filter(node -> node.state() == Node.State.ready)
                 .map(node -> candidateFrom(node, false))
@@ -179,7 +179,7 @@ public class NodePrioritizer {
 
     /** Create a candidate from given pre-existing node */
     private NodeCandidate candidateFrom(Node node, boolean isSurplus) {
-        Optional<Node> optionalParent = allNodes.parentOf(node);
+        Optional<Node> optionalParent = allNodesAndHosts.parentOf(node);
         if (optionalParent.isPresent()) {
             Node parent = optionalParent.get();
             return NodeCandidate.createChild(node,
@@ -217,7 +217,7 @@ public class NodePrioritizer {
      */
     private boolean canStillAllocate(Node node) {
         if (node.type() != NodeType.tenant || node.parentHostname().isEmpty()) return true;
-        Optional<Node> parent = allNodes.parentOf(node);
+        Optional<Node> parent = allNodesAndHosts.parentOf(node);
         return parent.isPresent() ? Nodes.canAllocateTenantNodeTo(parent.get(), dynamicProvisioning) : null;
     }