Create persistence for infrastructure versions

5 files changed, 217 insertions, 3 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureVersions.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureVersions.java
new file mode 100644
index 00000000000..b90ec492e1a
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureVersions.java
@@ -0,0 +1,63 @@
+package com.yahoo.vespa.hosted.provision.maintenance;
+
+import com.yahoo.component.Version;
+import com.yahoo.config.provision.NodeType;
+import com.yahoo.vespa.curator.Lock;
+import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.Optional;
+import java.util.logging.Logger;
+
+/**
+ * Multithread safe class to see and set target versions for infrastructure node types.
+ * {@link InfrastructureProvisioner} maintainer will then allocate the nodes of given node type
+ * with a wantedVersionVersion equal to the given target version.
+ *
+ * @author freva
+ */
+public class InfrastructureVersions {
+
+    private static Logger logger = Logger.getLogger(InfrastructureVersions.class.getName());
+
+    private final CuratorDatabaseClient db;
+
+    public InfrastructureVersions(CuratorDatabaseClient db) {
+        this.db = db;
+    }
+
+    public void setTargetVersion(NodeType nodeType, Version version, boolean force) {
+        if (nodeType != NodeType.config && nodeType != NodeType.confighost && nodeType != NodeType.proxyhost) {
+            throw new IllegalArgumentException("Cannot set version for type " + nodeType);
+        }
+
+        try (Lock lock = db.lockInfrastructureVersions()) {
+            Map<NodeType, Version> infrastructureVersions = db.readInfrastructureVersions();
+            Optional<Version> currentWantedVersion = Optional.ofNullable(infrastructureVersions.get(nodeType));
+
+            // Trying to set the version to the current version, skip
+            if (currentWantedVersion.equals(Optional.of(version))) return;
+
+            // If we don't force the set, we must set the new version to higher than the already set version
+            if (!force) {
+                if (currentWantedVersion.map(curVersion -> curVersion.compareTo(version) > 0).orElse(false))
+                    throw new IllegalArgumentException(String.format("Cannot downgrade version without setting 'force'. " +
+                            "Current wanted version: %s, attempted to set wanted version: %s",
+                            currentWantedVersion.get().toFullString(), version.toFullString()));
+            }
+
+            infrastructureVersions.put(nodeType, version);
+            db.writeInfrastructureVersions(infrastructureVersions);
+            logger.info("Set target version for " + nodeType + " to " + version.toFullString());
+        }
+    }
+
+    public Optional<Version> getTargetVersionFor(NodeType nodeType) {
+        return Optional.ofNullable(db.readInfrastructureVersions().get(nodeType));
+    }
+
+    public Map<NodeType, Version> getTargetVersions() {
+        return Collections.unmodifiableMap(db.readInfrastructureVersions());
+    }
+}
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 7fc305d397f..44eda5e83ec 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
@@ -2,9 +2,11 @@
 package com.yahoo.vespa.hosted.provision.persistence;
 
 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.NodeFlavors;
+import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.Zone;
 import com.yahoo.log.LogLevel;
 import com.yahoo.path.Path;
@@ -22,6 +24,7 @@ import java.time.Clock;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -70,6 +73,7 @@ public class CuratorDatabaseClient {
         for (Node.State state : Node.State.values())
             curatorDatabase.create(toPath(state));
         curatorDatabase.create(inactiveJobsPath());
+        curatorDatabase.create(infrastructureVersionsPath());
     }
 
     /**
@@ -363,5 +367,27 @@ public class CuratorDatabaseClient {
     private Path inactiveJobsPath() {
         return root.append("inactiveJobs");
     }
-    
+
+
+    public Map<NodeType, Version> readInfrastructureVersions() {
+        byte[] data = curatorDatabase.getData(infrastructureVersionsPath()).get();
+        if (data.length == 0) return new HashMap<>(); // infrastructure versions have never been written
+        return InfrastructureVersionsSerializer.fromJson(data);
+    }
+
+    public void writeInfrastructureVersions(Map<NodeType, Version> infrastructureVersions) {
+        NestedTransaction transaction = new NestedTransaction();
+        CuratorTransaction curatorTransaction = curatorDatabase.newCuratorTransactionIn(transaction);
+        curatorTransaction.add(CuratorOperations.setData(infrastructureVersionsPath().getAbsolute(),
+                InfrastructureVersionsSerializer.toJson(infrastructureVersions)));
+        transaction.commit();
+    }
+
+    public Lock lockInfrastructureVersions() {
+        return lock(root.append("locks").append("infrastructureVersionsLock"), defaultLockTimeout);
+    }
+
+    private Path infrastructureVersionsPath() {
+        return root.append("infrastructureVersions");
+    }
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/InfrastructureVersionsSerializer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/InfrastructureVersionsSerializer.java
new file mode 100644
index 00000000000..a48888fb4f0
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/InfrastructureVersionsSerializer.java
@@ -0,0 +1,42 @@
+// Copyright 2018 Yahoo Holdings. 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.component.Version;
+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.util.HashMap;
+import java.util.Map;
+
+/**
+ * @author freva
+ */
+class InfrastructureVersionsSerializer {
+
+    private InfrastructureVersionsSerializer() {}
+
+    static byte[] toJson(Map<NodeType, Version> versionsByNodeType) {
+        try {
+            Slime slime = new Slime();
+            Cursor object = slime.setObject();
+            versionsByNodeType.forEach((nodeType, version) ->
+                    object.setString(NodeSerializer.toString(nodeType), version.toFullString()));
+            return SlimeUtils.toJsonBytes(slime);
+        } catch (IOException e) {
+            throw new RuntimeException("Serialization of a infrastructure version failed", e);
+        }
+    }
+
+    static Map<NodeType, Version> fromJson(byte[] data) {
+        Map<NodeType, Version> infrastructureVersions = new HashMap<>();
+        Inspector inspector = SlimeUtils.jsonToSlime(data).get();
+        inspector.traverse((ObjectTraverser) (key, value) ->
+                infrastructureVersions.put(NodeSerializer.nodeTypeFromString(key), Version.fromString(value.asString())));
+        return infrastructureVersions;
+    }
+}
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java
index d164252b018..4d753599c4e 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/NodeSerializer.java
@@ -317,7 +317,7 @@ public class NodeSerializer {
         throw new IllegalArgumentException("Serialized form of '" + agent + "' not defined");
     }
 
