Add http request content stream support

modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java

@@ -0,0 +1,124 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.http.netty4;
+
+import org.apache.http.entity.ByteArrayEntity;
+import org.apache.http.entity.ContentType;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.ESNetty4IntegTestCase;
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.internal.node.NodeClient;
+import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.settings.ClusterSettings;
+import org.elasticsearch.common.settings.IndexScopedSettings;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.settings.SettingsFilter;
+import org.elasticsearch.common.util.CollectionUtils;
+import org.elasticsearch.features.NodeFeature;
+import org.elasticsearch.http.HttpContent;
+import org.elasticsearch.plugins.ActionPlugin;
+import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.rest.BaseRestHandler;
+import org.elasticsearch.rest.RestController;
+import org.elasticsearch.rest.RestHandler;
+import org.elasticsearch.rest.RestRequest;
+import org.elasticsearch.rest.RestResponse;
+import org.elasticsearch.rest.RestStatus;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
+import java.util.function.Predicate;
+import java.util.function.Supplier;
+
+public class Netty4HttpRequestStreamIT extends ESNetty4IntegTestCase {
+
+    @Override
+    protected Collection<Class<? extends Plugin>> nodePlugins() {
+        return CollectionUtils.concatLists(List.of(RequestContentStreamPlugin.class), super.nodePlugins());
+    }
+
+    @Override
+    protected boolean addMockHttpTransport() {
+        return false; // enable http
+    }
+
+    public void testBasicStream() throws IOException {
+        var totalBytes = 1024 * 1024;
+        var request = new Request("POST", RequestContentStreamPlugin.ROUTE);
+        request.setEntity(new ByteArrayEntity(randomByteArrayOfLength(totalBytes), ContentType.APPLICATION_JSON));
+
+        var respose = getRestClient().performRequest(request);
+        assertEquals(200, respose.getStatusLine().getStatusCode());
+        var gotTotalBytes = new BytesArray(respose.getEntity().getContent().readAllBytes()).utf8ToString();
+        assertEquals("" + totalBytes, gotTotalBytes);
+    }
+
+    public static class RequestContentStreamPlugin extends Plugin implements ActionPlugin {
+
+        static final String ROUTE = "/_test/request-stream/basic";
+        private static final Logger LOGGER = LogManager.getLogger("StreamRestHandler");
+
+        @Override
+        public Collection<RestHandler> getRestHandlers(
+            Settings settings,
+            NamedWriteableRegistry namedWriteableRegistry,
+            RestController restController,
+            ClusterSettings clusterSettings,
+            IndexScopedSettings indexScopedSettings,
+            SettingsFilter settingsFilter,
+            IndexNameExpressionResolver indexNameExpressionResolver,
+            Supplier<DiscoveryNodes> nodesInCluster,
+            Predicate<NodeFeature> clusterSupportsFeature
+        ) {
+            return List.of(new BaseRestHandler() {
+                @Override
+                public String getName() {
+                    return ROUTE;
+                }
+
+                @Override
+                public List<Route> routes() {
+                    return List.of(new Route(RestRequest.Method.POST, ROUTE));
+                }
+
+                @Override
+                protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) {
+                    return channel -> {
+
+                        // Example of handling streamed content
+
+                        var contentStream = request.contentStream();
+                        var totalReceivedBytes = new AtomicInteger();
+
+                        Consumer<HttpContent.Chunk> contentConsumer = (chunk) -> {
+                            LOGGER.info("got next chunk with size {}", chunk.bytes().length());
+                            totalReceivedBytes.addAndGet(chunk.bytes().length());
+                            if (chunk.isLast() == false) {
+                                contentStream.request(1024 * 10); // ask for 10kb, will round up to 16kb
+                            } else {
+                                channel.sendResponse(new RestResponse(RestStatus.OK, Integer.toString(totalReceivedBytes.get())));
+                            }
+                        };
+
+                        contentStream.setHandler(contentConsumer); // must setup handler before first request
+                        contentStream.request(1024); // initiate first chunk, will round up to 8kb
+                    };
+                }
+            });
+        }
+    }
+
+}

modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpAggregator.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.http.netty4;
+
+import io.netty.handler.codec.http.HttpObjectAggregator;
+import io.netty.handler.codec.http.HttpRequest;
+
+import org.elasticsearch.http.HttpPreRequest;
+import org.elasticsearch.http.netty4.internal.HttpHeadersAuthenticatorUtils;
+
+import java.util.function.Predicate;
+
+public class Netty4HttpAggregator extends HttpObjectAggregator {
+    private static final Predicate<HttpPreRequest> IGNORE_TEST = (req) -> req.uri().startsWith("/_test/request-stream") == false;
+
+    private final Predicate<HttpPreRequest> decider;
+    private boolean shouldAggregate;
+
+    public Netty4HttpAggregator(int maxContentLength) {
+        this(maxContentLength, IGNORE_TEST);
+    }
+
+    public Netty4HttpAggregator(int maxContentLength, Predicate<HttpPreRequest> decider) {
+        super(maxContentLength);
+        this.decider = decider;
+    }
+
+    @Override
+    public boolean acceptInboundMessage(Object msg) throws Exception {
+        if (msg instanceof HttpRequest request) {
+            var preReq = HttpHeadersAuthenticatorUtils.asHttpPreRequest(request);
+            shouldAggregate = decider.test(preReq);
+        }
+        if (shouldAggregate) {
+            return super.acceptInboundMessage(msg);
+        } else {
+            return false;
+        }
+    }
+}

modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java

@@ -9,18 +9,24 @@
 package org.elasticsearch.http.netty4;
 
 import io.netty.buffer.ByteBuf;
+import io.netty.buffer.Unpooled;
 import io.netty.channel.Channel;
 import io.netty.channel.ChannelDuplexHandler;
 import io.netty.channel.ChannelFuture;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.channel.ChannelPromise;
 import io.netty.handler.codec.compression.JdkZlibEncoder;
+import io.netty.handler.codec.http.DefaultFullHttpRequest;
 import io.netty.handler.codec.http.DefaultHttpContent;
 import io.netty.handler.codec.http.DefaultHttpResponse;
 import io.netty.handler.codec.http.DefaultLastHttpContent;
+import io.netty.handler.codec.http.EmptyHttpHeaders;
 import io.netty.handler.codec.http.FullHttpRequest;
+import io.netty.handler.codec.http.HttpContent;
 import io.netty.handler.codec.http.HttpObject;
+import io.netty.handler.codec.http.HttpRequest;
 import io.netty.handler.codec.http.HttpResponse;
+import io.netty.handler.codec.http.LastHttpContent;
 import io.netty.handler.ssl.SslCloseCompletionEvent;
 import io.netty.util.ReferenceCountUtil;
 import io.netty.util.concurrent.Future;
