Merge pull request #27994 from vespa-engine/freva/download-image-in-reserved

Start downloading container image while in reserved

4 files changed, 28 insertions, 29 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/container/ContainerOperations.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/container/ContainerOperations.java
index 264035b86a1..fa933e9622a 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/container/ContainerOperations.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/container/ContainerOperations.java
@@ -36,7 +36,7 @@ public class ContainerOperations {
 
     public ContainerOperations(ContainerEngine containerEngine, Cgroup cgroup, FileSystem fileSystem, Timer timer) {
         this.containerEngine = Objects.requireNonNull(containerEngine);
-        this.imageDownloader = new ContainerImageDownloader(containerEngine);
+        this.imageDownloader = new ContainerImageDownloader(containerEngine, timer);
         this.imagePruner = new ContainerImagePruner(containerEngine, timer);
         this.containerStatsCollector = new ContainerStatsCollector(containerEngine, cgroup, fileSystem);
     }
@@ -62,8 +62,8 @@ public class ContainerOperations {
     }
 
     /** Pull image asynchronously. Returns true if image is still downloading and false if download is complete */
-    public boolean pullImageAsyncIfNeeded(TaskContext context, DockerImage dockerImage, RegistryCredentials registryCredentials) {
-        return !imageDownloader.get(context, dockerImage, registryCredentials);
+    public boolean pullImageAsyncIfNeeded(TaskContext context, DockerImage dockerImage, RegistryCredentialsProvider credentialsProvider) {
+        return !imageDownloader.get(context, dockerImage, credentialsProvider);
     }
 
     /** Executes a command inside container identified by given context. Does NOT throw on non-zero exit code */
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/container/image/ContainerImageDownloader.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/container/image/ContainerImageDownloader.java
index 1e37e080528..d3327bf5148 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/container/image/ContainerImageDownloader.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/container/image/ContainerImageDownloader.java
@@ -3,10 +3,13 @@ package com.yahoo.vespa.hosted.node.admin.container.image;
 
 import com.yahoo.concurrent.DaemonThreadFactory;
 import com.yahoo.config.provision.DockerImage;
+import com.yahoo.jdisc.Timer;
 import com.yahoo.vespa.hosted.node.admin.component.TaskContext;
 import com.yahoo.vespa.hosted.node.admin.container.ContainerEngine;
-import com.yahoo.vespa.hosted.node.admin.container.RegistryCredentials;
+import com.yahoo.vespa.hosted.node.admin.container.RegistryCredentialsProvider;
 
+import java.time.Duration;
+import java.time.Instant;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.Objects;
@@ -26,13 +29,15 @@ public class ContainerImageDownloader {
     private static final Logger LOG = Logger.getLogger(ContainerImageDownloader.class.getName());
 
     private final ContainerEngine containerEngine;
+    private final Timer timer;
 
     private final ExecutorService executorService = Executors.newSingleThreadExecutor(
             new DaemonThreadFactory("container-image-downloader")); // Download one image at a time
     private final Set<DockerImage> pendingDownloads = Collections.synchronizedSet(new HashSet<>());
 
-    public ContainerImageDownloader(ContainerEngine containerEngine) {
+    public ContainerImageDownloader(ContainerEngine containerEngine, Timer timer) {
         this.containerEngine = Objects.requireNonNull(containerEngine);
+        this.timer = Objects.requireNonNull(timer);
     }
 
     /**
@@ -40,12 +45,14 @@ public class ContainerImageDownloader {
      *
      * @return true if the image download has completed.
      */
-    public boolean get(TaskContext context, DockerImage image, RegistryCredentials registryCredentials) {
+    public boolean get(TaskContext context, DockerImage image, RegistryCredentialsProvider credentialsProvider) {
         if (pendingDownloads.contains(image)) return false;
         if (containerEngine.hasImage(context, image)) return true;
         executorService.submit(() -> {
             try {
-                containerEngine.pullImage(context, image, registryCredentials);
+                Instant start = timer.currentTime();
+                containerEngine.pullImage(context, image, credentialsProvider.get());
+                LOG.log(Level.INFO, "Downloaded container image " + image + " in " + Duration.between(start, timer.currentTime()));
             } catch (RuntimeException e) {
                 LOG.log(Level.SEVERE, "Failed to download container image " + image, e);
             } finally {
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 4c17bfbe039..c7c2686fde7 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
@@ -21,7 +21,6 @@ import com.yahoo.vespa.hosted.node.admin.configserver.orchestrator.OrchestratorE
 import com.yahoo.vespa.hosted.node.admin.container.Container;
 import com.yahoo.vespa.hosted.node.admin.container.ContainerOperations;
 import com.yahoo.vespa.hosted.node.admin.container.ContainerResources;
-import com.yahoo.vespa.hosted.node.admin.container.RegistryCredentials;
 import com.yahoo.vespa.hosted.node.admin.container.RegistryCredentialsProvider;
 import com.yahoo.vespa.hosted.node.admin.maintenance.ContainerWireguardTask;
 import com.yahoo.vespa.hosted.node.admin.maintenance.StorageMaintainer;
@@ -431,9 +430,8 @@ public class NodeAgentImpl implements NodeAgent {
         NodeSpec node = context.node();
         if (node.wantedDockerImage().equals(container.map(c -> c.image()))) return false;
 
-        RegistryCredentials credentials = registryCredentialsProvider.get();
         return node.wantedDockerImage()
-                   .map(image -> containerOperations.pullImageAsyncIfNeeded(context, image, credentials))
+                   .map(image -> containerOperations.pullImageAsyncIfNeeded(context, image, registryCredentialsProvider))
                    .orElse(false);
     }
 
@@ -487,17 +485,14 @@ public class NodeAgentImpl implements NodeAgent {
         }
 
         switch (node.state()) {
-            case ready:
-            case reserved:
-            case failed:
-            case inactive:
-            case parked:
+            case ready, reserved, failed, inactive, parked -> {
                 storageMaintainer.syncLogs(context, true);
+                if (node.state() == NodeState.reserved) downloadImageIfNeeded(context, container);
                 removeContainerIfNeededUpdateContainerState(context, container);
                 updateNodeRepoWithCurrentAttributes(context, Optional.empty());
                 stopServicesIfNeeded(context);
-                break;
-            case active:
+            }
+            case active -> {
                 storageMaintainer.syncLogs(context, true);
                 storageMaintainer.cleanDiskIfFull(context);
                 storageMaintainer.handleCoreDumpsForContainer(context, container, false);
@@ -550,11 +545,9 @@ public class NodeAgentImpl implements NodeAgent {
                     orchestrator.resume(context.hostname().value());
                     suspendedInOrchestrator = false;
                 }
-                break;
-            case provisioned:
-                nodeRepository.setNodeState(context.hostname().value(), NodeState.ready);
-                break;
-            case dirty:
+            }
+            case provisioned -> nodeRepository.setNodeState(context.hostname().value(), NodeState.ready);
+            case dirty -> {
                 removeContainerIfNeededUpdateContainerState(context, container);
                 context.log(logger, "State is " + node.state() + ", will delete application storage and mark node as ready");
                 credentialsMaintainers.forEach(maintainer -> maintainer.clearCredentials(context));
@@ -562,9 +555,7 @@ public class NodeAgentImpl implements NodeAgent {
                 storageMaintainer.archiveNodeStorage(context);
                 updateNodeRepoWithCurrentAttributes(context, Optional.empty());
                 nodeRepository.setNodeState(context.hostname().value(), NodeState.ready);
-                break;
-            default:
-                throw ConvergenceException.ofError("UNKNOWN STATE " + node.state().name());
+            }
         }
     }
 
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/container/image/ContainerImageDownloaderTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/container/image/ContainerImageDownloaderTest.java
index 9fd14e7e665..7f002eee315 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/container/image/ContainerImageDownloaderTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/container/image/ContainerImageDownloaderTest.java
@@ -2,6 +2,7 @@
 package com.yahoo.vespa.hosted.node.admin.container.image;
 
 import com.yahoo.config.provision.DockerImage;
+import com.yahoo.jdisc.test.TestTimer;
 import com.yahoo.vespa.hosted.node.admin.component.TaskContext;
 import com.yahoo.vespa.hosted.node.admin.component.TestTaskContext;
 import com.yahoo.vespa.hosted.node.admin.container.ContainerEngineMock;
@@ -21,15 +22,15 @@ public class ContainerImageDownloaderTest {
     @Timeout(5_000)
     void test_download() {
         ContainerEngineMock podman = new ContainerEngineMock().asyncImageDownload(true);
-        ContainerImageDownloader downloader = new ContainerImageDownloader(podman);
+        ContainerImageDownloader downloader = new ContainerImageDownloader(podman, new TestTimer());
         TaskContext context = new TestTaskContext();
         DockerImage image = DockerImage.fromString("registry.example.com/repo/vespa:7.42");
 
-        assertFalse(downloader.get(context, image, RegistryCredentials.none), "Download started");
-        assertFalse(downloader.get(context, image, RegistryCredentials.none), "Download pending");
+        assertFalse(downloader.get(context, image, () -> RegistryCredentials.none), "Download started");
+        assertFalse(downloader.get(context, image, () -> RegistryCredentials.none), "Download pending");
         podman.completeDownloadOf(image);
         boolean downloadCompleted;
-        while (!(downloadCompleted = downloader.get(context, image, RegistryCredentials.none))) ;
+        while (!(downloadCompleted = downloader.get(context, image, () -> RegistryCredentials.none))) ;
         assertTrue(downloadCompleted, "Download completed");
     }