Refactor nodeprioritizer

6 files changed, 186 insertions, 184 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerAllocator.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerAllocator.java
index fca15283e5d..ef74fb8c5c7 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerAllocator.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerAllocator.java
@@ -103,7 +103,7 @@ public class DockerAllocator {
             offers.add(node);
         }
 
-        return allocation.offer(offers, false);
+        return null;//allocation.offer(offers, false);
     }
 
     /**
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 ffddaa305c8..fd7b4bad7a1 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
@@ -19,12 +19,10 @@ import java.util.function.BiConsumer;
  * @author bratseth
  */
 class GroupPreparer {
-
+    private static final boolean canChangeGroup = true;
     private final NodeRepository nodeRepository;
     private final Clock clock;
 
-    private static final boolean canChangeGroup = true;
-
     public GroupPreparer(NodeRepository nodeRepository, Clock clock) {
         this.nodeRepository = nodeRepository;
         this.clock = clock;
@@ -54,6 +52,7 @@ class GroupPreparer {
             // Use new, ready nodes. Lock ready pool to ensure that nodes are not grabbed by others.
             try (Mutex readyLock = nodeRepository.lockUnallocated()) {
 
+                // Create a prioritized set of nodes for the cluster
                 NodePrioritizer prioritizer = new NodePrioritizer(
                         nodeRepository.getNodes(),
                         application,
@@ -62,14 +61,23 @@ class GroupPreparer {
                         nodeRepository.getAvailableFlavors(),
                         1,
                         nofSpares);
-                prioritizer.initNodes(surplusActiveNodes, nodeRepository.dynamicAllocationEnabled());
+
+                prioritizer.addApplicationNodes();
+                prioritizer.addSurplusNodes(surplusActiveNodes);
+                prioritizer.addReadyNodes();
+                if (nodeRepository.dynamicAllocationEnabled()) {
+                    prioritizer.addNewDockerNodes();
+                }
+
+                // Allocate from the prioritized list
                 NodeAllocation allocation = new NodeAllocation(application, cluster, requestedNodes, highestIndex, clock);
-                prioritizer.offer(allocation);
+                allocation.offer(prioritizer.prioritize());
 
+                // Book-keeping
                 if (allocation.fullfilled()) {
-                    nodeRepository.reserve(prioritizer.filterInactiveAndReadyNodes(allocation.getAcceptedNodes()));
-                    nodeRepository.addDockerNodes(prioritizer.filterNewNodes(allocation.getAcceptedNodes()));
-                    surplusActiveNodes.removeAll(prioritizer.filterSurplusNodes(allocation.getAcceptedNodes()));
+                    nodeRepository.reserve(allocation.getAcceptedInactiveAndReadyNodes());
+                    nodeRepository.addDockerNodes(allocation.getAcceptedNewNodes());
+                    surplusActiveNodes.removeAll(allocation.getAcceptedSurplusNodes());
                     List<Node> result = allocation.finalNodes(surplusActiveNodes);
                     return result;
                 } else {
@@ -89,4 +97,4 @@ class GroupPreparer {
         }
         return ".";
     }
-}
+}
\ No newline at end of file
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 16aa31f0840..02ecd182d54 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
@@ -34,7 +34,7 @@ public class NodeAllocation {
     private final NodeSpec requestedNodes;
 
     /** The nodes this has accepted so far */
-    private final Set<Node> nodes = new LinkedHashSet<>();
+    private final Set<NodePriority> nodes = new LinkedHashSet<>();
 
     /** The number of nodes in the accepted nodes which are of the requested flavor */
     private int acceptedOfRequestedFlavor = 0;
@@ -69,19 +69,21 @@ public class NodeAllocation {
      * Note that if unallocated nodes are offered before allocated nodes, this will unnecessarily
      * reject allocated nodes due to index duplicates.
      *
-     * @param offeredNodes the nodes which are potentially on offer. These may belong to a different application etc.
-     * @param canChangeGroup whether it is ok to change the group the offered node is to belong to if necessary
+     * @param nodesPrioritized the nodes which are potentially on offer. These may belong to a different application etc.
      * @return the subset of offeredNodes which was accepted, with the correct allocation assigned
      */
-    public List<Node> offer(List<Node> offeredNodes, boolean canChangeGroup) {
+    public List<Node> offer(List<NodePriority> nodesPrioritized) {
+
         List<Node> accepted = new ArrayList<>();
-        for (Node offered : offeredNodes) {
+        for (NodePriority offeredPriority : nodesPrioritized) {
+            Node offered = offeredPriority.node;
+
             if (offered.allocation().isPresent()) {
                 boolean wantToRetireNode = false;
                 ClusterMembership membership = offered.allocation().get().membership();
                 if ( ! offered.allocation().get().owner().equals(application)) continue; // wrong application
                 if ( ! membership.cluster().equalsIgnoringGroupAndVespaVersion(cluster)) continue; // wrong cluster id/type
-                if ((! canChangeGroup || saturated()) && ! membership.cluster().group().equals(cluster.group())) continue; // wrong group and we can't or have no reason to change it
+                if ((! offeredPriority.isSurplusNode || saturated()) && ! membership.cluster().group().equals(cluster.group())) continue; // wrong group and we can't or have no reason to change it
                 if ( offered.allocation().get().isRemovable()) continue; // don't accept; causes removal
                 if ( indexes.contains(membership.index())) continue; // duplicate index (just to be sure)
 
@@ -91,8 +93,9 @@ public class NodeAllocation {
                 if ( offered.flavor().isRetired()) wantToRetireNode = true;
                 if ( offered.status().wantToRetire()) wantToRetireNode = true;
 
-                if ((!saturated() && hasCompatibleFlavor(offered)) || acceptToRetire(offered) )
-                    accepted.add(acceptNode(offered, wantToRetireNode));
+                if ((!saturated() && hasCompatibleFlavor(offered)) || acceptToRetire(offered) ) {
+                    accepted.add(acceptNode(offeredPriority, wantToRetireNode));
+                }
             }
             else if (! saturated() && hasCompatibleFlavor(offered)) {
                 if ( offeredNodeHasParentHostnameAlreadyAccepted(this.nodes, offered)) {
@@ -106,17 +109,18 @@ public class NodeAllocation {
                     continue;
                 }
                 Node alloc = offered.allocate(application, ClusterMembership.from(cluster, highestIndex.add(1)), clock.instant());
-                accepted.add(acceptNode(alloc, false));
+                offeredPriority.node = alloc;
+                accepted.add(acceptNode(offeredPriority, false));
             }
         }
 
         return accepted;
     }
 
-    private boolean offeredNodeHasParentHostnameAlreadyAccepted(Collection<Node> accepted, Node offered) {
-        for (Node acceptedNode : accepted) {
-            if (acceptedNode.parentHostname().isPresent() && offered.parentHostname().isPresent() &&
-                    acceptedNode.parentHostname().get().equals(offered.parentHostname().get())) {
+    private boolean offeredNodeHasParentHostnameAlreadyAccepted(Collection<NodePriority> accepted, Node offered) {
+        for (NodePriority acceptedNode : accepted) {
+            if (acceptedNode.node.parentHostname().isPresent() && offered.parentHostname().isPresent() &&
+                    acceptedNode.node.parentHostname().get().equals(offered.parentHostname().get())) {
                 return true;
             }
         }
@@ -150,31 +154,29 @@ public class NodeAllocation {
         return requestedNodes.isCompatible(node.flavor());
     }
 
-    /** Updates the state of some existing nodes in this list by replacing them by id with the given instances. */
-    public void update(List<Node> updatedNodes) {
-        nodes.removeAll(updatedNodes);
-        nodes.addAll(updatedNodes);
-    }
-
-    private Node acceptNode(Node node, boolean wantToRetire) {
+    private Node acceptNode(NodePriority nodePriority, boolean wantToRetire) {
+        Node node = nodePriority.node;
         if (! wantToRetire) {
             if ( ! node.state().equals(Node.State.active)) {
                 // reactivated node - make sure its not retired
                 node = node.unretire();
+                nodePriority.node= node;
             }
             acceptedOfRequestedFlavor++;
         } else {
             ++wasRetiredJustNow;
             // Retire nodes which are of an unwanted flavor, retired flavor or have an overlapping parent host
             node = node.retire(clock.instant());
+            nodePriority.node= node;
         }
         if ( ! node.allocation().get().membership().cluster().equals(cluster)) {
             // group may be different
             node = setCluster(cluster, node);
+            nodePriority.node= node;
         }
         indexes.add(node.allocation().get().membership().index());
         highestIndex.set(Math.max(highestIndex.get(), node.allocation().get().membership().index()));
-        nodes.add(node);
+        nodes.add(nodePriority);
         return node;
     }
 
@@ -210,47 +212,62 @@ public class NodeAllocation {
      * @param surplusNodes this will add nodes not any longer needed by this group to this list
      * @return the final list of nodes
      */
-    public List<Node> finalNodes(List<Node> surplusNodes) {
-        long currentRetired = nodes.stream().filter(node -> node.allocation().get().membership().retired()).count();
+    List<Node> finalNodes(List<Node> surplusNodes) {
+        long currentRetired = nodes.stream().filter(node -> node.node.allocation().get().membership().retired()).count();
         long surplus = requestedNodes.surplusGiven(nodes.size()) - currentRetired;
 
-        List<Node> changedNodes = new ArrayList<>();
         if (surplus > 0) { // retire until surplus is 0, prefer to retire higher indexes to minimize redistribution
-            for (Node node : byDecreasingIndex(nodes)) {
-                if ( ! node.allocation().get().membership().retired() && node.state().equals(Node.State.active)) {
-                    Node retiredNode = node.retire(Agent.application, clock.instant());
-                    changedNodes.add(retiredNode);
-                    surplusNodes.add(retiredNode); // offer this node to other groups
+            for (NodePriority node : byDecreasingIndex(nodes)) {
+                if ( ! node.node.allocation().get().membership().retired() && node.node.state().equals(Node.State.active)) {
+                    Node retiredNode = node.node.retire(Agent.application, clock.instant());
+                    node.node = retiredNode;
+                    surplusNodes.add(node.node); // offer this node to other groups
                     if (--surplus == 0) break;
                 }
             }
         }
         else if (surplus < 0) { // unretire until surplus is 0
-            for (Node node : byIncreasingIndex(nodes)) {
-                if ( node.allocation().get().membership().retired() && hasCompatibleFlavor(node)) {
-                    changedNodes.add(node.unretire());
+            for (NodePriority node : byIncreasingIndex(nodes)) {
+                if ( node.node.allocation().get().membership().retired() && hasCompatibleFlavor(node.node)) {
+                    node.node = node.node.unretire();
                     if (++surplus == 0) break;
                 }
             }
         }
-        update(changedNodes);
-        return new ArrayList<>(nodes);
+
+        return nodes.stream().map(n -> n.node).collect(Collectors.toList());
+    }
+
+    List<Node> getAcceptedInactiveAndReadyNodes() {
+        return nodes.stream().map(n -> n.node)
+                .filter(n -> n.state().equals(Node.State.inactive) || n.state().equals(Node.State.ready))
+                .collect(Collectors.toList());
+    }
+
+    List<Node> getAcceptedSurplusNodes() {
+        return nodes.stream()
+                .filter(n -> n.isSurplusNode)
+                .map(n -> n.node)
+                .collect(Collectors.toList());
     }
 
-    public Set<Node> getAcceptedNodes() {
-        return nodes;
+    List<Node> getAcceptedNewNodes() {
+        return nodes.stream()
+                .filter(n -> n.isNewNode)
+                .map(n -> n.node)
+                .collect(Collectors.toList());
     }
 
-    private List<Node> byDecreasingIndex(Set<Node> nodes) {
+    private List<NodePriority> byDecreasingIndex(Set<NodePriority> nodes) {
         return nodes.stream().sorted(nodeIndexComparator().reversed()).collect(Collectors.toList());
     }
 
-    private List<Node> byIncreasingIndex(Set<Node> nodes) {
+    private List<NodePriority> byIncreasingIndex(Set<NodePriority> nodes) {
         return nodes.stream().sorted(nodeIndexComparator()).collect(Collectors.toList());
     }
 
-    private Comparator<Node> nodeIndexComparator() {
-        return Comparator.comparing((Node n) -> n.allocation().get().membership().index());
+    private Comparator<NodePriority> nodeIndexComparator() {
+        return Comparator.comparing((NodePriority n) -> n.node.allocation().get().membership().index());
     }
 
 }
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 854308e4bbd..f76e5d766bf 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
@@ -18,11 +18,10 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
 
 /**
  * Builds up a priority queue of which nodes should be offered to the allocation.
- *
+ * <p>
  * Builds up a list of NodePriority objects and sorts them according to the
  * NodePriority::compare method.
  *
@@ -38,9 +37,10 @@ public class NodePrioritizer {
     private final int maxRetires;
     private final ClusterSpec clusterSpec;
 
+    private final boolean isAllocatingForReplacement;
     private final List<Node> spareHosts;
     private final List<Node> headroomViolatedHosts;
-    private final long failedNodesInCluster;
+    private final boolean isDocker;
 
     int nofViolations = 0;
 
@@ -53,40 +53,44 @@ public class NodePrioritizer {
 
         // Add spare and headroom allocations
         spareHosts = DockerCapacityConstraints.findSpareHosts(allNodes, spares);
-        List<Node> nodesWithHeadroomAndSpares =
-                DockerCapacityConstraints.addHeadroomAndSpareNodes(allNodes, nodeFlavors, spares);
+        headroomViolatedHosts = new ArrayList<>();
 
-        this.capacity = new DockerHostCapacity(nodesWithHeadroomAndSpares);
+        this.capacity = new DockerHostCapacity(allNodes);
 
-        failedNodesInCluster = allNodes.stream()
+        long nofFailedNodes = allNodes.stream()
                 .filter(node -> node.state().equals(Node.State.failed))
                 .filter(node -> node.allocation().isPresent())
                 .filter(node -> node.allocation().get().owner().equals(appId))
                 .filter(node -> node.allocation().get().membership().cluster().id().equals(clusterSpec.id()))
                 .count();
 
-        // TODO Find hosts where we have headroom violations
-        headroomViolatedHosts = new ArrayList<>();
-
+        long nofNodesInCluster = allNodes.stream()
+                .filter(node -> node.allocation().isPresent())
+                .filter(node -> node.allocation().get().owner().equals(appId))
+                .filter(node -> node.allocation().get().membership().cluster().id().equals(clusterSpec.id()))
+                .count();
 
+        isAllocatingForReplacement = isReplacement(nofNodesInCluster, nofFailedNodes);
+        isDocker = isDocker();
     }
 
-    void initNodes(List<Node> surplusNodes, boolean dynamicAllocationEnabled) {
-        addApplicationNodes();
-        addSurplusNodes(surplusNodes);
-        addReadyNodes();
-        if (dynamicAllocationEnabled && getDockerFlavor() != null) {
-            addNewDockerNodes();
-        }
+    List<NodePriority> prioritize() {
+        List<NodePriority> priorityList = new ArrayList<>(nodes.values());
+        Collections.sort(priorityList, (a, b) -> NodePriority.compare(a, b));
+        return priorityList;
     }
 
-    private void addSurplusNodes(List<Node> surplusNodes) {
+    void addSurplusNodes(List<Node> surplusNodes) {
         for (Node node : surplusNodes) {
-            nodes.put(node, toNodePriority(node, true, false));
+            NodePriority nodePri = toNodePriority(node, true, false);
+            if (!nodePri.violatesSpares || isAllocatingForReplacement) {
+                nodes.put(node, nodePri);
+            }
         }
     }
 
-    private void addNewDockerNodes() {
+    void addNewDockerNodes() {
+        if (!isDocker) return;
         DockerHostCapacity capacity = new DockerHostCapacity(allNodes);
 
         for (Node node : allNodes) {
@@ -102,81 +106,42 @@ public class NodePrioritizer {
                     }
                 }
 
-                if (!conflictingCluster && capacity.hasCapacity(node, getDockerFlavor())) {
+                if (!conflictingCluster && capacity.hasCapacity(node, getFlavor())) {
                     Set<String> ipAddresses = DockerHostCapacity.findFreeIps(node, allNodes);
                     if (ipAddresses.isEmpty()) continue;
                     String ipAddress = ipAddresses.stream().findFirst().get();
                     String hostname = lookupHostname(ipAddress);
                     if (hostname == null) continue;
                     Node newNode = Node.createDockerNode("fake-" + hostname, Collections.singleton(ipAddress),
-                            Collections.emptySet(), hostname, Optional.of(node.hostname()), getDockerFlavor(), NodeType.tenant);
-                    nodes.put(newNode, toNodePriority(newNode, false, true));
+                            Collections.emptySet(), hostname, Optional.of(node.hostname()), getFlavor(), NodeType.tenant);
+                    NodePriority nodePri = toNodePriority(newNode, false, true);
+                    if (!nodePri.violatesSpares || isAllocatingForReplacement) {
+                        nodes.put(newNode, nodePri);
+                    }
                 }
             }
         }
     }
 
-    /**
-     * From ipAddress - get hostname
-     *
-     * @return hostname or null if not able to do the loopup
-     */
-    private static String lookupHostname(String ipAddress) {
-        try {
-            return InetAddress.getByName(ipAddress).getHostName();
-        } catch (UnknownHostException e) {
-            e.printStackTrace();
-        }
-        return null;
-    }
-
-    private void addReadyNodes() {
-        allNodes.stream()
-                .filter(node -> node.type().equals(requestedNodes.type()))
-                .filter(node -> node.state().equals(Node.State.ready))
-                .map(node -> toNodePriority(node, false, false))
-                .forEach(nodePriority -> nodes.put(nodePriority.node, nodePriority));
-    }
-
     void addApplicationNodes() {
-        List<Node.State> legalStates = Arrays.asList(Node.State.active,Node.State.inactive, Node.State.reserved);
+        List<Node.State> legalStates = Arrays.asList(Node.State.active, Node.State.inactive, Node.State.reserved);
         allNodes.stream()
                 .filter(node -> node.type().equals(requestedNodes.type()))
                 .filter(node -> legalStates.contains(node.state()))
                 .filter(node -> node.allocation().isPresent())
                 .filter(node -> node.allocation().get().owner().equals(appId))
                 .map(node -> toNodePriority(node, false, false))
+                .filter(n -> !n.violatesSpares || isAllocatingForReplacement)
                 .forEach(nodePriority -> nodes.put(nodePriority.node, nodePriority));
     }
 
-    List<Node> filterNewNodes(Set<Node> acceptedNodes) {
-        List<Node> newNodes = new ArrayList<>();
-        for (Node node : acceptedNodes) {
-            if (nodes.get(node).isNewNode) {
-                newNodes.add(node);
-            }
-        }
-        return newNodes;
-    }
-
-    List<Node> filterSurplusNodes(Set<Node> acceptedNodes) {
-        List<Node> surplusNodes = new ArrayList<>();
-        for (Node node : acceptedNodes) {
-            if (nodes.get(node).isSurplusNode) {
-                surplusNodes.add(node);
-            }
-        }
-        return surplusNodes;
-    }
-
-    List<Node> filterInactiveAndReadyNodes(Set<Node> acceptedNodes) {
-        List<Node> inactiveAndReady = new ArrayList<>();
-        for (Node node : acceptedNodes) {
-            if (node.state().equals(Node.State.inactive) || node.state().equals(Node.State.ready)) {
-                inactiveAndReady.add(node);
-            }
-        }
-        return inactiveAndReady;
+    void addReadyNodes() {
+        allNodes.stream()
+                .filter(node -> node.type().equals(requestedNodes.type()))
+                .filter(node -> node.state().equals(Node.State.ready))
+                .map(node -> toNodePriority(node, false, false))
+                .filter(n -> !n.violatesSpares || isAllocatingForReplacement)
+                .forEach(nodePriority -> nodes.put(nodePriority.node, nodePriority));
     }
 
     /**
@@ -184,62 +149,53 @@ public class NodePrioritizer {
      * parameters to the priority sorting procedure.
      */
     private NodePriority toNodePriority(Node node, boolean isSurplusNode, boolean isNewNode) {
-            NodePriority pri = new NodePriority();
-            pri.node = node;
-            pri.isSurplusNode = isSurplusNode;
-            pri.isNewNode = isNewNode;
-            pri.preferredOnFlavor = requestedNodes.specifiesNonStockFlavor() && node.flavor().equals(getDockerFlavor());
-            pri.parent = findParentNode(node);
-
-            if (pri.parent.isPresent()) {
-                Node parent = pri.parent.get();
-                pri.freeParentCapacity = capacity.freeCapacityOf(parent, true);
-
-                /**
-                 * To be conservative we have a restriction of how many nodes we can retire for each cluster,
-                 * pr. allocation iteration. TODO also account for previously retired nodes? (thus removing the pr iteration restriction)
-                 */
-                if (nofViolations <= maxRetires) {
-                    // Spare violation
-                    if (spareHosts.contains(parent)) {
-                        pri.violatesSpares = true;
-                        nofViolations++;
-                    }
+        NodePriority pri = new NodePriority();
+        pri.node = node;
+        pri.isSurplusNode = isSurplusNode;
+        pri.isNewNode = isNewNode;
+        pri.preferredOnFlavor = requestedNodes.specifiesNonStockFlavor() && node.flavor().equals(getFlavor());
+        pri.parent = findParentNode(node);
+
+        if (pri.parent.isPresent()) {
+            Node parent = pri.parent.get();
+            pri.freeParentCapacity = capacity.freeCapacityOf(parent, true);
+
+            /**
+             * To be conservative we have a restriction of how many nodes we can retire for each cluster,
+             * pr. allocation iteration. TODO also account for previously retired nodes? (thus removing the pr iteration restriction)
+             */
+            if (nofViolations <= maxRetires) {
+                if (spareHosts.contains(parent)) {
+                    pri.violatesSpares = true;
+                    nofViolations++;
+                }
 
-                    // Headroom violation
-                    if (headroomViolatedHosts.contains(parent)) {
-                        pri.violatesHeadroom = true;
-                        nofViolations++;
-                    }
+                // Headroom violation
+                if (headroomViolatedHosts.contains(parent)) {
+                    pri.violatesHeadroom = true;
+                    nofViolations++;
                 }
             }
-            return pri;
+        }
+        return pri;
     }
 
-    void offer(NodeAllocation allocation) {
-        List<NodePriority> prioritizedNodes = nodes.values().stream().collect(Collectors.toList());
-        Collections.sort(prioritizedNodes, (a,b) -> NodePriority.compare(a,b));
-
-        for (NodePriority nodePriority : prioritizedNodes) {
-
-            // The replacement heuristic assumes that new nodes are offered after already existing nodes
-            boolean isReplacement = isReplacement(allocation.getAcceptedNodes().size());
-
-            // Only add new allocations that violates the spare constraint if this is a replacement
-            if (!nodePriority.violatesSpares || isReplacement || !nodePriority.isNewNode) {
-                List<Node> acceptedNodes = allocation.offer(Collections.singletonList(nodePriority.node), nodePriority.isSurplusNode);
-                // Update with the potentially changed node (new object)
-                if (!acceptedNodes.isEmpty()) {
-                    nodePriority.node = acceptedNodes.get(0);
-                    nodes.remove(nodePriority.node);
-                    nodes.put(acceptedNodes.get(0), nodePriority);
-                }
-            }
+    /**
+     * From ipAddress - get hostname
+     *
+     * @return hostname or null if not able to do the loopup
+     */
+    private static String lookupHostname(String ipAddress) {
+        try {
+            return InetAddress.getByName(ipAddress).getHostName();
+        } catch (UnknownHostException e) {
+            e.printStackTrace();
         }
+        return null;
     }
 
-    private boolean isReplacement(int nodesAccepted) {
-        if (failedNodesInCluster == 0) return false;
+    private boolean isReplacement(long nofNodesInCluster, long nodeFailedNodes) {
+        if (nodeFailedNodes == 0) return false;
 
         int wantedCount = 0;
         if (requestedNodes instanceof NodeSpec.CountNodeSpec) {
@@ -247,10 +203,10 @@ public class NodePrioritizer {
             wantedCount = countSpec.getCount();
         }
 
-        return (wantedCount <= nodesAccepted + failedNodesInCluster);
+        return (wantedCount > nofNodesInCluster - nodeFailedNodes);
     }
 
-    private Flavor getDockerFlavor() {
+    private Flavor getFlavor() {
         if (requestedNodes instanceof NodeSpec.CountNodeSpec) {
             NodeSpec.CountNodeSpec countSpec = (NodeSpec.CountNodeSpec) requestedNodes;
             return countSpec.getFlavor();
@@ -258,6 +214,14 @@ public class NodePrioritizer {
         return null;
     }
 
+    private boolean isDocker() {
+        Flavor flavor = getFlavor();
+        if (flavor != null) {
+            return flavor.getType().equals(Flavor.Type.DOCKER_CONTAINER);
+        }
+        return false;
+    }
+
     private Optional<Node> findParentNode(Node node) {
         if (!node.parentHostname().isPresent()) return Optional.empty();
         return allNodes.stream()
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/Preparer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/Preparer.java
index 639d6ea17f4..a0f9452ef9f 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/Preparer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/Preparer.java
@@ -88,7 +88,12 @@ class Preparer {
             }
         }
     }
-    
+
+    /**
+     * Nodes are immutable so when changing attributes to the node we create a new instance.
+     *
+     * This method is used to both add new nodes and replaces old node references with the new references.
+     */
     private List<Node> replace(List<Node> list, List<Node> changed) {
         list.removeAll(changed);
         list.addAll(changed);
@@ -121,5 +126,4 @@ class Preparer {
         }
         return retired;
     }
-
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
index fb8fdee222c..0885b941401 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
@@ -51,6 +51,11 @@ public class MockNodeRepository extends NodeRepository {
         populate();
     }
 
+    @Override
+    public boolean dynamicAllocationEnabled() {
+        return true;
+    }
+
     private void populate() {
         NodeRepositoryProvisioner provisioner = new NodeRepositoryProvisioner(this, flavors, Zone.defaultZone());
 
@@ -68,14 +73,14 @@ public class MockNodeRepository extends NodeRepository {
         nodes.add(createNode("node2", "host2.yahoo.com", ipAddresses, Optional.empty(), flavors.getFlavorOrThrow("default"), NodeType.tenant));
         nodes.add(createNode("node3", "host3.yahoo.com", ipAddresses, Optional.empty(), flavors.getFlavorOrThrow("expensive"), NodeType.tenant));
 
-        // TODO: Use docker flavor
-        Node node4 = createNode("node4", "host4.yahoo.com", ipAddresses, Optional.of("dockerhost4"), flavors.getFlavorOrThrow("default"), NodeType.tenant);
+        Node node4 = createNode("node4", "host4.yahoo.com", ipAddresses, Optional.of("dockerhost1.yahoo.com"), flavors.getFlavorOrThrow("docker"), NodeType.tenant);
         node4 = node4.with(node4.status().withVespaVersion(new Version("6.41.0")));
         nodes.add(node4);
 
-        Node node5 = createNode("node5", "host5.yahoo.com", ipAddresses, Optional.of("parent1.yahoo.com"), flavors.getFlavorOrThrow("default"), NodeType.tenant);
+        Node node5 = createNode("node5", "host5.yahoo.com", ipAddresses, Optional.of("dockerhost2.yahoo.com"), flavors.getFlavorOrThrow("docker"), NodeType.tenant);
         nodes.add(node5.with(node5.status().withVespaVersion(new Version("1.2.3"))));
 
+
         nodes.add(createNode("node6", "host6.yahoo.com", ipAddresses, Optional.empty(), flavors.getFlavorOrThrow("default"), NodeType.tenant));
         nodes.add(createNode("node7", "host7.yahoo.com", ipAddresses, Optional.empty(), flavors.getFlavorOrThrow("default"), NodeType.tenant));
         // 8 and 9 are added by web service calls
@@ -87,8 +92,13 @@ public class MockNodeRepository extends NodeRepository {
         nodes.add(node10);
 
         nodes.add(createNode("node55", "host55.yahoo.com", ipAddresses, Optional.empty(), flavors.getFlavorOrThrow("default"), NodeType.tenant));
-        nodes.add(createNode("parent1", "parent1.yahoo.com", ipAddresses, additionalIpAddresses, Optional.empty(), flavors.getFlavorOrThrow("default"), NodeType.host));
-        nodes.add(createNode("parent2", "dockerhost4", ipAddresses, additionalIpAddresses, Optional.empty(), flavors.getFlavorOrThrow("default"), NodeType.host));
+
+        /** Setup docker hosts (two of these will be reserved for spares */
+        nodes.add(createNode("dockerhost1", "dockerhost1.yahoo.com", ipAddresses, additionalIpAddresses, Optional.empty(), flavors.getFlavorOrThrow("large"), NodeType.host));
+        nodes.add(createNode("dockerhost2", "dockerhost2.yahoo.com", ipAddresses, additionalIpAddresses, Optional.empty(), flavors.getFlavorOrThrow("large"), NodeType.host));
+        nodes.add(createNode("dockerhost3", "dockerhost3.yahoo.com", ipAddresses, additionalIpAddresses, Optional.empty(), flavors.getFlavorOrThrow("large"), NodeType.host));
+        nodes.add(createNode("dockerhost4", "dockerhost4.yahoo.com", ipAddresses, additionalIpAddresses, Optional.empty(), flavors.getFlavorOrThrow("large"), NodeType.host));
+        nodes.add(createNode("dockerhost5", "dockerhost5.yahoo.com", ipAddresses, additionalIpAddresses, Optional.empty(), flavors.getFlavorOrThrow("large"), NodeType.host));
 
         nodes = addNodes(nodes);
         nodes.remove(6);
@@ -114,7 +124,7 @@ public class MockNodeRepository extends NodeRepository {
 
         ApplicationId app3 = ApplicationId.from(TenantName.from("tenant3"), ApplicationName.from("application3"), InstanceName.from("instance3"));
         ClusterSpec cluster3 = ClusterSpec.request(ClusterSpec.Type.content, ClusterSpec.Id.from("id3"), Version.fromString("6.42"));
-        activate(provisioner.prepare(app3, cluster3, Capacity.fromNodeCount(2), 1, null), app3, provisioner);
+        activate(provisioner.prepare(app3, cluster3, Capacity.fromNodeCount(2, "docker"), 1, null), app3, provisioner);
     }
 
     private void activate(List<HostSpec> hosts, ApplicationId application, NodeRepositoryProvisioner provisioner) {
@@ -122,5 +132,4 @@ public class MockNodeRepository extends NodeRepository {
         provisioner.activate(transaction, application, hosts);
         transaction.commit();
     }
-
-}
+}
\ No newline at end of file