From 2ff3afa7741b143960b11c0b3597343ca68a57df Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Mon, 29 Jul 2024 17:57:37 -0700 Subject: [PATCH 01/14] add http request content stream support --- .../netty4/Netty4HttpRequestStreamIT.java | 121 ++++++++++++++++++ .../http/netty4/Netty4HttpAggregator.java | 46 +++++++ .../netty4/Netty4HttpPipeliningHandler.java | 61 +++++++-- .../http/netty4/Netty4HttpRequest.java | 28 ++-- .../Netty4HttpRequestContentStream.java | 117 +++++++++++++++++ .../netty4/Netty4HttpRequestFullContent.java | 34 +++++ .../netty4/Netty4HttpServerTransport.java | 4 +- .../Netty4HttpServerTransportTests.java | 2 +- .../http/HttpClientStatsTracker.java | 14 +- .../org/elasticsearch/http/HttpContent.java | 73 +++++++++++ .../org/elasticsearch/http/HttpRequest.java | 2 +- .../org/elasticsearch/rest/RestRequest.java | 25 +++- .../http/HttpClientStatsTrackerTests.java | 4 +- .../elasticsearch/http/TestHttpRequest.java | 5 +- .../rest/RestControllerTests.java | 7 +- .../elasticsearch/rest/RestRequestTests.java | 3 +- .../test/rest/FakeRestRequest.java | 15 ++- .../audit/logfile/LoggingAuditTrailTests.java | 5 +- 18 files changed, 520 insertions(+), 46 deletions(-) create mode 100644 modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java create mode 100644 modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpAggregator.java create mode 100644 modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java create mode 100644 modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestFullContent.java create mode 100644 server/src/main/java/org/elasticsearch/http/HttpContent.java diff --git a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java new file mode 100644 index 0000000000000..5f7a114b535e4 --- /dev/null +++ b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java @@ -0,0 +1,121 @@ +/* + * 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> 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 getRestHandlers( + Settings settings, + NamedWriteableRegistry namedWriteableRegistry, + RestController restController, + ClusterSettings clusterSettings, + IndexScopedSettings indexScopedSettings, + SettingsFilter settingsFilter, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier nodesInCluster, + Predicate clusterSupportsFeature + ) { + return List.of(new BaseRestHandler() { + @Override + public String getName() { + return ROUTE; + } + + @Override + public List routes() { + return List.of(new Route(RestRequest.Method.POST, ROUTE)); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { + return channel -> { + var contentStream = request.contentStream(); + var totalReceivedBytes = new AtomicInteger(); + + Consumer contentConsumer = (chunk) -> { + LOGGER.info("got next chunk with size {}", chunk.bytes().length()); + totalReceivedBytes.addAndGet(chunk.bytes().length()); + if (chunk.isLast() == false) { + contentStream.request(1024); + } else { + channel.sendResponse(new RestResponse(RestStatus.OK, Integer.toString(totalReceivedBytes.get()))); + } + }; + + contentStream.setHandler(contentConsumer); + contentStream.request(1024); + }; + } + }); + } + } + +} diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpAggregator.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpAggregator.java new file mode 100644 index 0000000000000..16f1c2bbd2e37 --- /dev/null +++ b/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 IGNORE_TEST = (req) -> req.uri().startsWith("/_test/request-stream") == false; + + private final Predicate decider; + private boolean shouldAggregate; + + public Netty4HttpAggregator(int maxContentLength) { + this(maxContentLength, IGNORE_TEST); + } + + public Netty4HttpAggregator(int maxContentLength, Predicate 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; + } + } +} diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java index b915011514d9a..1a12be4d64a38 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java +++ b/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(); } diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java index 1e35f084c87ec..9d08d95deef2a 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java +++ b/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> 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 diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java new file mode 100644 index 0000000000000..4ace76271c18c --- /dev/null +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java @@ -0,0 +1,117 @@ +/* + * 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.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; +import io.netty.channel.Channel; +import io.netty.handler.codec.http.LastHttpContent; + +import org.elasticsearch.common.bytes.BytesReference; +import org.elasticsearch.http.HttpContent; +import org.elasticsearch.transport.netty4.Netty4Utils; + +import java.util.function.Consumer; + +public class Netty4HttpRequestContentStream implements HttpContent.Stream { + + private final Channel channel; + private Consumer handler; + private CompositeByteBuf aggregate; + private int requestedBytes = 0; + private boolean cancelled = false; + + public Netty4HttpRequestContentStream(Channel channel) { + this.channel = channel; + channel.config().setAutoRead(false); + } + + @Override + public Consumer handler() { + return handler; + } + + @Override + public void setHandler(Consumer chunkHandler) { + this.handler = chunkHandler; + } + + @Override + public void request(int bytes) { + assert handler != null : "handler must be set before requesting next chunk"; + requestedBytes += bytes; + channel.read(); + } + + @Override + public void cancel() { + cancelled = true; + channel.config().setAutoRead(true); + } + + public void handleNettyContent(io.netty.handler.codec.http.HttpContent httpContent) { + if (cancelled) { + httpContent.release(); + return; + } + + assert handler != null : "handler must be set before processing http content"; + var isLast = httpContent instanceof LastHttpContent; + var content = httpContent.content(); + + if (aggregate == null && content.readableBytes() >= requestedBytes) { + // network chunk is large enough and there is no ongoing aggregation + handler.accept(new Netty4ContentChunk(content, isLast)); + } else if (content.readableBytes() >= requestedBytes) { + // there is ongoing aggregation, and we reached requestBytes limit, send aggregated chunk downstream + aggregate.addComponent(content); + handler.accept(new Netty4ContentChunk(aggregate, isLast)); + aggregate = null; + } else { + // accumulation step + if (aggregate == null) { + aggregate = channel.alloc().compositeHeapBuffer(); + } + aggregate.addComponent(content); + channel.read(); + } + requestedBytes = Math.max(0, requestedBytes - content.readableBytes()); + if (isLast) { + channel.config().setAutoRead(true); + } + } + + static class Netty4ContentChunk implements Chunk { + + private final ByteBuf nettyBuffer; + private final BytesReference bytes; + private final boolean isLast; + + Netty4ContentChunk(ByteBuf nettyBuffer, boolean isLast) { + this.nettyBuffer = nettyBuffer; + this.bytes = Netty4Utils.toBytesReference(nettyBuffer); + this.isLast = isLast; + } + + @Override + public BytesReference bytes() { + return bytes; + } + + @Override + public void release() { + nettyBuffer.release(); + } + + @Override + public boolean isLast() { + return isLast; + } + } +} diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestFullContent.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestFullContent.java new file mode 100644 index 0000000000000..b45e73b8dac88 --- /dev/null +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestFullContent.java @@ -0,0 +1,34 @@ +/* + * 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.elasticsearch.common.bytes.BytesReference; +import org.elasticsearch.http.HttpContent; +import org.elasticsearch.transport.netty4.Netty4Utils; + +public class Netty4HttpRequestFullContent implements HttpContent.Full { + + private final io.netty.handler.codec.http.HttpContent httpContent; + private final BytesReference bytes; + + public Netty4HttpRequestFullContent(io.netty.handler.codec.http.HttpContent httpContent) { + this.httpContent = httpContent; + this.bytes = Netty4Utils.toBytesReference(httpContent.content()); + } + + @Override + public BytesReference bytes() { + return bytes; + } + + @Override + public void release() { + httpContent.release(); + } +} diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java index f48a3143fd016..75a5d0022a1e9 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java @@ -28,6 +28,7 @@ import io.netty.handler.codec.http.HttpResponse; import io.netty.handler.codec.http.HttpResponseEncoder; import io.netty.handler.codec.http.HttpUtil; +import io.netty.handler.flow.FlowControlHandler; import io.netty.handler.ssl.SslHandler; import io.netty.handler.timeout.ReadTimeoutException; import io.netty.handler.timeout.ReadTimeoutHandler; @@ -351,6 +352,7 @@ protected HttpMessage createMessage(String[] initialLine) throws Exception { } decoder.setCumulator(ByteToMessageDecoder.COMPOSITE_CUMULATOR); ch.pipeline().addLast("decoder", decoder); // parses the HTTP bytes request into HTTP message pieces + ch.pipeline().addLast(new FlowControlHandler()); if (httpValidator != null) { // runs a validation function on the first HTTP message piece which contains all the headers // if validation passes, the pieces of that particular request are forwarded, otherwise they are discarded @@ -364,7 +366,7 @@ protected HttpMessage createMessage(String[] initialLine) throws Exception { ); } // combines the HTTP message pieces into a single full HTTP request (with headers and body) - final HttpObjectAggregator aggregator = new HttpObjectAggregator(handlingSettings.maxContentLength()); + final HttpObjectAggregator aggregator = new Netty4HttpAggregator(handlingSettings.maxContentLength()); aggregator.setMaxCumulationBufferComponents(transport.maxCompositeBufferComponents); ch.pipeline() .addLast("decoder_compress", new HttpContentDecompressor()) // this handles request body decompression diff --git a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpServerTransportTests.java b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpServerTransportTests.java index bc6e5fef834e8..911873ed07ab5 100644 --- a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpServerTransportTests.java +++ b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpServerTransportTests.java @@ -904,7 +904,7 @@ public void dispatchBadRequest(final RestChannel channel, final ThreadContext th assertThat(channel.request().getHttpRequest().header(headerReference.get()), is(headerValueReference.get())); assertThat(channel.request().getHttpRequest().method(), is(translateRequestMethod(httpMethodReference.get()))); // assert content is dropped - assertThat(channel.request().getHttpRequest().content().utf8ToString(), is("")); + assertThat(channel.request().getHttpRequest().content().asFull().bytes().utf8ToString(), is("")); try { channel.sendResponse(new RestResponse(channel, (Exception) ((ElasticsearchWrapperException) cause).getCause())); } catch (IOException e) { diff --git a/server/src/main/java/org/elasticsearch/http/HttpClientStatsTracker.java b/server/src/main/java/org/elasticsearch/http/HttpClientStatsTracker.java index e8c92c8540df2..3fe5029036397 100644 --- a/server/src/main/java/org/elasticsearch/http/HttpClientStatsTracker.java +++ b/server/src/main/java/org/elasticsearch/http/HttpClientStatsTracker.java @@ -24,6 +24,7 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.Semaphore; +import java.util.function.Consumer; import java.util.function.LongPredicate; import java.util.stream.Stream; @@ -226,7 +227,18 @@ synchronized void update(HttpRequest httpRequest, HttpChannel httpChannel, long lastRequestTimeMillis = currentTimeMillis; lastUri = httpRequest.uri(); requestCount += 1; - requestSizeBytes += httpRequest.content().length(); + if (httpRequest.content().isFull()) { + requestSizeBytes += httpRequest.content().asFull().bytes().length(); + } else { + var contentStream = httpRequest.content().asStream(); + var chunkHandler = contentStream.handler(); + Consumer bytesCountingHandler = (chunk) -> { + requestSizeBytes += chunk.bytes().length(); + chunkHandler.accept(chunk); + }; + contentStream.setHandler(bytesCountingHandler); + } + } private static String getFirstValueForHeader(final HttpRequest request, final String header) { diff --git a/server/src/main/java/org/elasticsearch/http/HttpContent.java b/server/src/main/java/org/elasticsearch/http/HttpContent.java new file mode 100644 index 0000000000000..4fb8f2548f7f9 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/http/HttpContent.java @@ -0,0 +1,73 @@ +/* + * 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; + +import org.elasticsearch.common.bytes.BytesArray; +import org.elasticsearch.common.bytes.BytesReference; + +import java.util.function.Consumer; + +public sealed interface HttpContent permits HttpContent.Full, HttpContent.Stream { + + static Full fromBytesReference(BytesReference bytesRef) { + return new ByteRefHttpContent(bytesRef); + } + + static Full empty() { + return new ByteRefHttpContent(BytesArray.EMPTY); + } + + default boolean isFull() { + return this instanceof Full; + } + + default boolean isStream() { + return this instanceof Stream; + } + + default Full asFull() { + assert this instanceof Full; + return (Full) this; + } + + default Stream asStream() { + assert this instanceof Stream; + return (Stream) this; + } + + non-sealed interface Full extends HttpContent, Chunk { + @Override + default boolean isLast() { + return true; + } + + @Override + default void release() {} + } + + non-sealed interface Stream extends HttpContent { + Consumer handler(); + + void setHandler(Consumer chunkHandler); + + void request(int bytes); + + void cancel(); + } + + interface Chunk { + BytesReference bytes(); + + void release(); + + boolean isLast(); + } + + record ByteRefHttpContent(BytesReference bytes) implements Full {} +} diff --git a/server/src/main/java/org/elasticsearch/http/HttpRequest.java b/server/src/main/java/org/elasticsearch/http/HttpRequest.java index 2757fa15ce477..99f5172772d51 100644 --- a/server/src/main/java/org/elasticsearch/http/HttpRequest.java +++ b/server/src/main/java/org/elasticsearch/http/HttpRequest.java @@ -27,7 +27,7 @@ enum HttpVersion { HTTP_1_1 } - BytesReference content(); + HttpContent content(); List strictCookies(); diff --git a/server/src/main/java/org/elasticsearch/rest/RestRequest.java b/server/src/main/java/org/elasticsearch/rest/RestRequest.java index 66ba0c743813e..ae80c789b3675 100644 --- a/server/src/main/java/org/elasticsearch/rest/RestRequest.java +++ b/server/src/main/java/org/elasticsearch/rest/RestRequest.java @@ -23,6 +23,7 @@ import org.elasticsearch.core.TimeValue; import org.elasticsearch.core.Tuple; import org.elasticsearch.http.HttpChannel; +import org.elasticsearch.http.HttpContent; import org.elasticsearch.http.HttpRequest; import org.elasticsearch.telemetry.tracing.Traceable; import org.elasticsearch.xcontent.ParsedMediaType; @@ -71,7 +72,7 @@ public class RestRequest implements ToXContent.Params, Traceable { private final long requestId; public boolean isContentConsumed() { - return contentConsumed; + return contentConsumed || isStreamedContent(); } @SuppressWarnings("this-escape") @@ -278,16 +279,32 @@ public final String path() { } public boolean hasContent() { - return contentLength() > 0; + return contentLength() > 0 || isStreamedContent(); } public int contentLength() { - return httpRequest.content().length(); + if (httpRequest.content().isFull()) { + return httpRequest.content().asFull().bytes().length(); + } else { + return 0; + } } public BytesReference content() { this.contentConsumed = true; - return httpRequest.content(); + if (httpRequest.content().isFull()) { + return httpRequest.content().asFull().bytes(); + } else { + return BytesArray.EMPTY; + } + } + + public boolean isStreamedContent() { + return httpRequest.content().isStream(); + } + + public HttpContent.Stream contentStream() { + return httpRequest.content().asStream(); } /** diff --git a/server/src/test/java/org/elasticsearch/http/HttpClientStatsTrackerTests.java b/server/src/test/java/org/elasticsearch/http/HttpClientStatsTrackerTests.java index 2dfaaf34bb1f1..5924382b708e3 100644 --- a/server/src/test/java/org/elasticsearch/http/HttpClientStatsTrackerTests.java +++ b/server/src/test/java/org/elasticsearch/http/HttpClientStatsTrackerTests.java @@ -118,7 +118,7 @@ public void testStatsCollection() { assertThat(clientStats.remoteAddress(), equalTo(NetworkAddress.format(httpChannel.getRemoteAddress()))); assertThat(clientStats.lastUri(), equalTo(httpRequest1.uri())); assertThat(clientStats.requestCount(), equalTo(1L)); - requestLength += httpRequest1.content().length(); + requestLength += httpRequest1.content().asFull().bytes().length(); assertThat(clientStats.requestSizeBytes(), equalTo(requestLength)); assertThat(clientStats.closedTimeMillis(), equalTo(-1L)); assertThat(clientStats.openedTimeMillis(), equalTo(openTimeMillis)); @@ -148,7 +148,7 @@ public void testStatsCollection() { assertThat(clientStats.remoteAddress(), equalTo(NetworkAddress.format(httpChannel.getRemoteAddress()))); assertThat(clientStats.lastUri(), equalTo(httpRequest2.uri())); assertThat(clientStats.requestCount(), equalTo(2L)); - requestLength += httpRequest2.content().length(); + requestLength += httpRequest2.content().asFull().bytes().length(); assertThat(clientStats.requestSizeBytes(), equalTo(requestLength)); assertThat(clientStats.closedTimeMillis(), equalTo(-1L)); assertThat(clientStats.openedTimeMillis(), equalTo(openTimeMillis)); diff --git a/server/src/test/java/org/elasticsearch/http/TestHttpRequest.java b/server/src/test/java/org/elasticsearch/http/TestHttpRequest.java index e7b0232afa245..4dda3252b94ee 100644 --- a/server/src/test/java/org/elasticsearch/http/TestHttpRequest.java +++ b/server/src/test/java/org/elasticsearch/http/TestHttpRequest.java @@ -8,7 +8,6 @@ package org.elasticsearch.http; -import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.rest.ChunkedRestResponseBodyPart; import org.elasticsearch.rest.RestRequest; @@ -48,8 +47,8 @@ public String uri() { } @Override - public BytesReference content() { - return BytesArray.EMPTY; + public HttpContent content() { + return HttpContent.empty(); } @Override diff --git a/server/src/test/java/org/elasticsearch/rest/RestControllerTests.java b/server/src/test/java/org/elasticsearch/rest/RestControllerTests.java index 67f42e6cf1808..2f03b8a84ede3 100644 --- a/server/src/test/java/org/elasticsearch/rest/RestControllerTests.java +++ b/server/src/test/java/org/elasticsearch/rest/RestControllerTests.java @@ -26,6 +26,7 @@ import org.elasticsearch.common.xcontent.XContentHelper; import org.elasticsearch.core.IOUtils; import org.elasticsearch.core.RestApiVersion; +import org.elasticsearch.http.HttpContent; import org.elasticsearch.http.HttpHeadersValidationException; import org.elasticsearch.http.HttpInfo; import org.elasticsearch.http.HttpRequest; @@ -830,11 +831,11 @@ public String uri() { } @Override - public BytesReference content() { + public HttpContent content() { if (hasContent) { - return new BytesArray("test"); + return HttpContent.fromBytesReference(new BytesArray("test")); } - return BytesArray.EMPTY; + return HttpContent.empty(); } @Override diff --git a/server/src/test/java/org/elasticsearch/rest/RestRequestTests.java b/server/src/test/java/org/elasticsearch/rest/RestRequestTests.java index bb06dbe5d09aa..86c4d478dc96f 100644 --- a/server/src/test/java/org/elasticsearch/rest/RestRequestTests.java +++ b/server/src/test/java/org/elasticsearch/rest/RestRequestTests.java @@ -13,6 +13,7 @@ import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.core.CheckedConsumer; import org.elasticsearch.http.HttpChannel; +import org.elasticsearch.http.HttpContent; import org.elasticsearch.http.HttpRequest; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.rest.FakeRestRequest; @@ -83,7 +84,7 @@ public void testContentLengthDoesNotConsumesContent() { private void runConsumesContentTest(final CheckedConsumer consumer, final boolean expected) { final HttpRequest httpRequest = mock(HttpRequest.class); when(httpRequest.uri()).thenReturn(""); - when(httpRequest.content()).thenReturn(new BytesArray(new byte[1])); + when(httpRequest.content()).thenReturn(HttpContent.fromBytesReference(new BytesArray(new byte[1]))); when(httpRequest.getHeaders()).thenReturn( Collections.singletonMap("Content-Type", Collections.singletonList(randomFrom("application/json", "application/x-ndjson"))) ); diff --git a/test/framework/src/main/java/org/elasticsearch/test/rest/FakeRestRequest.java b/test/framework/src/main/java/org/elasticsearch/test/rest/FakeRestRequest.java index 3a9c4b371c9da..7717f2e6971cd 100644 --- a/test/framework/src/main/java/org/elasticsearch/test/rest/FakeRestRequest.java +++ b/test/framework/src/main/java/org/elasticsearch/test/rest/FakeRestRequest.java @@ -14,6 +14,7 @@ import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.xcontent.LoggingDeprecationHandler; import org.elasticsearch.http.HttpChannel; +import org.elasticsearch.http.HttpContent; import org.elasticsearch.http.HttpRequest; import org.elasticsearch.http.HttpResponse; import org.elasticsearch.rest.ChunkedRestResponseBodyPart; @@ -53,24 +54,24 @@ public static class FakeHttpRequest implements HttpRequest { private final Method method; private final String uri; - private final BytesReference content; + private final HttpContent content; private final Map> headers; private final Exception inboundException; public FakeHttpRequest(Method method, String uri, BytesReference content, Map> headers) { - this(method, uri, content, headers, null); + this(method, uri, content == null ? HttpContent.empty() : HttpContent.fromBytesReference(content), headers, null); } private FakeHttpRequest( Method method, String uri, - BytesReference content, + HttpContent content, Map> headers, Exception inboundException ) { this.method = method; this.uri = uri; - this.content = content == null ? BytesArray.EMPTY : content; + this.content = content; this.headers = headers; this.inboundException = inboundException; } @@ -86,7 +87,7 @@ public String uri() { } @Override - public BytesReference content() { + public HttpContent content() { return content; } @@ -194,7 +195,7 @@ public static class Builder { private Map params = new HashMap<>(); - private BytesReference content = BytesArray.EMPTY; + private HttpContent content = HttpContent.empty(); private String path = "/"; @@ -220,7 +221,7 @@ public Builder withParams(Map params) { } public Builder withContent(BytesReference contentBytes, XContentType xContentType) { - this.content = contentBytes; + this.content = HttpContent.fromBytesReference(contentBytes); if (xContentType != null) { headers.put("Content-Type", Collections.singletonList(xContentType.mediaType())); } diff --git a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/audit/logfile/LoggingAuditTrailTests.java b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/audit/logfile/LoggingAuditTrailTests.java index 17bad90415e7c..9c909fbdf6082 100644 --- a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/audit/logfile/LoggingAuditTrailTests.java +++ b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/audit/logfile/LoggingAuditTrailTests.java @@ -2636,7 +2636,10 @@ public void testAuthenticationSuccessRest() throws Exception { checkedFields.put(LoggingAuditTrail.REQUEST_ID_FIELD_NAME, requestId); checkedFields.put(LoggingAuditTrail.URL_PATH_FIELD_NAME, "_uri"); if (includeRequestBody && Strings.hasLength(request.content())) { - checkedFields.put(LoggingAuditTrail.REQUEST_BODY_FIELD_NAME, request.getHttpRequest().content().utf8ToString()); + checkedFields.put( + LoggingAuditTrail.REQUEST_BODY_FIELD_NAME, + request.getHttpRequest().content().asFull().bytes().utf8ToString() + ); } if (params.isEmpty() == false) { checkedFields.put(LoggingAuditTrail.URL_QUERY_FIELD_NAME, "foo=bar&evac=true"); From 07f97ccb7a42413266b51a7243cda28ba164b457 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Mon, 29 Jul 2024 19:16:24 -0700 Subject: [PATCH 02/14] add comments and fix last chunk aggregation --- .../netty4/Netty4HttpRequestStreamIT.java | 9 ++++-- .../Netty4HttpRequestContentStream.java | 18 ++++++++++-- .../org/elasticsearch/http/HttpContent.java | 28 +++++++++++++++++++ 3 files changed, 49 insertions(+), 6 deletions(-) diff --git a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java index 5f7a114b535e4..b6fbb816985a2 100644 --- a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java +++ b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java @@ -97,6 +97,9 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { return channel -> { + + // Example of handling streamed content + var contentStream = request.contentStream(); var totalReceivedBytes = new AtomicInteger(); @@ -104,14 +107,14 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli LOGGER.info("got next chunk with size {}", chunk.bytes().length()); totalReceivedBytes.addAndGet(chunk.bytes().length()); if (chunk.isLast() == false) { - contentStream.request(1024); + 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); - contentStream.request(1024); + contentStream.setHandler(contentConsumer); // must setup handler before first request + contentStream.request(1024); // initiate first chunk, will round up to 8kb }; } }); diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java index 4ace76271c18c..0500e2dc779be 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java @@ -19,6 +19,13 @@ import java.util.function.Consumer; +/** + * Implementation of HTTP request content stream. This implementation relies on + * {@link io.netty.channel.ChannelConfig#setAutoRead(boolean)} to stop arrival of new chunks. + * Otherwise, downstream might receive much more data than requested. + * This can be achieved with {@link io.netty.handler.flow.FlowControlHandler} placed before stream + * constructor in {@link io.netty.channel.ChannelPipeline} and after {@link io.netty.handler.codec.http.HttpRequestDecoder}. + */ public class Netty4HttpRequestContentStream implements HttpContent.Stream { private final Channel channel; @@ -67,10 +74,12 @@ public void handleNettyContent(io.netty.handler.codec.http.HttpContent httpConte if (aggregate == null && content.readableBytes() >= requestedBytes) { // network chunk is large enough and there is no ongoing aggregation + requestedBytes = 0; handler.accept(new Netty4ContentChunk(content, isLast)); } else if (content.readableBytes() >= requestedBytes) { // there is ongoing aggregation, and we reached requestBytes limit, send aggregated chunk downstream - aggregate.addComponent(content); + aggregate.addComponent(true, content); + requestedBytes = 0; handler.accept(new Netty4ContentChunk(aggregate, isLast)); aggregate = null; } else { @@ -78,11 +87,14 @@ public void handleNettyContent(io.netty.handler.codec.http.HttpContent httpConte if (aggregate == null) { aggregate = channel.alloc().compositeHeapBuffer(); } - aggregate.addComponent(content); + aggregate.addComponent(true, content); + requestedBytes -= content.readableBytes(); channel.read(); } - requestedBytes = Math.max(0, requestedBytes - content.readableBytes()); if (isLast) { + if (aggregate != null) { + handler.accept(new Netty4ContentChunk(aggregate, true)); + } channel.config().setAutoRead(true); } } diff --git a/server/src/main/java/org/elasticsearch/http/HttpContent.java b/server/src/main/java/org/elasticsearch/http/HttpContent.java index 4fb8f2548f7f9..116dca75f4d1c 100644 --- a/server/src/main/java/org/elasticsearch/http/HttpContent.java +++ b/server/src/main/java/org/elasticsearch/http/HttpContent.java @@ -10,9 +10,13 @@ import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.bytes.BytesReference; +import org.elasticsearch.core.Nullable; import java.util.function.Consumer; +/** + * A super-interface for different HTTP content implementations + */ public sealed interface HttpContent permits HttpContent.Full, HttpContent.Stream { static Full fromBytesReference(BytesReference bytesRef) { @@ -41,6 +45,9 @@ default Stream asStream() { return (Stream) this; } + /** + * Full content represents a complete http body content that can be accessed immediately + */ non-sealed interface Full extends HttpContent, Chunk { @Override default boolean isLast() { @@ -51,13 +58,34 @@ default boolean isLast() { default void release() {} } + /** + * Stream is a lazy-loaded content. Stream supports only single handler, this handler must be set before requesting next chunk. + */ non-sealed interface Stream extends HttpContent { + /** + * Returns current handler + */ + @Nullable Consumer handler(); + /** + * Sets handler that can handle next chunk + */ void setHandler(Consumer chunkHandler); + /** + * Request next chunk of bytes. The size of next chunk might deffer from requested size in following cases: + *
    + *
  • Last content always less or equal requested size
    • + *
  • Implementation is free to round up chunk size to optimal network chunk size. For example, transport default chunk size + * is 8kb, then requesting 10kb might return 16kb chunk
    • + *
+ */ void request(int bytes); + /** + * Abort stream consumption, it's terminal operation and no more chunks can be received from stream + */ void cancel(); } From 5d4469898114f65eea05005fdf35b9f9a28f6cdd Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Mon, 29 Jul 2024 19:40:54 -0700 Subject: [PATCH 03/14] change composite buffer type --- .../http/netty4/Netty4HttpRequestContentStream.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java index 0500e2dc779be..37436e16b822a 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java @@ -85,7 +85,7 @@ public void handleNettyContent(io.netty.handler.codec.http.HttpContent httpConte } else { // accumulation step if (aggregate == null) { - aggregate = channel.alloc().compositeHeapBuffer(); + aggregate = channel.alloc().compositeBuffer(); } aggregate.addComponent(true, content); requestedBytes -= content.readableBytes(); From 839e5758f6cc27175c04280209f266e3decb1b92 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Wed, 31 Jul 2024 20:16:07 -0700 Subject: [PATCH 04/14] update stream interface --- .../netty4/Netty4HttpRequestStreamIT.java | 126 ++++++++++++----- .../netty4/Netty4HttpPipeliningHandler.java | 4 +- .../http/netty4/Netty4HttpRequest.java | 21 +-- .../netty4/Netty4HttpRequestBodyStream.java | 85 ++++++++++++ .../Netty4HttpRequestContentStream.java | 129 ------------------ .../netty4/Netty4HttpRequestFullContent.java | 34 ----- .../transport/netty4/Netty4Utils.java | 10 ++ .../Netty4HttpServerTransportTests.java | 2 +- .../http/{HttpContent.java => HttpBody.java} | 51 +++---- .../http/HttpClientStatsTracker.java | 14 +- .../org/elasticsearch/http/HttpRequest.java | 4 +- .../elasticsearch/rest/BaseRestHandler.java | 11 ++ .../org/elasticsearch/rest/RestRequest.java | 16 +-- .../http/HttpClientStatsTrackerTests.java | 4 +- .../elasticsearch/http/TestHttpRequest.java | 4 +- .../rest/RestControllerTests.java | 8 +- .../elasticsearch/rest/RestRequestTests.java | 4 +- .../test/rest/FakeRestRequest.java | 14 +- .../audit/logfile/LoggingAuditTrailTests.java | 5 +- 19 files changed, 255 insertions(+), 291 deletions(-) create mode 100644 modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java delete mode 100644 modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java delete mode 100644 modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestFullContent.java rename server/src/main/java/org/elasticsearch/http/{HttpContent.java => HttpBody.java} (53%) diff --git a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java index b6fbb816985a2..4f503202af5e2 100644 --- a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java +++ b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java @@ -8,16 +8,27 @@ 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 io.netty.bootstrap.Bootstrap; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelInitializer; +import io.netty.channel.SimpleChannelInboundHandler; +import io.netty.channel.nio.NioEventLoopGroup; +import io.netty.channel.socket.SocketChannel; +import io.netty.channel.socket.nio.NioSocketChannel; +import io.netty.handler.codec.http.DefaultFullHttpRequest; +import io.netty.handler.codec.http.FullHttpResponse; +import io.netty.handler.codec.http.HttpHeaderNames; +import io.netty.handler.codec.http.HttpMethod; +import io.netty.handler.codec.http.HttpObjectAggregator; +import io.netty.handler.codec.http.HttpRequestEncoder; +import io.netty.handler.codec.http.HttpResponseDecoder; +import io.netty.handler.codec.http.HttpVersion; + 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.bytes.ReleasableBytesReference; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.settings.ClusterSettings; import org.elasticsearch.common.settings.IndexScopedSettings; @@ -25,19 +36,22 @@ 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.http.HttpServerTransport; import org.elasticsearch.plugins.ActionPlugin; import org.elasticsearch.plugins.Plugin; import org.elasticsearch.rest.BaseRestHandler; +import org.elasticsearch.rest.RestChannel; 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.nio.charset.StandardCharsets; import java.util.Collection; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Consumer; import java.util.function.Predicate; @@ -45,6 +59,54 @@ public class Netty4HttpRequestStreamIT extends ESNetty4IntegTestCase { + // ensures content integrity, no loses + public void testReceiveAllChunks() throws Exception { + var sendBytes = randomByteArrayOfLength(1024 * 1024 * 10); + var consumedBytes = new AtomicInteger(); + + var latch = new CountDownLatch(1); + var client = nettyClient(internalCluster().getRandomNodeName(), (resp -> { + consumedBytes.set(Integer.parseInt(resp.content().toString(StandardCharsets.UTF_8))); + latch.countDown(); + })); + + var req = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.POST, RequestContentStreamPlugin.ROUTE); + req.headers().add(HttpHeaderNames.CONTENT_LENGTH, sendBytes.length); + req.content().writeBytes(sendBytes); + + client.connect().await().channel().writeAndFlush(req).get(); + safeAwait(latch); + assertEquals(sendBytes.length, consumedBytes.get()); + client.config().group().shutdownGracefully().get(10, TimeUnit.SECONDS); + } + + void setChunkConsumer(Consumer channelConsumer) { + internalCluster().getInstances(RequestContentStreamPlugin.class).forEach(p -> { p.channelConsumer = channelConsumer; }); + } + + Bootstrap nettyClient(String node, Consumer responseHandler) throws Exception { + var httpServer = internalCluster().getInstance(HttpServerTransport.class, node); + var remoteAddr = randomFrom(httpServer.boundAddress().boundAddresses()); + return new Bootstrap().group(new NioEventLoopGroup(1)) + .channel(NioSocketChannel.class) + .remoteAddress(remoteAddr.getAddress(), remoteAddr.getPort()) + .handler(new ChannelInitializer() { + @Override + protected void initChannel(SocketChannel ch) { + var p = ch.pipeline(); + p.addLast(new HttpRequestEncoder()); + p.addLast(new HttpResponseDecoder()); + p.addLast(new HttpObjectAggregator(4096)); + p.addLast(new SimpleChannelInboundHandler() { + @Override + protected void channelRead0(ChannelHandlerContext ctx, FullHttpResponse msg) { + responseHandler.accept(msg); + } + }); + } + }); + } + @Override protected Collection> nodePlugins() { return CollectionUtils.concatLists(List.of(RequestContentStreamPlugin.class), super.nodePlugins()); @@ -55,21 +117,10 @@ 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"); + Consumer channelConsumer; @Override public Collection getRestHandlers( @@ -96,25 +147,26 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { - return channel -> { - - // Example of handling streamed content - - var contentStream = request.contentStream(); - var totalReceivedBytes = new AtomicInteger(); - - Consumer 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()))); + return new RequestBodyChunkConsumer() { + + int totalBytes = 0; + + @Override + public void handleChunk(RestChannel channel, ReleasableBytesReference chunk, boolean isLast) { + try (chunk) { + totalBytes += chunk.length(); + if (isLast == false) { + request.contentStream().requestBytes(1024); + } else { + channel.sendResponse(new RestResponse(RestStatus.OK, Integer.toString(totalBytes))); + } } - }; + } - contentStream.setHandler(contentConsumer); // must setup handler before first request - contentStream.request(1024); // initiate first chunk, will round up to 8kb + @Override + public void accept(RestChannel channel) throws Exception { + request.contentStream().requestBytes(1024); // ask for first chunk + } }; } }); diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java index 1a12be4d64a38..9b90dba1eed98 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java @@ -80,7 +80,7 @@ private record ChunkedWrite(PromiseCombiner combiner, ChannelPromise onDone, Chu * 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; + private Netty4HttpRequestBodyStream currentRequestStream; /* * The current read and write sequence numbers. Read sequence numbers are attached to requests in the order they are read from the @@ -138,7 +138,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { currentRequestStream = null; netty4HttpRequest = new Netty4HttpRequest(readSequence++, fullHttpRequest); } else { - var contentStream = new Netty4HttpRequestContentStream(ctx.channel()); + var contentStream = new Netty4HttpRequestBodyStream(ctx.channel()); currentRequestStream = contentStream; netty4HttpRequest = new Netty4HttpRequest( readSequence++, diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java index 9d08d95deef2a..43ff436944ef4 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java @@ -20,12 +20,13 @@ import io.netty.handler.codec.http.cookie.ServerCookieEncoder; import org.elasticsearch.common.bytes.BytesReference; -import org.elasticsearch.http.HttpContent; +import org.elasticsearch.http.HttpBody; 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,26 +40,26 @@ public class Netty4HttpRequest implements HttpRequest { private final FullHttpRequest request; - private final HttpContent content; + private final HttpBody content; private final Map> headers; private final AtomicBoolean released; private final Exception inboundException; private final boolean pooled; private final int sequence; - Netty4HttpRequest(int sequence, FullHttpRequest request, Netty4HttpRequestContentStream contentStream) { + Netty4HttpRequest(int sequence, FullHttpRequest request, Netty4HttpRequestBodyStream contentStream) { this(sequence, request, new AtomicBoolean(false), false, contentStream, null); } Netty4HttpRequest(int sequence, FullHttpRequest request) { - this(sequence, request, new AtomicBoolean(false), true, new Netty4HttpRequestFullContent(request)); + this(sequence, request, new AtomicBoolean(false), true, Netty4Utils.fullHttpBodyFrom(request.content())); } Netty4HttpRequest(int sequence, FullHttpRequest request, Exception inboundException) { - this(sequence, request, new AtomicBoolean(false), true, new Netty4HttpRequestFullContent(request), inboundException); + this(sequence, request, new AtomicBoolean(false), true, Netty4Utils.fullHttpBodyFrom(request.content()), inboundException); } - private Netty4HttpRequest(int sequence, FullHttpRequest request, AtomicBoolean released, boolean pooled, HttpContent content) { + private Netty4HttpRequest(int sequence, FullHttpRequest request, AtomicBoolean released, boolean pooled, HttpBody content) { this(sequence, request, released, pooled, content, null); } @@ -67,7 +68,7 @@ private Netty4HttpRequest( FullHttpRequest request, AtomicBoolean released, boolean pooled, - HttpContent content, + HttpBody content, Exception inboundException ) { this.sequence = sequence; @@ -90,7 +91,7 @@ public String uri() { } @Override - public HttpContent content() { + public HttpBody body() { assert released.get() == false; return content; } @@ -171,9 +172,9 @@ public HttpRequest removeHeader(String header) { copiedHeadersWithout, copiedTrailingHeadersWithout ); - HttpContent copiedContent; + HttpBody copiedContent; if (content.isFull()) { - copiedContent = HttpContent.fromBytesReference(content.asFull().bytes()); + copiedContent = HttpBody.fromBytesReference(content.asFull().bytes()); } else { copiedContent = content; } diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java new file mode 100644 index 0000000000000..63132424af5e5 --- /dev/null +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java @@ -0,0 +1,85 @@ +/* + * 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.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; +import io.netty.channel.Channel; +import io.netty.handler.codec.http.HttpContent; +import io.netty.handler.codec.http.LastHttpContent; + +import org.elasticsearch.http.HttpBody; +import org.elasticsearch.transport.netty4.Netty4Utils; + +/** + * Implementation of HTTP request content stream. This implementation relies on + * {@link io.netty.channel.ChannelConfig#setAutoRead(boolean)} to stop arrival of new chunks. + * Otherwise, downstream might receive much more data than requested. + * This can be achieved with {@link io.netty.handler.flow.FlowControlHandler} placed before stream + * constructor in {@link io.netty.channel.ChannelPipeline} and after {@link io.netty.handler.codec.http.HttpRequestDecoder}. + */ +public class Netty4HttpRequestBodyStream implements HttpBody.Stream { + + private final Channel channel; + private HttpBody.ChunkHandler handler; + private CompositeByteBuf aggregate; + private int requestedBytes = 0; + + public Netty4HttpRequestBodyStream(Channel channel) { + this.channel = channel; + channel.config().setAutoRead(false); + } + + @Override + public ChunkHandler handler() { + return handler; + } + + @Override + public void setHandler(ChunkHandler chunkHandler) { + this.handler = chunkHandler; + } + + @Override + public void requestBytes(int bytes) { + assert handler != null : "handler must be set before requesting next chunk"; + requestedBytes += bytes; + channel.read(); + } + + public void handleNettyContent(HttpContent httpContent) { + assert handler != null : "handler must be set before processing http content"; + var isLast = httpContent instanceof LastHttpContent; + var content = httpContent.content(); + + if (content.readableBytes() >= requestedBytes || isLast) { + ByteBuf sendBuf; + if (aggregate == null) { + sendBuf = content; + } else { + aggregate.addComponent(true, content); + sendBuf = aggregate; + aggregate = null; + } + requestedBytes = 0; + handler.onNext(Netty4Utils.toReleasableBytesReference(sendBuf), isLast); + if (isLast) { + channel.config().setAutoRead(true); + } + } else { + if (aggregate == null) { + aggregate = channel.alloc().compositeBuffer(); + } + aggregate.addComponent(true, content); + requestedBytes -= content.readableBytes(); + channel.read(); + } + } + +} diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java deleted file mode 100644 index 37436e16b822a..0000000000000 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestContentStream.java +++ /dev/null @@ -1,129 +0,0 @@ -/* - * 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.buffer.ByteBuf; -import io.netty.buffer.CompositeByteBuf; -import io.netty.channel.Channel; -import io.netty.handler.codec.http.LastHttpContent; - -import org.elasticsearch.common.bytes.BytesReference; -import org.elasticsearch.http.HttpContent; -import org.elasticsearch.transport.netty4.Netty4Utils; - -import java.util.function.Consumer; - -/** - * Implementation of HTTP request content stream. This implementation relies on - * {@link io.netty.channel.ChannelConfig#setAutoRead(boolean)} to stop arrival of new chunks. - * Otherwise, downstream might receive much more data than requested. - * This can be achieved with {@link io.netty.handler.flow.FlowControlHandler} placed before stream - * constructor in {@link io.netty.channel.ChannelPipeline} and after {@link io.netty.handler.codec.http.HttpRequestDecoder}. - */ -public class Netty4HttpRequestContentStream implements HttpContent.Stream { - - private final Channel channel; - private Consumer handler; - private CompositeByteBuf aggregate; - private int requestedBytes = 0; - private boolean cancelled = false; - - public Netty4HttpRequestContentStream(Channel channel) { - this.channel = channel; - channel.config().setAutoRead(false); - } - - @Override - public Consumer handler() { - return handler; - } - - @Override - public void setHandler(Consumer chunkHandler) { - this.handler = chunkHandler; - } - - @Override - public void request(int bytes) { - assert handler != null : "handler must be set before requesting next chunk"; - requestedBytes += bytes; - channel.read(); - } - - @Override - public void cancel() { - cancelled = true; - channel.config().setAutoRead(true); - } - - public void handleNettyContent(io.netty.handler.codec.http.HttpContent httpContent) { - if (cancelled) { - httpContent.release(); - return; - } - - assert handler != null : "handler must be set before processing http content"; - var isLast = httpContent instanceof LastHttpContent; - var content = httpContent.content(); - - if (aggregate == null && content.readableBytes() >= requestedBytes) { - // network chunk is large enough and there is no ongoing aggregation - requestedBytes = 0; - handler.accept(new Netty4ContentChunk(content, isLast)); - } else if (content.readableBytes() >= requestedBytes) { - // there is ongoing aggregation, and we reached requestBytes limit, send aggregated chunk downstream - aggregate.addComponent(true, content); - requestedBytes = 0; - handler.accept(new Netty4ContentChunk(aggregate, isLast)); - aggregate = null; - } else { - // accumulation step - if (aggregate == null) { - aggregate = channel.alloc().compositeBuffer(); - } - aggregate.addComponent(true, content); - requestedBytes -= content.readableBytes(); - channel.read(); - } - if (isLast) { - if (aggregate != null) { - handler.accept(new Netty4ContentChunk(aggregate, true)); - } - channel.config().setAutoRead(true); - } - } - - static class Netty4ContentChunk implements Chunk { - - private final ByteBuf nettyBuffer; - private final BytesReference bytes; - private final boolean isLast; - - Netty4ContentChunk(ByteBuf nettyBuffer, boolean isLast) { - this.nettyBuffer = nettyBuffer; - this.bytes = Netty4Utils.toBytesReference(nettyBuffer); - this.isLast = isLast; - } - - @Override - public BytesReference bytes() { - return bytes; - } - - @Override - public void release() { - nettyBuffer.release(); - } - - @Override - public boolean isLast() { - return isLast; - } - } -} diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestFullContent.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestFullContent.java deleted file mode 100644 index b45e73b8dac88..0000000000000 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestFullContent.java +++ /dev/null @@ -1,34 +0,0 @@ -/* - * 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.elasticsearch.common.bytes.BytesReference; -import org.elasticsearch.http.HttpContent; -import org.elasticsearch.transport.netty4.Netty4Utils; - -public class Netty4HttpRequestFullContent implements HttpContent.Full { - - private final io.netty.handler.codec.http.HttpContent httpContent; - private final BytesReference bytes; - - public Netty4HttpRequestFullContent(io.netty.handler.codec.http.HttpContent httpContent) { - this.httpContent = httpContent; - this.bytes = Netty4Utils.toBytesReference(httpContent.content()); - } - - @Override - public BytesReference bytes() { - return bytes; - } - - @Override - public void release() { - httpContent.release(); - } -} diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4Utils.java b/modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4Utils.java index 1025ad11ba05e..7404226f069c5 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4Utils.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4Utils.java @@ -23,10 +23,12 @@ import org.elasticsearch.action.ActionListener; import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.bytes.BytesReference; +import org.elasticsearch.common.bytes.ReleasableBytesReference; import org.elasticsearch.common.recycler.Recycler; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.util.concurrent.EsExecutors; import org.elasticsearch.core.Booleans; +import org.elasticsearch.http.HttpBody; import org.elasticsearch.transport.TransportException; import java.io.IOException; @@ -123,6 +125,14 @@ public static BytesReference toBytesReference(final ByteBuf buffer) { } } + public static ReleasableBytesReference toReleasableBytesReference(final ByteBuf buffer) { + return new ReleasableBytesReference(toBytesReference(buffer), buffer::release); + } + + public static HttpBody.Full fullHttpBodyFrom(final ByteBuf buf) { + return new HttpBody.ByteRefHttpBody(toBytesReference(buf)); + } + public static Recycler createRecycler(Settings settings) { // If this method is called by super ctor the processors will not be set. Accessing NettyAllocator initializes netty's internals // setting the processors. We must do it ourselves first just in case. diff --git a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpServerTransportTests.java b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpServerTransportTests.java index 911873ed07ab5..9e213e6468356 100644 --- a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpServerTransportTests.java +++ b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpServerTransportTests.java @@ -904,7 +904,7 @@ public void dispatchBadRequest(final RestChannel channel, final ThreadContext th assertThat(channel.request().getHttpRequest().header(headerReference.get()), is(headerValueReference.get())); assertThat(channel.request().getHttpRequest().method(), is(translateRequestMethod(httpMethodReference.get()))); // assert content is dropped - assertThat(channel.request().getHttpRequest().content().asFull().bytes().utf8ToString(), is("")); + assertThat(channel.request().getHttpRequest().body().asFull().bytes().utf8ToString(), is("")); try { channel.sendResponse(new RestResponse(channel, (Exception) ((ElasticsearchWrapperException) cause).getCause())); } catch (IOException e) { diff --git a/server/src/main/java/org/elasticsearch/http/HttpContent.java b/server/src/main/java/org/elasticsearch/http/HttpBody.java similarity index 53% rename from server/src/main/java/org/elasticsearch/http/HttpContent.java rename to server/src/main/java/org/elasticsearch/http/HttpBody.java index 116dca75f4d1c..4eb72229b5516 100644 --- a/server/src/main/java/org/elasticsearch/http/HttpContent.java +++ b/server/src/main/java/org/elasticsearch/http/HttpBody.java @@ -10,21 +10,20 @@ import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.bytes.BytesReference; +import org.elasticsearch.common.bytes.ReleasableBytesReference; import org.elasticsearch.core.Nullable; -import java.util.function.Consumer; - /** * A super-interface for different HTTP content implementations */ -public sealed interface HttpContent permits HttpContent.Full, HttpContent.Stream { +public sealed interface HttpBody permits HttpBody.Full, HttpBody.Stream { static Full fromBytesReference(BytesReference bytesRef) { - return new ByteRefHttpContent(bytesRef); + return new ByteRefHttpBody(bytesRef); } static Full empty() { - return new ByteRefHttpContent(BytesArray.EMPTY); + return new ByteRefHttpBody(BytesArray.EMPTY); } default boolean isFull() { @@ -48,54 +47,36 @@ default Stream asStream() { /** * Full content represents a complete http body content that can be accessed immediately */ - non-sealed interface Full extends HttpContent, Chunk { - @Override - default boolean isLast() { - return true; - } - - @Override - default void release() {} + non-sealed interface Full extends HttpBody { + BytesReference bytes(); } /** * Stream is a lazy-loaded content. Stream supports only single handler, this handler must be set before requesting next chunk. */ - non-sealed interface Stream extends HttpContent { + non-sealed interface Stream extends HttpBody { /** * Returns current handler */ @Nullable - Consumer handler(); + ChunkHandler handler(); /** * Sets handler that can handle next chunk */ - void setHandler(Consumer chunkHandler); + void setHandler(ChunkHandler chunkHandler); /** - * Request next chunk of bytes. The size of next chunk might deffer from requested size in following cases: - *
    - *
  • Last content always less or equal requested size
    • - *
  • Implementation is free to round up chunk size to optimal network chunk size. For example, transport default chunk size - * is 8kb, then requesting 10kb might return 16kb chunk
    • - *
+ * Request next chunk of bytes. Implementation is free to round up chunk size to optimal network chunk size. + * For example, transport default chunk size is 8kb, request for 10kb might return 16kb chunk. */ - void request(int bytes); - - /** - * Abort stream consumption, it's terminal operation and no more chunks can be received from stream - */ - void cancel(); + void requestBytes(int bytes); } - interface Chunk { - BytesReference bytes(); - - void release(); - - boolean isLast(); + @FunctionalInterface + interface ChunkHandler { + void onNext(ReleasableBytesReference chunk, boolean isLast); } - record ByteRefHttpContent(BytesReference bytes) implements Full {} + record ByteRefHttpBody(BytesReference bytes) implements Full {} } diff --git a/server/src/main/java/org/elasticsearch/http/HttpClientStatsTracker.java b/server/src/main/java/org/elasticsearch/http/HttpClientStatsTracker.java index 3fe5029036397..ef05af8bb9ade 100644 --- a/server/src/main/java/org/elasticsearch/http/HttpClientStatsTracker.java +++ b/server/src/main/java/org/elasticsearch/http/HttpClientStatsTracker.java @@ -24,7 +24,6 @@ import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.Semaphore; -import java.util.function.Consumer; import java.util.function.LongPredicate; import java.util.stream.Stream; @@ -227,18 +226,9 @@ synchronized void update(HttpRequest httpRequest, HttpChannel httpChannel, long lastRequestTimeMillis = currentTimeMillis; lastUri = httpRequest.uri(); requestCount += 1; - if (httpRequest.content().isFull()) { - requestSizeBytes += httpRequest.content().asFull().bytes().length(); - } else { - var contentStream = httpRequest.content().asStream(); - var chunkHandler = contentStream.handler(); - Consumer bytesCountingHandler = (chunk) -> { - requestSizeBytes += chunk.bytes().length(); - chunkHandler.accept(chunk); - }; - contentStream.setHandler(bytesCountingHandler); + if (httpRequest.body().isFull()) { + requestSizeBytes += httpRequest.body().asFull().bytes().length(); } - } private static String getFirstValueForHeader(final HttpRequest request, final String header) { diff --git a/server/src/main/java/org/elasticsearch/http/HttpRequest.java b/server/src/main/java/org/elasticsearch/http/HttpRequest.java index 99f5172772d51..cce57d20be23f 100644 --- a/server/src/main/java/org/elasticsearch/http/HttpRequest.java +++ b/server/src/main/java/org/elasticsearch/http/HttpRequest.java @@ -27,7 +27,7 @@ enum HttpVersion { HTTP_1_1 } - HttpContent content(); + HttpBody body(); List strictCookies(); @@ -46,7 +46,7 @@ enum HttpVersion { Exception getInboundException(); /** - * Release any resources associated with this request. Implementations should be idempotent. The behavior of {@link #content()} + * Release any resources associated with this request. Implementations should be idempotent. The behavior of {@link #body()} * after this method has been invoked is undefined and implementation specific. */ void release(); diff --git a/server/src/main/java/org/elasticsearch/rest/BaseRestHandler.java b/server/src/main/java/org/elasticsearch/rest/BaseRestHandler.java index a17bc885f6b65..2d65b7d69724c 100644 --- a/server/src/main/java/org/elasticsearch/rest/BaseRestHandler.java +++ b/server/src/main/java/org/elasticsearch/rest/BaseRestHandler.java @@ -10,6 +10,7 @@ import org.apache.lucene.search.spell.LevenshteinDistance; import org.elasticsearch.client.internal.node.NodeClient; +import org.elasticsearch.common.bytes.ReleasableBytesReference; import org.elasticsearch.common.settings.Setting; import org.elasticsearch.common.settings.Setting.Property; import org.elasticsearch.common.util.set.Sets; @@ -115,6 +116,12 @@ public final void handleRequest(RestRequest request, RestChannel channel, NodeCl ); } + if (request.isStreamedContent()) { + assert action instanceof RequestBodyChunkConsumer; + var chunkConsumer = (RequestBodyChunkConsumer) action; + request.contentStream().setHandler((chunk, isLast) -> chunkConsumer.handleChunk(channel, chunk, isLast)); + } + usageCount.increment(); // execute the action action.accept(channel); @@ -172,6 +179,10 @@ protected interface RestChannelConsumer extends CheckedConsumer void runConsumesContentTest(final CheckedConsumer consumer, final boolean expected) { final HttpRequest httpRequest = mock(HttpRequest.class); when(httpRequest.uri()).thenReturn(""); - when(httpRequest.content()).thenReturn(HttpContent.fromBytesReference(new BytesArray(new byte[1]))); + when(httpRequest.body()).thenReturn(HttpBody.fromBytesReference(new BytesArray(new byte[1]))); when(httpRequest.getHeaders()).thenReturn( Collections.singletonMap("Content-Type", Collections.singletonList(randomFrom("application/json", "application/x-ndjson"))) ); diff --git a/test/framework/src/main/java/org/elasticsearch/test/rest/FakeRestRequest.java b/test/framework/src/main/java/org/elasticsearch/test/rest/FakeRestRequest.java index 7717f2e6971cd..bcea7f61fe158 100644 --- a/test/framework/src/main/java/org/elasticsearch/test/rest/FakeRestRequest.java +++ b/test/framework/src/main/java/org/elasticsearch/test/rest/FakeRestRequest.java @@ -13,8 +13,8 @@ import org.elasticsearch.common.bytes.BytesArray; import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.xcontent.LoggingDeprecationHandler; +import org.elasticsearch.http.HttpBody; import org.elasticsearch.http.HttpChannel; -import org.elasticsearch.http.HttpContent; import org.elasticsearch.http.HttpRequest; import org.elasticsearch.http.HttpResponse; import org.elasticsearch.rest.ChunkedRestResponseBodyPart; @@ -54,18 +54,18 @@ public static class FakeHttpRequest implements HttpRequest { private final Method method; private final String uri; - private final HttpContent content; + private final HttpBody content; private final Map> headers; private final Exception inboundException; public FakeHttpRequest(Method method, String uri, BytesReference content, Map> headers) { - this(method, uri, content == null ? HttpContent.empty() : HttpContent.fromBytesReference(content), headers, null); + this(method, uri, content == null ? HttpBody.empty() : HttpBody.fromBytesReference(content), headers, null); } private FakeHttpRequest( Method method, String uri, - HttpContent content, + HttpBody content, Map> headers, Exception inboundException ) { @@ -87,7 +87,7 @@ public String uri() { } @Override - public HttpContent content() { + public HttpBody body() { return content; } @@ -195,7 +195,7 @@ public static class Builder { private Map params = new HashMap<>(); - private HttpContent content = HttpContent.empty(); + private HttpBody content = HttpBody.empty(); private String path = "/"; @@ -221,7 +221,7 @@ public Builder withParams(Map params) { } public Builder withContent(BytesReference contentBytes, XContentType xContentType) { - this.content = HttpContent.fromBytesReference(contentBytes); + this.content = HttpBody.fromBytesReference(contentBytes); if (xContentType != null) { headers.put("Content-Type", Collections.singletonList(xContentType.mediaType())); } diff --git a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/audit/logfile/LoggingAuditTrailTests.java b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/audit/logfile/LoggingAuditTrailTests.java index 9c909fbdf6082..f56749330579f 100644 --- a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/audit/logfile/LoggingAuditTrailTests.java +++ b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/audit/logfile/LoggingAuditTrailTests.java @@ -2636,10 +2636,7 @@ public void testAuthenticationSuccessRest() throws Exception { checkedFields.put(LoggingAuditTrail.REQUEST_ID_FIELD_NAME, requestId); checkedFields.put(LoggingAuditTrail.URL_PATH_FIELD_NAME, "_uri"); if (includeRequestBody && Strings.hasLength(request.content())) { - checkedFields.put( - LoggingAuditTrail.REQUEST_BODY_FIELD_NAME, - request.getHttpRequest().content().asFull().bytes().utf8ToString() - ); + checkedFields.put(LoggingAuditTrail.REQUEST_BODY_FIELD_NAME, request.getHttpRequest().body().asFull().bytes().utf8ToString()); } if (params.isEmpty() == false) { checkedFields.put(LoggingAuditTrail.URL_QUERY_FIELD_NAME, "foo=bar&evac=true"); From 83bec150b2a1e749cbf3f5eec706178e23df86e9 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Fri, 2 Aug 2024 19:18:05 -0700 Subject: [PATCH 05/14] add backpressure integ test --- .../netty4/Netty4HttpRequestStreamIT.java | 176 --------- .../Netty4IncrementalRequestHandlingIT.java | 344 ++++++++++++++++++ .../netty4/Netty4HttpPipeliningHandler.java | 28 +- .../netty4/Netty4HttpRequestBodyStream.java | 18 +- .../netty4/Netty4HttpServerTransport.java | 2 +- 5 files changed, 368 insertions(+), 200 deletions(-) delete mode 100644 modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java create mode 100644 modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java diff --git a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java deleted file mode 100644 index 4f503202af5e2..0000000000000 --- a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4HttpRequestStreamIT.java +++ /dev/null @@ -1,176 +0,0 @@ -/* - * 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.bootstrap.Bootstrap; -import io.netty.channel.ChannelHandlerContext; -import io.netty.channel.ChannelInitializer; -import io.netty.channel.SimpleChannelInboundHandler; -import io.netty.channel.nio.NioEventLoopGroup; -import io.netty.channel.socket.SocketChannel; -import io.netty.channel.socket.nio.NioSocketChannel; -import io.netty.handler.codec.http.DefaultFullHttpRequest; -import io.netty.handler.codec.http.FullHttpResponse; -import io.netty.handler.codec.http.HttpHeaderNames; -import io.netty.handler.codec.http.HttpMethod; -import io.netty.handler.codec.http.HttpObjectAggregator; -import io.netty.handler.codec.http.HttpRequestEncoder; -import io.netty.handler.codec.http.HttpResponseDecoder; -import io.netty.handler.codec.http.HttpVersion; - -import org.elasticsearch.ESNetty4IntegTestCase; -import org.elasticsearch.client.internal.node.NodeClient; -import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver; -import org.elasticsearch.cluster.node.DiscoveryNodes; -import org.elasticsearch.common.bytes.ReleasableBytesReference; -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.HttpServerTransport; -import org.elasticsearch.plugins.ActionPlugin; -import org.elasticsearch.plugins.Plugin; -import org.elasticsearch.rest.BaseRestHandler; -import org.elasticsearch.rest.RestChannel; -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.nio.charset.StandardCharsets; -import java.util.Collection; -import java.util.List; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.TimeUnit; -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 { - - // ensures content integrity, no loses - public void testReceiveAllChunks() throws Exception { - var sendBytes = randomByteArrayOfLength(1024 * 1024 * 10); - var consumedBytes = new AtomicInteger(); - - var latch = new CountDownLatch(1); - var client = nettyClient(internalCluster().getRandomNodeName(), (resp -> { - consumedBytes.set(Integer.parseInt(resp.content().toString(StandardCharsets.UTF_8))); - latch.countDown(); - })); - - var req = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.POST, RequestContentStreamPlugin.ROUTE); - req.headers().add(HttpHeaderNames.CONTENT_LENGTH, sendBytes.length); - req.content().writeBytes(sendBytes); - - client.connect().await().channel().writeAndFlush(req).get(); - safeAwait(latch); - assertEquals(sendBytes.length, consumedBytes.get()); - client.config().group().shutdownGracefully().get(10, TimeUnit.SECONDS); - } - - void setChunkConsumer(Consumer channelConsumer) { - internalCluster().getInstances(RequestContentStreamPlugin.class).forEach(p -> { p.channelConsumer = channelConsumer; }); - } - - Bootstrap nettyClient(String node, Consumer responseHandler) throws Exception { - var httpServer = internalCluster().getInstance(HttpServerTransport.class, node); - var remoteAddr = randomFrom(httpServer.boundAddress().boundAddresses()); - return new Bootstrap().group(new NioEventLoopGroup(1)) - .channel(NioSocketChannel.class) - .remoteAddress(remoteAddr.getAddress(), remoteAddr.getPort()) - .handler(new ChannelInitializer() { - @Override - protected void initChannel(SocketChannel ch) { - var p = ch.pipeline(); - p.addLast(new HttpRequestEncoder()); - p.addLast(new HttpResponseDecoder()); - p.addLast(new HttpObjectAggregator(4096)); - p.addLast(new SimpleChannelInboundHandler() { - @Override - protected void channelRead0(ChannelHandlerContext ctx, FullHttpResponse msg) { - responseHandler.accept(msg); - } - }); - } - }); - } - - @Override - protected Collection> nodePlugins() { - return CollectionUtils.concatLists(List.of(RequestContentStreamPlugin.class), super.nodePlugins()); - } - - @Override - protected boolean addMockHttpTransport() { - return false; // enable http - } - - public static class RequestContentStreamPlugin extends Plugin implements ActionPlugin { - - static final String ROUTE = "/_test/request-stream/basic"; - Consumer channelConsumer; - - @Override - public Collection getRestHandlers( - Settings settings, - NamedWriteableRegistry namedWriteableRegistry, - RestController restController, - ClusterSettings clusterSettings, - IndexScopedSettings indexScopedSettings, - SettingsFilter settingsFilter, - IndexNameExpressionResolver indexNameExpressionResolver, - Supplier nodesInCluster, - Predicate clusterSupportsFeature - ) { - return List.of(new BaseRestHandler() { - @Override - public String getName() { - return ROUTE; - } - - @Override - public List routes() { - return List.of(new Route(RestRequest.Method.POST, ROUTE)); - } - - @Override - protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { - return new RequestBodyChunkConsumer() { - - int totalBytes = 0; - - @Override - public void handleChunk(RestChannel channel, ReleasableBytesReference chunk, boolean isLast) { - try (chunk) { - totalBytes += chunk.length(); - if (isLast == false) { - request.contentStream().requestBytes(1024); - } else { - channel.sendResponse(new RestResponse(RestStatus.OK, Integer.toString(totalBytes))); - } - } - } - - @Override - public void accept(RestChannel channel) throws Exception { - request.contentStream().requestBytes(1024); // ask for first chunk - } - }; - } - }); - } - } - -} diff --git a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java new file mode 100644 index 0000000000000..e78231bcc772e --- /dev/null +++ b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java @@ -0,0 +1,344 @@ +/* + * 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.bootstrap.Bootstrap; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.Unpooled; +import io.netty.channel.Channel; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelInitializer; +import io.netty.channel.SimpleChannelInboundHandler; +import io.netty.channel.nio.NioEventLoopGroup; +import io.netty.channel.socket.SocketChannel; +import io.netty.channel.socket.nio.NioSocketChannel; +import io.netty.handler.codec.http.DefaultFullHttpRequest; +import io.netty.handler.codec.http.DefaultHttpContent; +import io.netty.handler.codec.http.DefaultHttpRequest; +import io.netty.handler.codec.http.DefaultLastHttpContent; +import io.netty.handler.codec.http.FullHttpRequest; +import io.netty.handler.codec.http.FullHttpResponse; +import io.netty.handler.codec.http.HttpClientCodec; +import io.netty.handler.codec.http.HttpContent; +import io.netty.handler.codec.http.HttpObjectAggregator; +import io.netty.handler.codec.http.HttpRequest; + +import org.elasticsearch.ESNetty4IntegTestCase; +import org.elasticsearch.action.support.SubscribableListener; +import org.elasticsearch.client.internal.node.NodeClient; +import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver; +import org.elasticsearch.cluster.node.DiscoveryNodes; +import org.elasticsearch.common.bytes.BytesReference; +import org.elasticsearch.common.bytes.ReleasableBytesReference; +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.HttpBody; +import org.elasticsearch.http.HttpChannel; +import org.elasticsearch.http.HttpServerTransport; +import org.elasticsearch.plugins.ActionPlugin; +import org.elasticsearch.plugins.Plugin; +import org.elasticsearch.rest.BaseRestHandler; +import org.elasticsearch.rest.RestChannel; +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 org.elasticsearch.test.ESIntegTestCase; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.BlockingDeque; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingDeque; +import java.util.concurrent.TimeUnit; +import java.util.function.Predicate; +import java.util.function.Supplier; + +import static io.netty.handler.codec.http.HttpHeaderNames.CONTENT_LENGTH; +import static io.netty.handler.codec.http.HttpHeaderNames.CONTENT_TYPE; +import static io.netty.handler.codec.http.HttpHeaderValues.APPLICATION_JSON; +import static io.netty.handler.codec.http.HttpMethod.POST; +import static io.netty.handler.codec.http.HttpVersion.HTTP_1_1; + +@ESIntegTestCase.ClusterScope(numDataNodes = 1) +public class Netty4IncrementalRequestHandlingIT extends ESNetty4IntegTestCase { + + // ensure empty http content has single 0 size chunk + public void testEmptyContent() throws Exception { + try (var ctx = setupClientCtx()) { + ctx.clientChannel.writeAndFlush(fullRequest(Unpooled.EMPTY_BUFFER)); + var handler = ctx.awaitRestChannelAccepted(); + handler.stream.requestBytes(1); + var recvChunk = safePoll(handler.recvChunks); + assertTrue(recvChunk.isLast); + assertEquals(0, recvChunk.chunk.length()); + recvChunk.chunk.close(); + } + } + + // ensures content integrity, no loses and re-order + public void testReceiveAllChunks() throws Exception { + try (var ctx = setupClientCtx()) { + var dataSize = randomIntBetween(1024, 10 * 1024 * 1024); + var sendData = Unpooled.wrappedBuffer(randomByteArrayOfLength(dataSize)); + sendData.retain(); + ctx.clientChannel.writeAndFlush(fullRequest(sendData)); + + var handler = ctx.awaitRestChannelAccepted(); + + var gotAllChunks = false; + var recvData = Unpooled.buffer(dataSize); + while (gotAllChunks == false) { + handler.stream.requestBytes(randomIntBetween(1024, 64 * 1024)); + var recvChunk = safePoll(handler.recvChunks); + try (recvChunk.chunk) { + if (recvChunk.isLast) { + gotAllChunks = true; + } + recvData.writeBytes(BytesReference.toBytes(recvChunk.chunk)); + } + } + assertEquals(sendData, recvData); + } + } + + // ensure that client's socket is no longer writable if server not consuming data + public void testClientBackpressure() throws Exception { + try (var ctx = setupClientCtx()) { + + // a meaningful write payload that is larger than socket and channel buffers + final var writeChunkSize = 16 * 1024; + final var totalWriteChunks = randomIntBetween(1024, 2048); // 16-32MB + final var totalWriteSize = writeChunkSize * totalWriteChunks; + var writtenChunks = 0; + + ctx.clientChannel.writeAndFlush(notFullRequest(totalWriteSize)); + var handler = ctx.awaitRestChannelAccepted(); + + // write chunks until channel is no longer writable + while (writtenChunks <= totalWriteChunks && ctx.clientChannel.isWritable()) { + ctx.clientChannel.writeAndFlush(randomContent(writeChunkSize, false)); + writtenChunks += 1; + } + assertTrue("should not be able to write all chunks", writtenChunks < totalWriteChunks); + + // read all written chunks and check client socket is writable again + var totalReadSize = 0; + var mustReadBytes = writtenChunks * writeChunkSize; + while (mustReadBytes > 0) { + handler.stream.requestBytes(writeChunkSize); + var recvChunk = safePoll(handler.recvChunks); + mustReadBytes -= recvChunk.chunk.length(); + totalReadSize += recvChunk.chunk.length(); + recvChunk.chunk.close(); + } + + assertBusy(() -> assertTrue("channel must be writable again", ctx.clientChannel.isWritable())); + + // write remaining chunks + while (writtenChunks < totalWriteChunks) { // leave 1 for LastHttpContent + writtenChunks += 1; + ctx.clientChannel.writeAndFlush(randomContent(writeChunkSize, false)); + } + ctx.clientChannel.writeAndFlush(randomContent(writeChunkSize, true)); + + // consume remaining chunks + var allConsumed = false; + while (allConsumed == false) { + handler.stream.requestBytes(writeChunkSize); + var recvChunk = safePoll(handler.recvChunks); + totalReadSize += recvChunk.chunk.length(); + allConsumed = recvChunk.isLast; + recvChunk.chunk.close(); + } + assertEquals("did not received all bytes", totalWriteSize, totalReadSize); + + // send OK response back to wrap it up + handler.channel.sendResponse(new RestResponse(RestStatus.OK, "")); + var resp = (FullHttpResponse) safePoll(ctx.clientRespQueue); + assertEquals(200, resp.status().code()); + resp.release(); + } + } + + static T safePoll(BlockingDeque queue) { + try { + var t = queue.poll(SAFE_AWAIT_TIMEOUT.seconds(), TimeUnit.SECONDS); + assertNotNull("queue is empty", t); + return t; + } catch (InterruptedException e) { + throw new AssertionError(e); + } + } + + static FullHttpRequest fullRequest(ByteBuf content) { + var req = new DefaultFullHttpRequest(HTTP_1_1, POST, SingleRequestHandlerPlugin.ROUTE, Unpooled.wrappedBuffer(content)); + req.headers().add(CONTENT_LENGTH, content.readableBytes()); + req.headers().add(CONTENT_TYPE, APPLICATION_JSON); + return req; + } + + static HttpRequest notFullRequest(int contentLength) { + var req = new DefaultHttpRequest(HTTP_1_1, POST, SingleRequestHandlerPlugin.ROUTE); + req.headers().add(CONTENT_LENGTH, contentLength); + req.headers().add(CONTENT_TYPE, APPLICATION_JSON); + return req; + } + + static HttpContent randomContent(int size, boolean isLast) { + var buf = Unpooled.wrappedBuffer(randomByteArrayOfLength(size)); + if (isLast) { + return new DefaultLastHttpContent(buf); + } else { + return new DefaultHttpContent(buf); + } + } + + record Ctx(String nodeName, Bootstrap clientBootstrap, Channel clientChannel, BlockingDeque clientRespQueue) + implements + AutoCloseable { + + @Override + public void close() throws Exception { + safeGet(clientChannel.close()); + safeGet(clientBootstrap.config().group().shutdownGracefully()); + clientRespQueue.forEach(o -> { if (o instanceof FullHttpResponse resp) resp.release(); }); + for (var handler : SingleRequestHandlerPlugin.handlers) { + handler.recvChunks.forEach(c -> c.chunk.close()); + handler.httpChannel.close(); + } + SingleRequestHandlerPlugin.handlers.clear(); + } + + ServerRequestHandler awaitRestChannelAccepted() throws Exception { + var handlers = SingleRequestHandlerPlugin.handlers; + assertBusy(() -> { assertEquals(1, handlers.size()); }); + var handler = handlers.get(0); + safeAwait(handler.channelAccepted); + return handler; + } + } + + Ctx setupClientCtx() throws Exception { + var nodeName = internalCluster().getRandomNodeName(); + var clientRespQueue = new LinkedBlockingDeque<>(16); + var bootstrap = bootstrapClient(nodeName, clientRespQueue); + var channel = bootstrap.connect().sync().channel(); + return new Ctx(nodeName, bootstrap, channel, clientRespQueue); + } + + Bootstrap bootstrapClient(String node, BlockingQueue queue) { + var httpServer = internalCluster().getInstance(HttpServerTransport.class, node); + var remoteAddr = randomFrom(httpServer.boundAddress().boundAddresses()); + return new Bootstrap().group(new NioEventLoopGroup(1)) + .channel(NioSocketChannel.class) + .remoteAddress(remoteAddr.getAddress(), remoteAddr.getPort()) + .handler(new ChannelInitializer() { + @Override + protected void initChannel(SocketChannel ch) { + var p = ch.pipeline(); + p.addLast(new HttpClientCodec()); + p.addLast(new HttpObjectAggregator(4096)); + p.addLast(new SimpleChannelInboundHandler() { + @Override + protected void channelRead0(ChannelHandlerContext ctx, FullHttpResponse msg) { + msg.retain(); + queue.add(msg); + } + + @Override + public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { + queue.add(cause); + } + }); + } + }); + } + + @Override + protected Collection> nodePlugins() { + return CollectionUtils.concatLists(List.of(SingleRequestHandlerPlugin.class), super.nodePlugins()); + } + + @Override + protected boolean addMockHttpTransport() { + return false; // enable http + } + + static class ServerRequestHandler implements BaseRestHandler.RequestBodyChunkConsumer { + SubscribableListener channelAccepted = new SubscribableListener<>(); + RestChannel channel; + HttpChannel httpChannel; + HttpBody.Stream stream; + BlockingDeque recvChunks = new LinkedBlockingDeque<>(); + + record Chunk(ReleasableBytesReference chunk, boolean isLast) {} + + @Override + public void handleChunk(RestChannel channel, ReleasableBytesReference chunk, boolean isLast) { + recvChunks.add(new Chunk(chunk, isLast)); + } + + @Override + public void accept(RestChannel channel) throws Exception { + this.channel = channel; + this.httpChannel = channel.request().getHttpChannel(); + this.stream = channel.request().contentStream(); + channelAccepted.onResponse(null); + } + } + + public static class SingleRequestHandlerPlugin extends Plugin implements ActionPlugin { + + static final String ROUTE = "/_test/request-stream"; + + static final List handlers = Collections.synchronizedList(new ArrayList<>()); + + @Override + public Collection getRestHandlers( + Settings settings, + NamedWriteableRegistry namedWriteableRegistry, + RestController restController, + ClusterSettings clusterSettings, + IndexScopedSettings indexScopedSettings, + SettingsFilter settingsFilter, + IndexNameExpressionResolver indexNameExpressionResolver, + Supplier nodesInCluster, + Predicate clusterSupportsFeature + ) { + return List.of(new BaseRestHandler() { + @Override + public String getName() { + return ROUTE; + } + + @Override + public List routes() { + return List.of(new Route(RestRequest.Method.POST, ROUTE)); + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { + var handler = new ServerRequestHandler(); + handlers.add(handler); + return handler; + } + }); + } + } + +} diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java index 9b90dba1eed98..af4d0498e1dea 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java @@ -119,10 +119,10 @@ public Netty4HttpPipeliningHandler( @Override public void channelRead(final ChannelHandlerContext ctx, final Object msg) { - activityTracker.startActivity(); - try { - final Netty4HttpRequest netty4HttpRequest; - if (msg instanceof HttpRequest request) { + if (msg instanceof HttpRequest request) { + try { + activityTracker.startActivity(); + final Netty4HttpRequest netty4HttpRequest; if (request.decoderResult().isFailure()) { final Throwable cause = request.decoderResult().cause(); final Exception nonError; @@ -138,7 +138,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { currentRequestStream = null; netty4HttpRequest = new Netty4HttpRequest(readSequence++, fullHttpRequest); } else { - var contentStream = new Netty4HttpRequestBodyStream(ctx.channel()); + var contentStream = new Netty4HttpRequestBodyStream(ctx); currentRequestStream = contentStream; netty4HttpRequest = new Netty4HttpRequest( readSequence++, @@ -155,16 +155,16 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { } } handlePipelinedRequest(ctx, netty4HttpRequest); - } else { - assert msg instanceof HttpContent; - assert currentRequestStream != null; - currentRequestStream.handleNettyContent((HttpContent) msg); - if (msg instanceof LastHttpContent) { - currentRequestStream = null; - } + } finally { + activityTracker.stopActivity(); + } + } else { + assert msg instanceof HttpContent : "expect HttpContent got " + msg; + assert currentRequestStream != null : "current stream must exists before handling http content"; + currentRequestStream.handleNettyContent((HttpContent) msg); + if (msg instanceof LastHttpContent) { + currentRequestStream = null; } - } finally { - activityTracker.stopActivity(); } } diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java index 63132424af5e5..a9c7de81a58a2 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java @@ -10,7 +10,7 @@ import io.netty.buffer.ByteBuf; import io.netty.buffer.CompositeByteBuf; -import io.netty.channel.Channel; +import io.netty.channel.ChannelHandlerContext; import io.netty.handler.codec.http.HttpContent; import io.netty.handler.codec.http.LastHttpContent; @@ -26,14 +26,14 @@ */ public class Netty4HttpRequestBodyStream implements HttpBody.Stream { - private final Channel channel; + private final ChannelHandlerContext ctx; private HttpBody.ChunkHandler handler; private CompositeByteBuf aggregate; private int requestedBytes = 0; - public Netty4HttpRequestBodyStream(Channel channel) { - this.channel = channel; - channel.config().setAutoRead(false); + public Netty4HttpRequestBodyStream(ChannelHandlerContext ctx) { + this.ctx = ctx; + ctx.channel().config().setAutoRead(false); } @Override @@ -50,7 +50,7 @@ public void setHandler(ChunkHandler chunkHandler) { public void requestBytes(int bytes) { assert handler != null : "handler must be set before requesting next chunk"; requestedBytes += bytes; - channel.read(); + ctx.read(); } public void handleNettyContent(HttpContent httpContent) { @@ -70,15 +70,15 @@ public void handleNettyContent(HttpContent httpContent) { requestedBytes = 0; handler.onNext(Netty4Utils.toReleasableBytesReference(sendBuf), isLast); if (isLast) { - channel.config().setAutoRead(true); + ctx.channel().config().setAutoRead(true); } } else { if (aggregate == null) { - aggregate = channel.alloc().compositeBuffer(); + aggregate = ctx.alloc().compositeBuffer(); } aggregate.addComponent(true, content); requestedBytes -= content.readableBytes(); - channel.read(); + ctx.read(); } } diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java index 75a5d0022a1e9..e00373a450dc5 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java @@ -352,7 +352,6 @@ protected HttpMessage createMessage(String[] initialLine) throws Exception { } decoder.setCumulator(ByteToMessageDecoder.COMPOSITE_CUMULATOR); ch.pipeline().addLast("decoder", decoder); // parses the HTTP bytes request into HTTP message pieces - ch.pipeline().addLast(new FlowControlHandler()); if (httpValidator != null) { // runs a validation function on the first HTTP message piece which contains all the headers // if validation passes, the pieces of that particular request are forwarded, otherwise they are discarded @@ -386,6 +385,7 @@ protected boolean isContentAlwaysEmpty(HttpResponse msg) { if (handlingSettings.compression()) { ch.pipeline().addLast("encoder_compress", new HttpContentCompressor(handlingSettings.compressionLevel())); } + ch.pipeline().addLast(new FlowControlHandler()); ch.pipeline() .addLast( "pipelining", From 14a06cefb0c44237fdbeb4a6ae10571674c358a6 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Fri, 2 Aug 2024 19:36:43 -0700 Subject: [PATCH 06/14] poke ci From fbf5d2837915c83e6917ba14d1f33b3e07d401c2 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Sat, 3 Aug 2024 18:59:51 -0700 Subject: [PATCH 07/14] remove flow control handler --- .../netty4/Netty4HttpPipeliningHandler.java | 2 +- .../netty4/Netty4HttpRequestBodyStream.java | 175 ++++++++++++++---- .../netty4/Netty4HttpServerTransport.java | 2 - .../Netty4HttpRequestBodyStreamTests.java | 174 +++++++++++++++++ .../java/org/elasticsearch/http/HttpBody.java | 19 +- 5 files changed, 333 insertions(+), 39 deletions(-) create mode 100644 modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java index af4d0498e1dea..6572ba5b6843f 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java @@ -138,7 +138,7 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { currentRequestStream = null; netty4HttpRequest = new Netty4HttpRequest(readSequence++, fullHttpRequest); } else { - var contentStream = new Netty4HttpRequestBodyStream(ctx); + var contentStream = new Netty4HttpRequestBodyStream(ctx.channel()); currentRequestStream = contentStream; netty4HttpRequest = new Netty4HttpRequest( readSequence++, diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java index a9c7de81a58a2..108ee5cf2842f 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java @@ -9,31 +9,36 @@ package org.elasticsearch.http.netty4; import io.netty.buffer.ByteBuf; -import io.netty.buffer.CompositeByteBuf; -import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.Channel; import io.netty.handler.codec.http.HttpContent; import io.netty.handler.codec.http.LastHttpContent; import org.elasticsearch.http.HttpBody; import org.elasticsearch.transport.netty4.Netty4Utils; +import java.util.ArrayDeque; +import java.util.Queue; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.atomic.AtomicBoolean; + /** - * Implementation of HTTP request content stream. This implementation relies on - * {@link io.netty.channel.ChannelConfig#setAutoRead(boolean)} to stop arrival of new chunks. - * Otherwise, downstream might receive much more data than requested. - * This can be achieved with {@link io.netty.handler.flow.FlowControlHandler} placed before stream - * constructor in {@link io.netty.channel.ChannelPipeline} and after {@link io.netty.handler.codec.http.HttpRequestDecoder}. + * Netty based implementation of {@link HttpBody.Stream}. + * This implementation utilize {@link io.netty.channel.ChannelConfig#setAutoRead(boolean)} to prevent entire payload buffering. + * But sometimes upstream can send few chunks of data despite autoRead=off. + * In this case chunks will be queued until downstream {@link HttpBody.Stream#requestBytes} */ public class Netty4HttpRequestBodyStream implements HttpBody.Stream { - private final ChannelHandlerContext ctx; + private final Channel channel; + private final BlockingQueue requestQueue = new LinkedBlockingQueue<>(); + private final ChunkQueue chunkQueue = new ChunkQueue(); + private final AtomicBoolean isDone = new AtomicBoolean(false); private HttpBody.ChunkHandler handler; - private CompositeByteBuf aggregate; - private int requestedBytes = 0; - public Netty4HttpRequestBodyStream(ChannelHandlerContext ctx) { - this.ctx = ctx; - ctx.channel().config().setAutoRead(false); + public Netty4HttpRequestBodyStream(Channel channel) { + this.channel = channel; + channel.config().setAutoRead(false); } @Override @@ -49,37 +54,139 @@ public void setHandler(ChunkHandler chunkHandler) { @Override public void requestBytes(int bytes) { assert handler != null : "handler must be set before requesting next chunk"; - requestedBytes += bytes; - ctx.read(); + if (isDone.get()) { + return; + } + requestQueue.add(bytes); + channel.eventLoop().execute(this::flushQueued); // flush on channel's thread only + if (requestQueue.isEmpty() == false) { + channel.read(); + } } public void handleNettyContent(HttpContent httpContent) { assert handler != null : "handler must be set before processing http content"; - var isLast = httpContent instanceof LastHttpContent; + var content = httpContent.content(); + var isLast = httpContent instanceof LastHttpContent; + var currentRequest = requestQueue.peek(); + var availableBytes = chunkQueue.availableBytes() + content.readableBytes(); - if (content.readableBytes() >= requestedBytes || isLast) { - ByteBuf sendBuf; - if (aggregate == null) { - sendBuf = content; + if (currentRequest == null || (currentRequest > availableBytes && isLast == false)) { + // enqueue incoming chunk: + // 1. There is no ongoing request, it might happen that upstream send a chunk before request. + // 2. There are not enough accumulated bytes, unless it's a last chunk then we must flush + chunkQueue.offer(new ContentChunk(content, isLast)); + } else { + // flush queued chunks: + // 1. nothing in the queue and current chunk fulfills request, send it right a way + // 2. there are queued chunks, enqueue and flush + if (chunkQueue.isEmpty()) { + this.handler.onNext(Netty4Utils.toReleasableBytesReference(content), isLast); + requestQueue.poll(); } else { - aggregate.addComponent(true, content); - sendBuf = aggregate; - aggregate = null; - } - requestedBytes = 0; - handler.onNext(Netty4Utils.toReleasableBytesReference(sendBuf), isLast); - if (isLast) { - ctx.channel().config().setAutoRead(true); + chunkQueue.offer(new ContentChunk(content, isLast)); + flushQueued(); } - } else { - if (aggregate == null) { - aggregate = ctx.alloc().compositeBuffer(); + } + + if (isLast) { + isDone.set(true); + channel.config().setAutoRead(true); + } else if (requestQueue.peek() != null) { + channel.read(); + } + } + + // visible for tests + ChunkQueue contentChunkQueue() { + return chunkQueue; + } + + private boolean canFulfillCurrentRequest() { + return requestQueue.isEmpty() == false && chunkQueue.isEmpty() == false + // enough queued bytes + && (chunkQueue.availableBytes() >= requestQueue.peek() + // not enough queued bytes, but got last content + || (chunkQueue.availableBytes < requestQueue.peek() && chunkQueue.hasLast())); + } + + private void flushQueued() { + while (canFulfillCurrentRequest()) { + var currentRequest = requestQueue.poll(); + assert currentRequest != null; + ByteBuf buf; + boolean isLast = false; + // next chunk is large enough, no need to aggregate + if (chunkQueue.peekSize() >= currentRequest) { + var chunk = chunkQueue.poll(); + isLast = chunk.isLast; + buf = chunk.buf; + } else { + // aggregate multiple chunks + var compBuf = channel.alloc().compositeBuffer(); + while (compBuf.readableBytes() < currentRequest && chunkQueue.isEmpty() == false) { + var chunk = chunkQueue.poll(); + isLast = chunk.isLast; + compBuf.addComponent(true, chunk.buf()); + } + buf = compBuf; } - aggregate.addComponent(true, content); - requestedBytes -= content.readableBytes(); - ctx.read(); + handler.onNext(Netty4Utils.toReleasableBytesReference(buf), isLast); } } + record ContentChunk(ByteBuf buf, boolean isLast) {} + + // Wrapped queue that keeps track of total available bytes of all chunks in the queue + static class ChunkQueue { + private final Queue queue = new ArrayDeque<>(); + private int availableBytes = 0; + private boolean hasLast = false; + + void release() { + while (queue.isEmpty() == false) { + queue.poll().buf.release(); + } + } + + boolean isEmpty() { + return queue.isEmpty(); + } + + int availableBytes() { + return availableBytes; + } + + boolean hasLast() { + return hasLast; + } + + int size() { + return queue.size(); + } + + void offer(ContentChunk content) { + queue.offer(content); + availableBytes += content.buf.readableBytes(); + hasLast = content.isLast; + } + + // returns size of next chunk in queue + int peekSize() { + var chunk = queue.peek(); + if (chunk == null) { + return 0; + } else { + return chunk.buf.readableBytes(); + } + } + + ContentChunk poll() { + var chunk = queue.poll(); + assert chunk != null; + availableBytes -= chunk.buf().readableBytes(); + return chunk; + } + } } diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java index e00373a450dc5..71664431efa3f 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpServerTransport.java @@ -28,7 +28,6 @@ import io.netty.handler.codec.http.HttpResponse; import io.netty.handler.codec.http.HttpResponseEncoder; import io.netty.handler.codec.http.HttpUtil; -import io.netty.handler.flow.FlowControlHandler; import io.netty.handler.ssl.SslHandler; import io.netty.handler.timeout.ReadTimeoutException; import io.netty.handler.timeout.ReadTimeoutHandler; @@ -385,7 +384,6 @@ protected boolean isContentAlwaysEmpty(HttpResponse msg) { if (handlingSettings.compression()) { ch.pipeline().addLast("encoder_compress", new HttpContentCompressor(handlingSettings.compressionLevel())); } - ch.pipeline().addLast(new FlowControlHandler()); ch.pipeline() .addLast( "pipelining", diff --git a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java new file mode 100644 index 0000000000000..3ea541c519771 --- /dev/null +++ b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java @@ -0,0 +1,174 @@ +/* + * 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.buffer.Unpooled; +import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.SimpleChannelInboundHandler; +import io.netty.channel.embedded.EmbeddedChannel; +import io.netty.handler.codec.http.DefaultHttpContent; +import io.netty.handler.codec.http.DefaultLastHttpContent; +import io.netty.handler.codec.http.HttpContent; +import io.netty.handler.flow.FlowControlHandler; + +import org.elasticsearch.common.bytes.ReleasableBytesReference; +import org.elasticsearch.http.HttpBody; +import org.elasticsearch.test.ESTestCase; +import org.junit.Before; + +import java.util.ArrayList; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +public class Netty4HttpRequestBodyStreamTests extends ESTestCase { + + EmbeddedChannel channel; + Netty4HttpRequestBodyStream stream; + static HttpBody.ChunkHandler discardHandler = (chunk, isLast) -> chunk.close(); + + @Before + public void createStream() { + channel = new EmbeddedChannel(); + stream = new Netty4HttpRequestBodyStream(channel); + stream.setHandler(discardHandler); // set default handler, each test might override one + channel.pipeline().addLast(new SimpleChannelInboundHandler() { + @Override + protected void channelRead0(ChannelHandlerContext ctx, HttpContent msg) { + msg.retain(); + stream.handleNettyContent(msg); + } + }); + } + + public void testEnqueueChunksBeforeRequest() { + var chunkSize = 1024; + var totalChunks = randomIntBetween(1, 100); + for (int i = 0; i < totalChunks; i++) { + channel.writeInbound(randomContent(chunkSize)); + } + assertEquals(totalChunks, stream.contentChunkQueue().size()); + assertEquals(totalChunks * chunkSize, stream.contentChunkQueue().availableBytes()); + } + + // test that every small request consume at least one chunk + public void testFlushQueuedSmallRequests() { + var chunks = new ArrayList(); + var totalBytes = new AtomicInteger(); + stream.setHandler((chunk, isLast) -> { + chunks.add(chunk); + totalBytes.addAndGet(chunk.length()); + }); + // enqueue chunks + var chunkSize = 1024; + var totalChunks = randomIntBetween(1, 100); + for (int i = 0; i < totalChunks; i++) { + channel.writeInbound(randomContent(chunkSize)); + } + // consume all chunks + for (var i = 0; i < totalChunks; i++) { + stream.requestBytes(1); + } + assertEquals(totalChunks, chunks.size()); + assertEquals(chunkSize * totalChunks, totalBytes.get()); + } + + // test that one large request consume all chunks + public void testFlushQueuedLargeRequest() { + var requestedChunkSize = new AtomicInteger(); + stream.setHandler((chunk, isLast) -> { requestedChunkSize.addAndGet(chunk.length()); }); + var chunkSize = 1024; + var totalChunks = randomIntBetween(1, 100); + for (int i = 0; i < totalChunks; i++) { + channel.writeInbound(randomContent(chunkSize)); + } + stream.requestBytes(chunkSize * totalChunks); + assertEquals(chunkSize * totalChunks, requestedChunkSize.get()); + } + + // test that we flush queued chunks with last content when request is larger than availableBytes + public void testFlushQueuedLastContent() { + var gotLast = new AtomicBoolean(false); + var contentSize = new AtomicInteger(0); + stream.setHandler((chunk, isLast) -> { + contentSize.addAndGet(chunk.length()); + gotLast.set(isLast); + }); + channel.writeInbound(randomContent(1024)); + channel.writeInbound(randomLastContent(0)); + stream.requestBytes(1025); + assertTrue(gotLast.get()); + assertEquals(1024, contentSize.get()); + } + + // ensures that we enqueue chunk if cannot fulfill current request + public void testEnqueueWhenNotEnoughBytes() { + stream.requestBytes(1024); + var chunkSize = 32; + var totalChunks = randomIntBetween(1, 1024 / chunkSize - 1); + for (int i = 0; i < totalChunks; i++) { + channel.writeInbound(randomContent(chunkSize)); + } + assertEquals(chunkSize * totalChunks, stream.contentChunkQueue().availableBytes()); + } + + public void testFlushSingleLastContent() { + var flushed = new AtomicBoolean(false); + stream.setHandler((chunk, isLast) -> flushed.set(true)); + stream.requestBytes(1024); + channel.writeInbound(randomLastContent(0)); + assertTrue(flushed.get()); + } + + // test that we read from channel when not enough available bytes + public void testReadFromChannel() { + var gotChunks = new ArrayList(); + var gotLast = new AtomicBoolean(false); + stream.setHandler((chunk, isLast) -> { + gotChunks.add(chunk); + gotLast.set(isLast); + }); + channel.pipeline().addFirst(new FlowControlHandler()); // block all incoming messages, need explicit channel.read() + var chunkSize = 1024; + var totalChunks = randomIntBetween(1, 32); + for (int i = 0; i < totalChunks; i++) { + channel.writeInbound(randomContent(chunkSize)); + } + + assertEquals("should not read until requested", 0, stream.contentChunkQueue().size()); + stream.requestBytes(chunkSize); + assertEquals(1, gotChunks.size()); + + stream.requestBytes(Integer.MAX_VALUE); + assertEquals("should enqueue remaining chunks", totalChunks - 1, stream.contentChunkQueue().size()); + + // send last content and flush queued content + channel.writeInbound(randomLastContent(0)); + assertEquals(0, stream.contentChunkQueue().size()); + assertEquals(2, gotChunks.size()); + assertTrue(gotLast.get()); + } + + HttpContent randomContent(int size, boolean isLast) { + var buf = Unpooled.wrappedBuffer(randomByteArrayOfLength(size)); + if (isLast) { + return new DefaultLastHttpContent(buf); + } else { + return new DefaultHttpContent(buf); + } + } + + HttpContent randomContent(int size) { + return randomContent(size, false); + } + + HttpContent randomLastContent(int size) { + return randomContent(size, true); + } + +} diff --git a/server/src/main/java/org/elasticsearch/http/HttpBody.java b/server/src/main/java/org/elasticsearch/http/HttpBody.java index 4eb72229b5516..99984c3d1af1d 100644 --- a/server/src/main/java/org/elasticsearch/http/HttpBody.java +++ b/server/src/main/java/org/elasticsearch/http/HttpBody.java @@ -67,8 +67,23 @@ non-sealed interface Stream extends HttpBody { void setHandler(ChunkHandler chunkHandler); /** - * Request next chunk of bytes. Implementation is free to round up chunk size to optimal network chunk size. - * For example, transport default chunk size is 8kb, request for 10kb might return 16kb chunk. + * Request next chunk of bytes. For every request there will be at least one chunk. + * Size of the next chunk might vary, depending on following factors: + *
    + *
  • + * Size might round up to optimal network chunk size. + * For example, default HttpDecoder content size is 8kb. + * Request for 1 byte will produce 8kb chunk, 5 requests for 1 byte will produce 5 chunks of 8kb. + * Request for 10kb will produce 16kb chunk. + *
    • + *
  • + * Size will be lass or equal request size if its a last chunk. + * Request for Integer.MAX_VALUE(or other number larger than actual payload) will always produce single full content chunk. + *
    • + *
  • + * After last chunk there will be no more chunks. This method will not do anything. + *
    • + *
*/ void requestBytes(int bytes); } From 116e8021c01f77b5e7394cb5dfb9a9b403a7450e Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Sat, 3 Aug 2024 19:13:54 -0700 Subject: [PATCH 08/14] cleanup --- .../netty4/Netty4HttpRequestBodyStream.java | 23 ++++++++++--------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java index 108ee5cf2842f..a941d66fbcd40 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java @@ -58,10 +58,12 @@ public void requestBytes(int bytes) { return; } requestQueue.add(bytes); - channel.eventLoop().execute(this::flushQueued); // flush on channel's thread only - if (requestQueue.isEmpty() == false) { - channel.read(); - } + channel.eventLoop().execute(() -> { + flushQueued(); // flush on channel's thread only + if (requestQueue.isEmpty() == false) { + channel.read(); + } + }); } public void handleNettyContent(HttpContent httpContent) { @@ -105,10 +107,10 @@ ChunkQueue contentChunkQueue() { private boolean canFulfillCurrentRequest() { return requestQueue.isEmpty() == false && chunkQueue.isEmpty() == false - // enough queued bytes - && (chunkQueue.availableBytes() >= requestQueue.peek() - // not enough queued bytes, but got last content - || (chunkQueue.availableBytes < requestQueue.peek() && chunkQueue.hasLast())); + // enough queued bytes + && (chunkQueue.availableBytes() >= requestQueue.peek() + // not enough queued bytes, but got last content + || (chunkQueue.availableBytes < requestQueue.peek() && chunkQueue.hasLast())); } private void flushQueued() { @@ -118,7 +120,7 @@ private void flushQueued() { ByteBuf buf; boolean isLast = false; // next chunk is large enough, no need to aggregate - if (chunkQueue.peekSize() >= currentRequest) { + if (chunkQueue.nextChunkSize() >= currentRequest) { var chunk = chunkQueue.poll(); isLast = chunk.isLast; buf = chunk.buf; @@ -172,8 +174,7 @@ void offer(ContentChunk content) { hasLast = content.isLast; } - // returns size of next chunk in queue - int peekSize() { + int nextChunkSize() { var chunk = queue.peek(); if (chunk == null) { return 0; From d34d47f555f944d6c2a3f6153628c780cc4a4205 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Sat, 3 Aug 2024 19:28:58 -0700 Subject: [PATCH 09/14] update documentation --- .../java/org/elasticsearch/http/HttpBody.java | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/http/HttpBody.java b/server/src/main/java/org/elasticsearch/http/HttpBody.java index 99984c3d1af1d..6c575a8cb6040 100644 --- a/server/src/main/java/org/elasticsearch/http/HttpBody.java +++ b/server/src/main/java/org/elasticsearch/http/HttpBody.java @@ -34,25 +34,32 @@ default boolean isStream() { return this instanceof Stream; } + /** + * Assumes that HTTP body is a full content. If not sure, use {@link HttpBody#isFull()}. + */ default Full asFull() { assert this instanceof Full; return (Full) this; } + /** + * Assumes that HTTP body is a lazy-stream. If not sure, use {@link HttpBody#isStream()}. + */ default Stream asStream() { assert this instanceof Stream; return (Stream) this; } /** - * Full content represents a complete http body content that can be accessed immediately + * Full content represents a complete http body content that can be accessed immediately. */ non-sealed interface Full extends HttpBody { BytesReference bytes(); } /** - * Stream is a lazy-loaded content. Stream supports only single handler, this handler must be set before requesting next chunk. + * Stream is a lazy-loaded content. Stream supports only single handler, this handler must be + * set before requesting next chunk. */ non-sealed interface Stream extends HttpBody { /** @@ -67,18 +74,17 @@ non-sealed interface Stream extends HttpBody { void setHandler(ChunkHandler chunkHandler); /** - * Request next chunk of bytes. For every request there will be at least one chunk. - * Size of the next chunk might vary, depending on following factors: + * Request next chunk of bytes. For every request there will be at least one chunk. Size of + * the next chunk might vary, and it depends on following factors: *
    *
  • - * Size might round up to optimal network chunk size. - * For example, default HttpDecoder content size is 8kb. - * Request for 1 byte will produce 8kb chunk, 5 requests for 1 byte will produce 5 chunks of 8kb. - * Request for 10kb will produce 16kb chunk. + * Size might round up to optimal network chunk size. For example, default HttpDecoder + * content size is 8kb. Request for 1 byte will produce 8kb chunk, 5 requests for 1 byte + * will produce 5 chunks of 8kb each. Request for 10kb will produce 16kb chunk. *
    • *
  • - * Size will be lass or equal request size if its a last chunk. - * Request for Integer.MAX_VALUE(or other number larger than actual payload) will always produce single full content chunk. + * Last chunk always less or equal requested size. Request for Integer.MAX_VALUE or + * other number larger than actual payload, will always produce single full content chunk. *
    • *
  • * After last chunk there will be no more chunks. This method will not do anything. From 90c8b5d79306ab9635d8559b26601c377584f4e1 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Sat, 3 Aug 2024 19:59:31 -0700 Subject: [PATCH 10/14] simplify conditions around last content --- .../netty4/Netty4HttpRequestBodyStream.java | 44 ++++++++++++------- 1 file changed, 27 insertions(+), 17 deletions(-) diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java index a941d66fbcd40..4b5465eb708f5 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java @@ -54,16 +54,13 @@ public void setHandler(ChunkHandler chunkHandler) { @Override public void requestBytes(int bytes) { assert handler != null : "handler must be set before requesting next chunk"; - if (isDone.get()) { - return; - } + assert bytes != 0; requestQueue.add(bytes); - channel.eventLoop().execute(() -> { - flushQueued(); // flush on channel's thread only - if (requestQueue.isEmpty() == false) { - channel.read(); - } - }); + if (channel.eventLoop().inEventLoop()) { + flushAndRead(); + } else { + channel.eventLoop().submit(this::flushAndRead); + } } public void handleNettyContent(HttpContent httpContent) { @@ -93,7 +90,6 @@ public void handleNettyContent(HttpContent httpContent) { } if (isLast) { - isDone.set(true); channel.config().setAutoRead(true); } else if (requestQueue.peek() != null) { channel.read(); @@ -106,11 +102,9 @@ ChunkQueue contentChunkQueue() { } private boolean canFulfillCurrentRequest() { - return requestQueue.isEmpty() == false && chunkQueue.isEmpty() == false - // enough queued bytes - && (chunkQueue.availableBytes() >= requestQueue.peek() - // not enough queued bytes, but got last content - || (chunkQueue.availableBytes < requestQueue.peek() && chunkQueue.hasLast())); + return requestQueue.isEmpty() == false + && chunkQueue.isEmpty() == false + && (chunkQueue.availableBytes() >= requestQueue.peek() || chunkQueue.hasLast()); } private void flushQueued() { @@ -119,8 +113,8 @@ private void flushQueued() { assert currentRequest != null; ByteBuf buf; boolean isLast = false; - // next chunk is large enough, no need to aggregate - if (chunkQueue.nextChunkSize() >= currentRequest) { + // skip aggregation if next chunk is enough + if (chunkQueue.nextChunkSize() >= currentRequest || chunkQueue.nextIsLast()) { var chunk = chunkQueue.poll(); isLast = chunk.isLast; buf = chunk.buf; @@ -138,6 +132,13 @@ private void flushQueued() { } } + private void flushAndRead() { + flushQueued(); // flush on channel's thread only + if (requestQueue.isEmpty() == false) { + channel.read(); + } + } + record ContentChunk(ByteBuf buf, boolean isLast) {} // Wrapped queue that keeps track of total available bytes of all chunks in the queue @@ -183,6 +184,15 @@ int nextChunkSize() { } } + boolean nextIsLast() { + var chunk = queue.peek(); + if (chunk == null) { + return false; + } else { + return chunk.isLast; + } + } + ContentChunk poll() { var chunk = queue.poll(); assert chunk != null; From cd53cd735bd6bbf108356aaad8f0d5900111e6a0 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Tue, 6 Aug 2024 13:49:36 -0700 Subject: [PATCH 11/14] release queued buffers on channel close evt --- .../Netty4IncrementalRequestHandlingIT.java | 126 +++++++++++++----- .../netty4/Netty4HttpPipeliningHandler.java | 41 ++---- .../http/netty4/Netty4HttpRequest.java | 19 ++- .../netty4/Netty4HttpRequestBodyStream.java | 14 +- .../Netty4HttpRequestBodyStreamTests.java | 12 +- 5 files changed, 141 insertions(+), 71 deletions(-) diff --git a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java index e78231bcc772e..545288df02847 100644 --- a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java +++ b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java @@ -43,7 +43,6 @@ import org.elasticsearch.common.settings.SettingsFilter; import org.elasticsearch.common.util.CollectionUtils; import org.elasticsearch.features.NodeFeature; -import org.elasticsearch.http.HttpBody; import org.elasticsearch.http.HttpChannel; import org.elasticsearch.http.HttpServerTransport; import org.elasticsearch.plugins.ActionPlugin; @@ -80,39 +79,94 @@ public class Netty4IncrementalRequestHandlingIT extends ESNetty4IntegTestCase { // ensure empty http content has single 0 size chunk public void testEmptyContent() throws Exception { try (var ctx = setupClientCtx()) { - ctx.clientChannel.writeAndFlush(fullRequest(Unpooled.EMPTY_BUFFER)); - var handler = ctx.awaitRestChannelAccepted(); - handler.stream.requestBytes(1); - var recvChunk = safePoll(handler.recvChunks); - assertTrue(recvChunk.isLast); - assertEquals(0, recvChunk.chunk.length()); - recvChunk.chunk.close(); + var totalRequests = randomIntBetween(1, 10); + for (int reqNo = 0; reqNo < totalRequests; reqNo++) { + // send request with empty content + ctx.clientChannel.writeAndFlush(fullHttpRequest(Unpooled.EMPTY_BUFFER)); + var handler = ctx.awaitRestChannelAccepted(); + handler.stream.requestBytes(1); + + // should receive a single empty chunk + var recvChunk = safePoll(handler.recvChunks); + assertTrue(recvChunk.isLast); + assertEquals(0, recvChunk.chunk.length()); + recvChunk.chunk.close(); + + // send response to process following request + handler.sendResponse(new RestResponse(RestStatus.OK, "")); + } + assertBusy(() -> assertEquals("should receive all server responses", totalRequests, ctx.clientRespQueue.size())); } } // ensures content integrity, no loses and re-order public void testReceiveAllChunks() throws Exception { try (var ctx = setupClientCtx()) { - var dataSize = randomIntBetween(1024, 10 * 1024 * 1024); - var sendData = Unpooled.wrappedBuffer(randomByteArrayOfLength(dataSize)); - sendData.retain(); - ctx.clientChannel.writeAndFlush(fullRequest(sendData)); - - var handler = ctx.awaitRestChannelAccepted(); - - var gotAllChunks = false; - var recvData = Unpooled.buffer(dataSize); - while (gotAllChunks == false) { - handler.stream.requestBytes(randomIntBetween(1024, 64 * 1024)); - var recvChunk = safePoll(handler.recvChunks); - try (recvChunk.chunk) { - if (recvChunk.isLast) { - gotAllChunks = true; + var totalRequests = randomIntBetween(1, 10); + for (int reqNo = 0; reqNo < totalRequests; reqNo++) { + + // this dataset will be compared with one on server side + var dataSize = randomIntBetween(1024, 10 * 1024 * 1024); + var sendData = Unpooled.wrappedBuffer(randomByteArrayOfLength(dataSize)); + sendData.retain(); + ctx.clientChannel.writeAndFlush(fullHttpRequest(sendData)); + + var handler = ctx.awaitRestChannelAccepted(); + + // consume chunks at random pace + var gotAllChunks = false; + var recvData = Unpooled.buffer(dataSize); + while (gotAllChunks == false) { + handler.stream.requestBytes(randomIntBetween(1024, 64 * 1024)); + var recvChunk = safePoll(handler.recvChunks); + try (recvChunk.chunk) { + if (recvChunk.isLast) { + gotAllChunks = true; + } + recvData.writeBytes(BytesReference.toBytes(recvChunk.chunk)); } - recvData.writeBytes(BytesReference.toBytes(recvChunk.chunk)); } + + assertEquals("sent and received payloads are not the same", sendData, recvData); + handler.sendResponse(new RestResponse(RestStatus.OK, "")); } - assertEquals(sendData, recvData); + assertBusy(() -> assertEquals("should receive all server responses", totalRequests, ctx.clientRespQueue.size())); + } + } + + // ensures that all queued chunks are released when connection closed + public void testClientConnectionCloseMidStream() throws Exception { + try (var ctx = setupClientCtx()) { + // write half of http request + ctx.clientChannel.writeAndFlush(httpRequest(16 * 1024)); + ctx.clientChannel.writeAndFlush(randomContent(8 * 1024, false)); + + // await stream handler is ready and request full content + var handler = ctx.awaitRestChannelAccepted(); + handler.stream.requestBytes(Integer.MAX_VALUE); + assertBusy(() -> assertEquals(1, handler.stream.chunkQueue().size())); + + // terminate connection and wait resources are released + ctx.clientChannel.close(); + assertBusy(() -> assertEquals(0, handler.stream.chunkQueue().size())); + } + } + + // ensures that all queued chunks are released when server decides to close connection + public void testServerCloseConnectionMidStream() throws Exception { + try (var ctx = setupClientCtx()) { + // write half of http request + ctx.clientChannel.writeAndFlush(httpRequest(16 * 1024)); + ctx.clientChannel.writeAndFlush(randomContent(8 * 1024, false)); + + // await stream handler is ready and request full content + var handler = ctx.awaitRestChannelAccepted(); + handler.stream.requestBytes(Integer.MAX_VALUE); + assertBusy(() -> assertEquals(1, handler.stream.chunkQueue().size())); + + // terminate connection on server and wait resources are released + handler.httpChannel.close(); + assertBusy(() -> assertEquals(0, handler.stream.chunkQueue().size())); } } @@ -126,7 +180,7 @@ public void testClientBackpressure() throws Exception { final var totalWriteSize = writeChunkSize * totalWriteChunks; var writtenChunks = 0; - ctx.clientChannel.writeAndFlush(notFullRequest(totalWriteSize)); + ctx.clientChannel.writeAndFlush(httpRequest(totalWriteSize)); var handler = ctx.awaitRestChannelAccepted(); // write chunks until channel is no longer writable @@ -185,15 +239,19 @@ static T safePoll(BlockingDeque queue) { } } - static FullHttpRequest fullRequest(ByteBuf content) { + static FullHttpRequest fullHttpRequest(ByteBuf content) { var req = new DefaultFullHttpRequest(HTTP_1_1, POST, SingleRequestHandlerPlugin.ROUTE, Unpooled.wrappedBuffer(content)); req.headers().add(CONTENT_LENGTH, content.readableBytes()); req.headers().add(CONTENT_TYPE, APPLICATION_JSON); return req; } - static HttpRequest notFullRequest(int contentLength) { - var req = new DefaultHttpRequest(HTTP_1_1, POST, SingleRequestHandlerPlugin.ROUTE); + static HttpRequest httpRequest(int contentLength) { + return httpRequest(SingleRequestHandlerPlugin.ROUTE, contentLength); + } + + static HttpRequest httpRequest(String uri, int contentLength) { + var req = new DefaultHttpRequest(HTTP_1_1, POST, uri); req.headers().add(CONTENT_LENGTH, contentLength); req.headers().add(CONTENT_TYPE, APPLICATION_JSON); return req; @@ -283,7 +341,7 @@ static class ServerRequestHandler implements BaseRestHandler.RequestBodyChunkCon SubscribableListener channelAccepted = new SubscribableListener<>(); RestChannel channel; HttpChannel httpChannel; - HttpBody.Stream stream; + Netty4HttpRequestBodyStream stream; BlockingDeque recvChunks = new LinkedBlockingDeque<>(); record Chunk(ReleasableBytesReference chunk, boolean isLast) {} @@ -297,9 +355,15 @@ public void handleChunk(RestChannel channel, ReleasableBytesReference chunk, boo public void accept(RestChannel channel) throws Exception { this.channel = channel; this.httpChannel = channel.request().getHttpChannel(); - this.stream = channel.request().contentStream(); + this.stream = (Netty4HttpRequestBodyStream) channel.request().contentStream(); channelAccepted.onResponse(null); } + + void sendResponse(RestResponse response) { + assertEquals(SingleRequestHandlerPlugin.handlers.get(0), this); + SingleRequestHandlerPlugin.handlers.remove(0); + channel.sendResponse(response); + } } public static class SingleRequestHandlerPlugin extends Plugin implements ActionPlugin { diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java index 6572ba5b6843f..4136cc4ce4cab 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java @@ -9,18 +9,15 @@ 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; @@ -119,9 +116,9 @@ public Netty4HttpPipeliningHandler( @Override public void channelRead(final ChannelHandlerContext ctx, final Object msg) { - if (msg instanceof HttpRequest request) { - try { - activityTracker.startActivity(); + activityTracker.startActivity(); + try { + if (msg instanceof HttpRequest request) { final Netty4HttpRequest netty4HttpRequest; if (request.decoderResult().isFailure()) { final Throwable cause = request.decoderResult().cause(); @@ -135,36 +132,24 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { netty4HttpRequest = new Netty4HttpRequest(readSequence++, (FullHttpRequest) request, nonError); } else { if (request instanceof FullHttpRequest fullHttpRequest) { - currentRequestStream = null; netty4HttpRequest = new Netty4HttpRequest(readSequence++, fullHttpRequest); } else { var contentStream = new Netty4HttpRequestBodyStream(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++, request, contentStream); } } handlePipelinedRequest(ctx, netty4HttpRequest); - } finally { - activityTracker.stopActivity(); - } - } else { - assert msg instanceof HttpContent : "expect HttpContent got " + msg; - assert currentRequestStream != null : "current stream must exists before handling http content"; - currentRequestStream.handleNettyContent((HttpContent) msg); - if (msg instanceof LastHttpContent) { - currentRequestStream = null; + } else { + assert msg instanceof HttpContent : "expect HttpContent got " + msg; + assert currentRequestStream != null : "current stream must exists before handling http content"; + currentRequestStream.handleNettyContent((HttpContent) msg); + if (msg instanceof LastHttpContent) { + currentRequestStream = null; + } } + } finally { + activityTracker.stopActivity(); } } diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java index 43ff436944ef4..5fc2eb5efd54f 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java @@ -11,6 +11,7 @@ import io.netty.buffer.ByteBuf; import io.netty.buffer.Unpooled; import io.netty.handler.codec.http.DefaultFullHttpRequest; +import io.netty.handler.codec.http.EmptyHttpHeaders; import io.netty.handler.codec.http.FullHttpRequest; import io.netty.handler.codec.http.HttpHeaderNames; import io.netty.handler.codec.http.HttpHeaders; @@ -47,8 +48,22 @@ public class Netty4HttpRequest implements HttpRequest { private final boolean pooled; private final int sequence; - Netty4HttpRequest(int sequence, FullHttpRequest request, Netty4HttpRequestBodyStream contentStream) { - this(sequence, request, new AtomicBoolean(false), false, contentStream, null); + Netty4HttpRequest(int sequence, io.netty.handler.codec.http.HttpRequest request, Netty4HttpRequestBodyStream contentStream) { + this( + sequence, + new DefaultFullHttpRequest( + request.protocolVersion(), + request.method(), + request.uri(), + Unpooled.EMPTY_BUFFER, + request.headers(), + EmptyHttpHeaders.INSTANCE + ), + new AtomicBoolean(false), + false, + contentStream, + null + ); } Netty4HttpRequest(int sequence, FullHttpRequest request) { diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java index 4b5465eb708f5..ac61646fe91f1 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java @@ -20,7 +20,6 @@ import java.util.Queue; import java.util.concurrent.BlockingQueue; import java.util.concurrent.LinkedBlockingQueue; -import java.util.concurrent.atomic.AtomicBoolean; /** * Netty based implementation of {@link HttpBody.Stream}. @@ -33,11 +32,11 @@ public class Netty4HttpRequestBodyStream implements HttpBody.Stream { private final Channel channel; private final BlockingQueue requestQueue = new LinkedBlockingQueue<>(); private final ChunkQueue chunkQueue = new ChunkQueue(); - private final AtomicBoolean isDone = new AtomicBoolean(false); private HttpBody.ChunkHandler handler; public Netty4HttpRequestBodyStream(Channel channel) { this.channel = channel; + channel.closeFuture().addListener((f) -> releaseQueuedChunks()); channel.config().setAutoRead(false); } @@ -97,7 +96,7 @@ public void handleNettyContent(HttpContent httpContent) { } // visible for tests - ChunkQueue contentChunkQueue() { + ChunkQueue chunkQueue() { return chunkQueue; } @@ -133,12 +132,19 @@ private void flushQueued() { } private void flushAndRead() { - flushQueued(); // flush on channel's thread only + assert channel.eventLoop().inEventLoop(); + flushQueued(); if (requestQueue.isEmpty() == false) { channel.read(); } } + private void releaseQueuedChunks() { + while (chunkQueue.isEmpty() == false) { + chunkQueue.poll().buf.release(); + } + } + record ContentChunk(ByteBuf buf, boolean isLast) {} // Wrapped queue that keeps track of total available bytes of all chunks in the queue diff --git a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java index 3ea541c519771..a6623e2cb2ba3 100644 --- a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java +++ b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java @@ -52,8 +52,8 @@ public void testEnqueueChunksBeforeRequest() { for (int i = 0; i < totalChunks; i++) { channel.writeInbound(randomContent(chunkSize)); } - assertEquals(totalChunks, stream.contentChunkQueue().size()); - assertEquals(totalChunks * chunkSize, stream.contentChunkQueue().availableBytes()); + assertEquals(totalChunks, stream.chunkQueue().size()); + assertEquals(totalChunks * chunkSize, stream.chunkQueue().availableBytes()); } // test that every small request consume at least one chunk @@ -114,7 +114,7 @@ public void testEnqueueWhenNotEnoughBytes() { for (int i = 0; i < totalChunks; i++) { channel.writeInbound(randomContent(chunkSize)); } - assertEquals(chunkSize * totalChunks, stream.contentChunkQueue().availableBytes()); + assertEquals(chunkSize * totalChunks, stream.chunkQueue().availableBytes()); } public void testFlushSingleLastContent() { @@ -140,16 +140,16 @@ public void testReadFromChannel() { channel.writeInbound(randomContent(chunkSize)); } - assertEquals("should not read until requested", 0, stream.contentChunkQueue().size()); + assertEquals("should not read until requested", 0, stream.chunkQueue().size()); stream.requestBytes(chunkSize); assertEquals(1, gotChunks.size()); stream.requestBytes(Integer.MAX_VALUE); - assertEquals("should enqueue remaining chunks", totalChunks - 1, stream.contentChunkQueue().size()); + assertEquals("should enqueue remaining chunks", totalChunks - 1, stream.chunkQueue().size()); // send last content and flush queued content channel.writeInbound(randomLastContent(0)); - assertEquals(0, stream.contentChunkQueue().size()); + assertEquals(0, stream.chunkQueue().size()); assertEquals(2, gotChunks.size()); assertTrue(gotLast.get()); } From 7b66986453ae5789d81bf60625e53b737f7fb343 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Wed, 7 Aug 2024 18:40:07 -0700 Subject: [PATCH 12/14] change requestBytes(bytes) to next() --- .../Netty4IncrementalRequestHandlingIT.java | 118 ++++++------ .../http/netty4/Netty4HttpRequest.java | 8 +- .../netty4/Netty4HttpRequestBodyStream.java | 177 ++++-------------- .../Netty4HttpRequestBodyStreamTests.java | 74 +------- .../java/org/elasticsearch/http/HttpBody.java | 34 ++-- 5 files changed, 126 insertions(+), 285 deletions(-) diff --git a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java index 545288df02847..03b126727758b 100644 --- a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java +++ b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java @@ -28,13 +28,13 @@ import io.netty.handler.codec.http.HttpContent; import io.netty.handler.codec.http.HttpObjectAggregator; import io.netty.handler.codec.http.HttpRequest; +import io.netty.handler.codec.http.LastHttpContent; import org.elasticsearch.ESNetty4IntegTestCase; import org.elasticsearch.action.support.SubscribableListener; import org.elasticsearch.client.internal.node.NodeClient; import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver; import org.elasticsearch.cluster.node.DiscoveryNodes; -import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.bytes.ReleasableBytesReference; import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.common.settings.ClusterSettings; @@ -55,6 +55,7 @@ import org.elasticsearch.rest.RestResponse; import org.elasticsearch.rest.RestStatus; import org.elasticsearch.test.ESIntegTestCase; +import org.elasticsearch.transport.netty4.Netty4Utils; import java.util.ArrayList; import java.util.Collection; @@ -84,7 +85,7 @@ public void testEmptyContent() throws Exception { // send request with empty content ctx.clientChannel.writeAndFlush(fullHttpRequest(Unpooled.EMPTY_BUFFER)); var handler = ctx.awaitRestChannelAccepted(); - handler.stream.requestBytes(1); + handler.stream.next(); // should receive a single empty chunk var recvChunk = safePoll(handler.recvChunks); @@ -113,17 +114,16 @@ public void testReceiveAllChunks() throws Exception { var handler = ctx.awaitRestChannelAccepted(); - // consume chunks at random pace var gotAllChunks = false; var recvData = Unpooled.buffer(dataSize); while (gotAllChunks == false) { - handler.stream.requestBytes(randomIntBetween(1024, 64 * 1024)); + handler.stream.next(); var recvChunk = safePoll(handler.recvChunks); try (recvChunk.chunk) { if (recvChunk.isLast) { gotAllChunks = true; } - recvData.writeBytes(BytesReference.toBytes(recvChunk.chunk)); + recvData.writeBytes(Netty4Utils.toByteBuf(recvChunk.chunk)); } } @@ -138,14 +138,16 @@ public void testReceiveAllChunks() throws Exception { public void testClientConnectionCloseMidStream() throws Exception { try (var ctx = setupClientCtx()) { // write half of http request - ctx.clientChannel.writeAndFlush(httpRequest(16 * 1024)); - ctx.clientChannel.writeAndFlush(randomContent(8 * 1024, false)); + ctx.clientChannel.write(httpRequest(2 * 1024)); + ctx.clientChannel.writeAndFlush(randomContent(1024, false)); // await stream handler is ready and request full content var handler = ctx.awaitRestChannelAccepted(); - handler.stream.requestBytes(Integer.MAX_VALUE); assertBusy(() -> assertEquals(1, handler.stream.chunkQueue().size())); + // enable auto-read to receive channel close event + handler.stream.channel().config().setAutoRead(true); + // terminate connection and wait resources are released ctx.clientChannel.close(); assertBusy(() -> assertEquals(0, handler.stream.chunkQueue().size())); @@ -156,12 +158,11 @@ public void testClientConnectionCloseMidStream() throws Exception { public void testServerCloseConnectionMidStream() throws Exception { try (var ctx = setupClientCtx()) { // write half of http request - ctx.clientChannel.writeAndFlush(httpRequest(16 * 1024)); - ctx.clientChannel.writeAndFlush(randomContent(8 * 1024, false)); + ctx.clientChannel.write(httpRequest(2 * 1024)); + ctx.clientChannel.writeAndFlush(randomContent(1024, false)); // await stream handler is ready and request full content var handler = ctx.awaitRestChannelAccepted(); - handler.stream.requestBytes(Integer.MAX_VALUE); assertBusy(() -> assertEquals(1, handler.stream.chunkQueue().size())); // terminate connection on server and wait resources are released @@ -170,62 +171,40 @@ public void testServerCloseConnectionMidStream() throws Exception { } } - // ensure that client's socket is no longer writable if server not consuming data + static int MBytes(int m) { + return m * 1024 * 1024; + } + + // ensure that client's socket buffers data when server is not consuming data public void testClientBackpressure() throws Exception { try (var ctx = setupClientCtx()) { - - // a meaningful write payload that is larger than socket and channel buffers - final var writeChunkSize = 16 * 1024; - final var totalWriteChunks = randomIntBetween(1024, 2048); // 16-32MB - final var totalWriteSize = writeChunkSize * totalWriteChunks; - var writtenChunks = 0; - - ctx.clientChannel.writeAndFlush(httpRequest(totalWriteSize)); - var handler = ctx.awaitRestChannelAccepted(); - - // write chunks until channel is no longer writable - while (writtenChunks <= totalWriteChunks && ctx.clientChannel.isWritable()) { - ctx.clientChannel.writeAndFlush(randomContent(writeChunkSize, false)); - writtenChunks += 1; - } - assertTrue("should not be able to write all chunks", writtenChunks < totalWriteChunks); - - // read all written chunks and check client socket is writable again - var totalReadSize = 0; - var mustReadBytes = writtenChunks * writeChunkSize; - while (mustReadBytes > 0) { - handler.stream.requestBytes(writeChunkSize); - var recvChunk = safePoll(handler.recvChunks); - mustReadBytes -= recvChunk.chunk.length(); - totalReadSize += recvChunk.chunk.length(); - recvChunk.chunk.close(); + var payloadSize = MBytes(50); + ctx.clientChannel.writeAndFlush(httpRequest(payloadSize)); + for (int i = 0; i < 5; i++) { + ctx.clientChannel.writeAndFlush(randomContent(MBytes(10), false)); } + ctx.clientChannel.writeAndFlush(LastHttpContent.EMPTY_LAST_CONTENT); - assertBusy(() -> assertTrue("channel must be writable again", ctx.clientChannel.isWritable())); - - // write remaining chunks - while (writtenChunks < totalWriteChunks) { // leave 1 for LastHttpContent - writtenChunks += 1; - ctx.clientChannel.writeAndFlush(randomContent(writeChunkSize, false)); - } - ctx.clientChannel.writeAndFlush(randomContent(writeChunkSize, true)); + var handler = ctx.awaitRestChannelAccepted(); - // consume remaining chunks - var allConsumed = false; - while (allConsumed == false) { - handler.stream.requestBytes(writeChunkSize); - var recvChunk = safePoll(handler.recvChunks); - totalReadSize += recvChunk.chunk.length(); - allConsumed = recvChunk.isLast; - recvChunk.chunk.close(); + // Read buffers for socket and channel usually within few MBytes range all together. + // This test assumes that buffers will not exceed 10 MBytes, in other words there should + // be less than 10 MBytes in fly between http client's socket and rest handler. This + // loop ensures that reading 10 MBytes of content on server side should free almost + // same size in client's channel write buffer. + for (int mb = 0; mb <= 50; mb += 10) { + var minBufSize = payloadSize - MBytes(10 + mb); + var maxBufSize = payloadSize - MBytes(mb); + assertBusy(() -> { + var bufSize = ctx.clientChannel.bytesBeforeWritable(); + assertTrue( + "client's channel buffer should be in range [" + minBufSize + "," + maxBufSize + "], got " + bufSize, + bufSize >= minBufSize && bufSize <= maxBufSize + ); + }); + handler.consumeBytes(MBytes(10)); } - assertEquals("did not received all bytes", totalWriteSize, totalReadSize); - - // send OK response back to wrap it up - handler.channel.sendResponse(new RestResponse(RestStatus.OK, "")); - var resp = (FullHttpResponse) safePoll(ctx.clientRespQueue); - assertEquals(200, resp.status().code()); - resp.release(); + assertTrue(handler.stream.hasLast()); } } @@ -342,6 +321,7 @@ static class ServerRequestHandler implements BaseRestHandler.RequestBodyChunkCon RestChannel channel; HttpChannel httpChannel; Netty4HttpRequestBodyStream stream; + boolean recvLast = false; BlockingDeque recvChunks = new LinkedBlockingDeque<>(); record Chunk(ReleasableBytesReference chunk, boolean isLast) {} @@ -364,6 +344,22 @@ void sendResponse(RestResponse response) { SingleRequestHandlerPlugin.handlers.remove(0); channel.sendResponse(response); } + + void consumeBytes(int bytes) { + if (recvLast) { + return; + } + while (bytes > 0) { + stream.next(); + var recvChunk = safePoll(recvChunks); + bytes -= recvChunk.chunk.length(); + recvChunk.chunk.close(); + if (recvChunk.isLast) { + recvLast = true; + break; + } + } + } } public static class SingleRequestHandlerPlugin extends Plugin implements ActionPlugin { diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java index 5fc2eb5efd54f..2d1caba3c477e 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequest.java @@ -187,13 +187,7 @@ public HttpRequest removeHeader(String header) { copiedHeadersWithout, copiedTrailingHeadersWithout ); - HttpBody copiedContent; - if (content.isFull()) { - copiedContent = HttpBody.fromBytesReference(content.asFull().bytes()); - } else { - copiedContent = content; - } - return new Netty4HttpRequest(sequence, requestWithoutHeader, released, pooled, copiedContent); + return new Netty4HttpRequest(sequence, requestWithoutHeader, released, pooled, content); } @Override diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java index ac61646fe91f1..8497e3ee8a40d 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStream.java @@ -8,7 +8,6 @@ package org.elasticsearch.http.netty4; -import io.netty.buffer.ByteBuf; import io.netty.channel.Channel; import io.netty.handler.codec.http.HttpContent; import io.netty.handler.codec.http.LastHttpContent; @@ -18,20 +17,19 @@ import java.util.ArrayDeque; import java.util.Queue; -import java.util.concurrent.BlockingQueue; -import java.util.concurrent.LinkedBlockingQueue; /** * Netty based implementation of {@link HttpBody.Stream}. - * This implementation utilize {@link io.netty.channel.ChannelConfig#setAutoRead(boolean)} to prevent entire payload buffering. - * But sometimes upstream can send few chunks of data despite autoRead=off. - * In this case chunks will be queued until downstream {@link HttpBody.Stream#requestBytes} + * This implementation utilize {@link io.netty.channel.ChannelConfig#setAutoRead(boolean)} + * to prevent entire payload buffering. But sometimes upstream can send few chunks of data despite + * autoRead=off. In this case chunks will be queued until downstream calls {@link Stream#next()} */ public class Netty4HttpRequestBodyStream implements HttpBody.Stream { private final Channel channel; - private final BlockingQueue requestQueue = new LinkedBlockingQueue<>(); - private final ChunkQueue chunkQueue = new ChunkQueue(); + private final Queue chunkQueue = new ArrayDeque<>(); + private boolean requested = false; + private boolean hasLast = false; private HttpBody.ChunkHandler handler; public Netty4HttpRequestBodyStream(Channel channel) { @@ -50,160 +48,67 @@ public void setHandler(ChunkHandler chunkHandler) { this.handler = chunkHandler; } + private void sendQueuedOrRead() { + assert channel.eventLoop().inEventLoop(); + requested = true; + var chunk = chunkQueue.poll(); + if (chunk == null) { + channel.read(); + } else { + sendChunk(chunk); + } + } + @Override - public void requestBytes(int bytes) { + public void next() { assert handler != null : "handler must be set before requesting next chunk"; - assert bytes != 0; - requestQueue.add(bytes); if (channel.eventLoop().inEventLoop()) { - flushAndRead(); + sendQueuedOrRead(); } else { - channel.eventLoop().submit(this::flushAndRead); + channel.eventLoop().submit(this::sendQueuedOrRead); } } public void handleNettyContent(HttpContent httpContent) { assert handler != null : "handler must be set before processing http content"; - - var content = httpContent.content(); - var isLast = httpContent instanceof LastHttpContent; - var currentRequest = requestQueue.peek(); - var availableBytes = chunkQueue.availableBytes() + content.readableBytes(); - - if (currentRequest == null || (currentRequest > availableBytes && isLast == false)) { - // enqueue incoming chunk: - // 1. There is no ongoing request, it might happen that upstream send a chunk before request. - // 2. There are not enough accumulated bytes, unless it's a last chunk then we must flush - chunkQueue.offer(new ContentChunk(content, isLast)); + if (requested && chunkQueue.isEmpty()) { + sendChunk(httpContent); } else { - // flush queued chunks: - // 1. nothing in the queue and current chunk fulfills request, send it right a way - // 2. there are queued chunks, enqueue and flush - if (chunkQueue.isEmpty()) { - this.handler.onNext(Netty4Utils.toReleasableBytesReference(content), isLast); - requestQueue.poll(); - } else { - chunkQueue.offer(new ContentChunk(content, isLast)); - flushQueued(); - } + chunkQueue.add(httpContent); } - - if (isLast) { + if (httpContent instanceof LastHttpContent) { + hasLast = true; channel.config().setAutoRead(true); - } else if (requestQueue.peek() != null) { - channel.read(); } } - // visible for tests - ChunkQueue chunkQueue() { - return chunkQueue; + // visible for test + Channel channel() { + return channel; } - private boolean canFulfillCurrentRequest() { - return requestQueue.isEmpty() == false - && chunkQueue.isEmpty() == false - && (chunkQueue.availableBytes() >= requestQueue.peek() || chunkQueue.hasLast()); + // visible for test + Queue chunkQueue() { + return chunkQueue; } - private void flushQueued() { - while (canFulfillCurrentRequest()) { - var currentRequest = requestQueue.poll(); - assert currentRequest != null; - ByteBuf buf; - boolean isLast = false; - // skip aggregation if next chunk is enough - if (chunkQueue.nextChunkSize() >= currentRequest || chunkQueue.nextIsLast()) { - var chunk = chunkQueue.poll(); - isLast = chunk.isLast; - buf = chunk.buf; - } else { - // aggregate multiple chunks - var compBuf = channel.alloc().compositeBuffer(); - while (compBuf.readableBytes() < currentRequest && chunkQueue.isEmpty() == false) { - var chunk = chunkQueue.poll(); - isLast = chunk.isLast; - compBuf.addComponent(true, chunk.buf()); - } - buf = compBuf; - } - handler.onNext(Netty4Utils.toReleasableBytesReference(buf), isLast); - } + // visible for test + boolean hasLast() { + return hasLast; } - private void flushAndRead() { - assert channel.eventLoop().inEventLoop(); - flushQueued(); - if (requestQueue.isEmpty() == false) { - channel.read(); - } + private void sendChunk(HttpContent httpContent) { + assert requested; + requested = false; + var bytesRef = Netty4Utils.toReleasableBytesReference(httpContent.content()); + var isLast = httpContent instanceof LastHttpContent; + handler.onNext(bytesRef, isLast); } private void releaseQueuedChunks() { while (chunkQueue.isEmpty() == false) { - chunkQueue.poll().buf.release(); + chunkQueue.poll().release(); } } - record ContentChunk(ByteBuf buf, boolean isLast) {} - - // Wrapped queue that keeps track of total available bytes of all chunks in the queue - static class ChunkQueue { - private final Queue queue = new ArrayDeque<>(); - private int availableBytes = 0; - private boolean hasLast = false; - - void release() { - while (queue.isEmpty() == false) { - queue.poll().buf.release(); - } - } - - boolean isEmpty() { - return queue.isEmpty(); - } - - int availableBytes() { - return availableBytes; - } - - boolean hasLast() { - return hasLast; - } - - int size() { - return queue.size(); - } - - void offer(ContentChunk content) { - queue.offer(content); - availableBytes += content.buf.readableBytes(); - hasLast = content.isLast; - } - - int nextChunkSize() { - var chunk = queue.peek(); - if (chunk == null) { - return 0; - } else { - return chunk.buf.readableBytes(); - } - } - - boolean nextIsLast() { - var chunk = queue.peek(); - if (chunk == null) { - return false; - } else { - return chunk.isLast; - } - } - - ContentChunk poll() { - var chunk = queue.poll(); - assert chunk != null; - availableBytes -= chunk.buf().readableBytes(); - return chunk; - } - } } diff --git a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java index a6623e2cb2ba3..65e56dd82f8fc 100644 --- a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java +++ b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java @@ -47,13 +47,11 @@ protected void channelRead0(ChannelHandlerContext ctx, HttpContent msg) { } public void testEnqueueChunksBeforeRequest() { - var chunkSize = 1024; var totalChunks = randomIntBetween(1, 100); for (int i = 0; i < totalChunks; i++) { - channel.writeInbound(randomContent(chunkSize)); + channel.writeInbound(randomContent(1024)); } assertEquals(totalChunks, stream.chunkQueue().size()); - assertEquals(totalChunks * chunkSize, stream.chunkQueue().availableBytes()); } // test that every small request consume at least one chunk @@ -72,59 +70,12 @@ public void testFlushQueuedSmallRequests() { } // consume all chunks for (var i = 0; i < totalChunks; i++) { - stream.requestBytes(1); + stream.next(); } assertEquals(totalChunks, chunks.size()); assertEquals(chunkSize * totalChunks, totalBytes.get()); } - // test that one large request consume all chunks - public void testFlushQueuedLargeRequest() { - var requestedChunkSize = new AtomicInteger(); - stream.setHandler((chunk, isLast) -> { requestedChunkSize.addAndGet(chunk.length()); }); - var chunkSize = 1024; - var totalChunks = randomIntBetween(1, 100); - for (int i = 0; i < totalChunks; i++) { - channel.writeInbound(randomContent(chunkSize)); - } - stream.requestBytes(chunkSize * totalChunks); - assertEquals(chunkSize * totalChunks, requestedChunkSize.get()); - } - - // test that we flush queued chunks with last content when request is larger than availableBytes - public void testFlushQueuedLastContent() { - var gotLast = new AtomicBoolean(false); - var contentSize = new AtomicInteger(0); - stream.setHandler((chunk, isLast) -> { - contentSize.addAndGet(chunk.length()); - gotLast.set(isLast); - }); - channel.writeInbound(randomContent(1024)); - channel.writeInbound(randomLastContent(0)); - stream.requestBytes(1025); - assertTrue(gotLast.get()); - assertEquals(1024, contentSize.get()); - } - - // ensures that we enqueue chunk if cannot fulfill current request - public void testEnqueueWhenNotEnoughBytes() { - stream.requestBytes(1024); - var chunkSize = 32; - var totalChunks = randomIntBetween(1, 1024 / chunkSize - 1); - for (int i = 0; i < totalChunks; i++) { - channel.writeInbound(randomContent(chunkSize)); - } - assertEquals(chunkSize * totalChunks, stream.chunkQueue().availableBytes()); - } - - public void testFlushSingleLastContent() { - var flushed = new AtomicBoolean(false); - stream.setHandler((chunk, isLast) -> flushed.set(true)); - stream.requestBytes(1024); - channel.writeInbound(randomLastContent(0)); - assertTrue(flushed.get()); - } - // test that we read from channel when not enough available bytes public void testReadFromChannel() { var gotChunks = new ArrayList(); @@ -136,22 +87,17 @@ public void testReadFromChannel() { channel.pipeline().addFirst(new FlowControlHandler()); // block all incoming messages, need explicit channel.read() var chunkSize = 1024; var totalChunks = randomIntBetween(1, 32); - for (int i = 0; i < totalChunks; i++) { + for (int i = 0; i < totalChunks - 1; i++) { channel.writeInbound(randomContent(chunkSize)); } + channel.writeInbound(randomLastContent(chunkSize)); - assertEquals("should not read until requested", 0, stream.chunkQueue().size()); - stream.requestBytes(chunkSize); - assertEquals(1, gotChunks.size()); - - stream.requestBytes(Integer.MAX_VALUE); - assertEquals("should enqueue remaining chunks", totalChunks - 1, stream.chunkQueue().size()); - - // send last content and flush queued content - channel.writeInbound(randomLastContent(0)); - assertEquals(0, stream.chunkQueue().size()); - assertEquals(2, gotChunks.size()); - assertTrue(gotLast.get()); + for (int i = 0; i < totalChunks; i++) { + assertEquals("should not enqueue chunks", 0, stream.chunkQueue().size()); + stream.next(); + assertEquals("each next() should produce single chunk", i + 1, gotChunks.size()); + } + assertTrue("should receive last content", gotLast.get()); } HttpContent randomContent(int size, boolean isLast) { diff --git a/server/src/main/java/org/elasticsearch/http/HttpBody.java b/server/src/main/java/org/elasticsearch/http/HttpBody.java index 6c575a8cb6040..b4d88b837b117 100644 --- a/server/src/main/java/org/elasticsearch/http/HttpBody.java +++ b/server/src/main/java/org/elasticsearch/http/HttpBody.java @@ -74,24 +74,24 @@ non-sealed interface Stream extends HttpBody { void setHandler(ChunkHandler chunkHandler); /** - * Request next chunk of bytes. For every request there will be at least one chunk. Size of - * the next chunk might vary, and it depends on following factors: - *
      - *
    • - * Size might round up to optimal network chunk size. For example, default HttpDecoder - * content size is 8kb. Request for 1 byte will produce 8kb chunk, 5 requests for 1 byte - * will produce 5 chunks of 8kb each. Request for 10kb will produce 16kb chunk. - *
      • - *
    • - * Last chunk always less or equal requested size. Request for Integer.MAX_VALUE or - * other number larger than actual payload, will always produce single full content chunk. - *
      • - *
    • - * After last chunk there will be no more chunks. This method will not do anything. - *
      • - *
    + * Request next chunk of data from the network. The size of the chunk depends on following + * factors. If request is not compressed then chunk size will be up to + * {@link HttpTransportSettings#SETTING_HTTP_MAX_CHUNK_SIZE}. If request is compressed then + * chunk size will be up to max_chunk_size * compression_ratio. Multiple calls can be + * deduplicated when next chunk is not yet available. It's recommended to call "next" once + * for every chunk. + * 
    +         * {@code
+         *     stream.setHandler((chunk, isLast) -> {
+         *         processChunk(chunk);
+         *         if (isLast == false) {
+         *             stream.next();
+         *         }
+         *     });
+         * }
+         *
    */ - void requestBytes(int bytes); + void next(); } @FunctionalInterface From f140b7d06c6b8de9c6d7e28f78ab3a4907405ba8 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Fri, 9 Aug 2024 16:26:14 -0700 Subject: [PATCH 13/14] address comments and test changes --- .../Netty4IncrementalRequestHandlingIT.java | 159 ++++++++++-------- .../Netty4HttpRequestBodyStreamTests.java | 8 +- .../elasticsearch/rest/BaseRestHandler.java | 2 +- .../elasticsearch/rest/RestController.java | 6 +- .../org/elasticsearch/rest/RestRequest.java | 20 +-- 5 files changed, 110 insertions(+), 85 deletions(-) diff --git a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java index 03b126727758b..dcfe411f3a07e 100644 --- a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java +++ b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java @@ -43,7 +43,6 @@ import org.elasticsearch.common.settings.SettingsFilter; import org.elasticsearch.common.util.CollectionUtils; import org.elasticsearch.features.NodeFeature; -import org.elasticsearch.http.HttpChannel; import org.elasticsearch.http.HttpServerTransport; import org.elasticsearch.plugins.ActionPlugin; import org.elasticsearch.plugins.Plugin; @@ -54,15 +53,17 @@ import org.elasticsearch.rest.RestRequest; import org.elasticsearch.rest.RestResponse; import org.elasticsearch.rest.RestStatus; +import org.elasticsearch.tasks.Task; import org.elasticsearch.test.ESIntegTestCase; import org.elasticsearch.transport.netty4.Netty4Utils; -import java.util.ArrayList; import java.util.Collection; -import java.util.Collections; import java.util.List; import java.util.concurrent.BlockingDeque; import java.util.concurrent.BlockingQueue; +import java.util.concurrent.Callable; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingDeque; import java.util.concurrent.TimeUnit; import java.util.function.Predicate; @@ -82,9 +83,11 @@ public void testEmptyContent() throws Exception { try (var ctx = setupClientCtx()) { var totalRequests = randomIntBetween(1, 10); for (int reqNo = 0; reqNo < totalRequests; reqNo++) { + var opaqueId = opaqueId(reqNo); + // send request with empty content - ctx.clientChannel.writeAndFlush(fullHttpRequest(Unpooled.EMPTY_BUFFER)); - var handler = ctx.awaitRestChannelAccepted(); + ctx.clientChannel.writeAndFlush(fullHttpRequest(opaqueId, Unpooled.EMPTY_BUFFER)); + var handler = ctx.awaitRestChannelAccepted(opaqueId); handler.stream.next(); // should receive a single empty chunk @@ -105,25 +108,25 @@ public void testReceiveAllChunks() throws Exception { try (var ctx = setupClientCtx()) { var totalRequests = randomIntBetween(1, 10); for (int reqNo = 0; reqNo < totalRequests; reqNo++) { + var opaqueId = opaqueId(reqNo); // this dataset will be compared with one on server side var dataSize = randomIntBetween(1024, 10 * 1024 * 1024); var sendData = Unpooled.wrappedBuffer(randomByteArrayOfLength(dataSize)); sendData.retain(); - ctx.clientChannel.writeAndFlush(fullHttpRequest(sendData)); + ctx.clientChannel.writeAndFlush(fullHttpRequest(opaqueId, sendData)); - var handler = ctx.awaitRestChannelAccepted(); + var handler = ctx.awaitRestChannelAccepted(opaqueId); - var gotAllChunks = false; var recvData = Unpooled.buffer(dataSize); - while (gotAllChunks == false) { + while (true) { handler.stream.next(); var recvChunk = safePoll(handler.recvChunks); try (recvChunk.chunk) { + recvData.writeBytes(Netty4Utils.toByteBuf(recvChunk.chunk)); if (recvChunk.isLast) { - gotAllChunks = true; + break; } - recvData.writeBytes(Netty4Utils.toByteBuf(recvChunk.chunk)); } } @@ -137,12 +140,14 @@ public void testReceiveAllChunks() throws Exception { // ensures that all queued chunks are released when connection closed public void testClientConnectionCloseMidStream() throws Exception { try (var ctx = setupClientCtx()) { + var opaqueId = opaqueId(0); + // write half of http request - ctx.clientChannel.write(httpRequest(2 * 1024)); + ctx.clientChannel.write(httpRequest(opaqueId, 2 * 1024)); ctx.clientChannel.writeAndFlush(randomContent(1024, false)); // await stream handler is ready and request full content - var handler = ctx.awaitRestChannelAccepted(); + var handler = ctx.awaitRestChannelAccepted(opaqueId); assertBusy(() -> assertEquals(1, handler.stream.chunkQueue().size())); // enable auto-read to receive channel close event @@ -157,35 +162,34 @@ public void testClientConnectionCloseMidStream() throws Exception { // ensures that all queued chunks are released when server decides to close connection public void testServerCloseConnectionMidStream() throws Exception { try (var ctx = setupClientCtx()) { + var opaqueId = opaqueId(0); + // write half of http request - ctx.clientChannel.write(httpRequest(2 * 1024)); + ctx.clientChannel.write(httpRequest(opaqueId, 2 * 1024)); ctx.clientChannel.writeAndFlush(randomContent(1024, false)); // await stream handler is ready and request full content - var handler = ctx.awaitRestChannelAccepted(); + var handler = ctx.awaitRestChannelAccepted(opaqueId); assertBusy(() -> assertEquals(1, handler.stream.chunkQueue().size())); // terminate connection on server and wait resources are released - handler.httpChannel.close(); + handler.channel.request().getHttpChannel().close(); assertBusy(() -> assertEquals(0, handler.stream.chunkQueue().size())); } } - static int MBytes(int m) { - return m * 1024 * 1024; - } - // ensure that client's socket buffers data when server is not consuming data public void testClientBackpressure() throws Exception { try (var ctx = setupClientCtx()) { + var opaqueId = opaqueId(0); var payloadSize = MBytes(50); - ctx.clientChannel.writeAndFlush(httpRequest(payloadSize)); + ctx.clientChannel.writeAndFlush(httpRequest(opaqueId, payloadSize)); for (int i = 0; i < 5; i++) { ctx.clientChannel.writeAndFlush(randomContent(MBytes(10), false)); } ctx.clientChannel.writeAndFlush(LastHttpContent.EMPTY_LAST_CONTENT); - var handler = ctx.awaitRestChannelAccepted(); + var handler = ctx.awaitRestChannelAccepted(opaqueId); // Read buffers for socket and channel usually within few MBytes range all together. // This test assumes that buffers will not exceed 10 MBytes, in other words there should @@ -208,6 +212,14 @@ public void testClientBackpressure() throws Exception { } } + private String opaqueId(int reqNo) { + return getTestName() + "-" + reqNo; + } + + static int MBytes(int m) { + return m * 1024 * 1024; + } + static T safePoll(BlockingDeque queue) { try { var t = queue.poll(SAFE_AWAIT_TIMEOUT.seconds(), TimeUnit.SECONDS); @@ -218,21 +230,23 @@ static T safePoll(BlockingDeque queue) { } } - static FullHttpRequest fullHttpRequest(ByteBuf content) { - var req = new DefaultFullHttpRequest(HTTP_1_1, POST, SingleRequestHandlerPlugin.ROUTE, Unpooled.wrappedBuffer(content)); + static FullHttpRequest fullHttpRequest(String opaqueId, ByteBuf content) { + var req = new DefaultFullHttpRequest(HTTP_1_1, POST, ControlServerRequestPlugin.ROUTE, Unpooled.wrappedBuffer(content)); req.headers().add(CONTENT_LENGTH, content.readableBytes()); req.headers().add(CONTENT_TYPE, APPLICATION_JSON); + req.headers().add(Task.X_OPAQUE_ID_HTTP_HEADER, opaqueId); return req; } - static HttpRequest httpRequest(int contentLength) { - return httpRequest(SingleRequestHandlerPlugin.ROUTE, contentLength); + static HttpRequest httpRequest(String opaqueId, int contentLength) { + return httpRequest(ControlServerRequestPlugin.ROUTE, opaqueId, contentLength); } - static HttpRequest httpRequest(String uri, int contentLength) { + static HttpRequest httpRequest(String uri, String opaqueId, int contentLength) { var req = new DefaultHttpRequest(HTTP_1_1, POST, uri); req.headers().add(CONTENT_LENGTH, contentLength); req.headers().add(CONTENT_TYPE, APPLICATION_JSON); + req.headers().add(Task.X_OPAQUE_ID_HTTP_HEADER, opaqueId); return req; } @@ -245,37 +259,12 @@ static HttpContent randomContent(int size, boolean isLast) { } } - record Ctx(String nodeName, Bootstrap clientBootstrap, Channel clientChannel, BlockingDeque clientRespQueue) - implements - AutoCloseable { - - @Override - public void close() throws Exception { - safeGet(clientChannel.close()); - safeGet(clientBootstrap.config().group().shutdownGracefully()); - clientRespQueue.forEach(o -> { if (o instanceof FullHttpResponse resp) resp.release(); }); - for (var handler : SingleRequestHandlerPlugin.handlers) { - handler.recvChunks.forEach(c -> c.chunk.close()); - handler.httpChannel.close(); - } - SingleRequestHandlerPlugin.handlers.clear(); - } - - ServerRequestHandler awaitRestChannelAccepted() throws Exception { - var handlers = SingleRequestHandlerPlugin.handlers; - assertBusy(() -> { assertEquals(1, handlers.size()); }); - var handler = handlers.get(0); - safeAwait(handler.channelAccepted); - return handler; - } - } - Ctx setupClientCtx() throws Exception { var nodeName = internalCluster().getRandomNodeName(); var clientRespQueue = new LinkedBlockingDeque<>(16); var bootstrap = bootstrapClient(nodeName, clientRespQueue); var channel = bootstrap.connect().sync().channel(); - return new Ctx(nodeName, bootstrap, channel, clientRespQueue); + return new Ctx(getTestName(), nodeName, bootstrap, channel, clientRespQueue); } Bootstrap bootstrapClient(String node, BlockingQueue queue) { @@ -308,7 +297,7 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { @Override protected Collection> nodePlugins() { - return CollectionUtils.concatLists(List.of(SingleRequestHandlerPlugin.class), super.nodePlugins()); + return CollectionUtils.concatLists(List.of(ControlServerRequestPlugin.class), super.nodePlugins()); } @Override @@ -316,15 +305,46 @@ protected boolean addMockHttpTransport() { return false; // enable http } + record Ctx(String testName, String nodeName, Bootstrap clientBootstrap, Channel clientChannel, BlockingDeque clientRespQueue) + implements + AutoCloseable { + + @Override + public void close() throws Exception { + safeGet(clientChannel.close()); + safeGet(clientBootstrap.config().group().shutdownGracefully()); + clientRespQueue.forEach(o -> { if (o instanceof FullHttpResponse resp) resp.release(); }); + for (var opaqueId : ControlServerRequestPlugin.handlers.keySet()) { + if (opaqueId.startsWith(testName)) { + var handler = ControlServerRequestPlugin.handlers.get(opaqueId); + handler.recvChunks.forEach(c -> c.chunk.close()); + handler.channel.request().getHttpChannel().close(); + ControlServerRequestPlugin.handlers.remove(opaqueId); + } + } + } + + ServerRequestHandler awaitRestChannelAccepted(String opaqueId) throws Exception { + assertBusy(() -> assertTrue(ControlServerRequestPlugin.handlers.containsKey(opaqueId))); + var handler = ControlServerRequestPlugin.handlers.get(opaqueId); + safeAwait(handler.channelAccepted); + return handler; + } + } + static class ServerRequestHandler implements BaseRestHandler.RequestBodyChunkConsumer { - SubscribableListener channelAccepted = new SubscribableListener<>(); + final SubscribableListener channelAccepted = new SubscribableListener<>(); + final String opaqueId; + final BlockingDeque recvChunks = new LinkedBlockingDeque<>(); + final Netty4HttpRequestBodyStream stream; RestChannel channel; - HttpChannel httpChannel; - Netty4HttpRequestBodyStream stream; + boolean recvLast = false; - BlockingDeque recvChunks = new LinkedBlockingDeque<>(); - record Chunk(ReleasableBytesReference chunk, boolean isLast) {} + ServerRequestHandler(String opaqueId, Netty4HttpRequestBodyStream stream) { + this.opaqueId = opaqueId; + this.stream = stream; + } @Override public void handleChunk(RestChannel channel, ReleasableBytesReference chunk, boolean isLast) { @@ -334,14 +354,10 @@ public void handleChunk(RestChannel channel, ReleasableBytesReference chunk, boo @Override public void accept(RestChannel channel) throws Exception { this.channel = channel; - this.httpChannel = channel.request().getHttpChannel(); - this.stream = (Netty4HttpRequestBodyStream) channel.request().contentStream(); channelAccepted.onResponse(null); } void sendResponse(RestResponse response) { - assertEquals(SingleRequestHandlerPlugin.handlers.get(0), this); - SingleRequestHandlerPlugin.handlers.remove(0); channel.sendResponse(response); } @@ -360,13 +376,20 @@ void consumeBytes(int bytes) { } } } + + Future onChannelThread(Callable task) { + return this.stream.channel().eventLoop().submit(task); + } + + record Chunk(ReleasableBytesReference chunk, boolean isLast) {} } - public static class SingleRequestHandlerPlugin extends Plugin implements ActionPlugin { + // takes full control of rest handler from the outside + public static class ControlServerRequestPlugin extends Plugin implements ActionPlugin { static final String ROUTE = "/_test/request-stream"; - static final List handlers = Collections.synchronizedList(new ArrayList<>()); + static final ConcurrentHashMap handlers = new ConcurrentHashMap<>(); @Override public Collection getRestHandlers( @@ -393,8 +416,10 @@ public List routes() { @Override protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) { - var handler = new ServerRequestHandler(); - handlers.add(handler); + var stream = (Netty4HttpRequestBodyStream) request.contentStream(); + var opaqueId = request.getHeaders().get(Task.X_OPAQUE_ID_HTTP_HEADER).get(0); + var handler = new ServerRequestHandler(opaqueId, stream); + handlers.put(opaqueId, handler); return handler; } }); diff --git a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java index 65e56dd82f8fc..00066ffaf0201 100644 --- a/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java +++ b/modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpRequestBodyStreamTests.java @@ -46,6 +46,7 @@ protected void channelRead0(ChannelHandlerContext ctx, HttpContent msg) { }); } + // ensures that no chunks are sent downstream without request public void testEnqueueChunksBeforeRequest() { var totalChunks = randomIntBetween(1, 100); for (int i = 0; i < totalChunks; i++) { @@ -54,8 +55,8 @@ public void testEnqueueChunksBeforeRequest() { assertEquals(totalChunks, stream.chunkQueue().size()); } - // test that every small request consume at least one chunk - public void testFlushQueuedSmallRequests() { + // ensures all queued chunks can be flushed downstream + public void testFlushQueued() { var chunks = new ArrayList(); var totalBytes = new AtomicInteger(); stream.setHandler((chunk, isLast) -> { @@ -76,7 +77,8 @@ public void testFlushQueuedSmallRequests() { assertEquals(chunkSize * totalChunks, totalBytes.get()); } - // test that we read from channel when not enough available bytes + // ensures that we read from channel when chunks queue is empty + // and pass next chunk downstream without queuing public void testReadFromChannel() { var gotChunks = new ArrayList(); var gotLast = new AtomicBoolean(false); diff --git a/server/src/main/java/org/elasticsearch/rest/BaseRestHandler.java b/server/src/main/java/org/elasticsearch/rest/BaseRestHandler.java index 2d65b7d69724c..d73b679917978 100644 --- a/server/src/main/java/org/elasticsearch/rest/BaseRestHandler.java +++ b/server/src/main/java/org/elasticsearch/rest/BaseRestHandler.java @@ -110,7 +110,7 @@ public final void handleRequest(RestRequest request, RestChannel channel, NodeCl throw new IllegalArgumentException(unrecognized(request, unconsumedParams, candidateParams, "parameter")); } - if (request.hasContent() && request.isContentConsumed() == false) { + if (request.hasContent() && (request.isContentConsumed() == false && request.isFullContent())) { throw new IllegalArgumentException( "request [" + request.method() + " " + request.path() + "] does not support having a body" ); diff --git a/server/src/main/java/org/elasticsearch/rest/RestController.java b/server/src/main/java/org/elasticsearch/rest/RestController.java index 8592888d2dd03..bba1d87f8a38e 100644 --- a/server/src/main/java/org/elasticsearch/rest/RestController.java +++ b/server/src/main/java/org/elasticsearch/rest/RestController.java @@ -424,8 +424,7 @@ private void dispatchRequest( MethodHandlers methodHandlers, ThreadContext threadContext ) throws Exception { - final int contentLength = request.contentLength(); - if (contentLength > 0) { + if (request.hasContent()) { if (isContentTypeDisallowed(request) || handler.mediaTypesValid(request) == false) { sendContentTypeErrorMessage(request.getAllHeaderValues("Content-Type"), channel); return; @@ -453,6 +452,9 @@ private void dispatchRequest( return; } } + // TODO: estimate streamed content size for circuit breaker, + // something like http_max_chunk_size * avg_compression_ratio(for compressed content) + final int contentLength = request.isFullContent() ? request.contentLength() : 0; try { if (handler.canTripCircuitBreaker()) { inFlightRequestsBreaker(circuitBreakerService).addEstimateBytesAndMaybeBreak(contentLength, ""); diff --git a/server/src/main/java/org/elasticsearch/rest/RestRequest.java b/server/src/main/java/org/elasticsearch/rest/RestRequest.java index dbd098ac3a618..7d3544e156bdc 100644 --- a/server/src/main/java/org/elasticsearch/rest/RestRequest.java +++ b/server/src/main/java/org/elasticsearch/rest/RestRequest.java @@ -72,7 +72,7 @@ public class RestRequest implements ToXContent.Params, Traceable { private final long requestId; public boolean isContentConsumed() { - return contentConsumed || isStreamedContent(); + return contentConsumed; } @SuppressWarnings("this-escape") @@ -279,24 +279,20 @@ public final String path() { } public boolean hasContent() { - return contentLength() > 0 || isStreamedContent(); + return isStreamedContent() || contentLength() > 0; } public int contentLength() { - if (httpRequest.body().isFull()) { - return httpRequest.body().asFull().bytes().length(); - } else { - return 0; - } + return httpRequest.body().asFull().bytes().length(); + } + + public boolean isFullContent() { + return httpRequest.body().isFull(); } public BytesReference content() { this.contentConsumed = true; - if (httpRequest.body().isFull()) { - return httpRequest.body().asFull().bytes(); - } else { - return BytesArray.EMPTY; - } + return httpRequest.body().asFull().bytes(); } public boolean isStreamedContent() { From 660e4ca9823b1342a25b54d77322728b88ddab93 Mon Sep 17 00:00:00 2001 From: Mikhail Berezovskiy Date: Mon, 12 Aug 2024 13:38:33 -0700 Subject: [PATCH 14/14] add assertions and comments --- .../http/netty4/Netty4IncrementalRequestHandlingIT.java | 9 ++++++++- .../http/netty4/Netty4HttpPipeliningHandler.java | 2 ++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java index dcfe411f3a07e..4de7ca97ed51b 100644 --- a/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java +++ b/modules/transport-netty4/src/internalClusterTest/java/org/elasticsearch/http/netty4/Netty4IncrementalRequestHandlingIT.java @@ -187,7 +187,10 @@ public void testClientBackpressure() throws Exception { for (int i = 0; i < 5; i++) { ctx.clientChannel.writeAndFlush(randomContent(MBytes(10), false)); } - ctx.clientChannel.writeAndFlush(LastHttpContent.EMPTY_LAST_CONTENT); + assertFalse( + "should not flush last content immediately", + ctx.clientChannel.writeAndFlush(LastHttpContent.EMPTY_LAST_CONTENT).isDone() + ); var handler = ctx.awaitRestChannelAccepted(opaqueId); @@ -199,6 +202,9 @@ public void testClientBackpressure() throws Exception { for (int mb = 0; mb <= 50; mb += 10) { var minBufSize = payloadSize - MBytes(10 + mb); var maxBufSize = payloadSize - MBytes(mb); + // it is hard to tell that client's channel is no logger flushing data + // it might take a few busy-iterations before channel buffer flush to kernel + // and bytesBeforeWritable will stop changing assertBusy(() -> { var bufSize = ctx.clientChannel.bytesBeforeWritable(); assertTrue( @@ -226,6 +232,7 @@ static T safePoll(BlockingDeque queue) { assertNotNull("queue is empty", t); return t; } catch (InterruptedException e) { + Thread.currentThread().interrupt(); throw new AssertionError(e); } } diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java index 4136cc4ce4cab..18e76a51efa17 100644 --- a/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java +++ b/modules/transport-netty4/src/main/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandler.java @@ -131,8 +131,10 @@ public void channelRead(final ChannelHandlerContext ctx, final Object msg) { } netty4HttpRequest = new Netty4HttpRequest(readSequence++, (FullHttpRequest) request, nonError); } else { + assert currentRequestStream == null : "current stream must be null for new request"; if (request instanceof FullHttpRequest fullHttpRequest) { netty4HttpRequest = new Netty4HttpRequest(readSequence++, fullHttpRequest); + currentRequestStream = null; } else { var contentStream = new Netty4HttpRequestBodyStream(ctx.channel()); currentRequestStream = contentStream;