Change connector private ip enabled to boolean type

zane-neo · zane-neo · commit cff9dd52d5a3 · 2025-10-23T10:50:32.000+08:00
Signed-off-by: zane-neo &lt;zaniu@amazon.com&gt;
diff --git a/common/src/main/java/org/opensearch/ml/common/httpclient/MLHttpClientFactory.java b/common/src/main/java/org/opensearch/ml/common/httpclient/MLHttpClientFactory.java
@@ -8,7 +8,6 @@
 import static org.opensearch.secure_sm.AccessController.doPrivileged;
 
 import java.time.Duration;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import lombok.extern.log4j.Log4j2;
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
@@ -21,7 +20,7 @@ public static SdkAsyncHttpClient getAsyncHttpClient(
         Duration connectionTimeout,
         Duration readTimeout,
         int maxConnections,
-        AtomicBoolean connectorPrivateIpEnabled
+        boolean connectorPrivateIpEnabled
     ) {
         return doPrivileged(() -> {
             log
@@ -37,7 +36,7 @@ public static SdkAsyncHttpClient getAsyncHttpClient(
                 .readTimeout(readTimeout)
                 .maxConcurrency(maxConnections)
                 .build();
-            return new ValidatingHttpClient(delegate, connectorPrivateIpEnabled);
+            return new MLValidatableAsyncHttpClient(delegate, connectorPrivateIpEnabled);
         });
     }
 }
diff --git a/common/src/main/java/org/opensearch/ml/common/httpclient/MLValidatableAsyncHttpClient.java b/common/src/main/java/org/opensearch/ml/common/httpclient/MLValidatableAsyncHttpClient.java
@@ -11,18 +11,17 @@
 import java.util.Arrays;
 import java.util.Locale;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import lombok.extern.log4j.Log4j2;
 import software.amazon.awssdk.http.async.AsyncExecuteRequest;
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
 
 @Log4j2
