Use image for child node type if set

3 files changed, 36 insertions, 4 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerImages.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerImages.java
index 4beab8a3e80..c17af890b40 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerImages.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerImages.java
@@ -16,7 +16,7 @@ import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
 
 /**
- * Multithread safe class to get and set docker images for given host types.
+ * Multithread safe class to get and set docker images for given node types.
  *
  * @author freva
  */
@@ -57,9 +57,30 @@ public class DockerImages {
         return dockerImages.get();
     }
 
-    /** Returns the current docker image for given node type, or default */
+    /** Returns the current docker image for given node type, or the type for corresponding child nodes
+     * if it is a Docker host, or default */
     public DockerImage dockerImageFor(NodeType type) {
-        return getDockerImages().getOrDefault(type, defaultImage);
+        NodeType typeToUseForLookup = type;
+        if (type.isDockerHost()) {
+            switch (type) {
+                case confighost:
+                    typeToUseForLookup = NodeType.config;
+                    break;
+                case controllerhost:
+                    typeToUseForLookup = NodeType.controller;
+                    break;
+                case host:
+                    typeToUseForLookup = NodeType.tenant;
+                    break;
+                case proxyhost:
+                    typeToUseForLookup = NodeType.proxy;
+                    break;
+                default:
+                    break; // do not change
+            }
+        }
+
+        return getDockerImages().getOrDefault(typeToUseForLookup, defaultImage);
     }
 
     /** Set the docker image for nodes of given type */
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesResponse.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesResponse.java
index 248cfbec662..a4412a502aa 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesResponse.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesResponse.java
@@ -219,7 +219,7 @@ class NodesResponse extends HttpResponse {
         object.setString("storageType", serializer.toString(resources.storageType()));
     }
 
-    // Hack: For non-docker noder, return current docker image as default prefix + current Vespa version
+    // Hack: For non-docker nodes, return current docker image as default prefix + current Vespa version
     // TODO: Remove current + wanted docker image from response for non-docker types
     private Optional<DockerImage> currentDockerImage(Node node) {
         return node.status().dockerImage()
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DockerImagesTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DockerImagesTest.java
index 2eca8998931..cd6ae587b04 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DockerImagesTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DockerImagesTest.java
@@ -7,6 +7,8 @@ import com.yahoo.config.provision.NodeType;
 import com.yahoo.vespa.flags.InMemoryFlagSource;
 import org.junit.Test;
 
+import java.util.Optional;
+
 import static org.junit.Assert.assertEquals;
 
 /**
@@ -19,6 +21,9 @@ public class DockerImagesTest {
         var flagSource = new InMemoryFlagSource();
         var tester = new ProvisioningTester.Builder().flagSource(flagSource).build();
 
+        var proxyImage = DockerImage.fromString("docker-registry.domain.tld:8080/dist/proxy");
+        tester.nodeRepository().dockerImages().setDockerImage(NodeType.proxy, Optional.of(proxyImage));
+
         // Host uses tenant default image (for preload purposes)
         var defaultImage = DockerImage.fromString("docker-registry.domain.tld:8080/dist/vespa");
         var hosts = tester.makeReadyNodes(2, "default", NodeType.host);
@@ -35,6 +40,12 @@ public class DockerImagesTest {
                 assertEquals(defaultImage, tester.nodeRepository().dockerImages().dockerImageFor(node.type()));
             }
         }
+
+        // Proxy host uses image used by child nodes (proxy nodes), which is overridden in this case (for preload purposes)
+        var proxyHosts = tester.makeReadyNodes(2, "default", NodeType.proxyhost);
+        for (var host : proxyHosts) {
+            assertEquals(proxyImage, tester.nodeRepository().dockerImages().dockerImageFor(host.type()));
+        }
     }
 
 }