Pausable chunked HTTP responses

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

@@ -10,6 +10,7 @@
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.ESNetty4IntegTestCase;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.internal.node.NodeClient;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
@@ -161,6 +162,16 @@ public boolean isDone() {
                     return chunkIterator.hasNext() == false;
                 }
 
+                @Override
+                public boolean isEndOfResponse() {
+                    return true;
+                }
+
+                @Override
+                public void getContinuation(ActionListener<ChunkedRestResponseBody> listener) {
+                    assert false : "no continuations";
+                }
+
                 @Override
                 public ReleasableBytesReference encodeChunk(int sizeHint, Recycler<BytesRef> recycler) {
                     final var page = recycler.obtain(); // just to ensure nothing is leaked

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

@@ -0,0 +1,38 @@
+/*
+ * 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.util.concurrent.PromiseCombiner;
+
+import org.elasticsearch.rest.ChunkedRestResponseBody;
+
+final class Netty4ChunkedHttpContinuation implements Netty4HttpResponse {
+    private final int sequence;
+    private final ChunkedRestResponseBody body;
+    private final PromiseCombiner combiner;
+
+    Netty4ChunkedHttpContinuation(int sequence, ChunkedRestResponseBody body, PromiseCombiner combiner) {
+        this.sequence = sequence;
+        this.body = body;
+        this.combiner = combiner;
+    }
+
+    @Override
+    public int getSequence() {
+        return sequence;
+    }
+
+    public ChunkedRestResponseBody body() {
+        return body;
+    }
+
+    public PromiseCombiner combiner() {
+        return combiner;
+    }
+}

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

@@ -19,7 +19,7 @@
 /**
  * A http response that will be transferred via chunked encoding when handled by {@link Netty4HttpPipeliningHandler}.
  */
-public final class Netty4ChunkedHttpResponse extends DefaultHttpResponse implements Netty4HttpResponse, HttpResponse {
+final class Netty4ChunkedHttpResponse extends DefaultHttpResponse implements Netty4HttpResponse, HttpResponse {
 
     private final int sequence;
 

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

@@ -28,6 +28,8 @@
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.ExceptionsHelper;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.bytes.ReleasableBytesReference;
 import org.elasticsearch.core.Booleans;
 import org.elasticsearch.core.Nullable;
@@ -140,6 +142,8 @@ public void write(final ChannelHandlerContext ctx, final Object msg, final Chann
         try {
             final Netty4HttpResponse restResponse = (Netty4HttpResponse) msg;
             if (restResponse.getSequence() != writeSequence) {
+                assert restResponse instanceof Netty4ChunkedHttpContinuation == false
+                    : "received out-of-order continuation at [" + restResponse.getSequence() + "], expecting [" + writeSequence + "]";
                 assert restResponse.getSequence() > writeSequence
                     : "response sequence [" + restResponse.getSequence() + "] we below write sequence [" + writeSequence + "]";
                 if (outboundHoldingQueue.size() >= maxEventsHeld) {
@@ -198,6 +202,8 @@ private void doWrite(ChannelHandlerContext ctx, Netty4HttpResponse readyResponse
             doWriteFullResponse(ctx, fullResponse, promise);
         } else if (readyResponse instanceof Netty4ChunkedHttpResponse chunkedResponse) {
             doWriteChunkedResponse(ctx, chunkedResponse, promise);
+        } else if (readyResponse instanceof Netty4ChunkedHttpContinuation chunkedContinuation) {
+            doWriteChunkedContinuation(ctx, chunkedContinuation, promise);
         } else {
             assert false : readyResponse.getClass().getCanonicalName();
             throw new IllegalStateException("illegal message type: " + readyResponse.getClass().getCanonicalName());
@@ -236,13 +242,54 @@ private void doWriteChunkedResponse(ChannelHandlerContext ctx, Netty4ChunkedHttp
         }
     }
 
+    private void doWriteChunkedContinuation(ChannelHandlerContext ctx, Netty4ChunkedHttpContinuation continuation, ChannelPromise promise)
+        throws IOException {
+        final PromiseCombiner combiner = continuation.combiner();
+        assert currentChunkedWrite == null;
+        final var responseBody = continuation.body();
+        currentChunkedWrite = new ChunkedWrite(combiner, promise, responseBody);
+        // NB "writable" means there's space in the downstream ChannelOutboundBuffer, we aren't trying to saturate the physical channel.
+        while (ctx.channel().isWritable()) {
+            if (writeChunk(ctx, combiner, responseBody)) {
+                finishChunkedWrite();
+                return;
+            }
+        }
+    }
+
     private void finishChunkedWrite() {
         assert currentChunkedWrite != null;
         assert currentChunkedWrite.responseBody().isDone();
         final var finishingWrite = currentChunkedWrite;
+        final var endOfResponse = finishingWrite.responseBody().isEndOfResponse();
         currentChunkedWrite = null;
-        writeSequence++;
-        finishingWrite.combiner.finish(finishingWrite.onDone());
+        if (endOfResponse) {
+            writeSequence++;
+            finishingWrite.combiner.finish(finishingWrite.onDone());
+        } else {
+            final var channel = finishingWrite.onDone().channel();
+            ActionListener.run(ActionListener.assertOnce(new ActionListener<>() {
+                @Override
+                public void onResponse(ChunkedRestResponseBody continuation) {
+                    channel.writeAndFlush(
+                        new Netty4ChunkedHttpContinuation(writeSequence, continuation, finishingWrite.combiner()),
+                        finishingWrite.onDone() // pass the terminal listener/promise along the line
+                    );
+                }
+
+                @Override
+                public void onFailure(Exception e) {
+                    logger.error(
+                        Strings.format("failed to get continuation of HTTP response body for [%s], closing connection", channel),
+                        e
+                    );
+                    channel.close().addListener(ignored -> {
+                        finishingWrite.combiner().add(channel.newFailedFuture(e));
+                        finishingWrite.combiner().finish(finishingWrite.onDone());
+                    });
+                }
+            }), finishingWrite.responseBody()::getContinuation);
+        }
     }
 
     private void splitAndWrite(ChannelHandlerContext ctx, Netty4FullHttpResponse msg, ChannelPromise promise) {
@@ -326,7 +373,8 @@ private boolean writeChunk(ChannelHandlerContext ctx, PromiseCombiner combiner,
         );
         final ByteBuf content = Netty4Utils.toByteBuf(bytes);
         final boolean done = body.isDone();
-        final ChannelFuture f = ctx.write(done ? new DefaultLastHttpContent(content) : new DefaultHttpContent(content));
+        final boolean lastChunk = done && body.isEndOfResponse();
+        final ChannelFuture f = ctx.write(lastChunk ? new DefaultLastHttpContent(content) : new DefaultHttpContent(content));
         f.addListener(ignored -> bytes.close());
         combiner.add(f);
         return done;
@@ -337,6 +385,11 @@ private void failQueuedWrites() {
         while ((queuedWrite = queuedWrites.poll()) != null) {
             queuedWrite.failAsClosedChannel();
         }
+        if (currentChunkedWrite != null) {
+            final var finishingWrite = currentChunkedWrite;
+            currentChunkedWrite = null;
+            finishingWrite.combiner.finish(finishingWrite.onDone());
+        }
     }
 
     @Override

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

@@ -11,7 +11,7 @@
 /**
  * Super-interface for responses handled by the Netty4 HTTP transport.
  */
-public sealed interface Netty4HttpResponse permits Netty4FullHttpResponse, Netty4ChunkedHttpResponse {
+sealed interface Netty4HttpResponse permits Netty4FullHttpResponse, Netty4ChunkedHttpResponse, Netty4ChunkedHttpContinuation {
     /**
      * @return The sequence number for the request which corresponds with this response, for making sure that we send responses to pipelined
      * requests in the correct order.

modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpClient.java

@@ -35,6 +35,7 @@
 import org.elasticsearch.common.unit.ByteSizeValue;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.transport.netty4.NettyAllocator;
 
 import java.io.Closeable;
@@ -137,9 +138,20 @@ private synchronized List<FullHttpResponse> sendRequests(final SocketAddress rem
             channelFuture = clientBootstrap.connect(remoteAddress);
             channelFuture.sync();
 
+            boolean needsFinalFlush = false;
             for (HttpRequest request : requests) {
-                channelFuture.channel().writeAndFlush(request);
+                if (ESTestCase.randomBoolean()) {
+                    channelFuture.channel().writeAndFlush(request);
+                    needsFinalFlush = false;
+                } else {
+                    channelFuture.channel().write(request);
+                    needsFinalFlush = true;
+                }
+            }
+            if (needsFinalFlush) {
+                channelFuture.channel().flush();
             }
+
             if (latch.await(30L, TimeUnit.SECONDS) == false) {
                 fail("Failed to get all expected responses.");
             }

modules/transport-netty4/src/test/java/org/elasticsearch/http/netty4/Netty4HttpPipeliningHandlerTests.java

@@ -28,6 +28,7 @@
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.Randomness;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.bytes.BytesReference;
@@ -511,6 +512,16 @@ public boolean isDone() {
                 return remaining == 0;
             }
 
+            @Override
+            public boolean isEndOfResponse() {
+                return true;
+            }
+
+            @Override
+            public void getContinuation(ActionListener<ChunkedRestResponseBody> listener) {
+                fail("no continuations here");
+            }
+
             @Override
             public ReleasableBytesReference encodeChunk(int sizeHint, Recycler<BytesRef> recycler) {
                 assertThat(remaining, greaterThan(0));

server/src/main/java/org/elasticsearch/rest/ChunkedRestResponseBody.java

@@ -8,6 +8,8 @@
 package org.elasticsearch.rest;
 
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.common.bytes.ReleasableBytesReference;
 import org.elasticsearch.common.io.stream.BytesStream;
 import org.elasticsearch.common.io.stream.RecyclerBytesStreamOutput;
@@ -20,6 +22,8 @@
 import org.elasticsearch.core.Streams;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
+import org.elasticsearch.tasks.CancellableTask;
+import org.elasticsearch.tasks.Task;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentBuilder;
 
@@ -43,6 +47,30 @@ public interface ChunkedRestResponseBody {
      */
     boolean isDone();
 
+    /**
+     * @return {@code true} if this is the last chunked body in the response, or {@code false} if the REST layer should request further
+     * chunked bodies by calling {@link #getContinuation}.
+     */
+    boolean isEndOfResponse();
+
+    /**
+     * <p>Asynchronously retrieves the next part of the body. Called if {@link #isEndOfResponse} returns {@code false}.</p>
+     *
+     * <p>Note that this is called on a transport thread, so implementations must take care to dispatch any nontrivial work elsewhere.</p>
+
+     * <p>Note that the {@link Task} corresponding to any invocation of {@link Client#execute} completes as soon as the client action
+     * returns its response, so it no longer exists when this method is called and cannot be used to receive cancellation notifications.
+     * Instead, if the HTTP channel is closed while sending a response (including while waiting for a continuation) then the REST layer
+     * will invoke {@link RestResponse#close}. Implementations will typically explicitly create a {@link CancellableTask} to represent the
+     * computation and transmission of the entire {@link RestResponse}, and will cancel this task if the {@link RestResponse} is closed
+     * prematurely.</p>
+     *
+     * @param listener Listener to complete with the next part of the body. By the point this is called we have already started to send
+     *                 the body of the response, so there's no good ways to handle an exception here. Completing the listener exceptionally
+     *                 will log an error, abort sending the response, and close the HTTP connection.
+     */
+    void getContinuation(ActionListener<ChunkedRestResponseBody> listener);
+
     /**
      * Serializes approximately as many bytes of the response as request by {@code sizeHint} to a {@link ReleasableBytesReference} that
      * is created from buffers backed by the given {@code recycler}.
@@ -102,6 +130,16 @@ public boolean isDone() {
                 return serialization.hasNext() == false;
             }
 
+            @Override
+            public boolean isEndOfResponse() {
+                return true;
+            }
+
+            @Override
+            public void getContinuation(ActionListener<ChunkedRestResponseBody> listener) {
+                assert false : "no continuations";
+            }
+
             @Override
             public ReleasableBytesReference encodeChunk(int sizeHint, Recycler<BytesRef> recycler) throws IOException {
                 try {
@@ -180,6 +218,16 @@ public boolean isDone() {
                 return chunkIterator.hasNext() == false;
             }
 
+            @Override
+            public boolean isEndOfResponse() {
+                return true;
+            }
+
+            @Override
+            public void getContinuation(ActionListener<ChunkedRestResponseBody> listener) {
+                assert false : "no continuations";
+            }
+
             @Override
             public ReleasableBytesReference encodeChunk(int sizeHint, Recycler<BytesRef> recycler) throws IOException {
                 try {

server/src/main/java/org/elasticsearch/rest/LoggingChunkedRestResponseBody.java

@@ -9,6 +9,7 @@
 package org.elasticsearch.rest;
 
 import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.common.bytes.ReleasableBytesReference;
 import org.elasticsearch.common.recycler.Recycler;
 
@@ -30,6 +31,16 @@ public boolean isDone() {
         return inner.isDone();
     }
 
+    @Override
+    public boolean isEndOfResponse() {
+        return inner.isEndOfResponse();
+    }
+
+    @Override
+    public void getContinuation(ActionListener<ChunkedRestResponseBody> listener) {
+        inner.getContinuation(listener.map(continuation -> new LoggingChunkedRestResponseBody(continuation, loggerStream)));
+    }
+
     @Override
     public ReleasableBytesReference encodeChunk(int sizeHint, Recycler<BytesRef> recycler) throws IOException {
         var chunk = inner.encodeChunk(sizeHint, recycler);

server/src/main/java/org/elasticsearch/rest/RestController.java

@@ -907,11 +907,23 @@ public boolean isDone() {
             return delegate.isDone();
         }
 
+        @Override
+        public boolean isEndOfResponse() {
+            return delegate.isEndOfResponse();
+        }
+
+        @Override
+        public void getContinuation(ActionListener<ChunkedRestResponseBody> listener) {
+            delegate.getContinuation(
+                listener.map(continuation -> new EncodedLengthTrackingChunkedRestResponseBody(continuation, responseLengthRecorder))
+            );
+        }
+
         @Override
         public ReleasableBytesReference encodeChunk(int sizeHint, Recycler<BytesRef> recycler) throws IOException {
             final ReleasableBytesReference bytesReference = delegate.encodeChunk(sizeHint, recycler);
             responseLengthRecorder.addChunkLength(bytesReference.length());
-            if (isDone()) {
+            if (isDone() && isEndOfResponse()) {
                 responseLengthRecorder.close();
             }
             return bytesReference;