Lock node when updating IP config

6 files changed, 81 insertions, 36 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/HostResumeProvisioner.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/HostResumeProvisioner.java
index c606ede05d1..86c5a926900 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/HostResumeProvisioner.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/HostResumeProvisioner.java
@@ -9,17 +9,14 @@ import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.node.IP;
 import com.yahoo.vespa.hosted.provision.provisioning.FatalProvisioningException;
+import com.yahoo.vespa.hosted.provision.provisioning.HostIpConfig;
 import com.yahoo.vespa.hosted.provision.provisioning.HostProvisioner;
 import com.yahoo.yolean.Exceptions;
 
 import javax.naming.NamingException;
 import java.time.Duration;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
 import java.util.logging.Level;
 import java.util.logging.Logger;
-import java.util.stream.Collectors;
 
 /**
  * Resumes provisioning (requests additional IP addresses, updates DNS when IPs are ready) of hosts in state provisioned
@@ -41,23 +38,13 @@ public class HostResumeProvisioner extends NodeRepositoryMaintainer {
     @Override
     protected double maintain() {
         NodeList allNodes = nodeRepository().nodes().list();
-        Map<String, Set<Node>> nodesByProvisionedParentHostname =
-                allNodes.nodeType(NodeType.tenant, NodeType.config, NodeType.controller)
-                     .asList()
-                     .stream()
-                     .filter(node -> node.parentHostname().isPresent())
-                     .collect(Collectors.groupingBy(node -> node.parentHostname().get(), Collectors.toSet()));
-
         NodeList hosts = allNodes.state(Node.State.provisioned).nodeType(NodeType.host, NodeType.confighost, NodeType.controllerhost);
         int failures = 0;
         for (Node host : hosts) {
-            Set<Node> children = nodesByProvisionedParentHostname.getOrDefault(host.hostname(), Set.of());
-            // This doesn't actually require unallocated lock, but that is much easier than simultaneously holding
-            // the application locks of the host and all it's children.
-            try (var lock = nodeRepository().nodes().lockUnallocated()) {
-                List<Node> updatedNodes = hostProvisioner.provision(host, children);
-                verifyDns(updatedNodes);
-                nodeRepository().nodes().write(updatedNodes, lock);
+            NodeList children = allNodes.childrenOf(host);
+            try {
+                HostIpConfig hostIpConfig = hostProvisioner.provision(host, children.asSet());
+                setIpConfig(host, children, hostIpConfig);
             } catch (IllegalArgumentException | IllegalStateException e) {
                 log.log(Level.INFO, "Could not provision " + host.hostname() + " with " + children.size() + " children, will retry in " +
                                     interval() + ": " + Exceptions.toMessageString(e));
@@ -79,13 +66,21 @@ public class HostResumeProvisioner extends NodeRepositoryMaintainer {
         return asSuccessFactor(hosts.size(), failures);
     }
 
-    /** Verify DNS configuration of given nodes */
-    private void verifyDns(List<Node> nodes) {
+    private void setIpConfig(Node host, NodeList children, HostIpConfig hostIpConfig) {
+        if (hostIpConfig.isEmpty()) return;
+        NodeList nodes = NodeList.of(host).and(children);
         for (var node : nodes) {
-            boolean enclave = node.cloudAccount().isEnclave(nodeRepository().zone());
-            for (var ipAddress : node.ipConfig().primary()) {
-                IP.verifyDns(node.hostname(), ipAddress, nodeRepository().nameResolver(), !enclave);
-            }
+            verifyDns(node, hostIpConfig.require(node.hostname()));
+        }
+        nodeRepository().nodes().setIpConfig(hostIpConfig);
+    }
+
+    /** Verify DNS configuration of given node */
+    private void verifyDns(Node node, IP.Config ipConfig) {
+        boolean enclave = node.cloudAccount().isEnclave(nodeRepository().zone());
+        for (var ipAddress : ipConfig.primary()) {
+            IP.verifyDns(node.hostname(), ipAddress, nodeRepository().nameResolver(), !enclave);
         }
     }
+
 }
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 10f0c8aa554..6dc3b2b3193 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
@@ -21,6 +21,7 @@ import com.yahoo.vespa.hosted.provision.applications.Applications;
 import com.yahoo.vespa.hosted.provision.maintenance.NodeFailer;
 import com.yahoo.vespa.hosted.provision.node.filter.NodeFilter;
 import com.yahoo.vespa.hosted.provision.persistence.CuratorDb;
+import com.yahoo.vespa.hosted.provision.provisioning.HostIpConfig;
 import com.yahoo.vespa.orchestrator.HostNameNotFoundException;
 import com.yahoo.vespa.orchestrator.Orchestrator;
 
@@ -349,6 +350,15 @@ public class Nodes {
         }
     }
 
