Merge pull request #15966 from vespa-engine/hakonhall/ensure-fresh-node-with-lock

Acquire locks for children when patching

11 files changed, 205 insertions, 87 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeMutex.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeMutex.java
new file mode 100644
index 00000000000..814b40a73b7
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeMutex.java
@@ -0,0 +1,31 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision;
+
+import com.yahoo.transaction.Mutex;
+
+/**
+ * Holds a node and mutex pair, with the mutex closed by {@link #close()}.
+ *
+ * @author hakon
+ */
+public class NodeMutex implements Mutex {
+    private final Node node;
+    private final Mutex mutex;
+
+    public NodeMutex(Node node, Mutex mutex) {
+        this.node = node;
+        this.mutex = mutex;
+    }
+
+    public Node node() { return node; }
+    @Override public void close() { mutex.close(); }
+
+    /** Returns a node mutex with the same mutex as this, but the given node.  Be sure to close only one. */
+    public NodeMutex with(Node updatedNode) {
+        if (!node.equals(updatedNode)) {
+            throw new IllegalArgumentException("Updated node not equal to current");
+        }
+
+        return new NodeMutex(updatedNode, mutex);
+    }
+}
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
index b0e7c0bd61b..4eb38fa650e 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
@@ -28,6 +28,7 @@ import com.yahoo.vespa.hosted.provision.maintenance.InfrastructureVersions;
 import com.yahoo.vespa.hosted.provision.maintenance.NodeFailer;
 import com.yahoo.vespa.hosted.provision.maintenance.PeriodicApplicationMaintainer;
 import com.yahoo.vespa.hosted.provision.node.Agent;
+import com.yahoo.vespa.hosted.provision.node.Allocation;
 import com.yahoo.vespa.hosted.provision.node.IP;
 import com.yahoo.vespa.hosted.provision.node.NodeAcl;
 import com.yahoo.vespa.hosted.provision.node.filter.NodeFilter;
