[ML] Adding asynchronous start up logic for the inference API internals

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSender.java

@@ -7,7 +7,9 @@
 
 package org.elasticsearch.xpack.inference.external.http.sender;
 
+import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.cluster.service.ClusterService;
@@ -77,6 +79,7 @@ public Sender createSender() {
         }
     }
 
+    private static final Logger logger = LogManager.getLogger(HttpRequestSender.class);
     private static final TimeValue START_COMPLETED_WAIT_TIME = TimeValue.timeValueSeconds(5);
 
     private final ThreadPool threadPool;
@@ -133,8 +136,14 @@ private void startInternal(ActionListener<Void> listener) {
     @Override
     public void startSynchronously() {
         if (started.compareAndSet(false, true)) {
-            startInternal(ActionListener.noop());
+            ActionListener<Void> listener = ActionListener.wrap(unused -> {}, exception -> {
+                logger.error("Http sender failed to start", exception);
+                ExceptionsHelper.maybeDieOnAnotherThread(exception);
+            });
+            startInternal(listener);
         }
+        // Handle the case where start*() was already called and this would return immediately because the started flag is already true
+        waitForStartToComplete();
     }
 
     private void waitForStartToComplete() {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/client/AmazonBedrockRequestSender.java

@@ -99,6 +99,7 @@ protected AmazonBedrockRequestSender(
 
     @Override
     public void startAsynchronously(ActionListener<Void> listener) {
+
         throw new UnsupportedOperationException("not implemented");
     }
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderTests.java

@@ -36,6 +36,7 @@
 import org.junit.Before;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.EnumSet;
 import java.util.List;
 import java.util.Locale;
@@ -61,6 +62,7 @@
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -111,6 +113,91 @@ public void testCreateSender_CanCallStartMultipleTimes() throws Exception {
         }
     }
 
+    public void testStart_ThrowsException_WhenAnErrorOccurs() throws IOException {
+        var mockManager = mock(HttpClientManager.class);
+        when(mockManager.getHttpClient()).thenReturn(mock(HttpClient.class));
+        doThrow(new Error("failed")).when(mockManager).start();
+
+        var senderFactory = new HttpRequestSender.Factory(
+            ServiceComponentsTests.createWithEmptySettings(threadPool),
+            mockManager,
+            mockClusterServiceEmpty()
+        );
+
+        try (var sender = senderFactory.createSender()) {
+            // Checking for both exception types because there's a race condition between the Error being thrown on a separate thread
+            // and the startCompleted latch timing out waiting for the start to complete
+            var exception = expectThrowsAnyOf(List.of(Error.class, IllegalStateException.class), sender::startSynchronously);
+
+            if (exception instanceof Error) {
+                assertThat(exception.getMessage(), is("failed"));
+            } else {
+                // IllegalStateException can be thrown if the startCompleted latch times out waiting for the start to complete
+                assertThat(exception.getMessage(), is("Http sender startup did not complete in time"));
+            }
+        }
+    }
+
+    public void testStart_ThrowsExceptionWaitingForStartToComplete() throws IOException {
+        var mockManager = mock(HttpClientManager.class);
+        when(mockManager.getHttpClient()).thenReturn(mock(HttpClient.class));
+        // This won't get rethrown because it is not an Error
+        doThrow(new IllegalArgumentException("failed")).when(mockManager).start();
+
+        var senderFactory = new HttpRequestSender.Factory(
+            ServiceComponentsTests.createWithEmptySettings(threadPool),
+            mockManager,
+            mockClusterServiceEmpty()
+        );
+
+        try (var sender = senderFactory.createSender()) {
+            var exception = expectThrows(IllegalStateException.class, sender::startSynchronously);
+
+            assertThat(exception.getMessage(), is("Http sender startup did not complete in time"));
+        }
+    }
+
+    public void testCreateSender_CanCallStartAsyncMultipleTimes() throws Exception {
+        var asyncCalls = 3;
+        var senderFactory = new HttpRequestSender.Factory(createWithEmptySettings(threadPool), clientManager, mockClusterServiceEmpty());
+
+        try (var sender = createSender(senderFactory)) {
+            var listenerList = new ArrayList<PlainActionFuture<Void>>();
+
+            for (int i = 0; i < asyncCalls; i++) {
+                PlainActionFuture<Void> listener = new PlainActionFuture<>();
+                listenerList.add(listener);
+                sender.startAsynchronously(listener);
+            }
+
+            for (int i = 0; i < asyncCalls; i++) {
+                PlainActionFuture<Void> listener = listenerList.get(i);
+                assertNull(listener.actionGet(TIMEOUT));
+            }
+        }
+    }
+
+    public void testCreateSender_CanCallStartAsyncAndSyncMultipleTimes() throws Exception {
+        var asyncCalls = 3;
+        var senderFactory = new HttpRequestSender.Factory(createWithEmptySettings(threadPool), clientManager, mockClusterServiceEmpty());
+
+        try (var sender = createSender(senderFactory)) {
+            var listenerList = new ArrayList<PlainActionFuture<Void>>();
+
+            for (int i = 0; i < asyncCalls; i++) {
+                PlainActionFuture<Void> listener = new PlainActionFuture<>();
+                listenerList.add(listener);
+                sender.startAsynchronously(listener);
+                sender.startSynchronously();
+            }
+
+            for (int i = 0; i < asyncCalls; i++) {
+                PlainActionFuture<Void> listener = listenerList.get(i);
+                assertNull(listener.actionGet(TIMEOUT));
+            }
+        }
+    }
+
     public void testCreateSender_SendsRequestAndReceivesResponse() throws Exception {
         var senderFactory = new HttpRequestSender.Factory(createWithEmptySettings(threadPool), clientManager, mockClusterServiceEmpty());
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/amazonbedrock/client/AmazonBedrockMockRequestSender.java

@@ -71,7 +71,7 @@ public void startSynchronously() {
 
     @Override
     public void startAsynchronously(ActionListener<Void> listener) {
-        listener.onResponse(null);
+        throw new UnsupportedOperationException("not supported");
     }
 
     @Override

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/elastic/authorization/ElasticInferenceServiceAuthorizationRequestHandlerTests.java

@@ -197,6 +197,43 @@ public void testGetAuthorization_ReturnsAValidResponse() throws IOException {
         }
     }
 
+    public void testGetAuthorization_OnResponseCalledOnce() throws IOException {
+        var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager);
+        var eisGatewayUrl = getUrl(webServer);
+        var logger = mock(Logger.class);
+        var authHandler = new ElasticInferenceServiceAuthorizationRequestHandler(eisGatewayUrl, threadPool, logger);
+
+        PlainActionFuture<ElasticInferenceServiceAuthorizationModel> listener = new PlainActionFuture<>();
+        ActionListener<ElasticInferenceServiceAuthorizationModel> onlyOnceListener = ActionListener.assertOnce(listener);
+        String responseJson = """
+                {
+                    "models": [
+                        {
+                          "model_name": "model-a",
+                          "task_types": ["embed/text/sparse", "chat"]
+                        }
+                    ]
+                }
+            """;
+        webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson));
+
+        try (var sender = senderFactory.createSender()) {
+            authHandler.getAuthorization(onlyOnceListener, sender);
+            authHandler.waitForAuthRequestCompletion(TIMEOUT);
+
+            var authResponse = listener.actionGet(TIMEOUT);
+            assertThat(authResponse.getAuthorizedTaskTypes(), is(EnumSet.of(TaskType.SPARSE_EMBEDDING, TaskType.CHAT_COMPLETION)));
+            assertThat(authResponse.getAuthorizedModelIds(), is(Set.of("model-a")));
+            assertTrue(authResponse.isAuthorized());
+
+            var loggerArgsCaptor = ArgumentCaptor.forClass(String.class);
+            verify(logger, times(1)).debug(loggerArgsCaptor.capture());
+
+            var message = loggerArgsCaptor.getValue();
+            assertThat(message, is("Retrieving authorization information from the Elastic Inference Service."));
+        }
+    }
+
     public void testGetAuthorization_InvalidResponse() throws IOException {
         var senderMock = createMockSender();
         var senderFactory = mock(HttpRequestSender.Factory.class);