Chunk + Throttle Netty Writes

modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4MessageChannelHandler.java

@@ -22,12 +22,17 @@
 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.util.Attribute;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.transport.Transports;
 
+import java.nio.channels.ClosedChannelException;
+import java.util.ArrayDeque;
+import java.util.Queue;
 
 /**
  * A handler (must be the last one!) that does size based frame decoding and forwards the actual message
@@ -37,13 +42,17 @@ final class Netty4MessageChannelHandler extends ChannelDuplexHandler {
 
     private final Netty4Transport transport;
 
+    private final Queue<WriteOperation> queuedWrites = new ArrayDeque<>();
+
+    private WriteOperation currentWrite;
+
     Netty4MessageChannelHandler(Netty4Transport transport) {
         this.transport = transport;
     }
 
     @Override
-    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
-        Transports.assertTransportThread();
+    public void channelRead(ChannelHandlerContext ctx, Object msg) {
+        assert Transports.assertTransportThread();
         assert msg instanceof ByteBuf : "Expected message type ByteBuf, found: " + msg.getClass();
 
         final ByteBuf buffer = (ByteBuf) msg;
@@ -57,7 +66,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception
     }
 
     @Override
-    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
+    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
         ExceptionsHelper.maybeDieOnAnotherThread(cause);
         final Throwable unwrapped = ExceptionsHelper.unwrap(cause, ElasticsearchException.class);
         final Throwable newCause = unwrapped != null ? unwrapped : cause;
@@ -68,4 +77,113 @@ public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws E
             transport.onException(tcpChannel, (Exception) newCause);
         }
     }
+
+    @Override
+    public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) {
+        assert msg instanceof ByteBuf;
+        final boolean queued = queuedWrites.offer(new WriteOperation((ByteBuf) msg, promise));
+        assert queued;
+    }
+
+    @Override
+    public void channelWritabilityChanged(ChannelHandlerContext ctx) {
+        if (ctx.channel().isWritable()) {
+            doFlush(ctx);
+        }
+        ctx.fireChannelWritabilityChanged();
+    }
+
+    @Override
+    public void flush(ChannelHandlerContext ctx) {
+        Channel channel = ctx.channel();
+        if (channel.isWritable() || channel.isActive() == false) {
+            doFlush(ctx);
+        }
+    }
+
+    @Override
+    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
+        doFlush(ctx);
+        super.channelInactive(ctx);
+    }
+
+    private void doFlush(ChannelHandlerContext ctx) {
+        assert ctx.executor().inEventLoop();
+        final Channel channel = ctx.channel();
+        if (channel.isActive() == false) {
+            if (currentWrite != null) {
+                currentWrite.promise.tryFailure(new ClosedChannelException());
+            }
+            failQueuedWrites();
+            return;
+        }
+        while (channel.isWritable()) {
+            if (currentWrite == null) {
+                currentWrite = queuedWrites.poll();
+            }
+            if (currentWrite == null) {
+                break;
+            }
+            final WriteOperation write = currentWrite;
+            if (write.buf.readableBytes() == 0) {
+                write.promise.trySuccess();
+                currentWrite = null;
+                continue;
+            }
+            final int readableBytes = write.buf.readableBytes();
+            final int bufferSize = Math.min(readableBytes, 1 << 18);
+            final int readerIndex = write.buf.readerIndex();
+            final boolean sliced = readableBytes != bufferSize;
+            final ByteBuf writeBuffer;
+            if (sliced) {
+                writeBuffer = write.buf.retainedSlice(readerIndex, bufferSize);
+                write.buf.readerIndex(readerIndex + bufferSize);
+            } else {
+                writeBuffer = write.buf;
+            }
+            final ChannelFuture writeFuture = ctx.write(writeBuffer);
+            if (sliced == false || write.buf.readableBytes() == 0) {
+                currentWrite = null;
+                writeFuture.addListener(future -> {
+                    assert ctx.executor().inEventLoop();
+                    if (future.isSuccess()) {
+                        write.promise.trySuccess();
+                    } else {
+                        write.promise.tryFailure(future.cause());
+                    }
+                });
+            } else {
+                writeFuture.addListener(future -> {
+                    assert ctx.executor().inEventLoop();
+                    if (future.isSuccess() == false) {
+                        write.promise.tryFailure(future.cause());
+                    }
+                });
+            }
+            ctx.flush();
+            if (channel.isActive() == false) {
+                failQueuedWrites();
+                return;
+            }
+        }
+    }
+
+    private void failQueuedWrites() {
+        WriteOperation queuedWrite;
+        while ((queuedWrite = queuedWrites.poll()) != null) {
+            queuedWrite.promise.tryFailure(new ClosedChannelException());
+        }
+    }
+
+    private static final class WriteOperation {
+
+        private final ByteBuf buf;
+
+        private final ChannelPromise promise;
+
+        WriteOperation(ByteBuf buf, ChannelPromise promise) {
+            this.buf = buf;
+            this.promise = promise;
+        }
+    }
 }