spark 2.1 Encoders and UploadBlock

/** Provides a canonical set of Encoders for simple types. */public class Encoders {  /** Strings are encoded with their length followed by UTF-8 bytes. */  public static class Strings {    public static int encodedLength(String s) {      return 4 + s.getBytes(StandardCharsets.UTF_8).length;    }    public static void encode(ByteBuf buf, String s) {      byte[] bytes = s.getBytes(StandardCharsets.UTF_8);      buf.writeInt(bytes.length);      buf.writeBytes(bytes);    }    public static String decode(ByteBuf buf) {      int length = buf.readInt();      byte[] bytes = new byte[length];      buf.readBytes(bytes);      return new String(bytes, StandardCharsets.UTF_8);    }  }  /** Byte arrays are encoded with their length followed by bytes. */  public static class ByteArrays {    public static int encodedLength(byte[] arr) {      return 4 + arr.length;    }    public static void encode(ByteBuf buf, byte[] arr) {      buf.writeInt(arr.length);      buf.writeBytes(arr);    }    public static byte[] decode(ByteBuf buf) {      int length = buf.readInt();      byte[] bytes = new byte[length];      buf.readBytes(bytes);      return bytes;    }  }  /** String arrays are encoded with the number of strings followed by per-String encoding. */  public static class StringArrays {    public static int encodedLength(String[] strings) {      int totalLength = 4;      for (String s : strings) {        totalLength += Strings.encodedLength(s);      }      return totalLength;    }    public static void encode(ByteBuf buf, String[] strings) {      buf.writeInt(strings.length);      for (String s : strings) {        Strings.encode(buf, s);      }    }    public static String[] decode(ByteBuf buf) {      int numStrings = buf.readInt();      String[] strings = new String[numStrings];      for (int i = 0; i < strings.length; i ++) {        strings[i] = Strings.decode(buf);      }      return strings;    }  }}

Encodable

/** * Interface for an object which can be encoded into a ByteBuf. Multiple Encodable objects are * stored in a single, pre-allocated ByteBuf, so Encodables must also provide their length. * * Encodable objects should provide a static "decode(ByteBuf)" method which is invoked by * {@link MessageDecoder}. During decoding, if the object uses the ByteBuf as its data (rather than * just copying data from it), then you must retain() the ByteBuf. * * Additionally, when adding a new Encodable Message, add it to {@link Message.Type}. */public interface Encodable {  /** Number of bytes of the encoded form of this object. */  int encodedLength();  /**   * Serializes this object by writing into the given ByteBuf.   * This method must write exactly encodedLength() bytes.   */  void encode(ByteBuf buf);}

BlockTransferMessage

/** * Messages handled by the {@link org.apache.spark.network.shuffle.ExternalShuffleBlockHandler}, or * by Spark's NettyBlockTransferService. * * At a high level: *   - OpenBlock is handled by both services, but only services shuffle files for the external *     shuffle service. It returns a StreamHandle. *   - UploadBlock is only handled by the NettyBlockTransferService. *   - RegisterExecutor is only handled by the external shuffle service. */public abstract class BlockTransferMessage implements Encodable {  protected abstract Type type();  /** Preceding every serialized message is its type, which allows us to deserialize it. */  public enum Type {    OPEN_BLOCKS(0), UPLOAD_BLOCK(1), REGISTER_EXECUTOR(2), STREAM_HANDLE(3), REGISTER_DRIVER(4),    HEARTBEAT(5);    private final byte id;    Type(int id) {      assert id < 128 : "Cannot have more than 128 message types";      this.id = (byte) id;    }    public byte id() { return id; }  }  // NB: Java does not support static methods in interfaces, so we must put this in a static class.  public static class Decoder {    /** Deserializes the 'type' byte followed by the message itself. */    public static BlockTransferMessage fromByteBuffer(ByteBuffer msg) {      ByteBuf buf = Unpooled.wrappedBuffer(msg);      byte type = buf.readByte();      switch (type) {        case 0: return OpenBlocks.decode(buf);        case 1: return UploadBlock.decode(buf);        case 2: return RegisterExecutor.decode(buf);        case 3: return StreamHandle.decode(buf);        case 4: return RegisterDriver.decode(buf);        case 5: return ShuffleServiceHeartbeat.decode(buf);        default: throw new IllegalArgumentException("Unknown message type: " + type);      }    }  }  /** Serializes the 'type' byte followed by the message itself. */  public ByteBuffer toByteBuffer() {    // 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();    return buf.nioBuffer();  }}

UploadBlock

/** Request to upload a block with a certain StorageLevel. Returns nothing (empty byte array). */public class UploadBlock extends BlockTransferMessage {  public final String appId;  public final String execId;  public final String blockId;  // TODO: StorageLevel is serialized separately in here because StorageLevel is not available in  // this package. We should avoid this hack.  public final byte[] metadata;  public final byte[] blockData;  /**   * @param metadata Meta-information about block, typically StorageLevel.   * @param blockData The actual block's bytes.   */  public UploadBlock(      String appId,      String execId,      String blockId,      byte[] metadata,      byte[] blockData) {    this.appId = appId;    this.execId = execId;    this.blockId = blockId;    this.metadata = metadata;    this.blockData = blockData;  }  @Override  protected Type type() { return Type.UPLOAD_BLOCK; }  @Override  public int hashCode() {    int objectsHashCode = Objects.hashCode(appId, execId, blockId);    return (objectsHashCode * 41 + Arrays.hashCode(metadata)) * 41 + Arrays.hashCode(blockData);  }  @Override  public String toString() {    return Objects.toStringHelper(this)      .add("appId", appId)      .add("execId", execId)      .add("blockId", blockId)      .add("metadata size", metadata.length)      .add("block size", blockData.length)      .toString();  }  @Override  public boolean equals(Object other) {    if (other != null && other instanceof UploadBlock) {      UploadBlock o = (UploadBlock) other;      return Objects.equal(appId, o.appId)        && Objects.equal(execId, o.execId)        && Objects.equal(blockId, o.blockId)        && Arrays.equals(metadata, o.metadata)        && Arrays.equals(blockData, o.blockData);    }    return false;  }  @Override  public int encodedLength() {    return Encoders.Strings.encodedLength(appId)      + Encoders.Strings.encodedLength(execId)      + Encoders.Strings.encodedLength(blockId)      + Encoders.ByteArrays.encodedLength(metadata)      + Encoders.ByteArrays.encodedLength(blockData);  }  @Override  public void encode(ByteBuf buf) {    Encoders.Strings.encode(buf, appId);    Encoders.Strings.encode(buf, execId);    Encoders.Strings.encode(buf, blockId);    Encoders.ByteArrays.encode(buf, metadata);    Encoders.ByteArrays.encode(buf, blockData);  }  public static UploadBlock decode(ByteBuf buf) {    String appId = Encoders.Strings.decode(buf);    String execId = Encoders.Strings.decode(buf);    String blockId = Encoders.Strings.decode(buf);    byte[] metadata = Encoders.ByteArrays.decode(buf);    byte[] blockData = Encoders.ByteArrays.decode(buf);    return new UploadBlock(appId, execId, blockId, metadata, blockData);  }}