Acquire lock for children when patching

4 files changed, 79 insertions, 38 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
index 18ccdf0dd92..8e606a4af3d 100644
--- 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
@@ -19,4 +19,9 @@ public class NodeMutex implements 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 newNode) {
+        return new NodeMutex(newNode, 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 899e8c12c06..76cae7c0e50 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
@@ -924,13 +924,12 @@ public class NodeRepository extends AbstractComponent {
     /** Create a lock which provides exclusive rights to modifying unallocated nodes */
     public Mutex lockUnallocated() { return db.lockInactive(); }
 
-    /** Returns a lock, and an up to date node fetched under an appropriate lock, if it exists. */
+    /** Returns the unallocated/application lock, and the node acquired under that lock. */
     public Optional<NodeMutex> lockAndGet(Node node) {
         Node staleNode = node;
 
         for (int i = 0; i < 4; ++i) {
-            Mutex lock = lock(staleNode);
-            Optional<Mutex> lockToClose = Optional.of(lock);
+            Mutex lockToClose = lock(staleNode);
             try {
                 Optional<Node> freshNode = getNode(staleNode.hostname(), staleNode.state());
                 if (freshNode.isEmpty()) {
@@ -942,30 +941,32 @@ public class NodeRepository extends AbstractComponent {
 
                 if (Objects.equals(freshNode.get().allocation().map(Allocation::owner),
                                    staleNode.allocation().map(Allocation::owner))) {
-                    lockToClose = Optional.empty();
-                    return Optional.of(new NodeMutex(freshNode.get(), lock));
+                    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 {
-                lockToClose.ifPresent(Mutex::close);
+                if (lockToClose != null) lockToClose.close();
             }
         }
 
         throw new IllegalStateException("Giving up trying to fetch an up to date node under lock: " + node.hostname());
     }
 
-    /** Returns a lock, and an up to date node fetched under an appropriate lock, if it exists. */
+    /** 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 a lock, and an up to date node fetched under an appropriate lock. */
+    /** 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 a lock, and an up to date node fetched under an appropriate lock. */
+    /** 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));
     }
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 570ebdf767d..15bf9dee805 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,20 +91,27 @@ public class NodePatcher {
         inspector.traverse((String name, Inspector value) -> {
             try {
                 patchedNodes.update(applyField(patchedNodes.node(), name, value, inspector, false));
-            } 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));
+                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);
             }
-        } );
+        });
 
         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" :
@@ -272,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 List.copyOf(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 b5454d97cfd..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
@@ -160,12 +160,11 @@ public class NodesV2ApiHandler extends LoggingRequestHandler {
     private HttpResponse handlePATCH(HttpRequest request) {
         String path = request.getUri().getPath();
         if (path.startsWith("/nodes/v2/node/")) {
-            try (NodeMutex lock = nodeRepository.lockAndGetRequired(nodeFromRequest(request))) {
-                var patchedNodes = new NodePatcher(nodeFlavors, request.getData(), lock.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 " + lock.node().hostname());
+                return new MessageResponse("Updated " + patcher.nodeMutexOfHost().node().hostname());
             }
         }
         else if (path.startsWith("/nodes/v2/upgrade/")) {