Adding context preserving fix

Author: jonathan-buttner

x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/rest/ServerSentEventsRestActionListenerTests.java

@@ -17,6 +17,7 @@
 import org.apache.http.nio.util.SimpleInputBuffer;
 import org.apache.http.protocol.HttpContext;
 import org.apache.http.util.EntityUtils;
+import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.RequestOptions;
 import org.elasticsearch.client.Response;
@@ -43,6 +44,7 @@
 import org.elasticsearch.rest.RestHandler;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.test.ESIntegTestCase;
+import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.inference.external.response.streaming.ServerSentEvent;
@@ -52,6 +54,7 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Deque;
 import java.util.Iterator;
 import java.util.List;
@@ -96,6 +99,14 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
     }
 
     public static class StreamingPlugin extends Plugin implements ActionPlugin {
+        private final SetOnce<ThreadPool> threadPool = new SetOnce<>();
+
+        @Override
+        public Collection<?> createComponents(PluginServices services) {
+            threadPool.set(services.threadPool());
+            return Collections.emptyList();
+        }
+
         @Override
         public Collection<RestHandler> getRestHandlers(
             Settings settings,
@@ -122,7 +133,7 @@ public void handleRequest(RestRequest request, RestChannel channel, NodeClient c
                     var publisher = new RandomPublisher(requestCount, withError);
                     var inferenceServiceResults = new StreamingInferenceServiceResults(publisher);
                     var inferenceResponse = new InferenceAction.Response(inferenceServiceResults, inferenceServiceResults.publisher());
-                    new ServerSentEventsRestActionListener(channel).onResponse(inferenceResponse);
+                    new ServerSentEventsRestActionListener(channel, threadPool.get()).onResponse(inferenceResponse);
                 }
             }, new RestHandler() {
                 @Override
@@ -132,7 +143,7 @@ public List<Route> routes() {
 
                 @Override
                 public void handleRequest(RestRequest request, RestChannel channel, NodeClient client) {
-                    new ServerSentEventsRestActionListener(channel).onFailure(expectedException);
+                    new ServerSentEventsRestActionListener(channel, threadPool.get()).onFailure(expectedException);
                 }
             }, new RestHandler() {
                 @Override
@@ -143,7 +154,7 @@ public List<Route> routes() {
                 @Override
                 public void handleRequest(RestRequest request, RestChannel channel, NodeClient client) {
                     var inferenceResponse = new InferenceAction.Response(new SingleInferenceServiceResults());
-                    new ServerSentEventsRestActionListener(channel).onResponse(inferenceResponse);
+                    new ServerSentEventsRestActionListener(channel, threadPool.get()).onResponse(inferenceResponse);
                 }
             });
         }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java

@@ -197,9 +197,11 @@ public List<RestHandler> getRestHandlers(
         Supplier<DiscoveryNodes> nodesInCluster,
         Predicate<NodeFeature> clusterSupportsFeature
     ) {
+        assert serviceComponents.get() != null : "serviceComponents must be set before retrieving the rest handlers";
+
         var availableRestActions = List.of(
             new RestInferenceAction(),
-            new RestStreamInferenceAction(),
+            new RestStreamInferenceAction(serviceComponents.get().threadPool()),
             new RestGetInferenceModelAction(),
             new RestPutInferenceModelAction(),
             new RestUpdateInferenceModelAction(),
@@ -208,7 +210,7 @@ public List<RestHandler> getRestHandlers(
             new RestGetInferenceServicesAction()
         );
         List<RestHandler> conditionalRestActions = UnifiedCompletionFeature.UNIFIED_COMPLETION_FEATURE_FLAG.isEnabled()
-            ? List.of(new RestUnifiedCompletionInferenceAction())
+            ? List.of(new RestUnifiedCompletionInferenceAction(serviceComponents.get().threadPool()))
             : List.of();
 
         return Stream.concat(availableRestActions.stream(), conditionalRestActions.stream()).toList();

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rest/RestStreamInferenceAction.java

@@ -11,16 +11,25 @@
 import org.elasticsearch.rest.RestChannel;
 import org.elasticsearch.rest.Scope;
 import org.elasticsearch.rest.ServerlessScope;
+import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 
 import java.util.List;
+import java.util.Objects;
 
 import static org.elasticsearch.rest.RestRequest.Method.POST;
 import static org.elasticsearch.xpack.inference.rest.Paths.STREAM_INFERENCE_ID_PATH;
 import static org.elasticsearch.xpack.inference.rest.Paths.STREAM_TASK_TYPE_INFERENCE_ID_PATH;
 
 @ServerlessScope(Scope.PUBLIC)
 public class RestStreamInferenceAction extends BaseInferenceAction {
+    private final ThreadPool threadPool;
+
+    public RestStreamInferenceAction(ThreadPool threadPool) {
+        super();
+        this.threadPool = Objects.requireNonNull(threadPool);
+    }
+
     @Override
     public String getName() {
         return "stream_inference_action";
@@ -38,6 +47,6 @@ protected InferenceAction.Request prepareInferenceRequest(InferenceAction.Reques
 
     @Override
     protected ActionListener<InferenceAction.Response> listener(RestChannel channel) {
-        return new ServerSentEventsRestActionListener(channel);
+        return new ServerSentEventsRestActionListener(channel, threadPool);
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rest/RestUnifiedCompletionInferenceAction.java

@@ -12,17 +12,26 @@
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.Scope;
 import org.elasticsearch.rest.ServerlessScope;
+import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.inference.action.UnifiedCompletionAction;
 
 import java.io.IOException;
 import java.util.List;
+import java.util.Objects;
 
 import static org.elasticsearch.rest.RestRequest.Method.POST;
 import static org.elasticsearch.xpack.inference.rest.Paths.UNIFIED_INFERENCE_ID_PATH;
 import static org.elasticsearch.xpack.inference.rest.Paths.UNIFIED_TASK_TYPE_INFERENCE_ID_PATH;
 
 @ServerlessScope(Scope.PUBLIC)
 public class RestUnifiedCompletionInferenceAction extends BaseRestHandler {
+    private final ThreadPool threadPool;
+
+    public RestUnifiedCompletionInferenceAction(ThreadPool threadPool) {
+        super();
+        this.threadPool = Objects.requireNonNull(threadPool);
+    }
+
     @Override
     public String getName() {
         return "unified_inference_action";
@@ -44,6 +53,10 @@ protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient
             request = UnifiedCompletionAction.Request.parseRequest(params.inferenceEntityId(), params.taskType(), inferTimeout, parser);
         }
 
-        return channel -> client.execute(UnifiedCompletionAction.INSTANCE, request, new ServerSentEventsRestActionListener(channel));
+        return channel -> client.execute(
+            UnifiedCompletionAction.INSTANCE,
+            request,
+            new ServerSentEventsRestActionListener(channel, threadPool)
+        );
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rest/ServerSentEventsRestActionListener.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.ContextPreservingActionListener;
 import org.elasticsearch.common.bytes.ReleasableBytesReference;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.BytesStream;
@@ -29,6 +30,7 @@
 import org.elasticsearch.rest.RestResponse;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.tasks.TaskCancelledException;
+import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
@@ -38,6 +40,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.Flow;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -55,6 +58,7 @@ public class ServerSentEventsRestActionListener implements ActionListener<Infere
     private final AtomicBoolean isLastPart = new AtomicBoolean(false);
     private final RestChannel channel;
     private final ToXContent.Params params;
+    private final ThreadPool threadPool;
 
     /**
      * A listener for the first part of the next entry to become available for transmission.
@@ -66,13 +70,14 @@ public class ServerSentEventsRestActionListener implements ActionListener<Infere
      */
     private ActionListener<ChunkedRestResponseBodyPart> nextBodyPartListener;
 
-    public ServerSentEventsRestActionListener(RestChannel channel) {
-        this(channel, channel.request());
+    public ServerSentEventsRestActionListener(RestChannel channel, ThreadPool threadPool) {
+        this(channel, channel.request(), threadPool);
     }
 
-    public ServerSentEventsRestActionListener(RestChannel channel, ToXContent.Params params) {
+    public ServerSentEventsRestActionListener(RestChannel channel, ToXContent.Params params, ThreadPool threadPool) {
         this.channel = channel;
         this.params = params;
+        this.threadPool = Objects.requireNonNull(threadPool);
     }
 
     @Override
@@ -99,7 +104,7 @@ protected void ensureOpen() {
     }
 
     private void initializeStream(InferenceAction.Response response) {
-        nextBodyPartListener = ActionListener.wrap(bodyPart -> {
+        ActionListener<ChunkedRestResponseBodyPart> chunkedResponseBodyActionListener = ActionListener.wrap(bodyPart -> {
             // this is the first response, so we need to send the RestResponse to open the stream
             // all subsequent bytes will be delivered through the nextBodyPartListener
             channel.sendResponse(RestResponse.chunked(RestStatus.OK, bodyPart, this::release));
@@ -115,6 +120,12 @@ private void initializeStream(InferenceAction.Response response) {
                 )
             );
         });
+
+        nextBodyPartListener = ContextPreservingActionListener.wrapPreservingContext(
+            chunkedResponseBodyActionListener,
+            threadPool.getThreadContext()
+        );
+
         // subscribe will call onSubscribe, which requests the first chunk
         response.publisher().subscribe(subscriber);
     }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rest/RestStreamInferenceActionTests.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.test.rest.FakeRestRequest;
 import org.elasticsearch.test.rest.RestActionTestCase;
+import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.junit.Before;
@@ -25,7 +26,7 @@ public class RestStreamInferenceActionTests extends RestActionTestCase {
 
     @Before
     public void setUpAction() {
-        controller().registerHandler(new RestStreamInferenceAction());
+        controller().registerHandler(new RestStreamInferenceAction(new TestThreadPool(getTestName())));
     }
 
     public void testStreamIsTrue() {

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rest/RestUnifiedCompletionInferenceActionTests.java

@@ -17,6 +17,7 @@
 import org.elasticsearch.rest.RestResponse;
 import org.elasticsearch.test.rest.FakeRestRequest;
 import org.elasticsearch.test.rest.RestActionTestCase;
+import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.core.inference.action.UnifiedCompletionAction;
 import org.junit.Before;
@@ -30,7 +31,7 @@ public class RestUnifiedCompletionInferenceActionTests extends RestActionTestCas
 
     @Before
     public void setUpAction() {
-        controller().registerHandler(new RestUnifiedCompletionInferenceAction());
+        controller().registerHandler(new RestUnifiedCompletionInferenceAction(new TestThreadPool(getTestName())));
     }
 
     public void testStreamIsTrue() {