Require locked node list in some more APIs

10 files changed, 32 insertions, 18 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/LockedNodeList.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/LockedNodeList.java
index dc86daf2c67..9bc18533ddf 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/LockedNodeList.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/LockedNodeList.java
@@ -17,9 +17,16 @@ import java.util.Objects;
  */
 public final class LockedNodeList extends NodeList {
 
+    private final Mutex lock;
+
     public LockedNodeList(List<Node> nodes, Mutex lock) {
         super(nodes, false);
-        Objects.requireNonNull(lock, "lock must be non-null");
+        this.lock = Objects.requireNonNull(lock, "lock must be non-null");
+    }
+
+    /** Returns a new LockedNodeList with the for the same lock. */
+    public LockedNodeList childList(List<Node> nodes) {
+        return new LockedNodeList(nodes, lock);
     }
 
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/FailedExpirer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/FailedExpirer.java
index fa3f9435c70..c594f7f43ae 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/FailedExpirer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/FailedExpirer.java
@@ -115,6 +115,7 @@ public class FailedExpirer extends NodeRepositoryMaintainer {
                 nodesToRecycle.add(candidate);
             }
         }
+        // TODO: consider locked nodes, and perform on directly
         nodeRepository.nodes().deallocate(nodesToRecycle, Agent.FailedExpirer, "Expired by FailedExpirer");
     }
 
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java
index cc7db3c138a..1ff6d2b300d 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/node/IP.java
@@ -113,7 +113,7 @@ public record IP() {
          *
          * @throws IllegalArgumentException if there are IP conflicts with existing nodes
          */
-        public static List<Node> verify(List<Node> nodes, LockedNodeList allNodes) {
+        public static LockedNodeList verify(List<Node> nodes, LockedNodeList allNodes) {
             NodeList sortedNodes = allNodes.sortedBy(Comparator.comparing(Node::hostname));
             for (var node : nodes) {
                 for (var other : sortedNodes) {
@@ -135,7 +135,7 @@ public record IP() {
                                                            other.hostname());
                 }
             }
-            return nodes;
+            return allNodes.childList(nodes);
         }
 
         /** Returns whether IP address of existing node can be assigned to node */
@@ -152,7 +152,7 @@ public record IP() {
         }
 
         public static Node verify(Node node, LockedNodeList allNodes) {
-            return verify(List.of(node), allNodes).get(0);
+            return verify(List.of(node), allNodes).asList().get(0);
         }
 
     }
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 9b5bbaea39e..eb388026315 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
@@ -154,7 +154,7 @@ public class Nodes {
             if (existing.isPresent())
                 throw new IllegalStateException("Cannot add " + node + ": A node with this name already exists");
         }
-        return db.addNodesInState(nodes.asList(), Node.State.reserved, Agent.system);
+        return db.addNodesInState(nodes, Node.State.reserved, Agent.system);
     }
 
     /**
@@ -163,7 +163,8 @@ public class Nodes {
      * with the history of that node.
      */
     public List<Node> addNodes(List<Node> nodes, Agent agent) {
-        try (Mutex lock = lockUnallocated()) {
+        try (NodeMutexes existingNodesLocks = lockAndGetAll(nodes, Optional.empty());
+             Mutex allocationLock = lockUnallocated()) {
             List<Node> nodesToAdd =  new ArrayList<>();
             List<Node> nodesToRemove = new ArrayList<>();
             for (int i = 0; i < nodes.size(); i++) {
@@ -200,7 +201,7 @@ public class Nodes {
             }
             NestedTransaction transaction = new NestedTransaction();
             db.removeNodes(nodesToRemove, transaction);
-            List<Node> resultingNodes = db.addNodesInState(IP.Config.verify(nodesToAdd, list(lock)), Node.State.provisioned, agent, transaction);
+            List<Node> resultingNodes = db.addNodesInState(IP.Config.verify(nodesToAdd, list(allocationLock)), Node.State.provisioned, agent, transaction);
             transaction.commit();
             return resultingNodes;
         }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDb.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDb.java
index d7c4fc1b3d5..28598710ea8 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDb.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDb.java
@@ -18,6 +18,7 @@ import com.yahoo.vespa.curator.Lock;
 import com.yahoo.vespa.curator.recipes.CuratorCounter;
 import com.yahoo.vespa.curator.transaction.CuratorOperations;
 import com.yahoo.vespa.curator.transaction.CuratorTransaction;
+import com.yahoo.vespa.hosted.provision.LockedNodeList;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.applications.Application;
 import com.yahoo.vespa.hosted.provision.archive.ArchiveUris;
@@ -105,7 +106,7 @@ public class CuratorDb {
     }
 
     /** Adds a set of nodes. Rollbacks/fails transaction if any node is not in the expected state. */
-    public List<Node> addNodesInState(List<Node> nodes, Node.State expectedState, Agent agent, NestedTransaction transaction) {
+    public List<Node> addNodesInState(LockedNodeList nodes, Node.State expectedState, Agent agent, NestedTransaction transaction) {
         CuratorTransaction curatorTransaction = db.newCuratorTransactionIn(transaction);
         for (Node node : nodes) {
             if (node.state() != expectedState)
@@ -116,10 +117,10 @@ public class CuratorDb {
             curatorTransaction.add(CuratorOperations.create(nodePath(node).getAbsolute(), serialized));
         }
         transaction.onCommitted(() -> nodes.forEach(node -> log.log(Level.INFO, "Added " + node)));
-        return nodes;
+        return new ArrayList<>(nodes.asList());
     }
 
-    public List<Node> addNodesInState(List<Node> nodes, Node.State expectedState, Agent agent) {
+    public List<Node> addNodesInState(LockedNodeList nodes, Node.State expectedState, Agent agent) {
         NestedTransaction transaction = new NestedTransaction();
         List<Node> writtenNodes = addNodesInState(nodes, expectedState, agent, transaction);
         transaction.commit();
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/RealDataScenarioTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/RealDataScenarioTest.java
index 29ebf1789c0..9c843b3eb01 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/RealDataScenarioTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/RealDataScenarioTest.java
@@ -141,7 +141,7 @@ public class RealDataScenarioTest {
             if (nodeNext.get()) {
                 String json = input.substring(input.indexOf("{\""), input.lastIndexOf('}') + 1);
                 Node node = nodeSerializer.fromJson(json.getBytes(UTF_8));
-                nodeRepository.database().addNodesInState(List.of(node), node.state(), Agent.system);
+                nodeRepository.database().addNodesInState(new LockedNodeList(List.of(node), () -> { }), node.state(), Agent.system);
                 nodeNext.set(false);
             } else {
                 if (!zkNodePathPattern.matcher(input).matches()) return;
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/CapacityCheckerTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/CapacityCheckerTester.java
index f8ec271ce5f..523feeeb303 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/CapacityCheckerTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/CapacityCheckerTester.java
@@ -23,6 +23,7 @@ import com.yahoo.test.ManualClock;
 import com.yahoo.vespa.curator.Curator;
 import com.yahoo.vespa.curator.mock.MockCurator;
 import com.yahoo.vespa.flags.InMemoryFlagSource;
+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.autoscale.MemoryMetricsDb;
@@ -201,7 +202,7 @@ public class CapacityCheckerTester {
         nodeRepository.nodes().addNodes(hostsWithChildren.getOrDefault(tenantHostApp, List.of()), Agent.system);
         hostsWithChildren.forEach((applicationId, nodes) -> {
             if (applicationId.equals(tenantHostApp)) return;
-            nodeRepository.database().addNodesInState(nodes, Node.State.active, Agent.system);
+            nodeRepository.database().addNodesInState(new LockedNodeList(nodes, () -> { }), Node.State.active, Agent.system);
         });
         nodeRepository.nodes().addNodes(createEmptyHosts(numHosts, numEmptyHosts, emptyHostExcessCapacity, emptyHostExcessIps), Agent.system);
 
@@ -322,9 +323,9 @@ public class CapacityCheckerTester {
             }
         }
 
-        nodeRepository.database().addNodesInState(hosts, Node.State.active, Agent.system);
+        nodeRepository.database().addNodesInState(new LockedNodeList(hosts, () -> { }), Node.State.active, Agent.system);
         nodes.forEach((application, applicationNodes) -> {
-            nodeRepository.database().addNodesInState(applicationNodes, Node.State.active, Agent.system);
+            nodeRepository.database().addNodesInState(new LockedNodeList(applicationNodes, () -> { }), Node.State.active, Agent.system);
         });
         updateCapacityChecker();
     }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/DirtyExpirerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/DirtyExpirerTest.java
index ddd7413567a..262616d5eac 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/DirtyExpirerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/DirtyExpirerTest.java
@@ -6,6 +6,7 @@ import com.yahoo.config.provision.ClusterMembership;
 import com.yahoo.config.provision.Flavor;
 import com.yahoo.config.provision.NodeResources;
 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.node.Agent;
 import com.yahoo.vespa.hosted.provision.node.Allocation;
@@ -45,7 +46,7 @@ public class DirtyExpirerTest {
                         false))
                 .build();
 
-        tester.nodeRepository().database().addNodesInState(List.of(node), node.state(), Agent.system);
+        tester.nodeRepository().database().addNodesInState(new LockedNodeList(List.of(node), () -> { }), node.state(), Agent.system);
 
         Duration expiryTimeout = Duration.ofMinutes(30);
         DirtyExpirer expirer = new DirtyExpirer(tester.nodeRepository(), expiryTimeout, new TestMetric());
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/HostCapacityMaintainerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/HostCapacityMaintainerTest.java
index 66d4b67c7c2..e1e183c9362 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/HostCapacityMaintainerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/HostCapacityMaintainerTest.java
@@ -27,6 +27,7 @@ import com.yahoo.test.ManualClock;
 import com.yahoo.vespa.flags.InMemoryFlagSource;
 import com.yahoo.vespa.flags.PermanentFlags;
 import com.yahoo.vespa.flags.custom.ClusterCapacity;
+import com.yahoo.vespa.hosted.provision.LockedNodeList;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.Node.State;
 import com.yahoo.vespa.hosted.provision.NodeList;
@@ -722,7 +723,7 @@ public class HostCapacityMaintainerTest {
                     createNode("host4", Optional.empty(), NodeType.host, Node.State.provisioned, null),
                     createNode("host4-1", Optional.of("host4"), NodeType.tenant, Node.State.reserved, tenantApp),
                     createNode("host4-2", Optional.of("host4"), NodeType.tenant, Node.State.reserved, tenantApp))
-                .forEach(node -> nodeRepository.database().addNodesInState(List.of(node), node.state(), Agent.system));
+                .forEach(node -> nodeRepository.database().addNodesInState(new LockedNodeList(List.of(node), () -> { }), node.state(), Agent.system));
             return this;
         }
 
@@ -744,7 +745,7 @@ public class HostCapacityMaintainerTest {
 
         private Node addNode(String hostname, Optional<String> parentHostname, NodeType nodeType, Node.State state, ApplicationId application, Duration hostTTL) {
             Node node = createNode(hostname, parentHostname, nodeType, state, application, hostTTL);
-            return nodeRepository.database().addNodesInState(List.of(node), node.state(), Agent.system).get(0);
+            return nodeRepository.database().addNodesInState(new LockedNodeList(List.of(node), () -> { }), node.state(), Agent.system).get(0);
         }
 
         private Node createNode(String hostname, Optional<String> parentHostname, NodeType nodeType,
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/ProvisionedExpirerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/ProvisionedExpirerTest.java
index 359f75c27ab..ac1e452d7a5 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/ProvisionedExpirerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/ProvisionedExpirerTest.java
@@ -9,6 +9,7 @@ import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.Zone;
+import com.yahoo.vespa.hosted.provision.LockedNodeList;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.provisioning.ProvisioningTester;
@@ -45,7 +46,7 @@ public class ProvisionedExpirerTest {
         var nodes = IntStream.range(0, 15)
                              .mapToObj(i -> Node.create("id-" + i, "host-" + i, new Flavor(NodeResources.unspecified()), Node.State.provisioned, NodeType.host).build())
                              .toList();
-        tester.nodeRepository().database().addNodesInState(nodes, Node.State.provisioned, Agent.system);
+        tester.nodeRepository().database().addNodesInState(new LockedNodeList(nodes, () -> { }), Node.State.provisioned, Agent.system);
     }
 
 }