Skip checking if a docker image is downloaded if we are already running it

4 files changed, 22 insertions, 19 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/DockerOperations.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/DockerOperations.java
index c8fff58badc..a012de36b62 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/DockerOperations.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/DockerOperations.java
@@ -22,7 +22,7 @@ public interface DockerOperations {
 
     Optional<Container> getContainer(ContainerName containerName);
 
-    void scheduleDownloadOfImage(ContainerName containerName, ContainerNodeSpec nodeSpec, Runnable callback);
+    void scheduleDownloadOfImage(ContainerName containerName, DockerImage dockerImage, Runnable callback);
 
     ProcessResult executeCommandInContainerAsRoot(ContainerName containerName, String... command);
 
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/DockerOperationsImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/DockerOperationsImpl.java
index 5da51b83d37..0e20e16ea48 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/DockerOperationsImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/docker/DockerOperationsImpl.java
@@ -221,17 +221,16 @@ public class DockerOperationsImpl implements DockerOperations {
     }
 
     @Override
-    public void scheduleDownloadOfImage(ContainerName containerName, final ContainerNodeSpec nodeSpec, Runnable callback) {
+    public void scheduleDownloadOfImage(ContainerName containerName, DockerImage dockerImage, Runnable callback) {
         PrefixLogger logger = PrefixLogger.getNodeAgentLogger(DockerOperationsImpl.class, containerName);
 
-        logger.info("Schedule async download of " + nodeSpec.wantedDockerImage.get());
-        final CompletableFuture<DockerImage> asyncPullResult = docker.pullImageAsync(nodeSpec.wantedDockerImage.get());
-        asyncPullResult.whenComplete((dockerImage, throwable) -> {
+        logger.info("Schedule async download of " + dockerImage);
+        final CompletableFuture<DockerImage> asyncPullResult = docker.pullImageAsync(dockerImage);
+        asyncPullResult.whenComplete((image, throwable) -> {
             if (throwable != null) {
-                logger.warning("Failed to pull " + nodeSpec.wantedDockerImage, throwable);
+                logger.warning("Failed to pull " + dockerImage, throwable);
                 return;
             }
-            assert nodeSpec.wantedDockerImage.get().equals(dockerImage);
             callback.run();
         });
     }
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 d8ee3550272..69ee290ea72 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
@@ -329,6 +329,8 @@ public class NodeAgentImpl implements NodeAgent {
 
 
     private void scheduleDownLoadIfNeeded(ContainerNodeSpec nodeSpec) {
+        if (nodeSpec.currentDockerImage.equals(nodeSpec.wantedDockerImage)) return;
+
         if (dockerOperations.shouldScheduleDownloadOfImage(nodeSpec.wantedDockerImage.get())) {
             if (nodeSpec.wantedDockerImage.get().equals(imageBeingDownloaded)) {
                 // Downloading already scheduled, but not done.
@@ -336,7 +338,7 @@ public class NodeAgentImpl implements NodeAgent {
             }
             imageBeingDownloaded = nodeSpec.wantedDockerImage.get();
             // Create a signalWorkToBeDone when download is finished.
-            dockerOperations.scheduleDownloadOfImage(containerName, nodeSpec, this::signalWorkToBeDone);
+            dockerOperations.scheduleDownloadOfImage(containerName, imageBeingDownloaded, this::signalWorkToBeDone);
         } else if (imageBeingDownloaded != null) { // Image was downloading, but now it's ready
             imageBeingDownloaded = null;
         }
@@ -427,7 +429,7 @@ public class NodeAgentImpl implements NodeAgent {
                     maintainer.handleCoreDumpsForContainer(containerName, nodeSpec, environment);
                 });
                 scheduleDownLoadIfNeeded(nodeSpec);
-                if (imageBeingDownloaded != null) {
+                if (isDownloadingImage()) {
                     addDebugMessage("Waiting for image to download " + imageBeingDownloaded.asString());
                     return;
                 }
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 12479db8ba3..82d38da6478 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
@@ -93,6 +93,7 @@ public class NodeAgentImplTest {
         final long rebootGeneration = 0;
         final ContainerNodeSpec nodeSpec = nodeSpecBuilder
                 .wantedDockerImage(dockerImage)
+                .currentDockerImage(dockerImage)
                 .nodeState(Node.State.active)
                 .wantedVespaVersion(vespaVersion)
                 .vespaVersion(vespaVersion)
@@ -103,7 +104,6 @@ public class NodeAgentImplTest {
 
         NodeAgentImpl nodeAgent = makeNodeAgent(dockerImage, true);
         when(nodeRepository.getContainerNodeSpec(hostName)).thenReturn(Optional.of(nodeSpec));
-        when(dockerOperations.shouldScheduleDownloadOfImage(any())).thenReturn(false);
 
         nodeAgent.converge();
 
@@ -133,7 +133,6 @@ public class NodeAgentImplTest {
                 .wantedDockerImage(dockerImage)
                 .nodeState(Node.State.active)
                 .wantedVespaVersion(vespaVersion)
-                .vespaVersion(vespaVersion)
                 .wantedRestartGeneration(restartGeneration)
                 .currentRestartGeneration(restartGeneration)
                 .wantedRebootGeneration(rebootGeneration)
@@ -142,7 +141,6 @@ public class NodeAgentImplTest {
         NodeAgentImpl nodeAgent = makeNodeAgent(null, false);
 
         when(nodeRepository.getContainerNodeSpec(hostName)).thenReturn(Optional.of(nodeSpec));
-        when(dockerOperations.shouldScheduleDownloadOfImage(any())).thenReturn(false);
         when(pathResolver.getApplicationStoragePathForNodeAdmin()).thenReturn(Files.createTempDirectory("foo"));
 
         nodeAgent.converge();
@@ -171,7 +169,9 @@ public class NodeAgentImplTest {
         final long currentRestartGeneration = 1;
         final ContainerNodeSpec nodeSpec = nodeSpecBuilder
                 .wantedDockerImage(newDockerImage)
+                .currentDockerImage(dockerImage)
                 .nodeState(Node.State.active)
+                .wantedVespaVersion(vespaVersion)
                 .vespaVersion(vespaVersion)
                 .wantedRestartGeneration(wantedRestartGeneration)
                 .currentRestartGeneration(currentRestartGeneration)
@@ -190,7 +190,7 @@ public class NodeAgentImplTest {
 
         final InOrder inOrder = inOrder(dockerOperations);
         inOrder.verify(dockerOperations, times(1)).shouldScheduleDownloadOfImage(eq(newDockerImage));
-        inOrder.verify(dockerOperations, times(1)).scheduleDownloadOfImage(eq(containerName), eq(nodeSpec), any());
+        inOrder.verify(dockerOperations, times(1)).scheduleDownloadOfImage(eq(containerName), eq(newDockerImage), any());
     }
 
     @Test
@@ -199,7 +199,9 @@ public class NodeAgentImplTest {
         final long currentRestartGeneration = 1;
         final ContainerNodeSpec nodeSpec = nodeSpecBuilder
                 .wantedDockerImage(dockerImage)
+                .currentDockerImage(dockerImage)
                 .nodeState(Node.State.active)
+                .wantedVespaVersion(vespaVersion)
                 .vespaVersion(vespaVersion)
                 .wantedRestartGeneration(wantedRestartGeneration)
                 .currentRestartGeneration(currentRestartGeneration)
@@ -207,8 +209,6 @@ public class NodeAgentImplTest {
 
         NodeAgentImpl nodeAgent = makeNodeAgent(dockerImage, true);
 
-        when(dockerOperations.shouldScheduleDownloadOfImage(any())).thenReturn(false);
-
         try {
             nodeAgent.converge();
             fail("Expected to throw an exception");
@@ -225,6 +225,7 @@ public class NodeAgentImplTest {
         final long rebootGeneration = 0;
         final ContainerNodeSpec nodeSpec = nodeSpecBuilder
                 .wantedDockerImage(dockerImage)
+                .currentDockerImage(dockerImage)
                 .nodeState(Node.State.failed)
                 .wantedVespaVersion(vespaVersion)
                 .vespaVersion(vespaVersion)
@@ -287,6 +288,7 @@ public class NodeAgentImplTest {
 
         final ContainerNodeSpec nodeSpec = nodeSpecBuilder
                 .wantedDockerImage(dockerImage)
+                .currentDockerImage(dockerImage)
                 .nodeState(Node.State.inactive)
                 .wantedVespaVersion(vespaVersion)
                 .vespaVersion(vespaVersion)
@@ -321,6 +323,7 @@ public class NodeAgentImplTest {
 
         final ContainerNodeSpec nodeSpec = nodeSpecBuilder
                 .wantedDockerImage(dockerImage)
+                .currentDockerImage(dockerImage)
                 .nodeState(nodeState)
                 .build();
 
@@ -375,6 +378,7 @@ public class NodeAgentImplTest {
     @Test
     public void testRestartDeadContainerAfterNodeAdminRestart() throws IOException {
         final ContainerNodeSpec nodeSpec = nodeSpecBuilder
+                .currentDockerImage(dockerImage)
                 .wantedDockerImage(dockerImage)
                 .nodeState(Node.State.active)
                 .vespaVersion(vespaVersion)
@@ -383,7 +387,6 @@ public class NodeAgentImplTest {
         NodeAgentImpl nodeAgent = makeNodeAgent(dockerImage, false);
 
         when(nodeRepository.getContainerNodeSpec(eq(hostName))).thenReturn(Optional.of(nodeSpec));
-        when(dockerOperations.shouldScheduleDownloadOfImage(eq(dockerImage))).thenReturn(false);
         when(pathResolver.getApplicationStoragePathForNodeAdmin()).thenReturn(Files.createTempDirectory("foo"));
 
         nodeAgent.tick();
@@ -397,6 +400,7 @@ public class NodeAgentImplTest {
         final long restartGeneration = 1;
         final ContainerNodeSpec nodeSpec = nodeSpecBuilder
                 .wantedDockerImage(dockerImage)
+                .currentDockerImage(dockerImage)
                 .nodeState(Node.State.active)
                 .vespaVersion(vespaVersion)
                 .wantedRestartGeneration(restartGeneration)
@@ -406,7 +410,6 @@ public class NodeAgentImplTest {
         NodeAgentImpl nodeAgent = makeNodeAgent(dockerImage, true);
 
         when(nodeRepository.getContainerNodeSpec(eq(hostName))).thenReturn(Optional.of(nodeSpec));
-        when(dockerOperations.shouldScheduleDownloadOfImage(eq(dockerImage))).thenReturn(false);
 
         final InOrder inOrder = inOrder(orchestrator, dockerOperations, nodeRepository);
         doThrow(new RuntimeException("Failed 1st time"))
@@ -479,6 +482,7 @@ public class NodeAgentImplTest {
         ContainerNodeSpec.Membership membership = new ContainerNodeSpec.Membership("clustType", "clustId", "grp", 3, false);
         final ContainerNodeSpec nodeSpec = nodeSpecBuilder
                 .wantedDockerImage(dockerImage)
+                .currentDockerImage(dockerImage)
                 .nodeState(Node.State.active)
                 .vespaVersion(vespaVersion)
                 .owner(owner)
@@ -489,7 +493,6 @@ public class NodeAgentImplTest {
 
         when(nodeRepository.getContainerNodeSpec(eq(hostName))).thenReturn(Optional.of(nodeSpec));
         when(storageMaintainer.updateIfNeededAndGetDiskMetricsFor(eq(containerName))).thenReturn(Optional.of(42547019776L));
-        when(dockerOperations.shouldScheduleDownloadOfImage(eq(dockerImage))).thenReturn(false);
         when(dockerOperations.getContainerStats(eq(containerName)))
                 .thenReturn(Optional.of(stats1))
                 .thenReturn(Optional.of(stats2));
@@ -521,7 +524,6 @@ public class NodeAgentImplTest {
         NodeAgentImpl nodeAgent = makeNodeAgent(null, false);
 
         when(nodeRepository.getContainerNodeSpec(eq(hostName))).thenReturn(Optional.of(nodeSpec));
-        when(dockerOperations.shouldScheduleDownloadOfImage(eq(dockerImage))).thenReturn(false);
         when(dockerOperations.getContainerStats(eq(containerName))).thenReturn(Optional.empty());
 
         nodeAgent.converge(); // Run the converge loop once to initialize lastNodeSpec