Address comments

Signed-off-by: zane-neo <[email protected]>
zane-neo · Jul 11, 2024 · 607de1b · 607de1b
1 parent 49788fa
commit 607de1b
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 17 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/breaker/DiskCircuitBreaker.java b/plugin/src/main/java/org/opensearch/ml/breaker/DiskCircuitBreaker.java
@@ -15,15 +15,16 @@
 
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.core.common.unit.ByteSizeUnit;
+import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.ml.common.exception.MLException;
 
 /**
  * A circuit breaker for disk usage.
  */
-public class DiskCircuitBreaker extends ThresholdCircuitBreaker<Integer> {
+public class DiskCircuitBreaker extends ThresholdCircuitBreaker<ByteSizeValue> {
     private static final String ML_DISK_CB = "Disk Circuit Breaker";
-    public static final int DEFAULT_DISK_SHORTAGE_THRESHOLD = 5;
-    private static final long GB = 1024 * 1024 * 1024;
+    public static final ByteSizeValue DEFAULT_DISK_SHORTAGE_THRESHOLD = new ByteSizeValue(5, ByteSizeUnit.GB);
     private final File diskDir;
 
     public DiskCircuitBreaker(Settings settings, ClusterService clusterService, File diskDir) {
@@ -42,7 +43,7 @@ public String getName() {
     public boolean isOpen() {
         try {
             return AccessController.doPrivileged((PrivilegedExceptionAction<Boolean>) () -> {
-                return diskDir.getFreeSpace() / GB < getThreshold();  // in GB
+                return new ByteSizeValue(diskDir.getFreeSpace(), ByteSizeUnit.BYTES).compareTo(getThreshold()) < 0;  // in GB
             });
         } catch (PrivilegedActionException e) {
             throw new MLException("Failed to run disk circuit breaker");

diff --git a/plugin/src/main/java/org/opensearch/ml/breaker/ThresholdCircuitBreaker.java b/plugin/src/main/java/org/opensearch/ml/breaker/ThresholdCircuitBreaker.java
@@ -14,7 +14,7 @@
 @Data
 public abstract class ThresholdCircuitBreaker<T> implements CircuitBreaker {
 
-    private T threshold;
+    private volatile T threshold;
 
     public ThresholdCircuitBreaker(T threshold) {
         this.threshold = threshold;

diff --git a/plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java b/plugin/src/main/java/org/opensearch/ml/settings/MLCommonsSettings.java
@@ -9,6 +9,8 @@
 import java.util.function.Function;
 
 import org.opensearch.common.settings.Setting;
+import org.opensearch.core.common.unit.ByteSizeUnit;
+import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.ml.common.conversation.ConversationalIndexConstants;
 import org.opensearch.searchpipelines.questionanswering.generative.GenerativeQAProcessorConstants;
 
@@ -77,12 +79,10 @@ private MLCommonsSettings() {}
     public static final Setting<Integer> ML_COMMONS_JVM_HEAP_MEM_THRESHOLD = Setting
         .intSetting("plugins.ml_commons.jvm_heap_memory_threshold", 85, 0, 100, Setting.Property.NodeScope, Setting.Property.Dynamic);
 
-    public static final Setting<Integer> ML_COMMONS_DISK_FREE_SPACE_THRESHOLD = Setting
-        .intSetting(
+    public static final Setting<ByteSizeValue> ML_COMMONS_DISK_FREE_SPACE_THRESHOLD = Setting
+        .byteSizeSetting(
             "plugins.ml_commons.disk_free_space_threshold",
-            5,
-            0,
-            Integer.MAX_VALUE,
+            new ByteSizeValue(5L, ByteSizeUnit.GB),
             Setting.Property.NodeScope,
             Setting.Property.Dynamic
         );

diff --git a/plugin/src/test/java/org/opensearch/ml/breaker/DiskCircuitBreakerTests.java b/plugin/src/test/java/org/opensearch/ml/breaker/DiskCircuitBreakerTests.java
@@ -1,8 +1,6 @@
 /*
- *
- *  * Copyright OpenSearch Contributors
- *  * SPDX-License-Identifier: Apache-2.0
- *
+ *  Copyright OpenSearch Contributors
+ *  SPDX-License-Identifier: Apache-2.0
  */
 
 package org.opensearch.ml.breaker;
@@ -22,6 +20,8 @@
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.core.common.unit.ByteSizeUnit;
+import org.opensearch.core.common.unit.ByteSizeValue;
 
 public class DiskCircuitBreakerTests {
     @Mock
@@ -40,7 +40,7 @@ public void setup() {
     @Test
     public void test_isOpen_whenDiskFreeSpaceIsHigherThanMinValue_breakerIsNotOpen() {
         CircuitBreaker breaker = new DiskCircuitBreaker(
-            Settings.builder().put(ML_COMMONS_DISK_FREE_SPACE_THRESHOLD.getKey(), 5).build(),
+            Settings.builder().put(ML_COMMONS_DISK_FREE_SPACE_THRESHOLD.getKey(), new ByteSizeValue(4L, ByteSizeUnit.GB)).build(),
             clusterService,
             file
         );
@@ -51,7 +51,7 @@ public void test_isOpen_whenDiskFreeSpaceIsHigherThanMinValue_breakerIsNotOpen()
     @Test
     public void test_isOpen_whenDiskFreeSpaceIsLessThanMinValue_breakerIsOpen() {
         CircuitBreaker breaker = new DiskCircuitBreaker(
-            Settings.builder().put(ML_COMMONS_DISK_FREE_SPACE_THRESHOLD.getKey(), 5).build(),
+            Settings.builder().put(ML_COMMONS_DISK_FREE_SPACE_THRESHOLD.getKey(), new ByteSizeValue(5L, ByteSizeUnit.GB)).build(),
             clusterService,
             file
         );
@@ -62,7 +62,7 @@ public void test_isOpen_whenDiskFreeSpaceIsLessThanMinValue_breakerIsOpen() {
     @Test
     public void test_isOpen_whenDiskFreeSpaceConfiguredToZero_breakerIsNotOpen() {
         CircuitBreaker breaker = new DiskCircuitBreaker(
-            Settings.builder().put(ML_COMMONS_DISK_FREE_SPACE_THRESHOLD.getKey(), 5).build(),
+            Settings.builder().put(ML_COMMONS_DISK_FREE_SPACE_THRESHOLD.getKey(), new ByteSizeValue(0L, ByteSizeUnit.KB)).build(),
             clusterService,
             file
         );