Add cluster type to ClusterCapacity and use it when provioning hosts in HostCapacityMaintainer

2 files changed, 41 insertions, 9 deletions

diff --git a/flags/src/main/java/com/yahoo/vespa/flags/custom/ClusterCapacity.java b/flags/src/main/java/com/yahoo/vespa/flags/custom/ClusterCapacity.java
index dcef85f9a0d..06c015e5206 100644
--- a/flags/src/main/java/com/yahoo/vespa/flags/custom/ClusterCapacity.java
+++ b/flags/src/main/java/com/yahoo/vespa/flags/custom/ClusterCapacity.java
@@ -12,6 +12,7 @@ import java.util.OptionalDouble;
 
 import static com.yahoo.vespa.flags.custom.Validation.requireNonNegative;
 import static com.yahoo.vespa.flags.custom.Validation.validArchitectures;
+import static com.yahoo.vespa.flags.custom.Validation.validClusterTypes;
 import static com.yahoo.vespa.flags.custom.Validation.validDiskSpeeds;
 import static com.yahoo.vespa.flags.custom.Validation.validStorageTypes;
 import static com.yahoo.vespa.flags.custom.Validation.validateEnum;
@@ -31,6 +32,7 @@ public class ClusterCapacity {
     private final String diskSpeed;
     private final String storageType;
     private final String architecture;
+    private final String clusterType;
 
     @JsonCreator
     public ClusterCapacity(@JsonProperty("count") int count,
@@ -40,7 +42,8 @@ public class ClusterCapacity {
                            @JsonProperty("bandwidthGbps") Double bandwidthGbps,
                            @JsonProperty("diskSpeed") String diskSpeed,
                            @JsonProperty("storageType") String storageType,
-                           @JsonProperty("architecture") String architecture) {
+                           @JsonProperty("architecture") String architecture,
+                           @JsonProperty("clusterType") String clusterType) {
         this.count = (int) requireNonNegative("count", count);
         this.vcpu = vcpu == null ? OptionalDouble.empty() : OptionalDouble.of(requireNonNegative("vcpu", vcpu));
         this.memoryGb = memoryGb == null ? OptionalDouble.empty() : OptionalDouble.of(requireNonNegative("memoryGb", memoryGb));
@@ -49,12 +52,13 @@ public class ClusterCapacity {
         this.diskSpeed = validateEnum("diskSpeed", validDiskSpeeds, diskSpeed == null ? "fast" : diskSpeed);
         this.storageType = validateEnum("storageType", validStorageTypes, storageType == null ? "any" : storageType);
         this.architecture = validateEnum("architecture", validArchitectures, architecture == null ? "x86_64" : architecture);
+        this.clusterType = clusterType == null ? null : validateEnum("clusterType", validClusterTypes, clusterType);
     }
 
     /** Returns a new ClusterCapacity equal to {@code this}, but with the given count. */
     public ClusterCapacity withCount(int count) {
         return new ClusterCapacity(count, vcpuOrNull(), memoryGbOrNull(), diskGbOrNull(), bandwidthGbpsOrNull(),
-                                   diskSpeed, storageType, architecture);
+                                   diskSpeed, storageType, architecture, clusterType);
     }
 
     @JsonGetter("count") public int count() { return count; }
@@ -73,6 +77,7 @@ public class ClusterCapacity {
     @JsonGetter("diskSpeed") public String diskSpeed() { return diskSpeed; }
     @JsonGetter("storageType") public String storageType() { return storageType; }
     @JsonGetter("architecture") public String architecture() { return architecture; }
+    @JsonGetter("clusterType") public String clusterType() { return clusterType; }
 
     @JsonIgnore public Double vcpu() { return vcpu.orElse(0.0); }
     @JsonIgnore public Double memoryGb() { return memoryGb.orElse(0.0); }
@@ -90,6 +95,7 @@ public class ClusterCapacity {
                 ", diskSpeed=" + diskSpeed +
                 ", storageType=" + storageType +
                 ", architecture=" + architecture +
+                ", clusterType=" + clusterType +
                 '}';
     }
 
@@ -105,12 +111,13 @@ public class ClusterCapacity {
                 bandwidthGbps.equals(that.bandwidthGbps) &&
                 diskSpeed.equals(that.diskSpeed) &&
                 storageType.equals(that.storageType) &&
-                architecture.equals(that.architecture);
+                architecture.equals(that.architecture) &&
+                clusterType.equals(that.clusterType);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(count, vcpu, memoryGb, diskGb, bandwidthGbps, diskSpeed, storageType, architecture);
+        return Objects.hash(count, vcpu, memoryGb, diskGb, bandwidthGbps, diskSpeed, storageType, architecture, clusterType);
     }
 
 }
diff --git a/flags/src/test/java/com/yahoo/vespa/flags/custom/ClusterCapacityTest.java b/flags/src/test/java/com/yahoo/vespa/flags/custom/ClusterCapacityTest.java
index 23ab3a48ffa..6322ad1a2e1 100644
--- a/flags/src/test/java/com/yahoo/vespa/flags/custom/ClusterCapacityTest.java
+++ b/flags/src/test/java/com/yahoo/vespa/flags/custom/ClusterCapacityTest.java
@@ -12,10 +12,12 @@ public class ClusterCapacityTest {
 
     @Test
     void serialization() throws IOException {
-        ClusterCapacity clusterCapacity = new ClusterCapacity(7, 1.2, 3.4, 5.6, null, "fast", "local", "x86_64");
+        ClusterCapacity clusterCapacity = new ClusterCapacity(7, 1.2, 3.4, 5.6, null, "fast", "local", "x86_64", null);
         ObjectMapper mapper = new ObjectMapper();
         String json = mapper.writeValueAsString(clusterCapacity);
-        assertEquals("{\"count\":7,\"vcpu\":1.2,\"memoryGb\":3.4,\"diskGb\":5.6,\"diskSpeed\":\"fast\",\"storageType\":\"local\",\"architecture\":\"x86_64\"}", json);
+        assertEquals("""
+                             {"count":7,"vcpu":1.2,"memoryGb":3.4,"diskGb":5.6,"diskSpeed":"fast","storageType":"local","architecture":"x86_64"}""",
+                     json);
 
         ClusterCapacity deserialized = mapper.readValue(json, ClusterCapacity.class);
         assertEquals(7, deserialized.count());
@@ -30,10 +32,12 @@ public class ClusterCapacityTest {
 
     @Test
     void serialization2() throws IOException {
-        ClusterCapacity clusterCapacity = new ClusterCapacity(7, 1.2, 3.4, 5.6, 2.3, "any", "remote", "arm64");
+        ClusterCapacity clusterCapacity = new ClusterCapacity(7, 1.2, 3.4, 5.6, 2.3, "any", "remote", "arm64", null);
         ObjectMapper mapper = new ObjectMapper();
         String json = mapper.writeValueAsString(clusterCapacity);
-        assertEquals("{\"count\":7,\"vcpu\":1.2,\"memoryGb\":3.4,\"diskGb\":5.6,\"bandwidthGbps\":2.3,\"diskSpeed\":\"any\",\"storageType\":\"remote\",\"architecture\":\"arm64\"}", json);
+        assertEquals("""
+                             {"count":7,"vcpu":1.2,"memoryGb":3.4,"diskGb":5.6,"bandwidthGbps":2.3,"diskSpeed":"any","storageType":"remote","architecture":"arm64"}""",
+                     json);
 
         ClusterCapacity deserialized = mapper.readValue(json, ClusterCapacity.class);
         assertEquals(7, deserialized.count());
@@ -47,8 +51,29 @@ public class ClusterCapacityTest {
     }
 
     @Test
+    void serialization3() throws IOException {
+        ClusterCapacity clusterCapacity = new ClusterCapacity(7, 1.2, 3.4, 5.6, 2.3, "any", "remote", "arm64", "admin");
+        ObjectMapper mapper = new ObjectMapper();
+        String json = mapper.writeValueAsString(clusterCapacity);
+        assertEquals("""
+                             {"count":7,"vcpu":1.2,"memoryGb":3.4,"diskGb":5.6,"bandwidthGbps":2.3,"diskSpeed":"any","storageType":"remote","architecture":"arm64","clusterType":"admin"}""",
+                     json);
+
+        ClusterCapacity deserialized = mapper.readValue(json, ClusterCapacity.class);
+        assertEquals(7, deserialized.count());
+        assertEquals(1.2, deserialized.vcpu(), 0.0001);
+        assertEquals(3.4, deserialized.memoryGb(), 0.0001);
+        assertEquals(5.6, deserialized.diskGb(), 0.0001);
+        assertEquals(2.3, deserialized.bandwidthGbps(), 0.0001);
+        assertEquals("any", deserialized.diskSpeed());
+        assertEquals("remote", deserialized.storageType());
+        assertEquals("arm64", deserialized.architecture());
+        assertEquals("admin", deserialized.clusterType());
+    }
+
+    @Test
     void serializationWithNoNodeResources() throws IOException {
-        ClusterCapacity clusterCapacity = new ClusterCapacity(7, null, null, null, null, null, null, null);
+        ClusterCapacity clusterCapacity = new ClusterCapacity(7, null, null, null, null, null, null, null, null);
         ObjectMapper mapper = new ObjectMapper();
         String json = mapper.writeValueAsString(clusterCapacity);
         assertEquals("{\"count\":7,\"diskSpeed\":\"fast\",\"storageType\":\"any\",\"architecture\":\"x86_64\"}", json);