[ML] Refactor Apache streaming (elastic#113726)

- Apache calls `consumeContent` once per batch, so it should be safe to
  reuse the bytebuffer by resetting it after each invocation. This saves
  on memory management a bit.
- Simplifying the initial response handling - we now invoke the listener
  once we get the first set of bytes, either for a successful response
  or a failure response, and we will now just verify the HttpResponse
  status line on the first response. This helps with some pausing when
  we received responses with just headers from some providers, and the
  response is now smoother.
- Response conversion now happens completely within the Flow processor
  rather than split between the processor and the response handler (so
  future code can all be in one place).

Co-authored-by: Elastic Machine <[email protected]>

Author: 2 people

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

@@ -42,6 +42,7 @@
 class StreamingHttpResultPublisher implements HttpAsyncResponseConsumer<HttpResponse>, Flow.Publisher<HttpResult> {
     private final HttpSettings settings;
     private final ActionListener<Flow.Publisher<HttpResult>> listener;
+    private final AtomicBoolean listenerCalled = new AtomicBoolean(false);
 
     // used to manage the HTTP response
     private volatile HttpResponse response;
@@ -57,6 +58,7 @@ class StreamingHttpResultPublisher implements HttpAsyncResponseConsumer<HttpResp
     private final Deque<Runnable> queue = new ConcurrentLinkedDeque<>();
 
     // used to control the flow of data from the Apache client, if we're producing more bytes than we can consume then we'll pause
+    private final SimpleInputBuffer inputBuffer = new SimpleInputBuffer(4096);
     private final AtomicLong bytesInQueue = new AtomicLong(0);
     private final Object ioLock = new Object();
     private volatile IOControl savedIoControl;
@@ -69,16 +71,8 @@ class StreamingHttpResultPublisher implements HttpAsyncResponseConsumer<HttpResp
     }
 
     @Override
-    public void responseReceived(HttpResponse httpResponse) throws IOException {
+    public void responseReceived(HttpResponse httpResponse) {
         this.response = httpResponse;
-        if (response.getEntity() == null || response.getEntity().getContentLength() <= 0) {
-            // on success, we may receive an empty content payload to initiate the stream
-            this.queue.offer(() -> subscriber.onNext(new HttpResult(response, new byte[0])));
-        } else {
-            var firstResponse = HttpResult.create(settings.getMaxResponseSize(), response);
-            this.queue.offer(() -> subscriber.onNext(firstResponse));
-        }
-        this.listener.onResponse(this);
     }
 
     @Override
@@ -100,32 +94,39 @@ public void consumeContent(ContentDecoder contentDecoder, IOControl ioControl) t
             return;
         }
 
-        var buffer = new SimpleInputBuffer(4096);
-        var consumed = buffer.consumeContent(contentDecoder);
-        var allBytes = new byte[consumed];
-        buffer.read(allBytes);
-
-        // we can have empty bytes, don't bother sending them
-        if (allBytes.length > 0) {
-            queue.offer(() -> {
-                subscriber.onNext(new HttpResult(response, allBytes));
-                var currentBytesInQueue = bytesInQueue.updateAndGet(current -> Long.max(0, current - allBytes.length));
-                if (savedIoControl != null) {
-                    var maxBytes = settings.getMaxResponseSize().getBytes() * 0.5;
-                    if (currentBytesInQueue <= maxBytes) {
-                        resumeProducer();
+        try {
+            var consumed = inputBuffer.consumeContent(contentDecoder);
+            var allBytes = new byte[consumed];
+            inputBuffer.read(allBytes);
+
+            // we can have empty bytes, don't bother sending them
+            if (allBytes.length > 0) {
+                queue.offer(() -> {
+                    subscriber.onNext(new HttpResult(response, allBytes));
+                    var currentBytesInQueue = bytesInQueue.updateAndGet(current -> Long.max(0, current - allBytes.length));
+                    if (savedIoControl != null) {
+                        var maxBytes = settings.getMaxResponseSize().getBytes() * 0.5;
+                        if (currentBytesInQueue <= maxBytes) {
+                            resumeProducer();
+                        }
                     }
-                }
-            });
-        }
+                });
+            }
 
-        // always check if totalByteSize > the configured setting in case the settings change
-        if (bytesInQueue.accumulateAndGet(allBytes.length, Long::sum) >= settings.getMaxResponseSize().getBytes()) {
-            pauseProducer(ioControl);
-        }
+            // always check if totalByteSize > the configured setting in case the settings change
+            if (bytesInQueue.accumulateAndGet(allBytes.length, Long::sum) >= settings.getMaxResponseSize().getBytes()) {
+                pauseProducer(ioControl);
+            }
+
+            // always run in case we're waking up from a pause and need to start a new thread
+            taskRunner.requestNextRun();
 
-        // always run in case we're waking up from a pause and need to start a new thread
-        taskRunner.requestNextRun();
+            if (listenerCalled.compareAndSet(false, true)) {
+                listener.onResponse(this);
+            }
+        } finally {
+            inputBuffer.reset();
+        }
     }
 
     private void pauseProducer(IOControl ioControl) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/ResponseHandler.java

@@ -64,12 +64,11 @@ default boolean canHandleStreamingResponses() {
      * HttpResults to the InferenceServiceResults.
      *
      * @param request The original request sent to the server
-     * @param result The first result that initiated the stream. If the result is HTTP 200, this result will not contain content bytes
-     * @param flow The remaining stream of results from the server.  If the result is HTTP 200, these results will contain content bytes
+     * @param flow    The remaining stream of results from the server.  If the result is HTTP 200, these results will contain content bytes
      * @return an inference results with {@link InferenceServiceResults#publisher()} set and {@link InferenceServiceResults#isStreaming()}
      * set to true.
      */
-    default InferenceServiceResults parseResult(Request request, HttpResult result, Flow.Publisher<HttpResult> flow) {
+    default InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpResult> flow) {
         assert canHandleStreamingResponses() == false : "This must be implemented when canHandleStreamingResponses() == true";
         throw new UnsupportedOperationException("This must be implemented when canHandleStreamingResponses() == true");
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RetryingHttpSender.java

@@ -116,7 +116,9 @@ public void tryAction(ActionListener<InferenceServiceResults> listener) {
             try {
                 if (request.isStreaming() && responseHandler.canHandleStreamingResponses()) {
                     httpClient.stream(request.createHttpRequest(), context, retryableListener.delegateFailure((l, r) -> {
-                        r.subscribe(new StreamingResponseHandler(throttlerManager, logger, request, responseHandler, l));
+                        var streamingResponseHandler = new StreamingResponseHandler(throttlerManager, logger, request, responseHandler);
+                        r.subscribe(streamingResponseHandler);
+                        l.onResponse(responseHandler.parseResult(request, streamingResponseHandler));
                     }));
                 } else {
                     httpClient.send(request.createHttpRequest(), context, retryableListener.delegateFailure((l, r) -> {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/StreamingResponseHandler.java

@@ -10,8 +10,6 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.ExceptionsHelper;
-import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
@@ -27,26 +25,18 @@ class StreamingResponseHandler implements Flow.Processor<HttpResult, HttpResult>
     private final Logger throttlerLogger;
     private final Request request;
     private final ResponseHandler responseHandler;
-    private final ActionListener<InferenceServiceResults> listener;
 
     private final AtomicBoolean upstreamIsClosed = new AtomicBoolean(false);
     private final AtomicBoolean processedFirstItem = new AtomicBoolean(false);
 
     private volatile Flow.Subscription upstream;
     private volatile Flow.Subscriber<? super HttpResult> downstream;
 
-    StreamingResponseHandler(
-        ThrottlerManager throttlerManager,
-        Logger throttlerLogger,
-        Request request,
-        ResponseHandler responseHandler,
-        ActionListener<InferenceServiceResults> listener
-    ) {
+    StreamingResponseHandler(ThrottlerManager throttlerManager, Logger throttlerLogger, Request request, ResponseHandler responseHandler) {
         this.throttlerManager = throttlerManager;
         this.throttlerLogger = throttlerLogger;
         this.request = request;
         this.responseHandler = responseHandler;
-        this.listener = listener;
     }
 
     @Override
@@ -90,27 +80,21 @@ public void cancel() {
     @Override
     public void onSubscribe(Flow.Subscription subscription) {
         upstream = subscription;
-        // start the first request, which will call onNext and validate the first HttpResult
-        upstream.request(1);
     }
 
     @Override
     public void onNext(HttpResult item) {
         if (processedFirstItem.compareAndSet(false, true)) {
             try {
                 responseHandler.validateResponse(throttlerManager, throttlerLogger, request, item);
-                var inferenceServiceResults = responseHandler.parseResult(request, item, this);
-                assert downstream != null : "the responseHandler must invoke the subscribe method";
-                listener.onResponse(inferenceServiceResults);
             } catch (Exception e) {
                 logException(throttlerLogger, request, item, responseHandler.getRequestType(), e);
-                listener.onFailure(e);
                 upstream.cancel();
                 onError(e);
+                return;
             }
-        } else {
-            downstream.onNext(item);
         }
+        downstream.onNext(item);
     }
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/openai/OpenAiChatCompletionResponseHandler.java

@@ -35,7 +35,7 @@ public boolean canHandleStreamingResponses() {
     }
 
     @Override
-    public InferenceServiceResults parseResult(Request request, HttpResult result, Flow.Publisher<HttpResult> flow) {
+    public InferenceServiceResults parseResult(Request request, Flow.Publisher<HttpResult> flow) {
         var serverSentEventProcessor = new ServerSentEventProcessor(new ServerSentEventParser());
         var openAiProcessor = new OpenAiStreamingProcessor();
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/amazonbedrock/completion/AmazonBedrockChatCompletionResponseListener.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.request.amazonbedrock.completion.AmazonBedrockChatCompletionRequest;
 import org.elasticsearch.xpack.inference.external.response.amazonbedrock.AmazonBedrockResponseHandler;
 import org.elasticsearch.xpack.inference.external.response.amazonbedrock.AmazonBedrockResponseListener;
@@ -29,7 +30,7 @@ public AmazonBedrockChatCompletionResponseListener(
     @Override
     public void onResponse(ConverseResult result) {
         ((AmazonBedrockChatCompletionResponseHandler) responseHandler).acceptChatCompletionResponseObject(result);
-        inferenceResultsListener.onResponse(responseHandler.parseResult(request, null));
+        inferenceResultsListener.onResponse(responseHandler.parseResult(request, (HttpResult) null));
     }
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/amazonbedrock/embeddings/AmazonBedrockEmbeddingsResponseListener.java

@@ -11,6 +11,7 @@
 
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.xpack.inference.external.http.HttpResult;
 import org.elasticsearch.xpack.inference.external.request.amazonbedrock.embeddings.AmazonBedrockEmbeddingsRequest;
 import org.elasticsearch.xpack.inference.external.response.amazonbedrock.AmazonBedrockResponseHandler;
 import org.elasticsearch.xpack.inference.external.response.amazonbedrock.AmazonBedrockResponseListener;
@@ -28,7 +29,7 @@ public AmazonBedrockEmbeddingsResponseListener(
     @Override
     public void onResponse(InvokeModelResult result) {
         ((AmazonBedrockEmbeddingsResponseHandler) responseHandler).acceptEmbeddingsResult(result);
-        inferenceResultsListener.onResponse(responseHandler.parseResult(request, null));
+        inferenceResultsListener.onResponse(responseHandler.parseResult(request, (HttpResult) null));
     }
 
     @Override

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

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.inference.external.http;
 
-import org.apache.http.HttpEntity;
 import org.apache.http.HttpResponse;
 import org.apache.http.nio.ContentDecoder;
 import org.apache.http.nio.IOControl;
@@ -19,7 +18,6 @@
 import org.elasticsearch.threadpool.ThreadPool;
 import org.junit.Before;
 
-import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
@@ -80,25 +78,7 @@ public void testFirstResponseCallsListener() throws IOException {
         publisher = new StreamingHttpResultPublisher(threadPool, settings, listener);
 
         publisher.responseReceived(mock(HttpResponse.class));
-
-        assertThat("Listener's onResponse should be called when we receive a response", latch.getCount(), equalTo(0L));
-    }
-
-    /**
-     * When we receive an http response with an entity with no content
-     * Then we call the listener
-     * And we queue the initial payload
-     */
-    public void testEmptyFirstResponseCallsListener() throws IOException {
-        var latch = new CountDownLatch(1);
-        var listener = ActionTestUtils.<Flow.Publisher<HttpResult>>assertNoFailureListener(r -> latch.countDown());
-        publisher = new StreamingHttpResultPublisher(threadPool, settings, listener);
-
-        var response = mock(HttpResponse.class);
-        var entity = mock(HttpEntity.class);
-        when(entity.getContentLength()).thenReturn(-1L);
-        when(response.getEntity()).thenReturn(entity);
-        publisher.responseReceived(response);
+        publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
 
         assertThat("Listener's onResponse should be called when we receive a response", latch.getCount(), equalTo(0L));
     }
@@ -114,12 +94,8 @@ public void testNonEmptyFirstResponseCallsListener() throws IOException {
         publisher = new StreamingHttpResultPublisher(threadPool, settings, listener);
 
         when(settings.getMaxResponseSize()).thenReturn(ByteSizeValue.ofBytes(9000));
-        var response = mock(HttpResponse.class);
-        var entity = mock(HttpEntity.class);
-        when(entity.getContentLength()).thenReturn(5L);
-        when(entity.getContent()).thenReturn(new ByteArrayInputStream(message));
-        when(response.getEntity()).thenReturn(entity);
-        publisher.responseReceived(response);
+        publisher.responseReceived(mock(HttpResponse.class));
+        publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
 
         assertThat("Listener's onResponse should be called when we receive a response", latch.getCount(), equalTo(0L));
     }
@@ -146,6 +122,7 @@ public void testNonEmptyFirstResponseCallsListener() throws IOException {
     public void testSubscriberAndPublisherExchange() throws IOException {
         var subscriber = new TestSubscriber();
         publisher.responseReceived(mock(HttpResponse.class));
+        publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
 
         // subscribe
         publisher.subscribe(subscriber);
@@ -155,7 +132,6 @@ public void testSubscriberAndPublisherExchange() throws IOException {
         // request the initial http response
         subscriber.requestData();
         assertThat("onNext was called with the initial HttpResponse", subscriber.httpResult, notNullValue());
-        assertTrue("HttpResponse has an empty body (because there is no HttpEntity)", subscriber.httpResult.isBodyEmpty());
         subscriber.httpResult = null; // reset test
 
         // subscriber requests data, publisher has not sent data yet
@@ -175,14 +151,14 @@ public void testNon200Response() throws IOException {
         var subscriber = new TestSubscriber();
         // Apache sends a single response and closes the consumer
         publisher.responseReceived(mock(HttpResponse.class));
+        publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
         publisher.close();
 
         // subscriber requests data
         publisher.subscribe(subscriber);
         assertThat("subscribe must call onSubscribe", subscriber.subscription, notNullValue());
         subscriber.requestData();
         assertThat("onNext was called with the initial HttpResponse", subscriber.httpResult, notNullValue());
-        assertTrue("HttpResponse has an empty body (because there is no HttpEntity)", subscriber.httpResult.isBodyEmpty());
         subscriber.requestData();
         assertTrue("Publisher has been closed", publisher.isDone());
         assertTrue("Subscriber has been completed", subscriber.completed);
@@ -233,6 +209,7 @@ public void testTotalBytesDecrement() throws IOException {
 
         var subscriber = new TestSubscriber();
         publisher.responseReceived(mock(HttpResponse.class));
+        publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
         publisher.subscribe(subscriber);
         subscriber.requestData();
         subscriber.httpResult = null;
@@ -476,18 +453,25 @@ public void testDoubleSubscribeFails() {
      * When a new request is processed
      * Then it should reuse that ML Utility thread
      */
-    public void testReuseMlThread() throws IOException, ExecutionException, InterruptedException, TimeoutException {
+    public void testReuseMlThread() throws ExecutionException, InterruptedException, TimeoutException {
         try {
             threadPool = spy(createThreadPool(inferenceUtilityPool()));
             publisher = new StreamingHttpResultPublisher(threadPool, settings, listener);
             var subscriber = new TestSubscriber();
             publisher.responseReceived(mock(HttpResponse.class));
             publisher.subscribe(subscriber);
 
-            CompletableFuture.runAsync(subscriber::requestData, threadPool.executor(UTILITY_THREAD_POOL_NAME)).get(5, TimeUnit.SECONDS);
+            CompletableFuture.runAsync(() -> {
+                try {
+                    publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
+                    subscriber.requestData();
+                } catch (IOException e) {
+                    throw new RuntimeException(e);
+                }
+            }, threadPool.executor(UTILITY_THREAD_POOL_NAME)).get(5, TimeUnit.SECONDS);
             verify(threadPool, times(1)).executor(UTILITY_THREAD_POOL_NAME);
             assertThat("onNext was called with the initial HttpResponse", subscriber.httpResult, notNullValue());
-            assertTrue("HttpResponse has an empty body (because there is no HttpEntity)", subscriber.httpResult.isBodyEmpty());
+            assertFalse("Expected HttpResult to have data", subscriber.httpResult.isBodyEmpty());
         } finally {
             terminate(threadPool);
         }
@@ -514,6 +498,7 @@ public void testCancelBreaksInfiniteLoop() throws Exception {
 
             publisher = new StreamingHttpResultPublisher(threadPool, settings, listener);
             publisher.responseReceived(mock(HttpResponse.class));
+            publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
             // create an infinitely running Subscriber
             var subscriber = new Flow.Subscriber<HttpResult>() {
                 Flow.Subscription subscription;
@@ -616,6 +601,7 @@ private TestSubscriber subscribe() {
 
     private TestSubscriber runBefore(Runnable runDuringOnNext) throws IOException {
         publisher.responseReceived(mock(HttpResponse.class));
+        publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
         TestSubscriber subscriber = new TestSubscriber() {
             public void onNext(HttpResult item) {
                 runDuringOnNext.run();
@@ -628,6 +614,7 @@ public void onNext(HttpResult item) {
 
     private TestSubscriber runAfter(Runnable runDuringOnNext) throws IOException {
         publisher.responseReceived(mock(HttpResponse.class));
+        publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
         TestSubscriber subscriber = new TestSubscriber() {
             public void onNext(HttpResult item) {
                 runDuringOnNext.run();