[ML] Send mid-stream errors to users (#114549) (#114746)

If apache sends an error mid stream, forward it to the user rather than
the now-ignored listener.

Author: prwhelan

docs/changelog/114549.yaml

@@ -0,0 +1,5 @@
+pr: 114549
+summary: Send mid-stream errors to users
+area: Machine Learning
+type: bug
+issues: []

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

@@ -153,44 +153,31 @@ public void stream(HttpRequest request, HttpContext context, ActionListener<Flow
         // The caller must call start() first before attempting to send a request
         assert status.get() == Status.STARTED : "call start() before attempting to send a request";
 
-        // apache can sometimes send us the same error in the consumer and the callback
-        // sometimes it sends us an error just on the callback
-        // notifyOnce will dedupe for us
-        var callOnceListener = ActionListener.notifyOnce(listener);
-
-        SocketAccess.doPrivileged(
-            () -> client.execute(
-                request.requestProducer(),
-                new StreamingHttpResultPublisher(threadPool, settings, callOnceListener),
-                context,
-                new FutureCallback<>() {
-                    @Override
-                    public void completed(HttpResponse response) {
-                        // StreamingHttpResultPublisher will publish results to the Flow.Publisher returned in the ActionListener
-                    }
-
-                    @Override
-                    public void failed(Exception ex) {
-                        throttlerManager.warn(
-                            logger,
-                            format("Request from inference entity id [%s] failed", request.inferenceEntityId()),
-                            ex
-                        );
-                        failUsingUtilityThread(ex, callOnceListener);
-                    }
-
-                    @Override
-                    public void cancelled() {
-                        failUsingUtilityThread(
+        var streamingProcessor = new StreamingHttpResultPublisher(threadPool, settings, listener);
+
+        SocketAccess.doPrivileged(() -> client.execute(request.requestProducer(), streamingProcessor, context, new FutureCallback<>() {
+            @Override
+            public void completed(HttpResponse response) {
+                streamingProcessor.close();
+            }
+
+            @Override
+            public void failed(Exception ex) {
+                threadPool.executor(UTILITY_THREAD_POOL_NAME).execute(() -> streamingProcessor.failed(ex));
+            }
+
+            @Override
+            public void cancelled() {
+                threadPool.executor(UTILITY_THREAD_POOL_NAME)
+                    .execute(
+                        () -> streamingProcessor.failed(
                             new CancellationException(
                                 format("Request from inference entity id [%s] was cancelled", request.inferenceEntityId())
-                            ),
-                            callOnceListener
-                        );
-                    }
-                }
-            )
-        );
+                            )
+                        )
+                    );
+            }
+        }));
     }
 
     @Override

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

@@ -65,7 +65,7 @@ class StreamingHttpResultPublisher implements HttpAsyncResponseConsumer<HttpResp
 
     StreamingHttpResultPublisher(ThreadPool threadPool, HttpSettings settings, ActionListener<Flow.Publisher<HttpResult>> listener) {
         this.settings = Objects.requireNonNull(settings);
-        this.listener = Objects.requireNonNull(listener);
+        this.listener = ActionListener.notifyOnce(Objects.requireNonNull(listener));
 
         this.taskRunner = new RequestBasedTaskRunner(new OffloadThread(), threadPool, UTILITY_THREAD_POOL_NAME);
     }
@@ -152,9 +152,13 @@ public void responseCompleted(HttpContext httpContext) {}
     @Override
     public void failed(Exception e) {
         if (this.isDone.compareAndSet(false, true)) {
-            ex = e;
-            queue.offer(() -> subscriber.onError(e));
-            taskRunner.requestNextRun();
+            if (listenerCalled.compareAndSet(false, true)) {
+                listener.onFailure(e);
+            } else {
+                ex = e;
+                queue.offer(() -> subscriber.onError(e));
+                taskRunner.requestNextRun();
+            }
         }
     }
 

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.junit.Before;
+import org.mockito.ArgumentCaptor;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -38,6 +39,7 @@
 import static org.hamcrest.Matchers.notNullValue;
 import static org.hamcrest.Matchers.nullValue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
