Support diskSpeed, storageType and architecture in ClusterCapacity

2 files changed, 52 insertions, 16 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 b24d5f62174..fc5f2456a77 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
@@ -7,8 +7,8 @@ import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
-
 import java.util.Objects;
+import java.util.Optional;
 import java.util.OptionalDouble;
 
 /**
@@ -23,23 +23,34 @@ public class ClusterCapacity {
     private final OptionalDouble memoryGb;
     private final OptionalDouble diskGb;
     private final OptionalDouble bandwidthGbps;
+    private final Optional<String> diskSpeed;
+    private final Optional<String> storageType;
+    private final Optional<String> architecture;
+
 
     @JsonCreator
     public ClusterCapacity(@JsonProperty("count") int count,
                            @JsonProperty("vcpu") Double vcpu,
                            @JsonProperty("memoryGb") Double memoryGb,
                            @JsonProperty("diskGb") Double diskGb,
-                           @JsonProperty("bandwidthGbps") Double bandwidthGbps) {
+                           @JsonProperty("bandwidthGbps") Double bandwidthGbps,
+                           @JsonProperty("diskSpeed") String diskSpeed,
+                           @JsonProperty("storageType") String storageType,
+                           @JsonProperty("architecture") String architecture) {
         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));
         this.diskGb = diskGb == null ? OptionalDouble.empty() : OptionalDouble.of(requireNonNegative("diskGb", diskGb));
         this.bandwidthGbps = bandwidthGbps == null ? OptionalDouble.empty() : OptionalDouble.of(bandwidthGbps);
+        this.diskSpeed = Optional.ofNullable(diskSpeed);
+        this.storageType = Optional.ofNullable(storageType);
+        this.architecture = Optional.ofNullable(architecture);
     }
 
     /** 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());
+        return new ClusterCapacity(count, vcpuOrNull(), memoryGbOrNull(), diskGbOrNull(), bandwidthGbpsOrNull(),
+                                   diskSpeedOrNull(), storageTypeOrNull(), architectureOrNull());
     }
 
     @JsonGetter("count") public int count() { return count; }
@@ -55,13 +66,23 @@ public class ClusterCapacity {
     @JsonGetter("bandwidthGbps") public Double bandwidthGbpsOrNull() {
         return bandwidthGbps.isPresent() ? bandwidthGbps.getAsDouble() : null;
     }
+    @JsonGetter("diskSpeed") public String diskSpeedOrNull() {
+        return diskSpeed.orElse(null);
+    }
+    @JsonGetter("storageType") public String storageTypeOrNull() {
+        return storageType.orElse(null);
+    }
+    @JsonGetter("architecture") public String architectureOrNull() {
+        return architecture.orElse(null);
+    }
 
     @JsonIgnore public Double vcpu() { return vcpu.orElse(0.0); }
     @JsonIgnore public Double memoryGb() { return memoryGb.orElse(0.0); }
     @JsonIgnore public Double diskGb() { return diskGb.orElse(0.0); }
-
-    @JsonIgnore
-    public double bandwidthGbps() { return bandwidthGbps.orElse(1.0); }
+    @JsonIgnore public double bandwidthGbps() { return bandwidthGbps.orElse(1.0); }
+    @JsonIgnore public String diskSpeed() { return diskSpeed.orElse("fast"); }
+    @JsonIgnore public String storageType() { return storageType.orElse("any"); }
+    @JsonIgnore public String architecture() { return architecture.orElse("x86_64"); }
 
     @Override
     public String toString() {
@@ -71,6 +92,9 @@ public class ClusterCapacity {
                 ", memoryGb=" + memoryGb +
                 ", diskGb=" + diskGb +
                 ", bandwidthGbps=" + bandwidthGbps +
+                ", diskSpeed=" + diskSpeed +
+                ", storageType=" + storageType +
+                ", architecture=" + architecture +
                 '}';
     }
 
@@ -83,12 +107,15 @@ public class ClusterCapacity {
                 vcpu.equals(that.vcpu) &&
                 memoryGb.equals(that.memoryGb) &&
                 diskGb.equals(that.diskGb) &&
-                bandwidthGbps.equals(that.bandwidthGbps);
+                bandwidthGbps.equals(that.bandwidthGbps) &&
+                diskSpeed.equals(that.diskSpeed) &&
+                storageType.equals(that.storageType) &&
+                architecture.equals(that.architecture);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(count, vcpu, memoryGb, diskGb, bandwidthGbps);
+        return Objects.hash(count, vcpu, memoryGb, diskGb, bandwidthGbps, diskSpeed, storageType, architecture);
     }
 
     private static double requireNonNegative(String name, double value) {
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 796cdbf30d1..471e458846b 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,47 +12,56 @@ public class ClusterCapacityTest {
 
     @Test
     void serialization() throws IOException {
-        ClusterCapacity clusterCapacity = new ClusterCapacity(7, 1.2, 3.4, 5.6, null);
+        ClusterCapacity clusterCapacity = new ClusterCapacity(7, 1.2, 3.4, 5.6, null, "fast", "local", "x86_64");
         ObjectMapper mapper = new ObjectMapper();
         String json = mapper.writeValueAsString(clusterCapacity);
-        assertEquals("{\"count\":7,\"vcpu\":1.2,\"memoryGb\":3.4,\"diskGb\":5.6}", 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());
         assertEquals(1.2, deserialized.vcpu(), 0.0001);
         assertEquals(3.4, deserialized.memoryGb(), 0.0001);
         assertEquals(5.6, deserialized.diskGb(), 0.0001);
         assertEquals(1.0, deserialized.bandwidthGbps(), 0.0001);
-        assertEquals(7, deserialized.count());
+        assertEquals("fast", deserialized.diskSpeed());
+        assertEquals("local", deserialized.storageType());
+        assertEquals("x86_64", deserialized.architecture());
     }
 
     @Test
     void serialization2() throws IOException {
-        ClusterCapacity clusterCapacity = new ClusterCapacity(7, 1.2, 3.4, 5.6, 2.3);
+        ClusterCapacity clusterCapacity = new ClusterCapacity(7, 1.2, 3.4, 5.6, 2.3, "any", "remote", "arm64");
         ObjectMapper mapper = new ObjectMapper();
         String json = mapper.writeValueAsString(clusterCapacity);
-        assertEquals("{\"count\":7,\"vcpu\":1.2,\"memoryGb\":3.4,\"diskGb\":5.6,\"bandwidthGbps\":2.3}", 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());
         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(7, deserialized.count());
+        assertEquals("any", deserialized.diskSpeed());
+        assertEquals("remote", deserialized.storageType());
+        assertEquals("arm64", deserialized.architecture());
     }
 
     @Test
     void serializationWithNoNodeResources() throws IOException {
-        ClusterCapacity clusterCapacity = new ClusterCapacity(7, null, null, null, null);
+        ClusterCapacity clusterCapacity = new ClusterCapacity(7, null, null, null, null, null, null, null);
         ObjectMapper mapper = new ObjectMapper();
         String json = mapper.writeValueAsString(clusterCapacity);
         assertEquals("{\"count\":7}", json);
 
         ClusterCapacity deserialized = mapper.readValue(json, ClusterCapacity.class);
+        assertEquals(7, deserialized.count());
         assertEquals(0.0, deserialized.vcpu(), 0.0001);
         assertEquals(0.0, deserialized.memoryGb(), 0.0001);
         assertEquals(0.0, deserialized.diskGb(), 0.0001);
         assertEquals(1.0, deserialized.bandwidthGbps(), 0.0001);  // 1.0 is used as fallback
-        assertEquals(7, deserialized.count());
+        assertEquals("fast", deserialized.diskSpeed());
+        assertEquals("any", deserialized.storageType());
+        assertEquals("x86_64", deserialized.architecture());
     }
 
 }
\ No newline at end of file