[#2646] feat(netty): Support converting CompositeByteBuf to a direct nio buffer

Author: wForget

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,24 @@ 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();
   }
 

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;
+  }
 }