Merge branch 'master' into bratseth/receive-wanted-version

3 files changed, 36 insertions, 14 deletions

diff --git a/config-provisioning/src/main/java/com/yahoo/config/provision/ProvisionInfo.java b/config-provisioning/src/main/java/com/yahoo/config/provision/ProvisionInfo.java
index e903332bd6f..c5ccb7d2f58 100644
--- a/config-provisioning/src/main/java/com/yahoo/config/provision/ProvisionInfo.java
+++ b/config-provisioning/src/main/java/com/yahoo/config/provision/ProvisionInfo.java
@@ -23,6 +23,7 @@ public class ProvisionInfo {
     private static final String hostSpecKey = "hostSpec";
     private static final String hostSpecHostName = "hostName";
     private static final String hostSpecMembership = "membership";
+    private static final String hostSpecFlavor = "flavor";
     private static final String dockerImage = "dockerImage";
 
     private final Set<HostSpec> hosts = new LinkedHashSet<>();
@@ -50,29 +51,33 @@ public class ProvisionInfo {
             if (host.membership().get().cluster().dockerImage().isPresent())
                 cursor.setString(dockerImage, host.membership().get().cluster().dockerImage().get());
         }
+        if (host.flavor().isPresent())
+            cursor.setString(hostSpecFlavor, host.flavor().get().name());
     }
 
     public Set<HostSpec> getHosts() {
         return Collections.unmodifiableSet(hosts);
     }
 
