diff --git a/common/src/main/java/org/apache/uniffle/common/netty/buffer/NettyManagedBuffer.java b/common/src/main/java/org/apache/uniffle/common/netty/buffer/NettyManagedBuffer.java index 547a4e891a..6d78d610fc 100644 --- a/common/src/main/java/org/apache/uniffle/common/netty/buffer/NettyManagedBuffer.java +++ b/common/src/main/java/org/apache/uniffle/common/netty/buffer/NettyManagedBuffer.java @@ -20,8 +20,11 @@ import java.nio.ByteBuffer; import io.netty.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; import io.netty.buffer.Unpooled; +import org.apache.uniffle.common.util.NettyUtils; + public class NettyManagedBuffer extends ManagedBuffer { public static final NettyManagedBuffer EMPTY_BUFFER = @@ -45,6 +48,22 @@ public ByteBuf byteBuf() { @Override public ByteBuffer nioByteBuffer() { + // CompositeByteBuf.nioBuffer will return a heap buffer if the composite buffer has more than + // one component, even if all components are direct buffers. In native client scenarios + // (like gluten), we prefer to use direct buffer to reduce data copying. + if (buf instanceof CompositeByteBuf + && buf.isDirect() + && buf.nioBufferCount() > 1 + && NettyUtils.preferDirectForCompositeBuffer()) { + int length = buf.readableBytes(); + ByteBuffer merged = ByteBuffer.allocateDirect(length).order(buf.order()); + for (ByteBuffer buf : buf.nioBuffers()) { + merged.put(buf); + } + merged.flip(); + return merged; + } + return buf.nioBuffer(); } diff --git a/common/src/main/java/org/apache/uniffle/common/util/NettyUtils.java b/common/src/main/java/org/apache/uniffle/common/util/NettyUtils.java index 59668b3b1b..d8545ba403 100644 --- a/common/src/main/java/org/apache/uniffle/common/util/NettyUtils.java +++ b/common/src/main/java/org/apache/uniffle/common/util/NettyUtils.java @@ -17,6 +17,7 @@ package org.apache.uniffle.common.util; +import java.util.Optional; import java.util.concurrent.ThreadFactory; import java.util.concurrent.atomic.AtomicReferenceArray; @@ -179,4 +180,25 @@ public static UnpooledByteBufAllocator createUnpooledByteBufAllocator(boolean pr public static long getMaxDirectMemory() { return MAX_DIRECT_MEMORY_IN_BYTES; } + + private static final String PREFER_DIRECT_FOR_COMPOSITE_BUFFER_PROPERTY_KEY = + "rss.netty.preferDirectForCompositeBuffer"; + private static final String PREFER_DIRECT_FOR_COMPOSITE_BUFFER_ENV_KEY = + "RSS_NETTY_PREFER_DIRECT_FOR_COMPOSITE_BUFFER"; + private static final boolean PREFER_DIRECT_FOR_COMPOSITE_BUFFER_DEFAULT = false; + private static final boolean _preferDirectForCompositeBuffer; + + static { + _preferDirectForCompositeBuffer = + Optional.ofNullable(System.getProperty(PREFER_DIRECT_FOR_COMPOSITE_BUFFER_PROPERTY_KEY)) + .map(Boolean::new) + .orElse( + Optional.ofNullable(System.getenv(PREFER_DIRECT_FOR_COMPOSITE_BUFFER_ENV_KEY)) + .map(Boolean::new) + .orElse(PREFER_DIRECT_FOR_COMPOSITE_BUFFER_DEFAULT)); + } + + public static boolean preferDirectForCompositeBuffer() { + return _preferDirectForCompositeBuffer; + } }