@@ -53,6 +54,7 @@ import java.util.EnumSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.TreeSet;
@@ -656,6 +658,7 @@ public class NodeRepository extends AbstractComponent {
         if (toState == Node.State.active && node.allocation().isEmpty())
             illegal("Could not set " + node + " active. It has no allocation.");
 
+        // TODO: Work out a safe lock acquisition strategy for moves, e.g. migrate to lockNode.
         try (Mutex lock = lock(node)) {
             if (toState == State.active) {
                 for (Node currentActive : getNodes(node.allocation().get().owner(), State.active)) {
@@ -921,13 +924,61 @@ public class NodeRepository extends AbstractComponent {
     /** Create a lock which provides exclusive rights to modifying unallocated nodes */
     public Mutex lockUnallocated() { return db.lockInactive(); }
 
-    /** Acquires the appropriate lock for this node */
-    public Mutex lock(Node node) {
+    /** Returns the unallocated/application lock, and the node acquired under that lock. */
+    public Optional<NodeMutex> lockAndGet(Node node) {
+        Node staleNode = node;
+
+        final int maxRetries = 4;
+        for (int i = 0; i < maxRetries; ++i) {
+            Mutex lockToClose = lock(staleNode);
+            try {
+                // As an optimization we first try finding the node in the same state
+                Optional<Node> freshNode = getNode(staleNode.hostname(), staleNode.state());
+                if (freshNode.isEmpty()) {
+                    freshNode = getNode(staleNode.hostname());
+                    if (freshNode.isEmpty()) {
+                        return Optional.empty();
+                    }
+                }
+
+                if (Objects.equals(freshNode.get().allocation().map(Allocation::owner),
+                                   staleNode.allocation().map(Allocation::owner))) {
+                    NodeMutex nodeMutex = new NodeMutex(freshNode.get(), lockToClose);
+                    lockToClose = null;
+                    return Optional.of(nodeMutex);
+                }
+
+                // The wrong lock was held when the fresh node was fetched, so try again
+                staleNode = freshNode.get();
+            } finally {
+                if (lockToClose != null) lockToClose.close();
+            }
+        }
+
+        throw new IllegalStateException("Giving up (after " + maxRetries + " attempts) " +
+                "fetching an up to date node under lock: " + node.hostname());
+    }
+
+    /** Returns the unallocated/application lock, and the node acquired under that lock. */
+    public Optional<NodeMutex> lockAndGet(String hostname) {
+        return getNode(hostname).flatMap(this::lockAndGet);
+    }
+
+    /** Returns the unallocated/application lock, and the node acquired under that lock. */
+    public NodeMutex lockAndGetRequired(Node node) {
+        return lockAndGet(node).orElseThrow(() -> new IllegalArgumentException("No such node: " + node.hostname()));
+    }
+
+    /** Returns the unallocated/application lock, and the node acquired under that lock. */
+    public NodeMutex lockAndGetRequired(String hostname) {
+        return lockAndGet(hostname).orElseThrow(() -> new IllegalArgumentException("No such node: " + hostname));
+    }
+
+    private Mutex lock(Node node) {
         return node.allocation().isPresent() ? lock(node.allocation().get().owner()) : lockUnallocated();
     }
 
     private void illegal(String message) {
         throw new IllegalArgumentException(message);
     }
-
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/MaintenanceDeployment.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/MaintenanceDeployment.java
index 62d042d5e8b..66ddf7f9ffe 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/MaintenanceDeployment.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/MaintenanceDeployment.java
@@ -9,6 +9,7 @@ import com.yahoo.config.provision.TransientException;
 import com.yahoo.jdisc.Metric;
 import com.yahoo.transaction.Mutex;
 import com.yahoo.vespa.hosted.provision.Node;
+import com.yahoo.vespa.hosted.provision.NodeMutex;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.yolean.Exceptions;
@@ -192,14 +193,15 @@ class MaintenanceDeployment implements Closeable {
 
         /** Returns true only if this operation changes the state of the wantToRetire flag */
         private boolean markWantToRetire(Node node, boolean wantToRetire, Agent agent, NodeRepository nodeRepository) {
-            try (Mutex lock = nodeRepository.lock(node)) {
-                Optional<Node> nodeToMove = nodeRepository.getNode(node.hostname());
-                if (nodeToMove.isEmpty()) return false;
-                if (nodeToMove.get().state() != Node.State.active) return false;
+            Optional<NodeMutex> nodeMutex = nodeRepository.lockAndGet(node);
+            if (nodeMutex.isEmpty()) return false;
 
-                if (nodeToMove.get().status().wantToRetire() == wantToRetire) return false;
+            try (var nodeLock = nodeMutex.get()) {
+                if (nodeLock.node().state() != Node.State.active) return false;
 
-                nodeRepository.write(nodeToMove.get().withWantToRetire(wantToRetire, agent, nodeRepository.clock().instant()), lock);
+                if (nodeLock.node().status().wantToRetire() == wantToRetire) return false;
+
+                nodeRepository.write(nodeLock.node().withWantToRetire(wantToRetire, agent, nodeRepository.clock().instant()), nodeLock);
                 return true;
             }
         }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringUpgrader.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringUpgrader.java
index 771f5e53e13..68ac96fdf1d 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringUpgrader.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringUpgrader.java
@@ -5,6 +5,7 @@ import com.yahoo.component.Version;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
+import com.yahoo.vespa.hosted.provision.NodeMutex;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.node.filter.NodeListFilter;
@@ -61,10 +62,10 @@ public class RetiringUpgrader implements Upgrader {
     /** Retire and deprovision given host and its children */
     private void retire(Node host, Version target, Instant now, NodeList allNodes) {
         if (!host.type().isHost()) throw new IllegalArgumentException("Cannot retire non-host " + host);
-        try (var lock = nodeRepository.lock(host)) {
-            Optional<Node> currentNode = nodeRepository.getNode(host.hostname());
-            if (currentNode.isEmpty()) return;
-            host = currentNode.get();
+        Optional<NodeMutex> nodeMutex = nodeRepository.lockAndGet(host);
+        if (nodeMutex.isEmpty()) return;
+        try (var lock = nodeMutex.get()) {
+            host = lock.node();
             NodeType nodeType = host.type();
 
             LOG.info("Retiring and deprovisioning " + host + ": On stale OS version " +
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 2b03a5cae8c..2ec3912181b 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
@@ -8,9 +8,9 @@ import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.config.provision.Flavor;
 import com.yahoo.config.provision.HostSpec;
 import com.yahoo.config.provision.ParentHostUnavailableException;
-import com.yahoo.transaction.Mutex;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
+import com.yahoo.vespa.hosted.provision.NodeMutex;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.applications.Application;
 import com.yahoo.vespa.hosted.provision.applications.ScalingEvent;
@@ -126,13 +126,16 @@ class Activator {
     private void unreserveParentsOf(List<Node> nodes) {
         for (Node node : nodes) {
             if ( node.parentHostname().isEmpty()) continue;
-            Optional<Node> parent = nodeRepository.getNode(node.parentHostname().get());
+            Optional<Node> parentNode = nodeRepository.getNode(node.parentHostname().get());
+            if (parentNode.isEmpty()) continue;
+            if (parentNode.get().reservedTo().isEmpty()) continue;
+
+            // Above is an optimization to avoid unnecessary locking - now repeat all conditions under lock
+            Optional<NodeMutex> parent = nodeRepository.lockAndGet(node.parentHostname().get());
             if (parent.isEmpty()) continue;
-            if (parent.get().reservedTo().isEmpty()) continue;
-            try (Mutex lock = nodeRepository.lock(parent.get())) {
-                Optional<Node> lockedParent = nodeRepository.getNode(parent.get().hostname());
-                if (lockedParent.isEmpty()) continue;
-                nodeRepository.write(lockedParent.get().withoutReservedTo(), lock);
+            try (var lock = parent.get()) {
+                if (lock.node().reservedTo().isEmpty()) continue;
+                nodeRepository.write(lock.node().withoutReservedTo(), lock);
             }
         }
     }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodePatcher.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodePatcher.java
index 832310cf2c9..47b0021874d 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodePatcher.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodePatcher.java
@@ -13,8 +13,11 @@ import com.yahoo.slime.Inspector;
 import com.yahoo.slime.ObjectTraverser;
 import com.yahoo.slime.SlimeUtils;
 import com.yahoo.slime.Type;
+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.NodeMutex;
+import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.node.Address;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.node.Allocation;
@@ -34,7 +37,6 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.TreeSet;
-import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 import static com.yahoo.config.provision.NodeResources.DiskSpeed.fast;
@@ -48,28 +50,36 @@ import static com.yahoo.config.provision.NodeResources.StorageType.remote;
  *
  * @author bratseth
  */
-public class NodePatcher {
+public class NodePatcher implements AutoCloseable {
 
     private static final String WANT_TO_RETIRE = "wantToRetire";
     private static final String WANT_TO_DEPROVISION = "wantToDeprovision";
     private static final Set<String> RECURSIVE_FIELDS = Set.of(WANT_TO_RETIRE);
 
+    private final NodeRepository nodeRepository;
     private final com.google.common.base.Supplier<LockedNodeList> memoizedNodes;
     private final PatchedNodes patchedNodes;
     private final NodeFlavors nodeFlavors;
     private final Inspector inspector;
     private final Clock clock;
 
-    public NodePatcher(NodeFlavors nodeFlavors, InputStream json, Node node, Supplier<LockedNodeList> nodes, Clock clock) {
-        this.memoizedNodes = Suppliers.memoize(nodes::get);
-        this.patchedNodes = new PatchedNodes(node);
+    public NodePatcher(NodeFlavors nodeFlavors, InputStream json, Node node, NodeRepository nodeRepository) {
+        this.nodeRepository = nodeRepository;
         this.nodeFlavors = nodeFlavors;
-        this.clock = clock;
+        this.clock = nodeRepository.clock();
         try {
             this.inspector = SlimeUtils.jsonToSlime(IOUtils.readBytes(json, 1000 * 1000)).get();
         } catch (IOException e) {
             throw new UncheckedIOException("Error reading request body", e);
         }
+
+        this.patchedNodes = new PatchedNodes(nodeRepository.lockAndGetRequired(node));
+        try {
+            this.memoizedNodes = Suppliers.memoize(() -> nodeRepository.list(patchedNodes.nodeMutex()));
+        } catch (RuntimeException e) {
+            patchedNodes.close();
+            throw e;
+        }
     }
 
     /**
@@ -81,19 +91,27 @@ public class NodePatcher {
         inspector.traverse((String name, Inspector value) -> {
             try {
                 patchedNodes.update(applyField(patchedNodes.node(), name, value, inspector, false));
+
+                if (RECURSIVE_FIELDS.contains(name)) {
+                    for (Node child: patchedNodes.children()) {
+                        patchedNodes.update(applyField(child, name, value, inspector, true));
+                    }
+                }
             } catch (IllegalArgumentException e) {
                 throw new IllegalArgumentException("Could not set field '" + name + "'", e);
             }
-
-            if (RECURSIVE_FIELDS.contains(name)) {
-                for (Node child: patchedNodes.children())
-                    patchedNodes.update(applyField(child, name, value, inspector, true));
-            }
-        } );
+        });
 
         return patchedNodes.nodes();
     }
 
+    public NodeMutex nodeMutexOfHost() { return patchedNodes.nodeMutex(); }
+
+    @Override
+    public void close() {
+        patchedNodes.close();
+    }
+
     private Node applyField(Node node, String name, Inspector value, Inspector root, boolean applyingAsChild) {
         switch (name) {
             case "currentRebootGeneration" :
@@ -271,36 +289,55 @@ public class NodePatcher {
         return Optional.of(field).filter(Inspector::valid).map(this::asBoolean);
     }
 
-    private class PatchedNodes {
-        private final Map<String, Node> nodes = new HashMap<>();
+    private class PatchedNodes implements Mutex {
+        private final Map<String, NodeMutex> nodes = new HashMap<>();
         private final String hostname;
         private boolean fetchedChildren;
 
-        private PatchedNodes(Node node) {
-            this.hostname = node.hostname();
+        private PatchedNodes(NodeMutex nodeMutex) {
+            this.hostname = nodeMutex.node().hostname();
+            nodes.put(hostname, nodeMutex);
+            fetchedChildren = !nodeMutex.node().type().isHost();
+        }
 
-            nodes.put(hostname, node);
-            fetchedChildren = !node.type().isHost();
+        public NodeMutex nodeMutex() {
+            return nodes.get(hostname);
         }
 
         public Node node() {
-            return nodes.get(hostname);
+            return nodeMutex().node();
         }
 
         public List<Node> children() {
             if (!fetchedChildren) {
-                memoizedNodes.get().childrenOf(hostname).forEach(this::update);
+                memoizedNodes.get()
+                        .childrenOf(hostname)
+                        .forEach(node -> nodeRepository.lockAndGet(node)
+                                .ifPresent(nodeMutex -> nodes.put(nodeMutex.node().hostname(), nodeMutex)));
                 fetchedChildren = true;
             }
-            return nodes.values().stream().filter(node -> !node.type().isHost()).collect(Collectors.toList());
+            return nodes.values().stream()
+                    .map(NodeMutex::node)
+                    .filter(node -> !node.type().isHost())
+                    .collect(Collectors.toList());
         }
 
         public void update(Node node) {
-            nodes.put(node.hostname(), node);
+            NodeMutex currentNodeMutex = nodes.get(node.hostname());
+            if (currentNodeMutex == null) {
+                throw new IllegalStateException("unable to update non-existing child: " + node.hostname());
+            }
+
+            nodes.put(node.hostname(), currentNodeMutex.with(node));
         }
 
         public List<Node> nodes() {
-            return List.copyOf(nodes.values());
+            return nodes.values().stream().map(NodeMutex::node).collect(Collectors.toList());
+        }
+
+        @Override
+        public void close() {
+            nodes.values().forEach(NodeMutex::close);
         }
     }
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
index c43629aeb09..9b5dac49e32 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
@@ -24,9 +24,9 @@ import com.yahoo.slime.Cursor;
 import com.yahoo.slime.Inspector;
 import com.yahoo.slime.Slime;
 import com.yahoo.slime.SlimeUtils;
-import com.yahoo.transaction.Mutex;
 import com.yahoo.vespa.hosted.provision.NoSuchNodeException;
 import com.yahoo.vespa.hosted.provision.Node;
+import com.yahoo.vespa.hosted.provision.NodeMutex;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.applications.Application;
 import com.yahoo.vespa.hosted.provision.node.Address;
@@ -160,14 +160,12 @@ public class NodesV2ApiHandler extends LoggingRequestHandler {
     private HttpResponse handlePATCH(HttpRequest request) {
         String path = request.getUri().getPath();
         if (path.startsWith("/nodes/v2/node/")) {
-            // TODO: Node is fetched outside of lock, might change after getting lock
-            Node node = nodeFromRequest(request);
-            try (var lock = nodeRepository.lock(node)) {
-                var patchedNodes = new NodePatcher(nodeFlavors, request.getData(), node, () -> nodeRepository.list(lock),
-                                                   nodeRepository.clock()).apply();
-                nodeRepository.write(patchedNodes, lock);
+            try (NodePatcher patcher = new NodePatcher(nodeFlavors, request.getData(), nodeFromRequest(request), nodeRepository)) {
+                var patchedNodes = patcher.apply();
+                nodeRepository.write(patchedNodes, patcher.nodeMutexOfHost());
+
+                return new MessageResponse("Updated " + patcher.nodeMutexOfHost().node().hostname());
             }
-            return new MessageResponse("Updated " + node.hostname());
         }
         else if (path.startsWith("/nodes/v2/upgrade/")) {
             return setTargetVersions(request);
@@ -210,13 +208,11 @@ public class NodesV2ApiHandler extends LoggingRequestHandler {
     }
 
     private HttpResponse deleteNode(String hostname) {
-        Optional<Node> node = nodeRepository.getNode(hostname);
-        if (node.isEmpty()) throw new NotFoundException("No node with hostname '" + hostname + "'");
-        try (Mutex lock = nodeRepository.lock(node.get())) {
-            node = nodeRepository.getNode(hostname);
-            if (node.isEmpty()) throw new NotFoundException("No node with hostname '" + hostname + "'");
-            if (node.get().state() == Node.State.deprovisioned) {
-                nodeRepository.forget(node.get());
+        Optional<NodeMutex> nodeMutex = nodeRepository.lockAndGet(hostname);
+        if (nodeMutex.isEmpty()) throw new NotFoundException("No node with hostname '" + hostname + "'");
+        try (var lock = nodeMutex.get()) {
+            if (lock.node().state() == Node.State.deprovisioned) {
+                nodeRepository.forget(lock.node());
                 return new MessageResponse("Permanently removed " + hostname);
             } else {
                 List<Node> removedNodes = nodeRepository.removeRecursively(hostname);
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockDeployer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockDeployer.java
index 8099b08ae89..9461cb00c5f 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockDeployer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockDeployer.java
@@ -11,9 +11,9 @@ import com.yahoo.config.provision.Deployer;
 import com.yahoo.config.provision.Deployment;
 import com.yahoo.config.provision.HostFilter;
 import com.yahoo.config.provision.HostSpec;
-import com.yahoo.transaction.Mutex;
 import com.yahoo.transaction.NestedTransaction;
 import com.yahoo.vespa.hosted.provision.Node;
+import com.yahoo.vespa.hosted.provision.NodeMutex;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.hosted.provision.provisioning.NodeRepositoryProvisioner;
 
@@ -196,8 +196,8 @@ public class MockDeployer implements Deployer {
             lastDeployTimes.put(applicationId, clock.instant());
 
             for (Node node : nodeRepository.list().owner(applicationId).state(Node.State.active).wantToRetire().asList()) {
-                try (Mutex lock = nodeRepository.lock(node)) {
-                    nodeRepository.write(node.retire(nodeRepository.clock().instant()), lock);
+                try (NodeMutex lock = nodeRepository.lockAndGetRequired(node)) {
+                    nodeRepository.write(lock.node().retire(nodeRepository.clock().instant()), lock);
                 }
             }
             return redeployments++;
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java
index e6e19d0407e..bcf53a07490 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTest.java
@@ -182,18 +182,17 @@ public class NodeRepositoryTest {
         assertFalse(tester.nodeRepository().getNode("host1").get().history().hasEventAfter(History.Event.Type.deprovisioned, testStart));
 
         // Set host 1 properties and deprovision it
-        Node host1 = tester.nodeRepository().getNode("host1").get();
-        host1 = host1.withWantToRetire(true, true, Agent.system, tester.nodeRepository().clock().instant());
-        host1 = host1.withFirmwareVerifiedAt(tester.clock().instant());
-        host1 = host1.with(host1.status().withIncreasedFailCount());
-        host1 = host1.with(host1.reports().withReport(Report.basicReport("id", Report.Type.HARD_FAIL, tester.clock().instant(), "Test report")));
-        try (var lock = tester.nodeRepository().lock(host1)) {
+        try (var lock = tester.nodeRepository().lockAndGetRequired("host1")) {
+            Node host1 = lock.node().withWantToRetire(true, true, Agent.system, tester.nodeRepository().clock().instant());
+            host1 = host1.withFirmwareVerifiedAt(tester.clock().instant());
+            host1 = host1.with(host1.status().withIncreasedFailCount());
+            host1 = host1.with(host1.reports().withReport(Report.basicReport("id", Report.Type.HARD_FAIL, tester.clock().instant(), "Test report")));
             tester.nodeRepository().write(host1, lock);
         }
         tester.nodeRepository().removeRecursively("host1");
 
         // Host 1 is deprovisioned and unwanted properties are cleared
-        host1 = tester.nodeRepository().getNode("host1").get();
+        Node host1 = tester.nodeRepository().getNode("host1").get();
         assertEquals(Node.State.deprovisioned, host1.state());
         assertTrue(host1.history().hasEventAfter(History.Event.Type.deprovisioned, testStart));
 
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/InPlaceResizeProvisionTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/InPlaceResizeProvisionTest.java
index 230195a92bd..62b153c7e12 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/InPlaceResizeProvisionTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/InPlaceResizeProvisionTest.java
@@ -10,10 +10,10 @@ import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.OutOfCapacityException;
 import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.Zone;
-import com.yahoo.transaction.Mutex;
 import com.yahoo.vespa.flags.InMemoryFlagSource;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
+import com.yahoo.vespa.hosted.provision.NodeMutex;
 import com.yahoo.vespa.hosted.provision.node.Agent;
 import org.junit.Test;
 
@@ -206,8 +206,8 @@ public class InPlaceResizeProvisionTest {
 
         // ... same with setting a node to want to retire
         Node nodeToWantoToRetire = listCluster(content1).not().retired().asList().get(0);
-        try (Mutex lock = tester.nodeRepository().lock(nodeToWantoToRetire)) {
-            tester.nodeRepository().write(nodeToWantoToRetire.withWantToRetire(true, Agent.system,
+        try (NodeMutex lock = tester.nodeRepository().lockAndGetRequired(nodeToWantoToRetire)) {
+            tester.nodeRepository().write(lock.node().withWantToRetire(true, Agent.system,
                     tester.clock().instant()), lock);
         }
         new PrepareHelper(tester, app).prepare(content1, 8, 1, halvedResources).activate();
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
index 91c9f7e50ac..ebd856e96a0 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
@@ -25,7 +25,6 @@ import com.yahoo.config.provision.TenantName;
 import com.yahoo.config.provision.Zone;
 import com.yahoo.config.provisioning.FlavorsConfig;
 import com.yahoo.test.ManualClock;
-import com.yahoo.transaction.Mutex;
 import com.yahoo.transaction.NestedTransaction;
 import com.yahoo.vespa.curator.Curator;
 import com.yahoo.vespa.curator.mock.MockCurator;
@@ -153,9 +152,8 @@ public class ProvisioningTester {
     public List<Node> patchNodes(List<Node> nodes, UnaryOperator<Node> patcher) {
         List<Node> updated = new ArrayList<>();
         for (var node : nodes) {
-            try (var lock = nodeRepository.lock(node)) {
-                node = nodeRepository.getNode(node.hostname()).get();
-                node = patcher.apply(node);
+            try (var lock = nodeRepository.lockAndGetRequired(node)) {
+                node = patcher.apply(lock.node());
                 nodeRepository.write(node, lock);
                 updated.add(node);
             }
@@ -188,21 +186,21 @@ public class ProvisioningTester {
 
         // Add ip addresses and activate parent host if necessary
         for (HostSpec prepared : preparedNodes) {
-            Node node = nodeRepository.getNode(prepared.hostname()).get();
-            if (node.ipConfig().primary().isEmpty()) {
-                node = node.with(new IP.Config(Set.of("::" + 0 + ":0"), Set.of()));
-                try (Mutex lock = nodeRepository.lock(node)) {
+            try (var lock = nodeRepository.lockAndGetRequired(prepared.hostname())) {
+                Node node = lock.node();
+                if (node.ipConfig().primary().isEmpty()) {
+                    node = node.with(new IP.Config(Set.of("::" + 0 + ":0"), Set.of()));
                     nodeRepository.write(node, lock);
                 }
+                if (node.parentHostname().isEmpty()) continue;
+                Node parent = nodeRepository.getNode(node.parentHostname().get()).get();
+                if (parent.state() == Node.State.active) continue;
+                NestedTransaction t = new NestedTransaction();
+                if (parent.ipConfig().primary().isEmpty())
+                    parent = parent.with(new IP.Config(Set.of("::" + 0 + ":0"), Set.of("::" + 0 + ":2")));
+                nodeRepository.activate(List.of(parent), t);
+                t.commit();
             }
-            if (node.parentHostname().isEmpty()) continue;
-            Node parent = nodeRepository.getNode(node.parentHostname().get()).get();
-            if (parent.state() == Node.State.active) continue;
-            NestedTransaction t = new NestedTransaction();
-            if (parent.ipConfig().primary().isEmpty())
-                parent = parent.with(new IP.Config(Set.of("::" + 0 + ":0"), Set.of("::" + 0 + ":2")));
-            nodeRepository.activate(List.of(parent), t);
-            t.commit();
         }
 
         return activate(application, preparedNodes);