Revamp TransportRequest handlers to support Writeable (#26315)

Backport of PR to 6.x, excluding the changes to SimulatePipelineRequest

Author: ywelsch

server/src/main/java/org/elasticsearch/action/ActionRequest.java

@@ -34,6 +34,10 @@ public ActionRequest() {
         // this.listenerThreaded = request.listenerThreaded();
     }
 
+    public ActionRequest(StreamInput in) throws IOException {
+        super(in);
+    }
+
     public abstract ActionRequestValidationException validate();
 
     /**

server/src/main/java/org/elasticsearch/action/support/HandledTransportAction.java

@@ -23,6 +23,7 @@
 import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.threadpool.ThreadPool;
@@ -43,6 +44,12 @@ protected HandledTransportAction(Settings settings, String actionName, ThreadPoo
         this(settings, actionName, true, threadPool, transportService, actionFilters, indexNameExpressionResolver, request);
     }
 
+    protected HandledTransportAction(Settings settings, String actionName, ThreadPool threadPool, TransportService transportService,
+                                     ActionFilters actionFilters, Writeable.Reader<Request> requestReader,
+                                     IndexNameExpressionResolver indexNameExpressionResolver) {
+        this(settings, actionName, true, threadPool, transportService, actionFilters, requestReader, indexNameExpressionResolver);
+    }
+
     protected HandledTransportAction(Settings settings, String actionName, boolean canTripCircuitBreaker, ThreadPool threadPool,
                                      TransportService transportService, ActionFilters actionFilters,
                                      IndexNameExpressionResolver indexNameExpressionResolver, Supplier<Request> request) {
@@ -51,6 +58,14 @@ protected HandledTransportAction(Settings settings, String actionName, boolean c
             new TransportHandler());
     }
 
+    protected HandledTransportAction(Settings settings, String actionName, boolean canTripCircuitBreaker, ThreadPool threadPool,
+                                     TransportService transportService, ActionFilters actionFilters,
+                                     Writeable.Reader<Request> requestReader, IndexNameExpressionResolver indexNameExpressionResolver) {
+        super(settings, actionName, threadPool, actionFilters, indexNameExpressionResolver, transportService.getTaskManager());
+        transportService.registerRequestHandler(actionName, ThreadPool.Names.SAME, false, canTripCircuitBreaker, requestReader,
+            new TransportHandler());
+    }
+
     class TransportHandler implements TransportRequestHandler<Request> {
 
         @Override

server/src/main/java/org/elasticsearch/common/io/stream/Streamable.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.common.io.stream;
 
 import java.io.IOException;
+import java.util.function.Supplier;
 
 /**
  * Implementers can be written to a {@linkplain StreamOutput} and read from a {@linkplain StreamInput}. This allows them to be "thrown
@@ -43,4 +44,12 @@ public interface Streamable {
      * Write this object's fields to a {@linkplain StreamOutput}.
      */
     void writeTo(StreamOutput out) throws IOException;
+
+    static <T extends Streamable> Writeable.Reader<T> newWriteableReader(Supplier<T> supplier) {
+        return (StreamInput in) -> {
+            T request = supplier.get();
+            request.readFrom(in);
+            return request;
+        };
+    }
 }

server/src/main/java/org/elasticsearch/transport/RequestHandlerRegistry.java

@@ -19,11 +19,12 @@
 
 package org.elasticsearch.transport;
 
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.tasks.Task;
 import org.elasticsearch.tasks.TaskManager;
 
 import java.io.IOException;
