HDDS-15438. [DiskBalancer] Validate persisted diskbalancer.info while reading YAML.

hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/diskbalancer/DiskBalancerYaml.java

@@ -76,24 +76,63 @@ public static DiskBalancerInfo readDiskBalancerInfoFile(File path)
         throw new IOException("Unable to parse yaml file.", e);
       }
 
-      // getContainerStates() may be null if the key is absent; isNotBlank(null) is false.
-      String cs = diskBalancerInfoYaml.getContainerStates();
-      String containerStates = StringUtils.isNotBlank(cs)
-          ? cs.trim() : DiskBalancerConfiguration.DEFAULT_CONTAINER_STATES;
+      validateRequiredFields(diskBalancerInfoYaml);
+      DiskBalancerVersion version = getValidatedVersion(diskBalancerInfoYaml);
+      String containerStates = getValidatedContainerStates(diskBalancerInfoYaml);
       diskBalancerInfo = new DiskBalancerInfo(
           diskBalancerInfoYaml.operationalState,
           diskBalancerInfoYaml.getThreshold(),
           diskBalancerInfoYaml.getBandwidthInMB(),
           diskBalancerInfoYaml.getParallelThread(),
           diskBalancerInfoYaml.isStopAfterDiskEven(),
           containerStates,
-          DiskBalancerVersion.getDiskBalancerVersion(
-              diskBalancerInfoYaml.version));
+          version);
+      validatePersistedConfiguration(diskBalancerInfo);
     }
 
     return diskBalancerInfo;
   }
 
+  private static void validateRequiredFields(
+      DiskBalancerInfoYaml diskBalancerInfoYaml) throws IOException {
+    if (diskBalancerInfoYaml.getOperationalState() == null) {
+      throw new IOException("DiskBalancer operationalState is missing from persisted info.");
+    }
+    if (diskBalancerInfoYaml.getVersion() == null) {
+      throw new IOException("DiskBalancer info version is missing from persisted info.");
+    }
+  }
+
+  private static DiskBalancerVersion getValidatedVersion(
+      DiskBalancerInfoYaml diskBalancerInfoYaml) throws IOException {
+    int rawVersion = diskBalancerInfoYaml.getVersion();
+    DiskBalancerVersion version =
+        DiskBalancerVersion.getDiskBalancerVersion(rawVersion);
+    if (version == null) {
+      throw new IOException("Unsupported DiskBalancer info version: " + rawVersion);
+    }
+    return version;
+  }
+
+  private static String getValidatedContainerStates(
+      DiskBalancerInfoYaml diskBalancerInfoYaml) {
+    // getContainerStates() may be null if the key is absent; isNotBlank(null) is false.
+    String containerStates = diskBalancerInfoYaml.getContainerStates();
+    return StringUtils.isNotBlank(containerStates)
+        ? containerStates.trim() : DiskBalancerConfiguration.DEFAULT_CONTAINER_STATES;
+  }
+
+  private static void validatePersistedConfiguration(
+      DiskBalancerInfo diskBalancerInfo) throws IOException {
+    try {
+      diskBalancerInfo.toConfiguration();
+    } catch (IllegalArgumentException ex) {
+      throw new IOException(
+          "Invalid DiskBalancer configuration in persisted info: "
+              + ex.getMessage(), ex);
+    }
+  }
+
   /**
    * Datanode DiskBalancer Info to be written to the yaml file.
    */
@@ -105,7 +144,7 @@ public static class DiskBalancerInfoYaml {
     private boolean stopAfterDiskEven;
     private String containerStates;
 
-    private int version;
+    private Integer version;
 
     public DiskBalancerInfoYaml() {
       // Needed for snake-yaml introspection.
@@ -163,11 +202,11 @@ public void setStopAfterDiskEven(boolean stopAfterDiskEven) {
       this.stopAfterDiskEven = stopAfterDiskEven;
     }
 
-    public void setVersion(int version) {
+    public void setVersion(Integer version) {
       this.version = version;
     }
 
-    public int getVersion() {
+    public Integer getVersion() {
       return this.version;
     }
 

hadoop-hdds/container-service/src/test/java/org/apache/hadoop/ozone/container/diskbalancer/TestDiskBalancerYaml.java

@@ -18,6 +18,7 @@
 package org.apache.hadoop.ozone.container.diskbalancer;
 
 import static org.apache.hadoop.ozone.OzoneConsts.OZONE_SCM_DATANODE_DISK_BALANCER_INFO_FILE_DEFAULT;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.io.File;
 import java.io.IOException;
@@ -101,4 +102,56 @@ public void testReadYamlNullContainerStatesUsesDefault() throws IOException {
     DiskBalancerInfo info = DiskBalancerYaml.readDiskBalancerInfoFile(file);
     Assertions.assertEquals(DiskBalancerConfiguration.DEFAULT_CONTAINER_STATES, info.getContainerStates());
   }
+
+  @ParameterizedTest
+  @MethodSource("invalidDiskBalancerYamlCases")
+  public void testReadYamlRejectsInvalidPersistedInfo(String yaml,
+      String expectedMessage) throws IOException {
+    File file = new File(tmpDir.toString(),
+        OZONE_SCM_DATANODE_DISK_BALANCER_INFO_FILE_DEFAULT);
+    Files.write(file.toPath(), yaml.getBytes(StandardCharsets.UTF_8));
+
+    IOException ex = assertThrows(IOException.class,
+        () -> DiskBalancerYaml.readDiskBalancerInfoFile(file));
+
+    Assertions.assertTrue(ex.getMessage().contains(expectedMessage),
+        () -> "Expected message to contain '" + expectedMessage + "': "
+            + ex.getMessage());
+  }
+
+  public static Stream<Arguments> invalidDiskBalancerYamlCases() {
+    return Stream.of(
+        Arguments.of(validYaml()
+            .replace("version: 1\n", "version: 99\n"),
+            "Unsupported DiskBalancer info version: 99"),
+        Arguments.of(validYaml()
+            .replace("version: 1\n", ""),
+            "DiskBalancer info version is missing"),
+        Arguments.of(validYaml()
+            .replace("operationalState: RUNNING\n", ""),
+            "DiskBalancer operationalState is missing"),
+        Arguments.of(validYaml()
+            .replace("threshold: 10.0\n", "threshold: 0.0\n"),
+            "Invalid DiskBalancer configuration in persisted info"),
+        Arguments.of(validYaml()
+            .replace("bandwidthInMB: 100\n", "bandwidthInMB: 0\n"),
+            "Invalid DiskBalancer configuration in persisted info"),
+        Arguments.of(validYaml()
+            .replace("parallelThread: 5\n", "parallelThread: 0\n"),
+            "Invalid DiskBalancer configuration in persisted info"),
+        Arguments.of(validYaml()
+            .replace("containerStates: CLOSED,QUASI_CLOSED\n",
+                "containerStates: OPEN\n"),
+            "Invalid DiskBalancer configuration in persisted info"));
+  }
+
+  private static String validYaml() {
+    return "operationalState: RUNNING\n"
+        + "threshold: 10.0\n"
+        + "bandwidthInMB: 100\n"
+        + "parallelThread: 5\n"
+        + "stopAfterDiskEven: true\n"
+        + "containerStates: CLOSED,QUASI_CLOSED\n"
+        + "version: 1\n";
+  }
 }