-    private static ProvisionInfo fromSlime(Inspector inspector) {
+    private static ProvisionInfo fromSlime(Inspector inspector, Optional<NodeFlavors> nodeFlavors) {
         Inspector array = inspector.field(mappingKey);
         final Set<HostSpec> hosts = new LinkedHashSet<>();
         array.traverse(new ArrayTraverser() {
             @Override
             public void entry(int i, Inspector inspector) {
-                hosts.add(createHostSpec(inspector.field(hostSpecKey)));
+                hosts.add(createHostSpec(inspector.field(hostSpecKey), nodeFlavors));
             }
         });
         return new ProvisionInfo(hosts);
     }
 
-    private static HostSpec createHostSpec(Inspector object) {
+    private static HostSpec createHostSpec(Inspector object, Optional<NodeFlavors> nodeFlavors) {
         Optional<ClusterMembership> membership =
-            object.field(hostSpecMembership).valid() ? Optional.of(readMembership(object)) : Optional.empty();
-        HostSpec h = new HostSpec(object.field(hostSpecHostName).asString(), Collections.emptyList(), Optional.empty(), membership);
-        return h;
+                object.field(hostSpecMembership).valid() ? Optional.of(readMembership(object)) : Optional.empty();
+        Optional<Flavor> flavor =
+                object.field(hostSpecFlavor).valid() ? readFlavor(object, nodeFlavors) : Optional.empty();
+
+        return new HostSpec(object.field(hostSpecHostName).asString(),Collections.emptyList(), flavor, membership);
     }
 
     private static ClusterMembership readMembership(Inspector object) {
@@ -80,14 +85,19 @@ public class ProvisionInfo {
                                       object.field(dockerImage).valid() ? Optional.of(object.field(dockerImage).asString()) : Optional.empty());
     }
 
+    private static Optional<Flavor> readFlavor(Inspector object, Optional<NodeFlavors> nodeFlavors) {
+        return nodeFlavors.map(flavorMapper ->  flavorMapper.getFlavor(object.field(hostSpecFlavor).asString()))
+                .orElse(Optional.empty());
+    }
+
     public byte[] toJson() throws IOException {
         Slime slime = new Slime();
         toSlime(slime.setObject());
         return SlimeUtils.toJsonBytes(slime);
     }
 
-    public static ProvisionInfo fromJson(byte[] json) {
-        return fromSlime(SlimeUtils.jsonToSlime(json).get());
+    public static ProvisionInfo fromJson(byte[] json, Optional<NodeFlavors> nodeFlavors) {
+        return fromSlime(SlimeUtils.jsonToSlime(json).get(), nodeFlavors);
     }
 
     public ProvisionInfo merge(ProvisionInfo provisionInfo) {
diff --git a/config-provisioning/src/main/java/com/yahoo/config/provision/Zone.java b/config-provisioning/src/main/java/com/yahoo/config/provision/Zone.java
index 6ff26d0aad7..6559e08d3ed 100644
--- a/config-provisioning/src/main/java/com/yahoo/config/provision/Zone.java
+++ b/config-provisioning/src/main/java/com/yahoo/config/provision/Zone.java
@@ -20,11 +20,15 @@ public class Zone {
     private final RegionName region;
     private final SystemName systemName;
     private final FlavorDefaults flavorDefaults;
+    private final Optional<NodeFlavors> nodeFlavors;
 
     @Inject
-    public Zone(ConfigserverConfig configserverConfig) {
-        this(SystemName.from(configserverConfig.system()), Environment.from(configserverConfig.environment()),
-                RegionName.from(configserverConfig.region()), new FlavorDefaults(configserverConfig));
+    public Zone(ConfigserverConfig configserverConfig, NodeFlavors nodeFlavors) {
+        this(SystemName.from(configserverConfig.system()),
+             Environment.from(configserverConfig.environment()),
+             RegionName.from(configserverConfig.region()),
+             new FlavorDefaults(configserverConfig),
+             nodeFlavors);
     }
 
     /** Create from environment and region */
@@ -39,14 +43,19 @@ public class Zone {
 
     /** Create from environment and region. Useful for testing. */
     public Zone(SystemName system, Environment environment, RegionName region, String defaultFlavor) {
-        this(system, environment, region, new FlavorDefaults(defaultFlavor));
+        this(system, environment, region, new FlavorDefaults(defaultFlavor), null);
     }
 
-    private Zone(SystemName systemName, Environment environment, RegionName region, FlavorDefaults flavorDefaults) {
+    private Zone(SystemName systemName,
+                 Environment environment,
+                 RegionName region,
+                 FlavorDefaults flavorDefaults,
+                 NodeFlavors nodeFlavors) {
         this.environment = environment;
         this.region = region;
         this.flavorDefaults = flavorDefaults;
         this.systemName = systemName;
+        this.nodeFlavors = Optional.ofNullable(nodeFlavors);
     }
 
     /** Returns the current environment */
@@ -61,6 +70,9 @@ public class Zone {
     /** Returns the default hardware flavor to assign in this zone */
     public String defaultFlavor(ClusterSpec.Type clusterType) { return flavorDefaults.flavor(clusterType); }
 
+    /** Returns all available node flavors for the zone, or empty if not set for this Zone. */
+    public Optional<NodeFlavors> nodeFlavors() { return nodeFlavors; }
+
     /** Do not use */
     public static Zone defaultZone() {
         return new Zone(SystemName.defaultSystem(), Environment.defaultEnvironment(), RegionName.defaultName());
diff --git a/config-provisioning/src/test/java/com/yahoo/config/provision/ProvisionInfoTest.java b/config-provisioning/src/test/java/com/yahoo/config/provision/ProvisionInfoTest.java
index 52b4f6853b8..a6ccdb8536f 100644
--- a/config-provisioning/src/test/java/com/yahoo/config/provision/ProvisionInfoTest.java
+++ b/config-provisioning/src/test/java/com/yahoo/config/provision/ProvisionInfoTest.java
@@ -49,7 +49,7 @@ public class ProvisionInfoTest {
     }
 
     private void assertProvisionInfo(ProvisionInfo info) throws IOException {
-        ProvisionInfo serializedInfo = ProvisionInfo.fromJson(info.toJson());
+        ProvisionInfo serializedInfo = ProvisionInfo.fromJson(info.toJson(), Optional.empty());
         assertEquals(info.getHosts().size(), serializedInfo.getHosts().size());
         assertTrue(serializedInfo.getHosts().contains(h1));
         assertTrue(serializedInfo.getHosts().contains(h2));