Store docker images by node type in ZK

2 files changed, 83 insertions, 1 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
index 0a0e5b7c265..3ed8ef9f64a 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
@@ -5,6 +5,7 @@ import com.google.common.util.concurrent.UncheckedTimeoutException;
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.ApplicationLockException;
+import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.HostName;
 import com.yahoo.config.provision.NodeFlavors;
 import com.yahoo.config.provision.NodeType;
@@ -93,6 +94,7 @@ public class CuratorDatabaseClient {
         curatorDatabase.create(inactiveJobsPath());
         curatorDatabase.create(infrastructureVersionsPath());
         curatorDatabase.create(osVersionsPath());
+        curatorDatabase.create(dockerImagesPath());
         curatorDatabase.create(firmwareCheckPath());
         curatorDatabase.create(loadBalancersRoot);
         curatorDatabase.create(flagsRoot);
@@ -364,6 +366,8 @@ public class CuratorDatabaseClient {
         return curatorDatabase.getData(path).filter(data -> data.length > 0).map(mapper);
     }
 
+
+    // Maintenance jobs
     public Set<String> readInactiveJobs() {
         try {
             return read(inactiveJobsPath(), stringSetSerializer::fromJson).orElseGet(HashSet::new);
@@ -391,6 +395,8 @@ public class CuratorDatabaseClient {
         return root.append("inactiveJobs");
     }
 
+
+    // Infrastructure versions
     public Map<NodeType, Version> readInfrastructureVersions() {
         return read(infrastructureVersionsPath(), NodeTypeVersionsSerializer::fromJson).orElseGet(TreeMap::new);
     }
@@ -411,6 +417,8 @@ public class CuratorDatabaseClient {
         return root.append("infrastructureVersions");
     }
 
+
+    // OS versions
     public Map<NodeType, Version> readOsVersions() {
         return read(osVersionsPath(), NodeTypeVersionsSerializer::fromJson).orElseGet(TreeMap::new);
     }
@@ -431,6 +439,30 @@ public class CuratorDatabaseClient {
         return root.append("osVersions");
     }
 
+
+    // Docker images
+    public Map<NodeType, DockerImage> readDockerImages() {
+        return read(dockerImagesPath(), NodeTypeDockerImagesSerializer::fromJson).orElseGet(TreeMap::new);
+    }
+
+    public void writeDockerImages(Map<NodeType, DockerImage> dockerImages) {
+        NestedTransaction transaction = new NestedTransaction();
+        CuratorTransaction curatorTransaction = curatorDatabase.newCuratorTransactionIn(transaction);
+        curatorTransaction.add(CuratorOperations.setData(dockerImagesPath().getAbsolute(),
+                NodeTypeDockerImagesSerializer.toJson(dockerImages)));
+        transaction.commit();
+    }
+
+    public Lock lockDockerImages() {
+        return lock(lockRoot.append("dockerImagesLock"), defaultLockTimeout);
+    }
+
+    private Path dockerImagesPath() {
+        return root.append("dockerImages");
+    }
+
+
+    // Firmware checks
     /** Stores the instant after which a firmware check is required, or clears any outstanding ones if empty is given. */
     public void writeFirmwareCheck(Optional<Instant> after) {
         byte[] data = after.map(instant -> Long.toString(instant.toEpochMilli()).getBytes())
@@ -450,6 +482,8 @@ public class CuratorDatabaseClient {
         return root.append("firmwareCheck");
     }
 
+
+    // Load balancers
     public Map<LoadBalancerId, LoadBalancer> readLoadBalancers() {
         return curatorDatabase.getChildren(loadBalancersRoot).stream()
                               .map(LoadBalancerId::fromSerializedForm)
@@ -460,7 +494,7 @@ public class CuratorDatabaseClient {
                                                          Collections::unmodifiableMap));
     }
 
-    public Optional<LoadBalancer> readLoadBalancer(LoadBalancerId id) {
+    private Optional<LoadBalancer> readLoadBalancer(LoadBalancerId id) {
         return read(loadBalancerPath(id), LoadBalancerSerializer::fromJson);
     }
 
@@ -501,6 +535,8 @@ public class CuratorDatabaseClient {
         transaction.commit();
     }
 
+
+    // Flags
     public Optional<Flag> readFlag(FlagId id) {
         return read(flagPath(id), FlagSerializer::fromJson);
     }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeTypeDockerImagesSerializer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeTypeDockerImagesSerializer.java
new file mode 100644
index 00000000000..37dc8a8a1ad
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeTypeDockerImagesSerializer.java
@@ -0,0 +1,46 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.persistence;
+
+import com.yahoo.config.provision.DockerImage;
+import com.yahoo.config.provision.NodeType;
+import com.yahoo.slime.Cursor;
+import com.yahoo.slime.Inspector;
+import com.yahoo.slime.ObjectTraverser;
+import com.yahoo.slime.Slime;
+import com.yahoo.vespa.config.SlimeUtils;
+
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.Map;
+import java.util.TreeMap;
+
+/**
+ * Serializer for docker images that are set per node type.
+ *
+ * @author freva
+ */
+public class NodeTypeDockerImagesSerializer {
+
+    private NodeTypeDockerImagesSerializer() {}
+
+    public static byte[] toJson(Map<NodeType, DockerImage> dockerImages) {
+        Slime slime = new Slime();
+        Cursor object = slime.setObject();
+        dockerImages.forEach((nodeType, dockerImage) ->
+                object.setString(NodeSerializer.toString(nodeType), dockerImage.asString()));
+        try {
+            return SlimeUtils.toJsonBytes(slime);
+        } catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
+    public static Map<NodeType, DockerImage> fromJson(byte[] data) {
+        Map<NodeType, DockerImage> dockerImages = new TreeMap<>(); // Use TreeMap to sort by node type
+        Inspector inspector = SlimeUtils.jsonToSlime(data).get();
+        inspector.traverse((ObjectTraverser) (key, value) ->
+                dockerImages.put(NodeSerializer.nodeTypeFromString(key), DockerImage.fromString(value.asString())));
+        return dockerImages;
+    }
+
+}