@@ -59,7 +61,7 @@ public void setUp() throws Exception {
         super.setUp();
         threadPool = mock(ThreadPool.class);
         settings = mock(HttpSettings.class);
-        listener = ActionListener.noop();
+        listener = spy(ActionListener.noop());
 
         when(threadPool.executor(UTILITY_THREAD_POOL_NAME)).thenReturn(EsExecutors.DIRECT_EXECUTOR_SERVICE);
         when(settings.getMaxResponseSize()).thenReturn(ByteSizeValue.ofBytes(maxBytes));
@@ -235,33 +237,13 @@ public void testTotalBytesDecrement() throws IOException {
     }
 
     /**
-     * Given an error from Apache
-     * When the subscriber requests the next set of data
-     * Then the subscriber receives the error from Apache
+     * When there is an error from Apache before the publisher invokes the listener
+     * Then the publisher will forward the call to the listener's onFailure
      */
     public void testErrorBeforeRequest() {
-        var subscriber = subscribe();
         var exception = new NullPointerException("test");
-
         publisher.failed(exception);
-        assertThat("subscriber receives exception on next request", subscriber.throwable, nullValue());
-
-        subscriber.requestData();
-        assertThat("subscriber receives exception", subscriber.throwable, is(exception));
-    }
-
-    /**
-     * Given the subscriber is waiting for data
-     * When Apache sends an error
-     * Then the subscriber immediately receives the error
-     */
-    public void testErrorAfterRequest() {
-        var subscriber = subscribe();
-        var exception = new NullPointerException("test");
-
-        subscriber.requestData();
-        publisher.failed(exception);
-        assertThat("subscriber receives exception", subscriber.throwable, is(exception));
+        verify(listener).onFailure(exception);
     }
 
     /**
@@ -375,6 +357,76 @@ public void testCancelAfterRequest() {
         assertTrue("onComplete should be called", subscriber.completed);
     }
 
+    /**
+     * When cancel is called
+     * Then we only send onComplete once
+     */
+    public void testCancelIsIdempotent() throws IOException {
+        Flow.Subscriber<HttpResult> subscriber = mock();
+
+        var subscription = ArgumentCaptor.forClass(Flow.Subscription.class);
+        publisher.subscribe(subscriber);
+        verify(subscriber).onSubscribe(subscription.capture());
+
+        publisher.responseReceived(mock());
+        publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
+        subscription.getValue().request(1);
+
+        subscription.getValue().request(1);
+        publisher.cancel();
+        verify(subscriber, times(1)).onComplete();
+        subscription.getValue().request(1);
+        publisher.cancel();
+        verify(subscriber, times(1)).onComplete();
+    }
+
+    /**
+     * When close is called
+     * Then we only send onComplete once
+     */
+    public void testCloseIsIdempotent() throws IOException {
+        Flow.Subscriber<HttpResult> subscriber = mock();
+
+        var subscription = ArgumentCaptor.forClass(Flow.Subscription.class);
+        publisher.subscribe(subscriber);
+        verify(subscriber).onSubscribe(subscription.capture());
+
+        publisher.responseReceived(mock());
+        publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
+        subscription.getValue().request(1);
+
+        subscription.getValue().request(1);
+        publisher.close();
+        verify(subscriber, times(1)).onComplete();
+        subscription.getValue().request(1);
+        publisher.close();
+        verify(subscriber, times(1)).onComplete();
+    }
+
+    /**
+     * When failed is called
+     * Then we only send onError once
+     */
+    public void testFailedIsIdempotent() throws IOException {
+        var expectedException = new IllegalStateException("wow");
+        Flow.Subscriber<HttpResult> subscriber = mock();
+
+        var subscription = ArgumentCaptor.forClass(Flow.Subscription.class);
+        publisher.subscribe(subscriber);
+        verify(subscriber).onSubscribe(subscription.capture());
+
+        publisher.responseReceived(mock());
+        publisher.consumeContent(contentDecoder(message), mock(IOControl.class));
+        subscription.getValue().request(1);
+
+        subscription.getValue().request(1);
+        publisher.failed(expectedException);
+        verify(subscriber, times(1)).onError(eq(expectedException));
+        subscription.getValue().request(1);
+        publisher.failed(expectedException);
+        verify(subscriber, times(1)).onError(eq(expectedException));
+    }
+
     /**
      * Given the queue is being processed
      * When Apache cancels the publisher