-import java.util.function.Supplier;
 
 public class RequestHandlerRegistry<Request extends TransportRequest> {
 
@@ -32,15 +33,14 @@ public class RequestHandlerRegistry<Request extends TransportRequest> {
     private final boolean forceExecution;
     private final boolean canTripCircuitBreaker;
     private final String executor;
-    private final Supplier<Request> requestFactory;
     private final TaskManager taskManager;
+    private final Writeable.Reader<Request> requestReader;
 
-    public RequestHandlerRegistry(String action, Supplier<Request> requestFactory, TaskManager taskManager,
+    public RequestHandlerRegistry(String action, Writeable.Reader<Request> requestReader, TaskManager taskManager,
                                   TransportRequestHandler<Request> handler, String executor, boolean forceExecution,
                                   boolean canTripCircuitBreaker) {
         this.action = action;
-        this.requestFactory = requestFactory;
-        assert newRequest() != null;
+        this.requestReader = requestReader;
         this.handler = handler;
         this.forceExecution = forceExecution;
         this.canTripCircuitBreaker = canTripCircuitBreaker;
@@ -52,8 +52,8 @@ public String getAction() {
         return action;
     }
 
-    public Request newRequest() {
-            return requestFactory.get();
+    public Request newRequest(StreamInput in) throws IOException {
+        return requestReader.read(in);
     }
 
     public void processMessageReceived(Request request, TransportChannel channel) throws Exception {

server/src/main/java/org/elasticsearch/transport/TcpTransport.java

@@ -1505,9 +1505,8 @@ protected String handleRequest(TcpChannel channel, String profileName, final Str
                 }
                 transportChannel = new TcpTransportChannel(this, channel, transportName, action, requestId, version, profileName,
                     messageLengthBytes);
-                final TransportRequest request = reg.newRequest();
+                final TransportRequest request = reg.newRequest(stream);
                 request.remoteAddress(new TransportAddress(remoteAddress));
-                request.readFrom(stream);
                 // in case we throw an exception, i.e. when the limit is hit, we don't want to verify
                 validateRequest(stream, requestId, action);
                 threadPool.executor(reg.getExecutor()).execute(new RequestHandler(reg, request, transportChannel));

server/src/main/java/org/elasticsearch/transport/TransportActionProxy.java

@@ -21,6 +21,7 @@
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.threadpool.ThreadPool;
 
 import java.io.IOException;
@@ -100,11 +101,11 @@ public String executor() {
 
     static class ProxyRequest<T extends TransportRequest> extends TransportRequest {
         T wrapped;
-        Supplier<T> supplier;
+        Writeable.Reader<T> reader;
         DiscoveryNode targetNode;
 
-        ProxyRequest(Supplier<T> supplier) {
-            this.supplier = supplier;
+        ProxyRequest(Writeable.Reader<T> reader) {
+            this.reader = reader;
         }
 
         ProxyRequest(T wrapped, DiscoveryNode targetNode) {
@@ -116,8 +117,7 @@ static class ProxyRequest<T extends TransportRequest> extends TransportRequest {
         public void readFrom(StreamInput in) throws IOException {
             super.readFrom(in);
             targetNode = new DiscoveryNode(in);
-            wrapped = supplier.get();
-            wrapped.readFrom(in);
+            wrapped = reader.read(in);
         }
 
         @Override

server/src/main/java/org/elasticsearch/transport/TransportMessage.java

@@ -22,11 +22,12 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Streamable;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.transport.TransportAddress;
 
 import java.io.IOException;
 
-public abstract class TransportMessage implements Streamable {
+public abstract class TransportMessage implements Streamable, Writeable {
 
     private TransportAddress remoteAddress;
 

server/src/main/java/org/elasticsearch/transport/TransportRequest.java

@@ -39,6 +39,10 @@ public static class Empty extends TransportRequest {
     public TransportRequest() {
     }
 
+    public TransportRequest(StreamInput in) throws IOException {
+        parentTaskId = TaskId.readFromStream(in);
+    }
+
     /**
      * Set a reference to task that created this request.
      */

server/src/main/java/org/elasticsearch/transport/TransportService.java

@@ -32,6 +32,8 @@
 import org.elasticsearch.common.component.AbstractLifecycleComponent;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Streamable;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.common.logging.Loggers;
 import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.common.settings.ClusterSettings;
@@ -713,7 +715,24 @@ public <Request extends TransportRequest> void registerRequestHandler(String act
                                                     String executor, TransportRequestHandler<Request> handler) {
         handler = interceptor.interceptHandler(action, executor, false, handler);
         RequestHandlerRegistry<Request> reg = new RequestHandlerRegistry<>(
-            action, requestFactory, taskManager, handler, executor, false, true);
+            action, Streamable.newWriteableReader(requestFactory), taskManager, handler, executor, false, true);
+        registerRequestHandler(reg);
+    }
+
+    /**
+     * Registers a new request handler
+     *
+     * @param action         The action the request handler is associated with
+     * @param requestReader  a callable to be used construct new instances for streaming
+     * @param executor       The executor the request handling will be executed on
+     * @param handler        The handler itself that implements the request handling
+     */
+    public <Request extends TransportRequest> void registerRequestHandler(String action, String executor,
+                                                                          Writeable.Reader<Request> requestReader,
+                                                                          TransportRequestHandler<Request> handler) {
+        handler = interceptor.interceptHandler(action, executor, false, handler);
+        RequestHandlerRegistry<Request> reg = new RequestHandlerRegistry<>(
+            action, requestReader, taskManager, handler, executor, false, true);
         registerRequestHandler(reg);
     }
 
@@ -733,7 +752,28 @@ public <Request extends TransportRequest> void registerRequestHandler(String act
                                                                           TransportRequestHandler<Request> handler) {
         handler = interceptor.interceptHandler(action, executor, forceExecution, handler);
         RequestHandlerRegistry<Request> reg = new RequestHandlerRegistry<>(
-            action, request, taskManager, handler, executor, forceExecution, canTripCircuitBreaker);
+            action, Streamable.newWriteableReader(request), taskManager, handler, executor, forceExecution, canTripCircuitBreaker);
+        registerRequestHandler(reg);
+    }
+
+    /**
+     * Registers a new request handler
+     *
+     * @param action                The action the request handler is associated with
+     * @param requestReader               The request class that will be used to construct new instances for streaming
+     * @param executor              The executor the request handling will be executed on
+     * @param forceExecution        Force execution on the executor queue and never reject it
+     * @param canTripCircuitBreaker Check the request size and raise an exception in case the limit is breached.
+     * @param handler               The handler itself that implements the request handling
+     */
+    public <Request extends TransportRequest> void registerRequestHandler(String action,
+                                                                          String executor, boolean forceExecution,
+                                                                          boolean canTripCircuitBreaker,
+                                                                          Writeable.Reader<Request> requestReader,
+                                                                          TransportRequestHandler<Request> handler) {
+        handler = interceptor.interceptHandler(action, executor, forceExecution, handler);
+        RequestHandlerRegistry<Request> reg = new RequestHandlerRegistry<>(
+            action, requestReader, taskManager, handler, executor, forceExecution, canTripCircuitBreaker);
         registerRequestHandler(reg);
     }
 

server/src/test/java/org/elasticsearch/transport/TransportActionProxyTests.java

@@ -267,7 +267,7 @@ public void testIsProxyAction() {
     }
 
     public void testIsProxyRequest() {
-        assertTrue(TransportActionProxy.isProxyRequest(new TransportActionProxy.ProxyRequest<>(() -> null)));
+        assertTrue(TransportActionProxy.isProxyRequest(new TransportActionProxy.ProxyRequest<>((in) -> null)));
         assertFalse(TransportActionProxy.isProxyRequest(TransportRequest.Empty.INSTANCE));
     }
 }