Set docker image prefix using /nodes/v2/upgrade API

5 files changed, 109 insertions, 8 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
index 41077e3b4ec..7cff20a20d1 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
@@ -28,6 +28,7 @@ import com.yahoo.vespa.hosted.provision.node.filter.StateFilter;
 import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
 import com.yahoo.vespa.hosted.provision.persistence.DnsNameResolver;
 import com.yahoo.vespa.hosted.provision.persistence.NameResolver;
+import com.yahoo.vespa.hosted.provision.provisioning.DockerImages;
 import com.yahoo.vespa.hosted.provision.provisioning.FirmwareChecks;
 import com.yahoo.vespa.hosted.provision.provisioning.OsVersions;
 import com.yahoo.vespa.hosted.provision.restapi.v2.NotFoundException;
@@ -83,9 +84,9 @@ public class NodeRepository extends AbstractComponent {
     private final Zone zone;
     private final NodeFlavors flavors;
     private final NameResolver nameResolver;
-    private final DockerImage dockerImage;
     private final OsVersions osVersions;
     private final FirmwareChecks firmwareChecks;
+    private final DockerImages dockerImages;
     private final Flags flags;
 
     /**
@@ -108,9 +109,9 @@ public class NodeRepository extends AbstractComponent {
         this.clock = clock;
         this.flavors = flavors;
         this.nameResolver = nameResolver;
-        this.dockerImage = dockerImage;
         this.osVersions = new OsVersions(this.db);
         this.firmwareChecks = new FirmwareChecks(db, clock);
+        this.dockerImages = new DockerImages(db, dockerImage);
         this.flags = new Flags(this.db);
 
         // read and write all nodes to make sure they are stored in the latest version of the serialized format
@@ -122,7 +123,7 @@ public class NodeRepository extends AbstractComponent {
     public CuratorDatabaseClient database() { return db; }
 
     /** Returns the Docker image to use for nodes in this */
-    public DockerImage dockerImage() { return dockerImage; }
+    public DockerImage dockerImage(NodeType nodeType) { return dockerImages.dockerImageFor(nodeType); }
 
     /** @return The name resolver used to resolve hostname and ip addresses */
     public NameResolver nameResolver() { return nameResolver; }
@@ -133,6 +134,9 @@ public class NodeRepository extends AbstractComponent {
     /** Returns the status of firmware checks for hosts managed by this. */
     public FirmwareChecks firmwareChecks() { return firmwareChecks; }
 
+    /** Returns the docker images to use for nodes in this. */
+    public DockerImages dockerImages() { return dockerImages; }
+
     /** Returns feature flags of this node repository */
     public Flags flags() {
         return flags;
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
new file mode 100644
index 00000000000..5e0e12dd62c
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerImages.java
@@ -0,0 +1,80 @@
+package com.yahoo.vespa.hosted.provision.provisioning;
+
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import com.google.common.collect.ImmutableMap;
+import com.yahoo.config.provision.DockerImage;
+import com.yahoo.config.provision.NodeType;
+import com.yahoo.vespa.curator.Lock;
+import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
+
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import java.util.logging.Logger;
+
+/**
+ * Multithread safe class to get and set docker images for given host types.
+ *
+ * @author freva
+ */
+public class DockerImages {
+
+    private static final Duration defaultCacheTtl = Duration.ofMinutes(1);
+    private static final Logger log = Logger.getLogger(DockerImages.class.getName());
+
+    private final CuratorDatabaseClient db;
+    private final DockerImage defaultImage;
+    private final Duration cacheTtl;
+
+    /**
+     * Docker image is read on every request to /nodes/v2/node/[fqdn]. Cache current getDockerImages to avoid
+     * unnecessary ZK reads. When getDockerImages change, some nodes may need to wait for TTL until they see the new target,
+     * this is fine.
+     */
+    private volatile Supplier<Map<NodeType, DockerImage>> dockerImages;
+
+    public DockerImages(CuratorDatabaseClient db, DockerImage defaultImage) {
+        this(db, defaultImage, defaultCacheTtl);
+    }
+
+    DockerImages(CuratorDatabaseClient db, DockerImage defaultImage, Duration cacheTtl) {
+        this.db = db;
+        this.defaultImage = defaultImage;
+        this.cacheTtl = cacheTtl;
+        createCache();
+    }
+
+    private void createCache() {
+        this.dockerImages = Suppliers.memoizeWithExpiration(() -> Collections.unmodifiableMap(db.readDockerImages()),
+                cacheTtl.toMillis(), TimeUnit.MILLISECONDS);
+    }
+
+    /** Returns the current docker images for each node type */
+    public Map<NodeType, DockerImage> getDockerImages() {
+        return dockerImages.get();
+    }
+
+    /** Returns the current docker image for given node type, or default */
+    public DockerImage dockerImageFor(NodeType type) {
+        return getDockerImages().getOrDefault(type, defaultImage);
+    }
+
+    /** Set the docker image for nodes of given type */
+    public void setDockerImage(NodeType nodeType, Optional<DockerImage> dockerImage) {
+        if (nodeType.isDockerHost()) {
+            throw new IllegalArgumentException("Setting docker image for " + nodeType + " nodes is unsupported");
+        }
+        try (Lock lock = db.lockDockerImages()) {
+            Map<NodeType, DockerImage> dockerImages = db.readDockerImages();
+
+            dockerImage.ifPresentOrElse(image -> dockerImages.put(nodeType, image), () -> dockerImages.remove(nodeType));
+            db.writeDockerImages(dockerImages);
+            createCache(); // Throw away current cache
+            log.info("Set docker image for " + nodeType + " nodes to " + dockerImage.map(DockerImage::asString).orElse(null));
+        }
+    }
+
+}
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesApiHandler.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesApiHandler.java
index 3ce24f73b2c..9c3e83b9f5a 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesApiHandler.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesApiHandler.java
@@ -3,6 +3,7 @@ package com.yahoo.vespa.hosted.provision.restapi.v2;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.HostFilter;
 import com.yahoo.config.provision.HostName;
 import com.yahoo.config.provision.NodeFlavors;
@@ -106,7 +107,7 @@ public class NodesApiHandler extends LoggingRequestHandler {
         if (path.startsWith("/nodes/v2/acl/")) return new NodeAclResponse(request, nodeRepository);
         if (path.equals(    "/nodes/v2/command/")) return ResourcesResponse.fromStrings(request.getUri(), "restart", "reboot");
         if (path.equals(    "/nodes/v2/maintenance/")) return new JobsResponse(maintenance.jobControl());
-        if (path.equals(    "/nodes/v2/upgrade/")) return new UpgradeResponse(maintenance.infrastructureVersions(), nodeRepository.osVersions());
+        if (path.equals(    "/nodes/v2/upgrade/")) return new UpgradeResponse(maintenance.infrastructureVersions(), nodeRepository.osVersions(), nodeRepository.dockerImages());
         if (path.equals(    "/nodes/v2/flags/")) return new FlagsResponse(nodeRepository.flags().list());
         throw new NotFoundException("Nothing at path '" + path + "'");
     }
@@ -308,6 +309,7 @@ public class NodesApiHandler extends LoggingRequestHandler {
         boolean force = inspector.field("force").asBool();
         Inspector versionField = inspector.field("version");
         Inspector osVersionField = inspector.field("osVersion");
+        Inspector dockerImageField = inspector.field("dockerImage");
 
         if (versionField.valid()) {
             Version version = Version.fromString(versionField.asString());
@@ -327,8 +329,16 @@ public class NodesApiHandler extends LoggingRequestHandler {
             }
         }
 
+        if (dockerImageField.valid()) {
+            Optional<DockerImage> dockerImage = Optional.of(dockerImageField.asString())
+                    .filter(s -> !s.isEmpty())
+                    .map(DockerImage::fromString);
+            nodeRepository.dockerImages().setDockerImage(nodeType, dockerImage);
+            messageParts.add("docker image to " + dockerImage.map(DockerImage::asString).orElse(null));
+        }
+
         if (messageParts.isEmpty()) {
-            throw new IllegalArgumentException("At least one of 'version' and 'osVersion' must be set");
+            throw new IllegalArgumentException("At least one of 'version', 'osVersion' or 'dockerImage' must be set");
         }
 
         return new MessageResponse("Set " + String.join(", ", messageParts) +
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 a0ecb063618..dfb6004e8c8 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
@@ -160,7 +160,7 @@ class NodesResponse extends HttpResponse {
             toSlime(allocation.membership(), object.setObject("membership"));
             object.setLong("restartGeneration", allocation.restartGeneration().wanted());
             object.setLong("currentRestartGeneration", allocation.restartGeneration().current());
-            object.setString("wantedDockerImage", nodeRepository.dockerImage().withTag(allocation.membership().cluster().vespaVersion()).asString());
+            object.setString("wantedDockerImage", nodeRepository.dockerImage(node.type()).withTag(allocation.membership().cluster().vespaVersion()).asString());
             object.setString("wantedVespaVersion", allocation.membership().cluster().vespaVersion().toFullString());
             allocation.networkPorts().ifPresent(ports -> NetworkPortsSerializer.toSlime(ports, object.setArray("networkPorts")));
             orchestrator.apply(new HostName(node.hostname()))
@@ -219,7 +219,7 @@ class NodesResponse extends HttpResponse {
                 .or(() -> Optional.of(node)
                         .filter(n -> n.type().isDockerHost())
                         .flatMap(n -> n.status().vespaVersion()
-                                .map(version -> nodeRepository.dockerImages().dockerImage().withTag(version))));
+                                .map(version -> nodeRepository.dockerImages().dockerImageFor(n.type()).withTag(version))));
     }
 
 
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/UpgradeResponse.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/UpgradeResponse.java
index 392cba7baa9..87d7944f040 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/UpgradeResponse.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/UpgradeResponse.java
@@ -6,6 +6,7 @@ import com.yahoo.slime.Cursor;
 import com.yahoo.slime.JsonFormat;
 import com.yahoo.slime.Slime;
 import com.yahoo.vespa.hosted.provision.maintenance.InfrastructureVersions;
+import com.yahoo.vespa.hosted.provision.provisioning.DockerImages;
 import com.yahoo.vespa.hosted.provision.provisioning.OsVersions;
 
 import java.io.IOException;
@@ -20,11 +21,13 @@ public class UpgradeResponse extends HttpResponse {
 
     private final InfrastructureVersions infrastructureVersions;
     private final OsVersions osVersions;
+    private final DockerImages dockerImages;
 
-    public UpgradeResponse(InfrastructureVersions infrastructureVersions, OsVersions osVersions) {
+    public UpgradeResponse(InfrastructureVersions infrastructureVersions, OsVersions osVersions, DockerImages dockerImages) {
         super(200);
         this.infrastructureVersions = infrastructureVersions;
         this.osVersions = osVersions;
+        this.dockerImages = dockerImages;
     }
 
     @Override
@@ -38,6 +41,10 @@ public class UpgradeResponse extends HttpResponse {
         Cursor osVersionsObject = root.setObject("osVersions");
         osVersions.targets().forEach((nodeType, version) -> osVersionsObject.setString(nodeType.name(), version.toFullString()));
 
+
+        Cursor dockerImagesObject = root.setObject("dockerImages");
+        dockerImages.getDockerImages().forEach((nodeType, image) -> dockerImagesObject.setString(nodeType.name(), image.asString()));
+
         new JsonFormat(true).encode(stream, slime);
     }