@@ -70,6 +76,12 @@ private record ChunkedWrite(PromiseCombiner combiner, ChannelPromise onDone, Chu
     @Nullable
     private ChunkedWrite currentChunkedWrite;
 
+    /**
+     * HTTP request content stream for current request, it's null if there is no current request or request is fully-aggregated
+     */
+    @Nullable
+    private Netty4HttpRequestContentStream currentRequestStream;
+
     /*
      * The current read and write sequence numbers. Read sequence numbers are attached to requests in the order they are read from the
      * channel, and then transferred to responses. A response is not written to the channel context until its sequence number matches the
@@ -109,23 +121,48 @@ public Netty4HttpPipeliningHandler(
     public void channelRead(final ChannelHandlerContext ctx, final Object msg) {
         activityTracker.startActivity();
         try {
-            assert msg instanceof FullHttpRequest : "Should have fully aggregated message already but saw [" + msg + "]";
-            final FullHttpRequest fullHttpRequest = (FullHttpRequest) msg;
             final Netty4HttpRequest netty4HttpRequest;
-            if (fullHttpRequest.decoderResult().isFailure()) {
-                final Throwable cause = fullHttpRequest.decoderResult().cause();
-                final Exception nonError;
-                if (cause instanceof Error) {
-                    ExceptionsHelper.maybeDieOnAnotherThread(cause);
-                    nonError = new Exception(cause);
+            if (msg instanceof HttpRequest request) {
+                if (request.decoderResult().isFailure()) {
+                    final Throwable cause = request.decoderResult().cause();
+                    final Exception nonError;
+                    if (cause instanceof Error) {
+                        ExceptionsHelper.maybeDieOnAnotherThread(cause);
+                        nonError = new Exception(cause);
+                    } else {
+                        nonError = (Exception) cause;
+                    }
+                    netty4HttpRequest = new Netty4HttpRequest(readSequence++, (FullHttpRequest) request, nonError);
                 } else {
-                    nonError = (Exception) cause;
+                    if (request instanceof FullHttpRequest fullHttpRequest) {
+                        currentRequestStream = null;
+                        netty4HttpRequest = new Netty4HttpRequest(readSequence++, fullHttpRequest);
+                    } else {
+                        var contentStream = new Netty4HttpRequestContentStream(ctx.channel());
+                        currentRequestStream = contentStream;
+                        netty4HttpRequest = new Netty4HttpRequest(
+                            readSequence++,
+                            new DefaultFullHttpRequest(
+                                request.protocolVersion(),
+                                request.method(),
+                                request.uri(),
+                                Unpooled.EMPTY_BUFFER,
+                                request.headers(),
+                                EmptyHttpHeaders.INSTANCE
+                            ),
+                            contentStream
+                        );
+                    }
                 }
-                netty4HttpRequest = new Netty4HttpRequest(readSequence++, fullHttpRequest, nonError);
+                handlePipelinedRequest(ctx, netty4HttpRequest);
             } else {
-                netty4HttpRequest = new Netty4HttpRequest(readSequence++, fullHttpRequest);
+                assert msg instanceof HttpContent;
+                assert currentRequestStream != null;
+                currentRequestStream.handleNettyContent((HttpContent) msg);
+                if (msg instanceof LastHttpContent) {
+                    currentRequestStream = null;
+                }
             }
-            handlePipelinedRequest(ctx, netty4HttpRequest);
         } finally {
             activityTracker.stopActivity();
         }

modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java

@@ -20,12 +20,12 @@
 import io.netty.handler.codec.http.cookie.ServerCookieEncoder;
 
 import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.http.HttpContent;
 import org.elasticsearch.http.HttpRequest;
 import org.elasticsearch.http.HttpResponse;
 import org.elasticsearch.rest.ChunkedRestResponseBodyPart;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.RestStatus;
-import org.elasticsearch.transport.netty4.Netty4Utils;
 
 import java.util.AbstractMap;
 import java.util.Collection;
@@ -39,22 +39,26 @@
 public class Netty4HttpRequest implements HttpRequest {
 
     private final FullHttpRequest request;
-    private final BytesReference content;
+    private final HttpContent content;
     private final Map<String, List<String>> headers;
     private final AtomicBoolean released;
     private final Exception inboundException;
     private final boolean pooled;
     private final int sequence;
 
+    Netty4HttpRequest(int sequence, FullHttpRequest request, Netty4HttpRequestContentStream contentStream) {
+        this(sequence, request, new AtomicBoolean(false), false, contentStream, null);
+    }
+
     Netty4HttpRequest(int sequence, FullHttpRequest request) {
-        this(sequence, request, new AtomicBoolean(false), true, Netty4Utils.toBytesReference(request.content()));
+        this(sequence, request, new AtomicBoolean(false), true, new Netty4HttpRequestFullContent(request));
     }
 
     Netty4HttpRequest(int sequence, FullHttpRequest request, Exception inboundException) {
-        this(sequence, request, new AtomicBoolean(false), true, Netty4Utils.toBytesReference(request.content()), inboundException);
+        this(sequence, request, new AtomicBoolean(false), true, new Netty4HttpRequestFullContent(request), inboundException);
     }
 
-    private Netty4HttpRequest(int sequence, FullHttpRequest request, AtomicBoolean released, boolean pooled, BytesReference content) {
+    private Netty4HttpRequest(int sequence, FullHttpRequest request, AtomicBoolean released, boolean pooled, HttpContent content) {
         this(sequence, request, released, pooled, content, null);
     }
 
@@ -63,7 +67,7 @@ private Netty4HttpRequest(
         FullHttpRequest request,
         AtomicBoolean released,
         boolean pooled,
-        BytesReference content,
+        HttpContent content,
         Exception inboundException
     ) {
         this.sequence = sequence;
@@ -86,7 +90,7 @@ public String uri() {
     }
 
     @Override
-    public BytesReference content() {
+    public HttpContent content() {
         assert released.get() == false;
         return content;
     }
@@ -118,7 +122,7 @@ public HttpRequest releaseAndCopy() {
                 ),
                 new AtomicBoolean(false),
                 false,
-                Netty4Utils.toBytesReference(copiedContent)
+                content
             );
         } finally {
             release();
@@ -167,7 +171,13 @@ public HttpRequest removeHeader(String header) {
             copiedHeadersWithout,
             copiedTrailingHeadersWithout
         );
-        return new Netty4HttpRequest(sequence, requestWithoutHeader, released, pooled, content);
+        HttpContent copiedContent;
+        if (content.isFull()) {
+            copiedContent = HttpContent.fromBytesReference(content.asFull().bytes());
+        } else {
+            copiedContent = content;
+        }
+        return new Netty4HttpRequest(sequence, requestWithoutHeader, released, pooled, copiedContent);
     }
 
     @Override