Fix SdkAsyncHttpClient resource leak in connector executors (#4716)

* fix: Close SdkAsyncHttpClient in connector executors to prevent resource leak

SdkAsyncHttpClient instances were created lazily in AwsConnectorExecutor
and HttpJsonConnectorExecutor but never closed. This leaked Netty event
loops and connection pools, causing connection pool exhaustion in CI
where 23+ test classes run in the same JVM.

- Make RemoteConnectorExecutor extend AutoCloseable
- Add close() to AwsConnectorExecutor and HttpJsonConnectorExecutor
- Fix RemoteModel.close() to close executor before nulling reference
- Close short-lived executors in ExecuteConnectorTransportAction,
  GetTaskTransportAction, CancelBatchJobTransportAction, and
  RemoteAgenticConversationMemory
- Add client_config with generous timeouts to integration test
  connectors as additional safety margin

Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>

* refactor: Move close() and httpClientRef to AbstractConnectorExecutor

Per review feedback, consolidate the duplicated close() method and
httpClientRef field from AwsConnectorExecutor and HttpJsonConnectorExecutor
into AbstractConnectorExecutor to avoid duplication.

Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>

---------

Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>

Author: dhrubo-os

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/AbstractConnectorExecutor.java

@@ -6,6 +6,7 @@
 package org.opensearch.ml.engine.algorithms.remote;
 
 import java.net.http.HttpRequest;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.logging.log4j.Logger;
 import org.opensearch.action.support.ThreadedActionListener;
@@ -20,13 +21,16 @@
 import lombok.Getter;
 import lombok.Setter;
 import lombok.extern.log4j.Log4j2;
+import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
 
 @Log4j2
 @Setter
 @Getter
 public abstract class AbstractConnectorExecutor implements RemoteConnectorExecutor {
     private ConnectorClientConfig connectorClientConfig;
 
+    protected final AtomicReference<SdkAsyncHttpClient> httpClientRef = new AtomicReference<>();
+
     public void initialize(Connector connector) {
         if (connector.getConnectorClientConfig() != null) {
             connectorClientConfig = connector.getConnectorClientConfig();
@@ -78,4 +82,16 @@ protected ThreadedActionListener<Tuple<Integer, ModelTensors>> createThreadedLis
     ) {
         return new ThreadedActionListener<>(logger, getClient().threadPool(), "opensearch_ml_predict_remote", actionListener, false);
     }
+
+    /**
+     * Closes the underlying HTTP client. Safe to call concurrently — NettyNioAsyncHttpClient.close()
+     * gracefully drains in-flight requests before shutting down the event loop group.
+     */
+    @Override
+    public void close() {
+        SdkAsyncHttpClient client = httpClientRef.getAndSet(null);
+        if (client != null) {
+            client.close();
+        }
+    }
 }

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutor.java

@@ -20,7 +20,6 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.commons.text.StringEscapeUtils;
 import org.apache.logging.log4j.Logger;
@@ -77,15 +76,13 @@ public class AwsConnectorExecutor extends AbstractConnectorExecutor {
     @Getter
     private MLGuard mlGuard;
 
-    private final AtomicReference<SdkAsyncHttpClient> httpClientRef = new AtomicReference<>();
+    @Setter
+    private volatile boolean connectorPrivateIpEnabled;
 
     @Setter
     @Getter
     private StreamTransportService streamTransportService;
 
-    @Setter
-    private volatile boolean connectorPrivateIpEnabled;
-
     public AwsConnectorExecutor(Connector connector) {
         super.initialize(connector);
         this.connector = (AwsConnector) connector;

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutor.java

@@ -19,7 +19,6 @@
 import java.util.Locale;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicReference;
 
 import org.apache.commons.text.StringEscapeUtils;
 import org.apache.logging.log4j.Logger;
@@ -79,8 +78,6 @@ public class HttpJsonConnectorExecutor extends AbstractConnectorExecutor {
     @Setter
     private volatile boolean connectorPrivateIpEnabled;
 
-    private final AtomicReference<SdkAsyncHttpClient> httpClientRef = new AtomicReference<>();
-
     @Setter
     @Getter
     private StreamTransportService streamTransportService;

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteConnectorExecutor.java

@@ -62,7 +62,10 @@
 
 import lombok.Builder;
 
-public interface RemoteConnectorExecutor {
+public interface RemoteConnectorExecutor extends AutoCloseable {
+
+    @Override
+    default void close() {}
 
     public String RETRY_EXECUTOR = "opensearch_ml_predict_remote";
     String SKIP_SSL_VERIFICATION = "skip_ssl_verification";

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/remote/RemoteModel.java

@@ -95,6 +95,13 @@ public void asyncPredict(MLInput mlInput, ActionListener<MLTaskResponse> actionL
 
     @Override
     public void close() {
+        if (this.connectorExecutor != null) {
+            try {
+                this.connectorExecutor.close();
+            } catch (Exception e) {
+                log.error("Failed to close connector executor", e);
+            }
+        }
         this.connectorExecutor = null;
     }
 

ml-algorithms/src/main/java/org/opensearch/ml/engine/memory/RemoteAgenticConversationMemory.java

@@ -1300,8 +1300,11 @@ private void createSessionInRemoteContainer(
                     }
                 } catch (Exception e) {
                     listener.onFailure(e);
+                } finally {
+                    executor.close();
                 }
             }, e -> {
+                executor.close();
                 log.error("Failed to create session via remote connector", e);
                 listener.onFailure(e);
             }));

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/AbstractConnectorExecutorTest.java

@@ -188,4 +188,24 @@ public void testGetMcpRequestHeaders_VerifyHeaderOrder() {
         assertEquals(MCP_HEADER_AWS_ACCESS_KEY_ID, headerNameCaptor.getAllValues().get(0));
         assertEquals(MCP_HEADER_OPENSEARCH_URL, headerNameCaptor.getAllValues().get(1));
     }
+
+    @Test
+    public void testClose_closesHttpClient() {
+        // Inject a mock SdkAsyncHttpClient into httpClientRef
+        software.amazon.awssdk.http.async.SdkAsyncHttpClient mockHttpClient = mock(
+            software.amazon.awssdk.http.async.SdkAsyncHttpClient.class
+        );
+        executor.httpClientRef.set(mockHttpClient);
+
+        executor.close();
+
+        verify(mockHttpClient).close();
+        assertNull(executor.httpClientRef.get());
+    }
+
+    @Test
+    public void testClose_whenNoClient_doesNotThrow() {
+        assertNull(executor.httpClientRef.get());
+        executor.close(); // should be a no-op
+    }
 }

plugin/src/main/java/org/opensearch/ml/action/connector/ExecuteConnectorTransportAction.java

@@ -106,9 +106,11 @@ protected void doExecute(Task task, ActionRequest request, ActionListener<MLTask
                         connectorExecutor.setXContentRegistry(xContentRegistry);
                         connectorExecutor
                             .executeAction(finalConnectorAction, executeConnectorRequest.getMlInput(), ActionListener.wrap(taskResponse -> {
+                                connectorExecutor.close();
                                 connector.removeCredential();
                                 actionListener.onResponse(taskResponse);
                             }, e -> {
+                                connectorExecutor.close();
                                 connector.removeCredential();
                                 actionListener.onFailure(e);
                             }));

plugin/src/main/java/org/opensearch/ml/action/tasks/CancelBatchJobTransportAction.java

@@ -266,8 +266,12 @@ private void executeConnector(Connector connector, MLInput mlInput, ActionListen
             connectorExecutor.setClient(client);
             connectorExecutor.setXContentRegistry(xContentRegistry);
             connectorExecutor.executeAction(CANCEL_BATCH_PREDICT.name(), mlInput, ActionListener.wrap(taskResponse -> {
+                connectorExecutor.close();
                 processTaskResponse(taskResponse, actionListener);
-            }, actionListener::onFailure));
+            }, e -> {
+                connectorExecutor.close();
+                actionListener.onFailure(e);
+            }));
         }, e -> {
             log.error("Failed to decrypt credentials in connector", e);
             actionListener.onFailure(e);

plugin/src/main/java/org/opensearch/ml/action/tasks/GetTaskTransportAction.java

@@ -474,8 +474,10 @@ private void executeConnector(
             connectorExecutor.setClient(client);
             connectorExecutor.setXContentRegistry(xContentRegistry);
             connectorExecutor.executeAction(BATCH_PREDICT_STATUS.name(), mlInput, ActionListener.wrap(taskResponse -> {
+                connectorExecutor.close();
                 processTaskResponse(mlTask, taskId, isUserInitiatedGetTaskRequest, taskResponse, remoteJob, r, actionListener);
             }, e -> {
+                connectorExecutor.close();
                 // When the request to remote service fails, we will retry the request for next 10 minutes (10 runs).
                 // If it fails even then, we mark it as unreachable in task index and send message to DLQ
                 if (!isUserInitiatedGetTaskRequest) {