Assign groups before storing reserved nodes

4 files changed, 5 insertions, 16 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/Activator.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/Activator.java
index 462f43ee7af..9adff9f9d7a 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/Activator.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/Activator.java
@@ -72,7 +72,6 @@ class Activator {
 
         NodeList reserved = applicationNodes.state(Node.State.reserved).matching(node -> hostnames.contains(node.hostname()));
         reserved = updatePortsFrom(hosts, reserved);
-        reserved = updateGroupFrom(hosts, reserved);
         nodeRepository.nodes().reserve(reserved.asList()); // Re-reserve nodes to avoid reservation expiry
 
         NodeList oldActive = applicationNodes.state(Node.State.active); // All nodes active now
@@ -240,18 +239,6 @@ class Activator {
         return NodeList.copyOf(updated);
     }
 
-    /**
-     * Reserved nodes are stored before they are assigned a group.
-     */
-    private NodeList updateGroupFrom(Collection<HostSpec> hosts, NodeList nodes) {
-        List<Node> updated = new ArrayList<>();
-        for (Node node : nodes) {
-            var membership = getHost(node.hostname(), hosts).membership().get();
-            updated.add(node.with(node.allocation().get().with(membership)));
-        }
-        return NodeList.copyOf(updated);
-    }
-
     private HostSpec getHost(String hostname, Collection<HostSpec> fromHosts) {
         for (HostSpec host : fromHosts)
             if (host.hostname().equals(hostname))
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 3c8f1b070e7..e6b47d38779 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
@@ -143,9 +143,9 @@ public class GroupPreparer {
                 throw new NodeAllocationException(allocation.allocationFailureDetails(), true);
 
             // Carry out and return allocation
+            List<Node> acceptedNodes = allocation.finalNodes();
             nodeRepository.nodes().reserve(allocation.reservableNodes());
             nodeRepository.nodes().addReservedNodes(new LockedNodeList(allocation.newNodes(), allocationLock));
-            List<Node> acceptedNodes = allocation.finalNodes();
             surplusActiveNodes.removeAll(acceptedNodes);
             return acceptedNodes;
         }
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 0ddcbe6f6d1..40e5909d4d9 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
@@ -413,6 +413,8 @@ class NodeAllocation {
 
         GroupIndices groupIndices = new GroupIndices(requestedNodes, allNodes, nodeRepository.clock());
         Collection<NodeCandidate> finalNodes = groupIndices.assignTo(nodes.values());
+        nodes.clear();
+        finalNodes.forEach(candidate -> nodes.put(candidate.toNode().hostname(), candidate));
         return finalNodes.stream().map(NodeCandidate::toNode).toList();
     }
 
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicAllocationTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicAllocationTest.java
index 8e9417008b9..7cf1b0d5177 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicAllocationTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicAllocationTest.java
@@ -341,8 +341,8 @@ public class DynamicAllocationTest {
         tester.activate(application, hosts);
 
         NodeList activeNodes = tester.nodeRepository().nodes().list().owner(application);
-        assertEquals(Set.of("127.0.127.2", "::2"), activeNodes.asList().get(0).ipConfig().primary());
-        assertEquals(Set.of("127.0.127.13", "::d"), activeNodes.asList().get(1).ipConfig().primary());
+        assertEquals(Set.of("127.0.127.2", "::2"), activeNodes.asList().get(1).ipConfig().primary());
+        assertEquals(Set.of("127.0.127.13", "::d"), activeNodes.asList().get(0).ipConfig().primary());
     }
 
     @Test