Review comments: simplify some code.

Author: Marcelo Vanzin

network/common/src/main/java/org/apache/spark/network/protocol/MessageEncoder.java

@@ -73,7 +73,7 @@ public void encode(ChannelHandlerContext ctx, Message in, List<Object> out) {
     assert header.writableBytes() == 0;
 
     if (body != null && bodyLength > 0) {
-      out.add(new MessageWithHeader(header, headerLength, body, bodyLength));
+      out.add(new MessageWithHeader(header, body, bodyLength));
     } else {
       out.add(header);
     }

network/common/src/main/java/org/apache/spark/network/protocol/MessageWithHeader.java

@@ -39,11 +39,11 @@ class MessageWithHeader extends AbstractReferenceCounted implements FileRegion {
   private final long bodyLength;
   private int totalBytesTransferred;
 
-  MessageWithHeader(ByteBuf header, int headerLength, Object body, long bodyLength) {
+  MessageWithHeader(ByteBuf header, Object body, long bodyLength) {
     Preconditions.checkArgument(body instanceof ByteBuf || body instanceof FileRegion,
       "Body must be a ByteBuf or a FileRegion.");
     this.header = header;
-    this.headerLength = headerLength;
+    this.headerLength = header.readableBytes();
     this.body = body;
     this.bodyLength = bodyLength;
   }
@@ -65,22 +65,26 @@ public long transfered() {
 
   @Override
   public long transferTo(WritableByteChannel target, long position) throws IOException {
-    Preconditions.checkArgument(position >= 0 && position < count(), "Invalid position.");
+    Preconditions.checkArgument(position == totalBytesTransferred, "Invalid position.");
     long written = 0;
 
     if (position < headerLength) {
-      written += copyByteBuf(header, target, position);
+      written += copyByteBuf(header, target);
       if (header.readableBytes() > 0) {
         totalBytesTransferred += written;
         return written;
       }
     }
 
     if (body instanceof FileRegion) {
+      // Adjust the position. If the write is happening as part of the same call where the header
+      // (or some part of it) is written, `position` will be less than the header size, so we want
+      // to start from position 0 in the FileRegion object. Otherwise, we start from the position
+      // requested by the caller.
       long bodyPos = position > headerLength ? position - headerLength : 0;
       written += ((FileRegion)body).transferTo(target, bodyPos);
     } else if (body instanceof ByteBuf) {
-      written += copyByteBuf((ByteBuf) body, target, position);
+      written += copyByteBuf((ByteBuf) body, target);
     }
 
     totalBytesTransferred += written;
@@ -93,12 +97,7 @@ protected void deallocate() {
     ReferenceCountUtil.release(body);
   }
 
-  private int copyByteBuf(ByteBuf buf, WritableByteChannel target, long position)
-    throws IOException {
-
-    if (position > totalBytesTransferred) {
-      buf.skipBytes(Ints.checkedCast(position - totalBytesTransferred));
-    }
+  private int copyByteBuf(ByteBuf buf, WritableByteChannel target) throws IOException {
     int written = target.write(buf.nioBuffer());
     buf.skipBytes(written);
     return written;

network/common/src/test/java/org/apache/spark/network/protocol/MessageWithHeaderSuite.java

@@ -47,8 +47,7 @@ public void testShortWrite() throws Exception {
   public void testByteBufBody() throws Exception {
     ByteBuf header = Unpooled.copyLong(42);
     ByteBuf body = Unpooled.copyLong(84);
-    MessageWithHeader msg = new MessageWithHeader(header, header.readableBytes(), body,
-      body.readableBytes());
+    MessageWithHeader msg = new MessageWithHeader(header, body, body.readableBytes());
 
     ByteBuf result = doWrite(msg, 1);
     assertEquals(msg.count(), result.readableBytes());
@@ -60,7 +59,7 @@ private void testFileRegionBody(int totalWrites, int writesPerCall) throws Excep
     ByteBuf header = Unpooled.copyLong(42);
     int headerLength = header.readableBytes();
     TestFileRegion region = new TestFileRegion(totalWrites, writesPerCall);
-    MessageWithHeader msg = new MessageWithHeader(header, headerLength, region, region.count());
+    MessageWithHeader msg = new MessageWithHeader(header, region, region.count());
 
     ByteBuf result = doWrite(msg, totalWrites / writesPerCall);
     assertEquals(headerLength + region.count(), result.readableBytes());