From 607de1b80d7d1bd8bebf505ed2dfd4988acea164 Mon Sep 17 00:00:00 2001 From: zane-neo Date: Thu, 11 Jul 2024 13:40:09 +0800 Subject: [PATCH] Address comments Signed-off-by: zane-neo --- .../opensearch/ml/breaker/DiskCircuitBreaker.java | 9 +++++---- .../ml/breaker/ThresholdCircuitBreaker.java | 2 +- .../opensearch/ml/settings/MLCommonsSettings.java | 10 +++++----- .../ml/breaker/DiskCircuitBreakerTests.java | 14 +++++++------- 4 files changed, 18 insertions(+), 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 index 0d87222d1e..1fb1ef36cf 100644 --- 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 { +public class DiskCircuitBreaker extends ThresholdCircuitBreaker { 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) () -> { - 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 index bab8b0558c..155a7ba637 100644 --- 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 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 index 07ccc196b4..1e7a569a09 100644 --- 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 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 ML_COMMONS_DISK_FREE_SPACE_THRESHOLD = Setting - .intSetting( + public static final Setting 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 index 84d32db8b0..7eab9d63e8 100644 --- 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 );