review commits

Author: witgo

common/network-common/src/main/java/org/apache/spark/network/sasl/SaslEncryption.java

@@ -31,12 +31,12 @@
 import io.netty.channel.ChannelOutboundHandlerAdapter;
 import io.netty.channel.ChannelPromise;
 import io.netty.channel.FileRegion;
-import io.netty.handler.codec.ByteToMessageDecoder;
-import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
 import io.netty.handler.codec.MessageToMessageDecoder;
 import io.netty.util.AbstractReferenceCounted;
 
 import org.apache.spark.network.util.ByteArrayWritableChannel;
+import org.apache.spark.network.util.TransportFrameDecoder;
+
 /**
  * Provides SASL-based encription for transport channels. The single method exposed by this
  * class installs the needed channel handlers on a connected channel.
@@ -61,21 +61,12 @@ static void addToChannel(
     channel.pipeline()
       .addFirst(ENCRYPTION_HANDLER_NAME, new EncryptionHandler(backend, maxOutboundBlockSize))
       .addFirst("saslDecryption", new DecryptionHandler(backend))
-      // Each frame does not exceed 8 + maxOutboundBlockSize bytes
       .addFirst("saslFrameDecoder", createFrameDecoder());
   }
 
-  /**
-   * Creates a LengthFieldBasedFrameDecoder where the first 8 bytes are the length of the frame.
-   * This is used before all decoders.
-   */
-  static ByteToMessageDecoder createFrameDecoder() {
-    // maxFrameLength = 2G
-    // lengthFieldOffset = 0
-    // lengthFieldLength = 8
-    // lengthAdjustment = -8, i.e. exclude the 8 byte length itself
-    // initialBytesToStrip = 8, i.e. strip out the length field itself
-    return new LengthFieldBasedFrameDecoder(Integer.MAX_VALUE, 0, 8, -8, 8);
+  // Each frame does not exceed 8 + maxOutboundBlockSize bytes
+  private static TransportFrameDecoder createFrameDecoder() {
+    return new TransportFrameDecoder(false);
   }
 
   private static class EncryptionHandler extends ChannelOutboundHandlerAdapter {

common/network-common/src/main/java/org/apache/spark/network/util/TransportFrameDecoder.java

@@ -46,15 +46,25 @@ public class TransportFrameDecoder extends ChannelInboundHandlerAdapter {
 
   public static final String HANDLER_NAME = "frameDecoder";
   private static final int LENGTH_SIZE = 8;
+  private static final int MAX_FRAME_SIZE = Integer.MAX_VALUE;
   private static final int UNKNOWN_FRAME_SIZE = -1;
 
   private final LinkedList<ByteBuf> buffers = new LinkedList<>();
   private final ByteBuf frameLenBuf = Unpooled.buffer(LENGTH_SIZE, LENGTH_SIZE);
+  private final boolean isSupportLargeData;
 
   private long totalSize = 0;
   private long nextFrameSize = UNKNOWN_FRAME_SIZE;
   private volatile Interceptor interceptor;
 
+  public TransportFrameDecoder() {
+    this(true);
+  }
+
+  public TransportFrameDecoder(boolean isSupportLargeData) {
+    this.isSupportLargeData = isSupportLargeData;
+  }
+
   @Override
   public void channelRead(ChannelHandlerContext ctx, Object data) throws Exception {
     ByteBuf in = (ByteBuf) data;
@@ -77,7 +87,13 @@ public void channelRead(ChannelHandlerContext ctx, Object data) throws Exception
         totalSize -= read;
       } else {
         // Interceptor is not active, so try to decode one frame.
-        LinkedList<ByteBuf> frame = decodeNext();
+        Object frame ;
+        if (isSupportLargeData) {
+          frame = decodeList();
+        } else {
+          frame = decodeByteBuf();
+        }
+
         if (frame == null) {
           break;
         }
@@ -120,7 +136,36 @@ private long decodeFrameSize() {
     return nextFrameSize;
   }
 
-  private LinkedList<ByteBuf> decodeNext() throws Exception {
+  private ByteBuf decodeByteBuf() throws Exception {
+    long frameSize = decodeFrameSize();
+    if (frameSize == UNKNOWN_FRAME_SIZE || totalSize < frameSize) {
+      return null;
+    }
+
+    // Reset size for next frame.
+    nextFrameSize = UNKNOWN_FRAME_SIZE;
+
+    Preconditions.checkArgument(frameSize < MAX_FRAME_SIZE, "Too large frame: %s", frameSize);
+    Preconditions.checkArgument(frameSize > 0, "Frame length should be positive: %s", frameSize);
+
+    // If the first buffer holds the entire frame, return it.
+    int remaining = (int) frameSize;
+    if (buffers.getFirst().readableBytes() >= remaining) {
+      return nextBufferForFrame(remaining);
+    }
+
+    // Otherwise, create a composite buffer.
+    CompositeByteBuf frame = buffers.getFirst().alloc().compositeBuffer(Integer.MAX_VALUE);
+    while (remaining > 0) {
+      ByteBuf next = nextBufferForFrame(remaining);
+      remaining -= next.readableBytes();
+      frame.addComponent(next).writerIndex(frame.writerIndex() + next.readableBytes());
+    }
+    assert remaining == 0;
+    return frame;
+  }
+
+  private LinkedList<ByteBuf> decodeList() throws Exception {
     long frameSize = decodeFrameSize();
     if (frameSize == UNKNOWN_FRAME_SIZE || totalSize < frameSize) {
       return null;