Avoid transport_worker thread in TransportBroadcastAction (#98001)

`TransportBroadcastAction` derivatives do work which scales as the
number of shards in the cluster, both during coordination and when
processing responses. We must therefore not do this work on
`transport_worker` threads.

Relates #97920

Author: DaveCTurner

docs/changelog/98001.yaml

@@ -0,0 +1,5 @@
+pr: 98001
+summary: Avoid `transport_worker` thread in `TransportBroadcastAction`
+area: Distributed
+type: bug
+issues: []

server/src/main/java/org/elasticsearch/action/support/broadcast/TransportBroadcastAction.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.action.ActionRunnable;
 import org.elasticsearch.action.NoShardAvailableActionException;
 import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.ChannelActionListener;
 import org.elasticsearch.action.support.HandledTransportAction;
 import org.elasticsearch.action.support.TransportActions;
 import org.elasticsearch.cluster.ClusterState;
@@ -26,14 +27,15 @@
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
-import org.elasticsearch.transport.TransportChannel;
-import org.elasticsearch.transport.TransportRequestHandler;
 import org.elasticsearch.transport.TransportService;
+import org.elasticsearch.transport.Transports;
 
 import java.io.IOException;
+import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReferenceArray;
 
@@ -49,31 +51,47 @@ public abstract class TransportBroadcastAction<
     protected final TransportService transportService;
     protected final IndexNameExpressionResolver indexNameExpressionResolver;
 
-    final String transportShardAction;
-    private final String shardExecutor;
+    private final String transportShardAction;
+    private final Executor executor;
 
     protected TransportBroadcastAction(
         String actionName,
         ClusterService clusterService,
         TransportService transportService,
         ActionFilters actionFilters,
         IndexNameExpressionResolver indexNameExpressionResolver,
-        Writeable.Reader<Request> request,
-        Writeable.Reader<ShardRequest> shardRequest,
-        String shardExecutor
+        Writeable.Reader<Request> requestReader,
+        Writeable.Reader<ShardRequest> shardRequestReader,
+        String executor
     ) {
-        super(actionName, transportService, actionFilters, request);
+        // TODO replace SAME when removing workaround for https://github.com/elastic/elasticsearch/issues/97916
+        super(actionName, transportService, actionFilters, requestReader, ThreadPool.Names.SAME);
         this.clusterService = clusterService;
         this.transportService = transportService;
         this.indexNameExpressionResolver = indexNameExpressionResolver;
         this.transportShardAction = actionName + "[s]";
-        this.shardExecutor = shardExecutor;
-
-        transportService.registerRequestHandler(transportShardAction, ThreadPool.Names.SAME, shardRequest, new ShardTransportHandler());
+        this.executor = transportService.getThreadPool().executor(executor);
+        assert this.executor != EsExecutors.DIRECT_EXECUTOR_SERVICE : "O(#shards) work must always fork to an appropriate executor";
+
+        transportService.registerRequestHandler(
+            transportShardAction,
+            executor,
+            shardRequestReader,
+            (request, channel, task) -> ActionListener.completeWith(
+                new ChannelActionListener<>(channel),
+                () -> shardOperation(request, task)
+            )
+        );
     }
 
     @Override
     protected void doExecute(Task task, Request request, ActionListener<Response> listener) {
+        // workaround for https://github.com/elastic/elasticsearch/issues/97916 - TODO remove this when we can
+        executor.execute(ActionRunnable.wrap(listener, l -> doExecuteForked(task, request, listener)));
+    }
+
+    protected void doExecuteForked(Task task, Request request, ActionListener<Response> listener) {
+        assert Transports.assertNotTransportThread("O(#shards) work must always fork to an appropriate executor");
         new AsyncBroadcastAction(task, request, listener).start();
     }
 
@@ -184,7 +202,7 @@ protected void sendShardRequest(DiscoveryNode node, ShardRequest shardRequest, A
                 node,
                 transportShardAction,
                 shardRequest,
-                new ActionListenerResponseHandler<>(listener, TransportBroadcastAction.this::readShardResponse)
+                new ActionListenerResponseHandler<>(listener, TransportBroadcastAction.this::readShardResponse, executor)
             );
         }
 
@@ -238,11 +256,8 @@ void onOperation(@Nullable ShardRouting shard, final ShardIterator shardIt, int
             }
         }
 
-        protected AtomicReferenceArray<Object> shardsResponses() {
-            return shardsResponses;
-        }
-
         protected void finishHim() {
+            assert Transports.assertNotTransportThread("O(#shards) work must always fork to an appropriate executor");
             ActionListener.completeWith(listener, () -> newResponse(request, shardsResponses, clusterState));
         }
 
@@ -269,24 +284,4 @@ void setFailure(ShardIterator shardIt, int shardIndex, Exception e) {
             }
         }
     }
-
-    class ShardTransportHandler implements TransportRequestHandler<ShardRequest> {
-
-        @Override
-        public void messageReceived(ShardRequest request, TransportChannel channel, Task task) throws Exception {
-            asyncShardOperation(request, task, ActionListener.wrap(channel::sendResponse, e -> {
-                try {
-                    channel.sendResponse(e);
-                } catch (Exception e1) {
-                    logger.warn(() -> format("Failed to send error response for action [%s] and request [%s]", actionName, request), e1);
-                }
-            }));
-        }
-    }
-
-    private void asyncShardOperation(ShardRequest request, Task task, ActionListener<ShardResponse> listener) {
-        transportService.getThreadPool()
-            .executor(shardExecutor)
-            .execute(ActionRunnable.supply(listener, () -> shardOperation(request, task)));
-    }
 }

server/src/test/java/org/elasticsearch/action/admin/indices/validate/query/TransportValidateQueryActionTests.java

@@ -9,11 +9,11 @@
 package org.elasticsearch.action.admin.indices.validate.query;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.index.IndexNotFoundException;
 import org.elasticsearch.test.ESSingleNodeTestCase;
 
-import java.util.concurrent.atomic.AtomicBoolean;
-
-import static org.hamcrest.Matchers.equalTo;
+import java.util.concurrent.TimeUnit;
 
 public class TransportValidateQueryActionTests extends ESSingleNodeTestCase {
 
@@ -24,24 +24,16 @@ public class TransportValidateQueryActionTests extends ESSingleNodeTestCase {
      * them garbled together, or trying to write one after the channel had closed, etc.
      */
     public void testListenerOnlyInvokedOnceWhenIndexDoesNotExist() {
-        final AtomicBoolean invoked = new AtomicBoolean();
-        final ActionListener<ValidateQueryResponse> listener = new ActionListener<>() {
-
-            @Override
-            public void onResponse(final ValidateQueryResponse validateQueryResponse) {
-                fail("onResponse should not be invoked in this failure case");
-            }
-
-            @Override
-            public void onFailure(final Exception e) {
-                if (invoked.compareAndSet(false, true) == false) {
-                    fail("onFailure invoked more than once");
-                }
-            }
-
-        };
-        client().admin().indices().validateQuery(new ValidateQueryRequest("non-existent-index"), listener);
-        assertThat(invoked.get(), equalTo(true)); // ensure that onFailure was invoked
+        expectThrows(
+            IndexNotFoundException.class,
+            () -> PlainActionFuture.<ValidateQueryResponse, RuntimeException>get(
+                future -> client().admin()
+                    .indices()
+                    .validateQuery(new ValidateQueryRequest("non-existent-index"), ActionListener.assertOnce(future)),
+                10,
+                TimeUnit.SECONDS
+            )
+        );
     }
 
 }