Fix restart generation bug

1 files changed, 7 insertions, 2 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
index 82802a4c6af..6b9a7d4021a 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImpl.java
@@ -470,10 +470,15 @@ public class NodeAgentImpl implements NodeAgent {
         if (!node.equals(lastNode)) {
             logChangesToNodeSpec(lastNode, node);
 
-            if (lastNode == null) {
+            // Current reboot generation uninitialized or incremented from outside to cancel reboot
+            if (currentRebootGeneration < node.getCurrentRebootGeneration())
                 currentRebootGeneration = node.getCurrentRebootGeneration();
+
+            // Either we have changed allocation status (restart gen. only available to allocated nodes), or
+            // restart generation has been incremented from outside to cancel restart
+            if (currentRestartGeneration.isPresent() != node.getCurrentRestartGeneration().isPresent() ||
+                    currentRestartGeneration.map(current -> current < node.getCurrentRestartGeneration().get()).orElse(false))
                 currentRestartGeneration = node.getCurrentRestartGeneration();
-            }
 
             // Every time the node spec changes, we should clear the metrics for this container as the dimensions
             // will change and we will be reporting duplicate metrics.