Avoid building large CompositeByteBuf when sending transport messages (#105137)

original-brownbear · web-flow · commit f8795088349b · 2024-02-05T14:35:15.000-05:00
We can avoid building composite byte buf instances on the transport
layer (they have quite a bit of overhead and make heap dumps more
complicated to read). There's no need to add another round of references
to the BytesReference components here. Just write these out as they come
in. This would allow for some efficiency improving follow-ups where we
can essentially release the pages that have passed the write pipeline.
To avoid having this explode the size of the queue for writes per
channel, I moved that to a linked list. The slowdown from a linked list
is irrelevant I believe. Mostly the queue is empty so it doesn't matter
or if it isn't empty, operations other than dequeuing are much more
important to performance in this logic anyway (+ Netty internally uses a
LL down the line anyway).

I would regard this as step-1 in making the serialisation here more lazy
like on the REST layer to avoid copying bytes to the outbound buffer
that we already have as `byte[]`.
diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4TcpChannel.java b/modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4TcpChannel.java
@@ -165,7 +165,7 @@ public void sendMessage(BytesReference reference, ActionListener<Void> listener)
         // We need to both guard against double resolving the listener and not resolving it in case of event loop shutdown so we need to
         // use #notifyOnce here until https://github.com/netty/netty/issues/8007 is resolved.
         var wrapped = ActionListener.notifyOnce(listener);
-        channel.writeAndFlush(Netty4Utils.toByteBuf(reference), addPromise(wrapped, channel));
+        channel.writeAndFlush(reference, addPromise(wrapped, channel));
         if (channel.eventLoop().isShutdown()) {
             wrapped.onFailure(new TransportException("Cannot send message, event loop is shutting down."));
         }
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
@@ -72,9 +72,13 @@ public static void setAvailableProcessors(final int availableProcessors) {
      * pages of the BytesReference. Don't free the bytes of reference before the ByteBuf goes out of scope.
      */
     public static ByteBuf toByteBuf(final BytesReference reference) {
-        if (reference.length() == 0) {
-            return Unpooled.EMPTY_BUFFER;
+        if (reference.hasArray()) {
+            return Unpooled.wrappedBuffer(reference.array(), reference.arrayOffset(), reference.length());
         }
+        return compositeReferenceToByteBuf(reference);
+    }
+
+    private static ByteBuf compositeReferenceToByteBuf(BytesReference reference) {
         final BytesRefIterator iterator = reference.iterator();
         // usually we have one, two, or three components from the header, the message, and a buffer
         final List<ByteBuf> buffers = new ArrayList<>(3);
diff --git a/modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4WriteThrottlingHandler.java b/modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4WriteThrottlingHandler.java
@@ -9,19 +9,25 @@
 package org.elasticsearch.transport.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.util.concurrent.Future;
 import io.netty.util.concurrent.GenericFutureListener;
+import io.netty.util.concurrent.PromiseCombiner;
 
+import org.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.BytesRefIterator;
+import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.transport.Transports;
 
+import java.io.IOException;
 import java.nio.channels.ClosedChannelException;
-import java.util.ArrayDeque;
+import java.util.LinkedList;
 import java.util.Queue;
 
 /**
@@ -32,7 +38,7 @@
 public final class Netty4WriteThrottlingHandler extends ChannelDuplexHandler {
 
     public static final int MAX_BYTES_PER_WRITE = 1 << 18;
-    private final Queue<WriteOperation> queuedWrites = new ArrayDeque<>();
+    private final Queue<WriteOperation> queuedWrites = new LinkedList<>();
 
     private final ThreadContext threadContext;
     private WriteOperation currentWrite;
@@ -42,17 +48,36 @@ public Netty4WriteThrottlingHandler(ThreadContext threadContext) {
     }
 
     @Override
-    public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) {
-        assert msg instanceof ByteBuf;
+    public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws IOException {
+        if (msg instanceof BytesReference reference) {
+            if (reference.hasArray()) {
+                writeSingleByteBuf(ctx, Unpooled.wrappedBuffer(reference.array(), reference.arrayOffset(), reference.length()), promise);
+            } else {
+                BytesRefIterator iter = reference.iterator();
+                final PromiseCombiner combiner = new PromiseCombiner(ctx.executor());
+                BytesRef next;
+                while ((next = iter.next()) != null) {
+                    final ChannelPromise chunkPromise = ctx.newPromise();
+                    combiner.add((Future<Void>) chunkPromise);
+                    writeSingleByteBuf(ctx, Unpooled.wrappedBuffer(next.bytes, next.offset, next.length), chunkPromise);
+                }
+                combiner.finish(promise);
+            }
+        } else {
+            assert msg instanceof ByteBuf;
+            writeSingleByteBuf(ctx, (ByteBuf) msg, promise);
+        }
+    }
+
+    private void writeSingleByteBuf(ChannelHandlerContext ctx, ByteBuf buf, ChannelPromise promise) {
         assert Transports.assertDefaultThreadContext(threadContext);
         assert Transports.assertTransportThread();
-        final ByteBuf buf = (ByteBuf) msg;
         if (ctx.channel().isWritable() && currentWrite == null && queuedWrites.isEmpty()) {
             // nothing is queued for writing and the channel is writable, just pass the write down the pipeline directly
             if (buf.readableBytes() > MAX_BYTES_PER_WRITE) {
                 writeInSlices(ctx, promise, buf);
             } else {
-                ctx.write(msg, promise);
+                ctx.write(buf, promise);
             }
         } else {
             queueWrite(buf, promise);
diff --git a/modules/transport-netty4/src/test/java/org/elasticsearch/transport/netty4/Netty4WriteThrottlingHandlerTests.java b/modules/transport-netty4/src/test/java/org/elasticsearch/transport/netty4/Netty4WriteThrottlingHandlerTests.java
@@ -14,6 +14,9 @@
 import io.netty.channel.ChannelPromise;
 import io.netty.channel.embedded.EmbeddedChannel;
 
+import org.elasticsearch.common.bytes.BytesArray;
+import org.elasticsearch.common.bytes.BytesReference;
+import org.elasticsearch.common.bytes.CompositeBytesReference;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.test.ESTestCase;
@@ -28,6 +31,7 @@
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.lessThan;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
+import static org.hamcrest.Matchers.oneOf;
 
 public class Netty4WriteThrottlingHandlerTests extends ESTestCase {
 
@@ -56,42 +60,76 @@ public void testThrottlesLargeMessage() throws ExecutionException, InterruptedEx
         assertThat(writeableBytes, lessThan(Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE));
         final int fullSizeChunks = randomIntBetween(2, 10);
         final int extraChunkSize = randomIntBetween(0, 10);
-        final ByteBuf message = Unpooled.wrappedBuffer(
-            randomByteArrayOfLength(Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE * fullSizeChunks + extraChunkSize)
+        final byte[] messageBytes = randomByteArrayOfLength(
+            Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE * fullSizeChunks + extraChunkSize
         );
+        final Object message = wrapAsNettyOrEsBuffer(messageBytes);
         final ChannelPromise promise = embeddedChannel.newPromise();
         transportGroup.getLowLevelGroup().submit(() -> embeddedChannel.write(message, promise)).get();
         assertThat(seen, hasSize(1));
-        assertEquals(message.slice(0, Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE), seen.get(0));
+        assertSliceEquals(seen.get(0), message, 0, Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE);
         assertFalse(promise.isDone());
         transportGroup.getLowLevelGroup().submit(embeddedChannel::flush).get();
         assertTrue(promise.isDone());
         assertThat(seen, hasSize(fullSizeChunks + (extraChunkSize == 0 ? 0 : 1)));
         assertTrue(capturingHandler.didWriteAfterThrottled);
         if (extraChunkSize != 0) {
-            assertEquals(
-                message.slice(Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE * fullSizeChunks, extraChunkSize),
-                seen.get(seen.size() - 1)
+            assertSliceEquals(
+                seen.get(seen.size() - 1),
+                message,
+                Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE * fullSizeChunks,
+                extraChunkSize
             );
         }
     }
 
-    public void testPassesSmallMessageDirectly() throws ExecutionException, InterruptedException {
+    public void testThrottleLargeCompositeMessage() throws ExecutionException, InterruptedException {
         final List<ByteBuf> seen = new CopyOnWriteArrayList<>();
         final CapturingHandler capturingHandler = new CapturingHandler(seen);
         final EmbeddedChannel embeddedChannel = new EmbeddedChannel(
             capturingHandler,
             new Netty4WriteThrottlingHandler(new ThreadContext(Settings.EMPTY))
         );
+        // we assume that the channel outbound buffer is smaller than Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE
         final int writeableBytes = Math.toIntExact(embeddedChannel.bytesBeforeUnwritable());
         assertThat(writeableBytes, lessThan(Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE));
-        final ByteBuf message = Unpooled.wrappedBuffer(
-            randomByteArrayOfLength(randomIntBetween(0, Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE))
+        final int fullSizeChunks = randomIntBetween(2, 10);
+        final int extraChunkSize = randomIntBetween(0, 10);
+        final byte[] messageBytes = randomByteArrayOfLength(
+            Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE * fullSizeChunks + extraChunkSize
+        );
+        int splitOffset = randomIntBetween(0, messageBytes.length);
+        final BytesReference message = CompositeBytesReference.of(
+            new BytesArray(messageBytes, 0, splitOffset),
+            new BytesArray(messageBytes, splitOffset, messageBytes.length - splitOffset)
+        );
+        final ChannelPromise promise = embeddedChannel.newPromise();
+        transportGroup.getLowLevelGroup().submit(() -> embeddedChannel.write(message, promise)).get();
+        assertThat(seen, hasSize(oneOf(1, 2)));
+        assertSliceEquals(seen.get(0), message, 0, seen.get(0).readableBytes());
+        assertFalse(promise.isDone());
+        transportGroup.getLowLevelGroup().submit(embeddedChannel::flush).get();
+        assertTrue(promise.isDone());
+        assertThat(seen, hasSize(oneOf(fullSizeChunks, fullSizeChunks + 1)));
+        assertTrue(capturingHandler.didWriteAfterThrottled);
+        assertBufferEquals(Unpooled.compositeBuffer().addComponents(true, seen), message);
+    }
+
+    public void testPassesSmallMessageDirectly() throws ExecutionException, InterruptedException {
+        final List<ByteBuf> seen = new CopyOnWriteArrayList<>();
+        final CapturingHandler capturingHandler = new CapturingHandler(seen);
+        final EmbeddedChannel embeddedChannel = new EmbeddedChannel(
+            capturingHandler,
+            new Netty4WriteThrottlingHandler(new ThreadContext(Settings.EMPTY))
         );
+        final int writeableBytes = Math.toIntExact(embeddedChannel.bytesBeforeUnwritable());
+        assertThat(writeableBytes, lessThan(Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE));
+        final byte[] messageBytes = randomByteArrayOfLength(randomIntBetween(0, Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE));
+        final Object message = wrapAsNettyOrEsBuffer(messageBytes);
         final ChannelPromise promise = embeddedChannel.newPromise();
         transportGroup.getLowLevelGroup().submit(() -> embeddedChannel.write(message, promise)).get();
         assertThat(seen, hasSize(1)); // first message should be passed through straight away
-        assertSame(message, seen.get(0));
+        assertBufferEquals(seen.get(0), message);
         assertFalse(promise.isDone());
         transportGroup.getLowLevelGroup().submit(embeddedChannel::flush).get();
         assertTrue(promise.isDone());
@@ -107,13 +145,14 @@ public void testThrottlesOnUnwritable() throws ExecutionException, InterruptedEx
         );
         final int writeableBytes = Math.toIntExact(embeddedChannel.bytesBeforeUnwritable());
         assertThat(writeableBytes, lessThan(Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE));
-        final ByteBuf message = Unpooled.wrappedBuffer(randomByteArrayOfLength(writeableBytes + randomIntBetween(0, 10)));
+        final byte[] messageBytes = randomByteArrayOfLength(writeableBytes + randomIntBetween(0, 10));
+        final Object message = wrapAsNettyOrEsBuffer(messageBytes);
         final ChannelPromise promise = embeddedChannel.newPromise();
         transportGroup.getLowLevelGroup().submit(() -> embeddedChannel.write(message, promise)).get();
         assertThat(seen, hasSize(1)); // first message should be passed through straight away
-        assertSame(message, seen.get(0));
+        assertBufferEquals(seen.get(0), message);
         assertFalse(promise.isDone());
-        final ByteBuf messageToQueue = Unpooled.wrappedBuffer(
+        final Object messageToQueue = wrapAsNettyOrEsBuffer(
             randomByteArrayOfLength(randomIntBetween(0, Netty4WriteThrottlingHandler.MAX_BYTES_PER_WRITE))
         );
         final ChannelPromise promiseForQueued = embeddedChannel.newPromise();
@@ -126,6 +165,31 @@ public void testThrottlesOnUnwritable() throws ExecutionException, InterruptedEx
         assertTrue(promiseForQueued.isDone());
     }
 
+    private static void assertBufferEquals(ByteBuf expected, Object message) {
+        if (message instanceof ByteBuf buf) {
+            assertSame(expected, buf);
+        } else {
+            assertEquals(expected, Netty4Utils.toByteBuf(asInstanceOf(BytesReference.class, message)));
+        }
+    }
+
+    private static void assertSliceEquals(ByteBuf expected, Object message, int index, int length) {
+        assertEquals(
+            (message instanceof ByteBuf buf ? buf : Netty4Utils.toByteBuf(asInstanceOf(BytesReference.class, message))).slice(
+                index,
+                length
+            ),
+            expected
+        );
+    }
+
+    private static Object wrapAsNettyOrEsBuffer(byte[] messageBytes) {
+        if (randomBoolean()) {
+            return Unpooled.wrappedBuffer(messageBytes);
+        }
+        return new BytesArray(messageBytes);
+    }
+
     private static class CapturingHandler extends ChannelOutboundHandlerAdapter {
         private final List<ByteBuf> seen;
 

Original file line number	Diff line number	Diff line change
`@@ -165,7 +165,7 @@ public void sendMessage(BytesReference reference, ActionListener<Void> listener)`
`165`	`165`	`// We need to both guard against double resolving the listener and not resolving it in case of event loop shutdown so we need to`
`166`	`166`	`// use #notifyOnce here until https://github.com/netty/netty/issues/8007 is resolved.`
`167`	`167`	`var wrapped = ActionListener.notifyOnce(listener);`
`168`		`- channel.writeAndFlush(Netty4Utils.toByteBuf(reference), addPromise(wrapped, channel));`
	`168`	`+ channel.writeAndFlush(reference, addPromise(wrapped, channel));`
`169`	`169`	`if (channel.eventLoop().isShutdown()) {`
`170`	`170`	`wrapped.onFailure(new TransportException("Cannot send message, event loop is shutting down."));`
`171`	`171`	`}`