SPARK-4805 [CORE] BlockTransferMessage.toByteArray() trips assertion

Allocate enough room for type byte as well as message, to avoid tripping assertion about capacity of the buffer

Author: Sean Owen <sowen@cloudera.com>

Closes #3650 from srowen/SPARK-4805 and squashes the following commits:

9e1d502 [Sean Owen] Allocate enough room for type byte as well as message, to avoid tripping assertion about capacity of the buffer

(cherry picked from commit d8f84f26e388055ca7459810e001d05ab60af15b)
Signed-off-by: Aaron Davidson <aaron@databricks.com>

1 files changed, 2 insertions, 1 deletions

diff --git a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/BlockTransferMessage.java b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/BlockTransferMessage.java
index b4b13b8a6e..6c1210b332 100644
--- a/network/shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/BlockTransferMessage.java
+++ b/network/shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/BlockTransferMessage.java
@@ -67,7 +67,8 @@ public abstract class BlockTransferMessage implements Encodable {
 
   /** Serializes the 'type' byte followed by the message itself. */
   public byte[] toByteArray() {
-    ByteBuf buf = Unpooled.buffer(encodedLength());
+    // Allow room for encoded message, plus the type byte
+    ByteBuf buf = Unpooled.buffer(encodedLength() + 1);
     buf.writeByte(type().id);
     encode(buf);
     assert buf.writableBytes() == 0 : "Writable bytes remain: " + buf.writableBytes();