+    /** Update IP config for nodes in given config */
+    public void setIpConfig(HostIpConfig hostIpConfig) {
+        Predicate<Node> nodeInConfig = (node) -> hostIpConfig.contains(node.hostname());
+        performOn(nodeInConfig, (node, lock) -> {
+            IP.Config ipConfig = hostIpConfig.require(node.hostname());
+            return write(node.with(ipConfig), lock);
+        });
+    }
+
     /**
      * Parks this node and returns it in its new state.
      *
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostIpConfig.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostIpConfig.java
new file mode 100644
index 00000000000..891251fc892
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostIpConfig.java
@@ -0,0 +1,40 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.provisioning;
+
+import com.yahoo.vespa.hosted.provision.node.IP;
+
+import java.util.Map;
+import java.util.Objects;
+
+/**
+ * IP config of a host and its children.
+ *
+ * @author mpolden
+ */
+public record HostIpConfig(Map<String, IP.Config> ipConfigByHostname) {
+
+    public static final HostIpConfig EMPTY = new HostIpConfig(Map.of());
+
+    public HostIpConfig(Map<String, IP.Config> ipConfigByHostname) {
+        this.ipConfigByHostname = Map.copyOf(Objects.requireNonNull(ipConfigByHostname));
+    }
+
+    public Map<String, IP.Config> asMap() {
+        return ipConfigByHostname;
+    }
+
+    public boolean contains(String hostname) {
+        return ipConfigByHostname.containsKey(hostname);
+    }
+
+    public IP.Config require(String hostname) {
+        IP.Config ipConfig = this.ipConfigByHostname.get(hostname);
+        if (ipConfig == null) throw new IllegalArgumentException("No IP config exists for node '" + hostname + "'");
+        return ipConfig;
+    }
+
+    public boolean isEmpty() {
+        return this.equals(EMPTY);
+    }
+
+}
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostProvisioner.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostProvisioner.java
index 3d2f5b6737a..ce48c5adab8 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostProvisioner.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/HostProvisioner.java
@@ -74,12 +74,11 @@ public interface HostProvisioner {
      *
      * @param host the host to provision
      * @param children list of all the nodes that run on the given host
-     * @return a subset of {@code host} and {@code children} where the values have been modified and should
-     * be written back to node-repository.
+     * @return IP config for the provisioned host and its children
      * @throws FatalProvisioningException if the provisioning has irrecoverably failed and the input nodes
      * should be deleted from node-repo.
      */
-    List<Node> provision(Node host, Set<Node> children) throws FatalProvisioningException;
+    HostIpConfig provision(Node host, Set<Node> children) throws FatalProvisioningException;
 
     /**
      * Deprovisions a given host and resources associated with it and its children (such as DNS entries).
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockHostProvisioner.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockHostProvisioner.java
index 1a180d593a4..24ea9361823 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockHostProvisioner.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockHostProvisioner.java
@@ -15,6 +15,7 @@ import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.node.IP;
 import com.yahoo.vespa.hosted.provision.provisioning.FatalProvisioningException;
+import com.yahoo.vespa.hosted.provision.provisioning.HostIpConfig;
 import com.yahoo.vespa.hosted.provision.provisioning.HostProvisioner;
 import com.yahoo.vespa.hosted.provision.provisioning.HostResourcesCalculator;
 import com.yahoo.vespa.hosted.provision.provisioning.ProvisionedHost;
@@ -94,16 +95,16 @@ public class MockHostProvisioner implements HostProvisioner {
     }
 
     @Override
-    public List<Node> provision(Node host, Set<Node> children) throws FatalProvisioningException {
+    public HostIpConfig provision(Node host, Set<Node> children) throws FatalProvisioningException {
         if (behaviours.contains(Behaviour.failProvisioning)) throw new FatalProvisioningException("Failed to provision node(s)");
         if (host.state() != Node.State.provisioned) throw new IllegalStateException("Host to provision must be in " + Node.State.provisioned);
-        List<Node> result = new ArrayList<>();
-        result.add(withIpAssigned(host));
+        Map<String, IP.Config> result = new HashMap<>();
+        result.put(host.hostname(), createIpConfig(host));
         for (var child : children) {
             if (child.state() != Node.State.reserved) throw new IllegalStateException("Child to provisioned must be in " + Node.State.reserved);
-            result.add(withIpAssigned(child));
+            result.put(child.hostname(), createIpConfig(child));
         }
-        return result;
+        return new HostIpConfig(result);
     }
 
     @Override
@@ -207,9 +208,9 @@ public class MockHostProvisioner implements HostProvisioner {
                         .toList();
     }
 
-    public Node withIpAssigned(Node node) {
+    public IP.Config createIpConfig(Node node) {
         if (!node.type().isHost()) {
-            return node.with(node.ipConfig().withPrimary(nameResolver.resolveAll(node.hostname())));
+            return node.ipConfig().withPrimary(nameResolver.resolveAll(node.hostname()));
         }
         int hostIndex = Integer.parseInt(node.hostname().replaceAll("^[a-z]+|-\\d+$", ""));
         Set<String> addresses = Set.of("::" + hostIndex + ":0");
@@ -223,7 +224,7 @@ public class MockHostProvisioner implements HostProvisioner {
             }
         }
         IP.Pool pool = node.ipConfig().pool().withIpAddresses(ipAddressPool);
-        return node.with(node.ipConfig().withPrimary(addresses).withPool(pool));
+        return node.ipConfig().withPrimary(addresses).withPool(pool);
     }
 
     public enum Behaviour {
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicProvisioningTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicProvisioningTest.java
index 0e19c48591d..382d2840377 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicProvisioningTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicProvisioningTest.java
@@ -453,7 +453,7 @@ public class DynamicProvisioningTest {
         if (!provisionedHosts.isEmpty()) {
             List<Node> hosts = provisionedHosts.asList()
                                                .stream()
-                                               .map(h -> ((MockHostProvisioner)tester.hostProvisioner()).withIpAssigned(h))
+                                               .map(h -> h.with(((MockHostProvisioner) tester.hostProvisioner()).createIpConfig(h)))
                                                .toList();
             tester.move(Node.State.ready, hosts);
             tester.activateTenantHosts();