-public class ValidatingHttpClient implements SdkAsyncHttpClient {
+public class MLValidatableAsyncHttpClient implements SdkAsyncHttpClient {
     private final SdkAsyncHttpClient delegate;
-    private final AtomicBoolean connectorPrivateIpEnabled;
+    private final boolean connectorPrivateIpEnabled;
 
-    protected ValidatingHttpClient(SdkAsyncHttpClient client, AtomicBoolean connectorPrivateIpEnabled) {
+    protected MLValidatableAsyncHttpClient(SdkAsyncHttpClient client, boolean connectorPrivateIpEnabled) {
         this.delegate = client;
         this.connectorPrivateIpEnabled = connectorPrivateIpEnabled;
     }
@@ -54,7 +53,7 @@ public void close() {
      * @param connectorPrivateIpEnabled The port number of the remote inference server, port number must be in range [0, 65536].
      * @throws UnknownHostException Allow to use private IP or not.
      */
-    public void validate(String protocol, String host, int port, AtomicBoolean connectorPrivateIpEnabled) throws UnknownHostException {
+    public void validate(String protocol, String host, int port, boolean connectorPrivateIpEnabled) throws UnknownHostException {
         if (protocol != null && !"http".equalsIgnoreCase(protocol) && !"https".equalsIgnoreCase(protocol)) {
             log.error("Remote inference protocol is not http or https: {}", protocol);
             throw new IllegalArgumentException("Protocol is not http or https: " + protocol);
@@ -74,9 +73,9 @@ public void validate(String protocol, String host, int port, AtomicBoolean conne
         validateIp(host, connectorPrivateIpEnabled);
     }
 
-    private void validateIp(String hostName, AtomicBoolean connectorPrivateIpEnabled) throws UnknownHostException {
+    private void validateIp(String hostName, boolean connectorPrivateIpEnabled) throws UnknownHostException {
         InetAddress[] addresses = InetAddress.getAllByName(hostName);
-        if (connectorPrivateIpEnabled != null && !connectorPrivateIpEnabled.get() && hasPrivateIpAddress(addresses)) {
+        if (!connectorPrivateIpEnabled && hasPrivateIpAddress(addresses)) {
             log.error("Remote inference host name has private ip address: {}", hostName);
             throw new IllegalArgumentException("Remote inference host name has private ip address: " + hostName);
         }
diff --git a/common/src/main/java/org/opensearch/ml/common/settings/MLFeatureEnabledSetting.java b/common/src/main/java/org/opensearch/ml/common/settings/MLFeatureEnabledSetting.java
@@ -25,7 +25,6 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.settings.Settings;
@@ -38,7 +37,7 @@ public class MLFeatureEnabledSetting {
     private volatile Boolean isAgentFrameworkEnabled;
 
     private volatile Boolean isLocalModelEnabled;
-    private volatile AtomicBoolean isConnectorPrivateIpEnabled;
+    private volatile Boolean isConnectorPrivateIpEnabled;
 
     private volatile Boolean isControllerEnabled;
     private volatile Boolean isBatchIngestionEnabled;
@@ -70,7 +69,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         isRemoteInferenceEnabled = ML_COMMONS_REMOTE_INFERENCE_ENABLED.get(settings);
         isAgentFrameworkEnabled = ML_COMMONS_AGENT_FRAMEWORK_ENABLED.get(settings);
         isLocalModelEnabled = ML_COMMONS_LOCAL_MODEL_ENABLED.get(settings);
-        isConnectorPrivateIpEnabled = new AtomicBoolean(ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED.get(settings));
+        isConnectorPrivateIpEnabled = ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED.get(settings);
         isControllerEnabled = ML_COMMONS_CONTROLLER_ENABLED.get(settings);
         isBatchIngestionEnabled = ML_COMMONS_OFFLINE_BATCH_INGESTION_ENABLED.get(settings);
         isBatchInferenceEnabled = ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED.get(settings);
@@ -94,7 +93,7 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_LOCAL_MODEL_ENABLED, it -> isLocalModelEnabled = it);
         clusterService
             .getClusterSettings()
-            .addSettingsUpdateConsumer(ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED, it -> isConnectorPrivateIpEnabled.set(it));
+            .addSettingsUpdateConsumer(ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED, it -> isConnectorPrivateIpEnabled = it);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_CONTROLLER_ENABLED, it -> isControllerEnabled = it);
         clusterService
             .getClusterSettings()
@@ -145,7 +144,7 @@ public boolean isLocalModelEnabled() {
         return isLocalModelEnabled;
     }
 
-    public AtomicBoolean isConnectorPrivateIpEnabled() {
+    public boolean isConnectorPrivateIpEnabled() {
         return isConnectorPrivateIpEnabled;
     }
 
diff --git a/common/src/test/java/org/opensearch/ml/common/httpclient/MLHttpClientFactoryTests.java b/common/src/test/java/org/opensearch/ml/common/httpclient/MLHttpClientFactoryTests.java
@@ -8,7 +8,6 @@
 import static org.junit.Assert.assertNotNull;
 
 import java.time.Duration;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.Test;
 
@@ -18,8 +17,7 @@ public class MLHttpClientFactoryTests {
 
     @Test
     public void test_getSdkAsyncHttpClient_success() {
-        SdkAsyncHttpClient client = MLHttpClientFactory
-            .getAsyncHttpClient(Duration.ofSeconds(100), Duration.ofSeconds(100), 100, new AtomicBoolean(false));
+        SdkAsyncHttpClient client = MLHttpClientFactory.getAsyncHttpClient(Duration.ofSeconds(100), Duration.ofSeconds(100), 100, false);
         assertNotNull(client);
     }
 
diff --git a/common/src/test/java/org/opensearch/ml/common/httpclient/MLValidatableAsyncHttpClientTests.java b/common/src/test/java/org/opensearch/ml/common/httpclient/MLValidatableAsyncHttpClientTests.java
@@ -10,22 +10,23 @@
 import static org.junit.Assert.assertThrows;
 import static org.mockito.Mockito.mock;
 
-import java.util.concurrent.atomic.AtomicBoolean;
-
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
 
-public class ValidatingHttpClientTests {
+public class MLValidatableAsyncHttpClientTests {
     private static final String TEST_HOST = "api.openai.com";
     private static final String HTTP = "http";
     private static final String HTTPS = "https";
-    private static final AtomicBoolean PRIVATE_IP_DISABLED = new AtomicBoolean(false);
-    private static final AtomicBoolean PRIVATE_IP_ENABLED = new AtomicBoolean(true);
+    private static final boolean PRIVATE_IP_DISABLED = false;
+    private static final boolean PRIVATE_IP_ENABLED = true;
 
-    private final ValidatingHttpClient validatingHttpClient = new ValidatingHttpClient(mock(SdkAsyncHttpClient.class), PRIVATE_IP_DISABLED);
+    private final MLValidatableAsyncHttpClient validatingHttpClient = new MLValidatableAsyncHttpClient(
+        mock(SdkAsyncHttpClient.class),
+        PRIVATE_IP_DISABLED
+    );
 
     @Rule
     public ExpectedException expectedException = ExpectedException.none();
diff --git a/common/src/test/java/org/opensearch/ml/common/settings/MLFeatureEnabledSettingTests.java b/common/src/test/java/org/opensearch/ml/common/settings/MLFeatureEnabledSettingTests.java
@@ -81,7 +81,7 @@ public void testDefaults_allFeaturesEnabled() {
         assertTrue(setting.isRemoteInferenceEnabled());
         assertTrue(setting.isAgentFrameworkEnabled());
         assertTrue(setting.isLocalModelEnabled());
-        assertTrue(setting.isConnectorPrivateIpEnabled().get());
+        assertTrue(setting.isConnectorPrivateIpEnabled());
         assertTrue(setting.isControllerEnabled());
         assertTrue(setting.isOfflineBatchIngestionEnabled());
         assertTrue(setting.isOfflineBatchInferenceEnabled());
@@ -122,7 +122,7 @@ public void testDefaults_someFeaturesDisabled() {
         assertFalse(setting.isRemoteInferenceEnabled());
         assertFalse(setting.isAgentFrameworkEnabled());
         assertFalse(setting.isLocalModelEnabled());
-        assertFalse(setting.isConnectorPrivateIpEnabled().get());
+        assertFalse(setting.isConnectorPrivateIpEnabled());
         assertFalse(setting.isControllerEnabled());
         assertFalse(setting.isOfflineBatchIngestionEnabled());
         assertFalse(setting.isOfflineBatchInferenceEnabled());
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutor.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutor.java
@@ -18,7 +18,6 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.commons.text.StringEscapeUtils;
@@ -81,7 +80,7 @@ public class AwsConnectorExecutor extends AbstractConnectorExecutor {
     private StreamTransportService streamTransportService;
 
     @Setter
-    private AtomicBoolean connectorPrivateIpEnabled;
+    private boolean connectorPrivateIpEnabled;
 
     public AwsConnectorExecutor(Connector connector) {
         super.initialize(connector);
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutor.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutor.java
@@ -17,7 +17,6 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.commons.text.StringEscapeUtils;
@@ -74,7 +73,7 @@ public class HttpJsonConnectorExecutor extends AbstractConnectorExecutor {
     @Getter
     private MLGuard mlGuard;
     @Setter
-    private volatile AtomicBoolean connectorPrivateIpEnabled;
+    private volatile boolean connectorPrivateIpEnabled;
 
     private final AtomicReference<SdkAsyncHttpClient> httpClientRef = new AtomicReference<>();
 
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java
@@ -19,7 +19,6 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.apache.logging.log4j.Logger;
 import org.opensearch.ExceptionsHelper;
@@ -183,7 +182,7 @@ default void setScriptService(ScriptService scriptService) {}
 
     default void setClient(Client client) {}
 
-    default void setConnectorPrivateIpEnabled(AtomicBoolean connectorPrivateIpEnabled) {}
+    default void setConnectorPrivateIpEnabled(boolean connectorPrivateIpEnabled) {}
 
     default void setXContentRegistry(NamedXContentRegistry xContentRegistry) {}
 
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteModel.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteModel.java
@@ -11,7 +11,6 @@
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.settings.Settings;
@@ -127,7 +126,7 @@ public CompletionStage<Boolean> initModelAsync(MLModel model, Map<String, Object
             this.connectorExecutor.setRateLimiter((TokenBucket) params.get(RATE_LIMITER));
             this.connectorExecutor.setUserRateLimiterMap((Map<String, TokenBucket>) params.get(USER_RATE_LIMITER_MAP));
             this.connectorExecutor.setMlGuard((MLGuard) params.get(GUARDRAILS));
-            this.connectorExecutor.setConnectorPrivateIpEnabled((AtomicBoolean) params.get(CONNECTOR_PRIVATE_IP_ENABLED));
+            this.connectorExecutor.setConnectorPrivateIpEnabled((boolean) params.getOrDefault(CONNECTOR_PRIVATE_IP_ENABLED, false));
             return CompletableFuture.completedStage(true);
         }).exceptionally(e -> {
             log.error("Failed to init remote model.", e);
diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutorTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutorTest.java
@@ -19,7 +19,6 @@
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicBoolean;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -97,7 +96,7 @@ public void invokeRemoteService_invalidIpAddress() {
             .actions(Arrays.asList(predictAction))
             .build();
         HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
-        executor.setConnectorPrivateIpEnabled(new AtomicBoolean(false));
+        executor.setConnectorPrivateIpEnabled(false);
         executor
             .invokeRemoteService(
                 PREDICT.name(),
@@ -130,7 +129,7 @@ public void invokeRemoteService_EnabledPrivateIpAddress() {
             .actions(Arrays.asList(predictAction))
             .build();
         HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
-        executor.setConnectorPrivateIpEnabled(new AtomicBoolean(true));
+        executor.setConnectorPrivateIpEnabled(true);
         executor
             .invokeRemoteService(
                 PREDICT.name(),
@@ -160,7 +159,7 @@ public void invokeRemoteService_DisabledPrivateIpAddress() {
             .actions(Arrays.asList(predictAction))
             .build();
         HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
-        executor.setConnectorPrivateIpEnabled(new AtomicBoolean(false));
+        executor.setConnectorPrivateIpEnabled(false);
         executor
             .invokeRemoteService(
                 PREDICT.name(),
diff --git a/plugin/src/test/java/org/opensearch/ml/model/MLModelManagerTests.java b/plugin/src/test/java/org/opensearch/ml/model/MLModelManagerTests.java
@@ -65,7 +65,6 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
@@ -285,7 +284,7 @@ public void setup() throws URISyntaxException {
         when(client.threadPool()).thenReturn(threadPool);
         when(threadPool.getThreadContext()).thenReturn(threadContext);
         when(mlFeatureEnabledSetting.isControllerEnabled()).thenReturn(true);
-        when(mlFeatureEnabledSetting.isConnectorPrivateIpEnabled()).thenReturn(new AtomicBoolean(false));
+        when(mlFeatureEnabledSetting.isConnectorPrivateIpEnabled()).thenReturn(false);
 
         modelManager = spy(
             new MLModelManager(
