Combine restarts with reboots

2 files changed, 5 insertions, 3 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 407814f3b4d..02c6e8ba398 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
@@ -197,6 +197,8 @@ public class NodeAgentImpl implements NodeAgent {
         dockerOperations.createContainer(context, containerData, getContainerResources(context));
         dockerOperations.startContainer(context);
 
+        currentRebootGeneration = context.node().wantedRebootGeneration();
+        currentRestartGeneration = context.node().wantedRestartGeneration();
         hasStartedServices = true; // Automatically started with the container
         hasResumedNode = false;
         context.log(logger, "Container successfully started, new containerState is " + containerState);
@@ -322,7 +324,6 @@ public class NodeAgentImpl implements NodeAgent {
 
         storageMaintainer.handleCoreDumpsForContainer(context, Optional.of(existingContainer));
         dockerOperations.removeContainer(context, existingContainer);
-        currentRebootGeneration = context.node().wantedRebootGeneration();
         containerState = ABSENT;
         context.log(logger, "Container successfully removed, new containerState is " + containerState);
     }
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
index 465a430b497..9e1f84f07a1 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeagent/NodeAgentImplTest.java
@@ -264,7 +264,7 @@ public class NodeAgentImplTest {
                 .state(NodeState.active)
                 .wantedDockerImage(dockerImage).currentDockerImage(dockerImage)
                 .wantedVespaVersion(vespaVersion).currentVespaVersion(vespaVersion)
-                .wantedRestartGeneration(1).currentRestartGeneration(1);
+                .wantedRestartGeneration(2).currentRestartGeneration(1);
 
         NodeAgentContext firstContext = createContext(specBuilder.build());
         NodeAgentImpl nodeAgent = makeNodeAgent(dockerImage, true);
@@ -278,11 +278,12 @@ public class NodeAgentImplTest {
         ContainerResources resourcesAfterThird = ContainerResources.from(0, 2, 20);
         mockGetContainer(dockerImage, resourcesAfterThird, true);
 
-        InOrder inOrder = inOrder(orchestrator, dockerOperations);
+        InOrder inOrder = inOrder(orchestrator, dockerOperations, nodeRepository);
         inOrder.verify(orchestrator).resume(any(String.class));
         inOrder.verify(dockerOperations).removeContainer(eq(secondContext), any());
         inOrder.verify(dockerOperations, never()).updateContainer(any(), any());
         inOrder.verify(dockerOperations, never()).restartVespa(any());
+        inOrder.verify(nodeRepository).updateNodeAttributes(eq(hostName), eq(new NodeAttributes().withRestartGeneration(2)));
 
         nodeAgent.doConverge(secondContext);
         inOrder.verify(orchestrator).resume(any(String.class));