Use chunked REST serialization for large REST responses

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

@@ -0,0 +1,48 @@
+/*
+ * 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.DefaultHttpResponse;
+import io.netty.handler.codec.http.HttpResponseStatus;
+import io.netty.handler.codec.http.HttpVersion;
+
+import org.elasticsearch.rest.ChunkedRestResponseBody;
+import org.elasticsearch.rest.RestStatus;
+
+public class Netty4ChunkedHttpResponse extends DefaultHttpResponse implements Netty4RestResponse {
+
+    private final int sequence;
+
+    private final ChunkedRestResponseBody body;
+
+    Netty4ChunkedHttpResponse(int sequence, HttpVersion version, RestStatus status, ChunkedRestResponseBody body) {
+        super(version, HttpResponseStatus.valueOf(status.getStatus()));
+        this.sequence = sequence;
+        this.body = body;
+    }
+
+    public ChunkedRestResponseBody body() {
+        return body;
+    }
+
+    @Override
+    public int getSequence() {
+        return sequence;
+    }
+
+    @Override
+    public void addHeader(String name, String value) {
+        headers().add(name, value);
+    }
+
+    @Override
+    public boolean containsHeader(String name) {
+        return headers().contains(name);
+    }
+}

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

@@ -11,6 +11,7 @@
 import io.netty.buffer.ByteBuf;
 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;
@@ -26,11 +27,17 @@
 
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.common.bytes.ReleasableBytesReference;
 import org.elasticsearch.core.Booleans;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Tuple;
+import org.elasticsearch.rest.ChunkedRestResponseBody;
 import org.elasticsearch.transport.Transports;
+import org.elasticsearch.transport.netty4.Netty4Utils;
+import org.elasticsearch.transport.netty4.Netty4WriteThrottlingHandler;
 import org.elasticsearch.transport.netty4.NettyAllocator;
 
+import java.io.IOException;
 import java.nio.channels.ClosedChannelException;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
@@ -41,13 +48,22 @@
 
 /**
  * Implements HTTP pipelining ordering, ensuring that responses are completely served in the same order as their corresponding requests.
+ * This handler also throttles write operations and will not pass any writes to the next handler so long as the channel is not writable.
  */
 public class Netty4HttpPipeliningHandler extends ChannelDuplexHandler {
 
     private final Logger logger;
 
     private final int maxEventsHeld;
-    private final PriorityQueue<Tuple<Netty4HttpResponse, ChannelPromise>> outboundHoldingQueue;
+    private final PriorityQueue<Tuple<? extends Netty4RestResponse, ChannelPromise>> outboundHoldingQueue;
+
+    private record ChunkedWrite(PromiseCombiner combiner, ChannelPromise onDone, Netty4ChunkedHttpResponse response) {}
+
+    /**
+     * The current {@link ChunkedWrite} if a chunked write is executed at the moment.
+     */
+    @Nullable
+    private ChunkedWrite currentChunkedWrite;
 
     /*
      * The current read and write sequence numbers. Read sequence numbers are attached to requests in the order they are read from the
@@ -119,36 +135,32 @@ protected void handlePipelinedRequest(ChannelHandlerContext ctx, Netty4HttpReque
     }
 
     @Override
-    public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise promise) {
-        assert msg instanceof Netty4HttpResponse : "Invalid message type: " + msg.getClass();
+    public void write(final ChannelHandlerContext ctx, final Object msg, final ChannelPromise promise) throws IOException {
+        assert msg instanceof Netty4RestResponse : "Invalid message type: " + msg.getClass();
         boolean success = false;
         try {
-            final Netty4HttpResponse response = (Netty4HttpResponse) msg;
-            if (response.getSequence() != writeSequence) {
-                assert response.getSequence() > writeSequence
-                    : "response sequence [" + response.getSequence() + "] we below write sequence [" + writeSequence + "]";
+            final Netty4RestResponse restResponse = (Netty4RestResponse) msg;
+            if (restResponse.getSequence() != writeSequence) {
+                assert restResponse.getSequence() > writeSequence
+                    : "response sequence [" + restResponse.getSequence() + "] we below write sequence [" + writeSequence + "]";
                 if (outboundHoldingQueue.size() >= maxEventsHeld) {
                     int eventCount = outboundHoldingQueue.size() + 1;
                     throw new IllegalStateException(
                         "Too many pipelined events [" + eventCount + "]. Max events allowed [" + maxEventsHeld + "]."
                     );
                 }
                 // response is not at the current sequence number so we add it to the outbound queue and return
-                outboundHoldingQueue.add(new Tuple<>(response, promise));
+                outboundHoldingQueue.add(new Tuple<>(restResponse, promise));
                 success = true;
                 return;
             }
 
             // response is at the current sequence number and does not need to wait for any other response to be written so we write
             // it out directly
-            doWrite(ctx, response, promise);
+            doWrite(ctx, restResponse, promise);
             success = true;
             // see if we have any queued up responses that became writeable due to the above write
-            while (outboundHoldingQueue.isEmpty() == false && outboundHoldingQueue.peek().v1().getSequence() == writeSequence) {
-                final Tuple<Netty4HttpResponse, ChannelPromise> top = outboundHoldingQueue.poll();
-                assert top != null : "we know the outbound holding queue to not be empty at this point";
-                doWrite(ctx, top.v1(), top.v2());
-            }
+            drainQueuedNowWritable(ctx);
         } catch (IllegalStateException e) {
             ctx.channel().close();
         } finally {
@@ -158,6 +170,14 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
         }
     }
 
+    private void drainQueuedNowWritable(ChannelHandlerContext ctx) throws IOException {
+        while (outboundHoldingQueue.isEmpty() == false && outboundHoldingQueue.peek().v1().getSequence() == writeSequence) {
+            final Tuple<? extends Netty4RestResponse, ChannelPromise> top = outboundHoldingQueue.poll();
+            assert top != null : "we know the outbound holding queue to not be empty at this point";
+            doWrite(ctx, top.v1(), top.v2());
+        }
+    }
+
     private static final String DO_NOT_SPLIT = "es.unsafe.do_not_split_http_responses";
 
     private static final boolean DO_NOT_SPLIT_HTTP_RESPONSES;
@@ -169,6 +189,14 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
         SPLIT_THRESHOLD = (int) (NettyAllocator.suggestedMaxAllocationSize() * 0.99);
     }
 
+    private void doWrite(ChannelHandlerContext ctx, Netty4RestResponse readyResponse, ChannelPromise promise) throws IOException {
+        if (readyResponse instanceof Netty4HttpResponse) {
+            doWrite(ctx, (Netty4HttpResponse) readyResponse, promise);
+        } else {
+            doWrite(ctx, (Netty4ChunkedHttpResponse) readyResponse, promise);
+        }
+    }
+
     /**
      * Split up large responses to prevent batch compression {@link JdkZlibEncoder} down the pipeline.
      */
@@ -181,6 +209,27 @@ private void doWrite(ChannelHandlerContext ctx, Netty4HttpResponse readyResponse
         writeSequence++;
     }
 
+    private void doWrite(ChannelHandlerContext ctx, Netty4ChunkedHttpResponse readyResponse, ChannelPromise promise) throws IOException {
+        assert currentChunkedWrite == null : "unexpected existing write [" + currentChunkedWrite + "]";
+        final PromiseCombiner combiner = new PromiseCombiner(ctx.executor());
+        currentChunkedWrite = new ChunkedWrite(combiner, promise, readyResponse);
+        final ChannelPromise first = ctx.newPromise();
+        combiner.add((Future<Void>) first);
+        enqueueWrite(ctx, readyResponse, first);
+        while (ctx.channel().isWritable()) {
+            if (writeChunk(ctx, combiner, readyResponse.body())) {
+                finishChunkedWrite();
+                return;
+            }
+        }
+    }
+
+    private void finishChunkedWrite() {
+        currentChunkedWrite.combiner.finish(currentChunkedWrite.onDone);
+        currentChunkedWrite = null;
+        writeSequence++;
+    }
+
     private void splitAndWrite(ChannelHandlerContext ctx, Netty4HttpResponse msg, ChannelPromise promise) {
         final PromiseCombiner combiner = new PromiseCombiner(ctx.executor());
         HttpResponse response = new DefaultHttpResponse(msg.protocolVersion(), msg.status(), msg.headers());
@@ -193,15 +242,15 @@ private void splitAndWrite(ChannelHandlerContext ctx, Netty4HttpResponse msg, Ch
         combiner.finish(promise);
     }
 
-    public void channelWritabilityChanged(ChannelHandlerContext ctx) {
+    public void channelWritabilityChanged(ChannelHandlerContext ctx) throws IOException {
         if (ctx.channel().isWritable()) {
             doFlush(ctx);
         }
         ctx.fireChannelWritabilityChanged();
     }
 
     @Override
-    public void flush(ChannelHandlerContext ctx) {
+    public void flush(ChannelHandlerContext ctx) throws IOException {
         if (doFlush(ctx) == false) {
             ctx.flush();
         }
@@ -218,7 +267,7 @@ public void channelInactive(ChannelHandlerContext ctx) throws Exception {
      *
      * @return true if a call to this method resulted in a call to {@link ChannelHandlerContext#flush()} on the given {@code ctx}
      */
-    private boolean doFlush(ChannelHandlerContext ctx) {
+    private boolean doFlush(ChannelHandlerContext ctx) throws IOException {
         assert ctx.executor().inEventLoop();
         final Channel channel = ctx.channel();
         if (channel.isActive() == false) {
@@ -227,8 +276,28 @@ private boolean doFlush(ChannelHandlerContext ctx) {
         }
         boolean needsFlush = true;
         while (channel.isWritable()) {
-            final WriteOperation currentWrite = queuedWrites.poll();
+            WriteOperation currentWrite = queuedWrites.poll();
             if (currentWrite == null) {
+                drainQueuedNowWritable(ctx);
+                if (channel.isWritable() == false) {
+                    // try flushing to make channel writable again, loop will only continue if channel becomes writable again
+                    ctx.flush();
+                    needsFlush = false;
+                }
+                if (channel.isWritable() == false) {
+                    break;
+                }
+                currentWrite = queuedWrites.poll();
+            }
+            if (currentWrite == null) {
+                // TODO: this is confusing, handle the outstanding first chunk of a chunked write nicer than by putting it in the
+                // queued writes
+                if (currentChunkedWrite != null) {
+                    if (writeChunk(ctx, currentChunkedWrite.combiner, currentChunkedWrite.response.body())) {
+                        finishChunkedWrite();
+                        continue;
+                    }
+                }
                 break;
             }
             ctx.write(currentWrite.msg, currentWrite.promise);
@@ -248,6 +317,19 @@ private boolean doFlush(ChannelHandlerContext ctx) {
         return true;
     }
 
+    private boolean writeChunk(ChannelHandlerContext ctx, PromiseCombiner combiner, ChunkedRestResponseBody body) throws IOException {
+        final ReleasableBytesReference bytes = body.encodeChunk(
+            Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE,
+            serverTransport.recycler()
+        );
+        final ByteBuf content = Netty4Utils.toByteBuf(bytes);
+        final boolean done = body.isDone();
+        final ChannelFuture f = ctx.write(done ? new DefaultLastHttpContent(content) : new DefaultHttpContent(content));
+        f.addListener(ignored -> bytes.close());
+        combiner.add(f);
+        return done;
+    }
+
     private void failQueuedWrites() {
         WriteOperation queuedWrite;
         while ((queuedWrite = queuedWrites.poll()) != null) {
@@ -257,21 +339,29 @@ private void failQueuedWrites() {
 
     @Override
     public void close(ChannelHandlerContext ctx, ChannelPromise promise) {
-        List<Tuple<Netty4HttpResponse, ChannelPromise>> inflightResponses = removeAllInflightResponses();
+        if (currentChunkedWrite != null) {
+            safeFailPromise(currentChunkedWrite.onDone, new ClosedChannelException());
+            currentChunkedWrite = null;
+        }
+        List<Tuple<? extends Netty4RestResponse, ChannelPromise>> inflightResponses = removeAllInflightResponses();
 
         if (inflightResponses.isEmpty() == false) {
             ClosedChannelException closedChannelException = new ClosedChannelException();
-            for (Tuple<Netty4HttpResponse, ChannelPromise> inflightResponse : inflightResponses) {
-                try {
-                    inflightResponse.v2().setFailure(closedChannelException);
-                } catch (RuntimeException e) {
-                    logger.error("unexpected error while releasing pipelined http responses", e);
-                }
+            for (Tuple<? extends Netty4RestResponse, ChannelPromise> inflightResponse : inflightResponses) {
+                safeFailPromise(inflightResponse.v2(), closedChannelException);
             }
         }
         ctx.close(promise);
     }
 
+    private void safeFailPromise(ChannelPromise promise, Exception ex) {
+        try {
+            promise.setFailure(ex);
+        } catch (RuntimeException e) {
+            logger.error("unexpected error while releasing pipelined http responses", e);
+        }
+    }
+
     private Future<Void> enqueueWrite(ChannelHandlerContext ctx, HttpObject msg) {
         final ChannelPromise p = ctx.newPromise();
         enqueueWrite(ctx, msg, p);
@@ -299,8 +389,8 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
         }
     }
 
-    private List<Tuple<Netty4HttpResponse, ChannelPromise>> removeAllInflightResponses() {
-        ArrayList<Tuple<Netty4HttpResponse, ChannelPromise>> responses = new ArrayList<>(outboundHoldingQueue);
+    private List<Tuple<? extends Netty4RestResponse, ChannelPromise>> removeAllInflightResponses() {
+        ArrayList<Tuple<? extends Netty4RestResponse, ChannelPromise>> responses = new ArrayList<>(outboundHoldingQueue);
         outboundHoldingQueue.clear();
         return responses;
     }

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

@@ -22,6 +22,8 @@
 
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.http.HttpRequest;
+import org.elasticsearch.http.HttpResponse;
+import org.elasticsearch.rest.ChunkedRestResponseBody;
 import org.elasticsearch.rest.RestRequest;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.transport.netty4.Netty4Utils;
@@ -237,6 +239,11 @@ public Netty4HttpResponse createResponse(RestStatus status, BytesReference conte
         return new Netty4HttpResponse(sequence, request.protocolVersion(), status, contentRef);
     }
 
+    @Override
+    public HttpResponse createResponse(RestStatus status, ChunkedRestResponseBody content) {
+        return new Netty4ChunkedHttpResponse(sequence, request.protocolVersion(), status, content);
+    }
+
     @Override
     public Exception getInboundException() {
         return inboundException;

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

@@ -13,11 +13,10 @@
 import io.netty.handler.codec.http.HttpVersion;
 
 import org.elasticsearch.common.bytes.BytesReference;
-import org.elasticsearch.http.HttpResponse;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.transport.netty4.Netty4Utils;
 
-public class Netty4HttpResponse extends DefaultFullHttpResponse implements HttpResponse {
+public class Netty4HttpResponse extends DefaultFullHttpResponse implements Netty4RestResponse {
 
     private final int sequence;
 
@@ -26,6 +25,7 @@ public class Netty4HttpResponse extends DefaultFullHttpResponse implements HttpR
         this.sequence = sequence;
     }
 
+    @Override
     public int getSequence() {
         return sequence;
     }

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

@@ -0,0 +1,16 @@
+/*
+ * 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.http.HttpResponse;
+
+public interface Netty4RestResponse extends HttpResponse {
+
+    int getSequence();
+}