-    private NodeType nodeTypeFromString(String typeString) {
+    static NodeType nodeTypeFromString(String typeString) {
         switch (typeString) {
             case "tenant" : return NodeType.tenant;
             case "host" : return NodeType.host;
@@ -329,7 +329,7 @@ public class NodeSerializer {
         }
     }
 
-    private String toString(NodeType type) {
+    static String toString(NodeType type) {
         switch (type) {
             case tenant: return "tenant";
             case host: return "host";
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureVersionsTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureVersionsTest.java
new file mode 100644
index 00000000000..37b0097b17c
--- /dev/null
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureVersionsTest.java
@@ -0,0 +1,83 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.maintenance;
+
+import com.yahoo.component.Version;
+import com.yahoo.config.provision.NodeType;
+import com.yahoo.vespa.hosted.provision.NodeRepositoryTester;
+import org.junit.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+
+import static junit.framework.TestCase.fail;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * @author freva
+ */
+public class InfrastructureVersionsTest {
+
+    private final NodeRepositoryTester tester = new NodeRepositoryTester();
+    private final InfrastructureVersions infrastructureVersions =
+            new InfrastructureVersions(tester.nodeRepository().database());
+
+    private final Version version = Version.fromString("6.123.456");
+
+    @Test
+    public void can_only_downgrade_with_force() {
+        assertTrue(infrastructureVersions.getTargetVersions().isEmpty());
+
+        assertEquals(Optional.empty(), infrastructureVersions.getTargetVersionFor(NodeType.config));
+        infrastructureVersions.setTargetVersion(NodeType.config, version, false);
+        assertEquals(Optional.of(version), infrastructureVersions.getTargetVersionFor(NodeType.config));
+
+        // Upgrading to new version without force is fine
+        Version new_version = Version.fromString("6.123.457"); // version + 1
+        infrastructureVersions.setTargetVersion(NodeType.config, new_version, false);
+        assertEquals(Optional.of(new_version), infrastructureVersions.getTargetVersionFor(NodeType.config));
+
+        // Downgrading to old version without force fails
+        try {
+            infrastructureVersions.setTargetVersion(NodeType.config, version, false);
+            fail("Should not be able to downgrade without force");
+        } catch (IllegalArgumentException ignored) { }
+
+        infrastructureVersions.setTargetVersion(NodeType.config, version, true);
+        assertEquals(Optional.of(version), infrastructureVersions.getTargetVersionFor(NodeType.config));
+    }
+
+    @Test
+    public void can_only_set_version_on_certain_node_types() {
+        // We can set version for config
+        infrastructureVersions.setTargetVersion(NodeType.config, version, false);
+
+        try {
+            infrastructureVersions.setTargetVersion(NodeType.tenant, version, false);
+            fail("Should not be able to set version for tenant nodes");
+        } catch (IllegalArgumentException ignored) { }
+
+        try {
+            // Using 'force' does not help, force only applies to version downgrade
+            infrastructureVersions.setTargetVersion(NodeType.tenant, version, true);
+            fail("Should not be able to set version for tenant nodes");
+        } catch (IllegalArgumentException ignored) { }
+    }
+
+    @Test
+    public void can_store_multiple_versions() {
+        Version version2 = Version.fromString("6.456.123");
+
+        infrastructureVersions.setTargetVersion(NodeType.config, version, false);
+        infrastructureVersions.setTargetVersion(NodeType.confighost, version2, false);
+        infrastructureVersions.setTargetVersion(NodeType.proxyhost, version, false);
+
+        Map<NodeType, Version> expected = new HashMap<>();
+        expected.put(NodeType.config, version);
+        expected.put(NodeType.confighost, version2);
+        expected.put(NodeType.proxyhost, version);
+
+        assertEquals(expected, infrastructureVersions.getTargetVersions());
+    }
+}
\ No newline at end of file