Change the Push block stream to not encode shuffle Id, map index, and…

… reduceId in a string

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/OneForOneBlockPusher.java

@@ -115,7 +115,14 @@ public void start() {
     for (int i = 0; i < blockIds.length; i++) {
       assert buffers.containsKey(blockIds[i]) : "Could not find the block buffer for block "
         + blockIds[i];
-      ByteBuffer header = new PushBlockStream(appId, blockIds[i], i).toByteBuffer();
+      String[] blockIdParts = blockIds[i].split("_");
+      if (blockIdParts.length != 4 || !blockIdParts[0].equals(
+        PushBlockStream.SHUFFLE_PUSH_BLOCK_PREFIX)) {
+        throw new IllegalArgumentException(
+          "Unexpected shuffle push block id format: " + blockIds[i]);
+      }
+      ByteBuffer header = new PushBlockStream(appId, Integer.parseInt(blockIdParts[1]),
+        Integer.parseInt(blockIdParts[2]), Integer.parseInt(blockIdParts[3]) , i).toByteBuffer();
       client.uploadStream(new NioManagedBuffer(header), buffers.get(blockIds[i]),
         new BlockPushCallback(i, blockIds[i]));
     }

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/RemoteBlockPushResolver.java

@@ -273,16 +273,9 @@ void deleteExecutorDirs(Path[] dirs) {
   @Override
   public StreamCallbackWithID receiveBlockDataAsStream(PushBlockStream msg) {
     // Retrieve merged shuffle file metadata
-    String[] blockIdParts = msg.blockId.split("_");
-    if (blockIdParts.length != 4 || !blockIdParts[0].equals("shuffle")) {
-      throw new IllegalArgumentException("Unexpected shuffle block id format: " + msg.blockId);
-    }
-    AppShuffleId appShuffleId = new AppShuffleId(msg.appId, Integer.parseInt(blockIdParts[1]));
-    int mapIndex = Integer.parseInt(blockIdParts[2]);
-    int reduceId = Integer.parseInt(blockIdParts[3]);
+    AppShuffleId appShuffleId = new AppShuffleId(msg.appId, msg.shuffleId);
     AppShufflePartitionInfo partitionInfoBeforeCheck =
-      getOrCreateAppShufflePartitionInfo(appShuffleId, reduceId);
-
+      getOrCreateAppShufflePartitionInfo(appShuffleId, msg.reduceId);
     // Here partitionInfo will be null in 2 cases:
     // 1) The request is received for a block that has already been merged, this is possible due
     // to the retry logic.
@@ -321,17 +314,18 @@ public StreamCallbackWithID receiveBlockDataAsStream(PushBlockStream msg) {
     final boolean isTooLate = partitionInfoBeforeCheck == null;
     // Check if the given block is already merged by checking the bitmap against the given map index
     final AppShufflePartitionInfo partitionInfo = partitionInfoBeforeCheck != null
-      && partitionInfoBeforeCheck.mapTracker.contains(mapIndex) ? null : partitionInfoBeforeCheck;
+      && partitionInfoBeforeCheck.mapTracker.contains(msg.mapIndex) ? null
+        : partitionInfoBeforeCheck;
     if (partitionInfo != null) {
       return new PushBlockStreamCallback(
-        this, msg, appShuffleId, reduceId, mapIndex, partitionInfo);
+        this, msg, appShuffleId, msg.reduceId, msg.mapIndex, partitionInfo);
     } else {
       // For a duplicate block or a block which is late, respond back with a callback that handles
       // them differently.
       return new StreamCallbackWithID() {
         @Override
         public String getID() {
-          return msg.blockId;
+          return msg.streamId;
         }
 
         @Override
@@ -345,7 +339,7 @@ public void onComplete(String streamId) {
           if (isTooLate) {
             // Throw an exception here so the block data is drained from channel and server
             // responds RpcFailure to the client.
-            throw new RuntimeException(String.format("Block %s %s", msg.blockId,
+            throw new RuntimeException(String.format("Block %s %s", msg.streamId,
               ErrorHandler.BlockPushErrorHandler.TOO_LATE_MESSAGE_SUFFIX));
           }
           // For duplicate block that is received before the shuffle merge finalizes, the
@@ -463,7 +457,7 @@ private PushBlockStreamCallback(
 
     @Override
     public String getID() {
-      return msg.blockId;
+      return msg.streamId;
     }
 
     /**
@@ -635,7 +629,7 @@ public void onComplete(String streamId) throws IOException {
         // however is already finalized. We should thus respond RpcFailure to the client.
         if (shufflePartitions == null || !shufflePartitions.containsKey(reduceId)) {
           deferredBufs = null;
-          throw new RuntimeException(String.format("Block %s %s", msg.blockId,
+          throw new RuntimeException(String.format("Block %s %s", msg.streamId,
             ErrorHandler.BlockPushErrorHandler.TOO_LATE_MESSAGE_SUFFIX));
         }
         // Check if we can commit this block
@@ -668,7 +662,7 @@ public void onComplete(String streamId) throws IOException {
           deferredBufs = null;
           throw new RuntimeException(String.format("%s %s to merged shuffle",
             ErrorHandler.BlockPushErrorHandler.BLOCK_APPEND_COLLISION_DETECTED_MSG_PREFIX,
-            msg.blockId));
+            msg.streamId));
         }
       }
       isWriting = false;

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/PushBlockStream.java

@@ -23,24 +23,31 @@
 import org.apache.spark.network.protocol.Encoders;
 
 // Needed by ScalaDoc. See SPARK-7726
-import static org.apache.spark.network.shuffle.protocol.BlockTransferMessage.Type;
 
 
 /**
  * Request to push a block to a remote shuffle service to be merged in push based shuffle.
  * The remote shuffle service will also include this message when responding the push requests.
  */
 public class PushBlockStream extends BlockTransferMessage {
+  public static final String SHUFFLE_PUSH_BLOCK_PREFIX = "shufflePush";
   public final String appId;
-  public final String blockId;
+  public final int shuffleId;
+  public final int mapIndex;
+  public final int reduceId;
   // Similar to the chunkIndex in StreamChunkId, indicating the index of a block in a batch of
   // blocks to be pushed.
   public final int index;
+  public final String streamId;
 
-  public PushBlockStream(String appId, String blockId, int index) {
+  public PushBlockStream(String appId, int shuffleId, int mapIndex, int reduceId, int index) {
     this.appId = appId;
-    this.blockId = blockId;
+    this.shuffleId = shuffleId;
+    this.mapIndex = mapIndex;
+    this.reduceId = reduceId;
     this.index = index;
+    this.streamId = String.format("%s_%d_%d_%d", SHUFFLE_PUSH_BLOCK_PREFIX,
+      shuffleId, mapIndex, reduceId);
   }
 
   @Override
@@ -50,14 +57,16 @@ protected Type type() {
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(appId, blockId, index);
+    return Objects.hashCode(appId, shuffleId, mapIndex , reduceId, index);
   }
 
   @Override
   public String toString() {
     return Objects.toStringHelper(this)
       .add("appId", appId)
-      .add("blockId", blockId)
+      .add("streamId", shuffleId)
+      .add("mapIndex", mapIndex)
+      .add("reduceId", reduceId)
       .add("index", index)
       .toString();
   }
@@ -67,29 +76,34 @@ public boolean equals(Object other) {
     if (other != null && other instanceof PushBlockStream) {
       PushBlockStream o = (PushBlockStream) other;
       return Objects.equal(appId, o.appId)
-        && Objects.equal(blockId, o.blockId)
+        && shuffleId == o.shuffleId
+        && mapIndex == o.mapIndex
+        && reduceId == o.reduceId
         && index == o.index;
     }
     return false;
   }
 
   @Override
   public int encodedLength() {
-    return Encoders.Strings.encodedLength(appId)
-      + Encoders.Strings.encodedLength(blockId) + 4;
+    return Encoders.Strings.encodedLength(appId) + 16;
   }
 
   @Override
   public void encode(ByteBuf buf) {
     Encoders.Strings.encode(buf, appId);
-    Encoders.Strings.encode(buf, blockId);
+    buf.writeInt(shuffleId);
+    buf.writeInt(mapIndex);
+    buf.writeInt(reduceId);
     buf.writeInt(index);
   }
 
   public static PushBlockStream decode(ByteBuf buf) {
     String appId = Encoders.Strings.decode(buf);
-    String blockId = Encoders.Strings.decode(buf);
+    int streamId = buf.readInt();
+    int mapIdx = buf.readInt();
+    int reduceId = buf.readInt();
     int index = buf.readInt();
-    return new PushBlockStream(appId, blockId, index);
+    return new PushBlockStream(appId, streamId, mapIdx, reduceId, index);
   }
 }

common/network-shuffle/src/test/java/org/apache/spark/network/shuffle/OneForOneBlockPusherSuite.java

@@ -45,77 +45,77 @@ public class OneForOneBlockPusherSuite {
   @Test
   public void testPushOne() {
     LinkedHashMap<String, ManagedBuffer> blocks = Maps.newLinkedHashMap();
-    blocks.put("shuffle_0_0_0", new NioManagedBuffer(ByteBuffer.wrap(new byte[1])));
+    blocks.put("shufflePush_0_0_0", new NioManagedBuffer(ByteBuffer.wrap(new byte[1])));
     String[] blockIds = blocks.keySet().toArray(new String[blocks.size()]);
 
     BlockFetchingListener listener = pushBlocks(
       blocks,
       blockIds,
-      Arrays.asList(new PushBlockStream("app-id", "shuffle_0_0_0", 0)));
+      Arrays.asList(new PushBlockStream("app-id", 0, 0, 0, 0)));
 
-    verify(listener).onBlockFetchSuccess(eq("shuffle_0_0_0"), any());
+    verify(listener).onBlockFetchSuccess(eq("shufflePush_0_0_0"), any());
   }
 
   @Test
   public void testPushThree() {
     LinkedHashMap<String, ManagedBuffer> blocks = Maps.newLinkedHashMap();
-    blocks.put("b0", new NioManagedBuffer(ByteBuffer.wrap(new byte[12])));
-    blocks.put("b1", new NioManagedBuffer(ByteBuffer.wrap(new byte[23])));
-    blocks.put("b2", new NettyManagedBuffer(Unpooled.wrappedBuffer(new byte[23])));
+    blocks.put("shufflePush_0_0_0", new NioManagedBuffer(ByteBuffer.wrap(new byte[12])));
+    blocks.put("shufflePush_0_1_0", new NioManagedBuffer(ByteBuffer.wrap(new byte[23])));
+    blocks.put("shufflePush_0_2_0", new NettyManagedBuffer(Unpooled.wrappedBuffer(new byte[23])));
     String[] blockIds = blocks.keySet().toArray(new String[blocks.size()]);
 
     BlockFetchingListener listener = pushBlocks(
       blocks,
       blockIds,
-      Arrays.asList(new PushBlockStream("app-id", "b0", 0),
-        new PushBlockStream("app-id", "b1", 1),
-        new PushBlockStream("app-id", "b2", 2)));
+      Arrays.asList(new PushBlockStream("app-id", 0, 0, 0, 0),
+        new PushBlockStream("app-id", 0, 1, 0, 1),
+        new PushBlockStream("app-id", 0, 2, 0, 2)));
 
-    for (int i = 0; i < 3; i ++) {
-      verify(listener, times(1)).onBlockFetchSuccess(eq("b" + i), any());
-    }
+    verify(listener, times(1)).onBlockFetchSuccess(eq("shufflePush_0_0_0"), any());
+    verify(listener, times(1)).onBlockFetchSuccess(eq("shufflePush_0_1_0"), any());
+    verify(listener, times(1)).onBlockFetchSuccess(eq("shufflePush_0_2_0"), any());
   }
 
   @Test
   public void testServerFailures() {
     LinkedHashMap<String, ManagedBuffer> blocks = Maps.newLinkedHashMap();
-    blocks.put("b0", new NioManagedBuffer(ByteBuffer.wrap(new byte[12])));
-    blocks.put("b1", new NioManagedBuffer(ByteBuffer.wrap(new byte[0])));
-    blocks.put("b2", new NioManagedBuffer(ByteBuffer.wrap(new byte[0])));
+    blocks.put("shufflePush_0_0_0", new NioManagedBuffer(ByteBuffer.wrap(new byte[12])));
+    blocks.put("shufflePush_0_1_0", new NioManagedBuffer(ByteBuffer.wrap(new byte[0])));
+    blocks.put("shufflePush_0_2_0", new NioManagedBuffer(ByteBuffer.wrap(new byte[0])));
     String[] blockIds = blocks.keySet().toArray(new String[blocks.size()]);
 
     BlockFetchingListener listener = pushBlocks(
       blocks,
       blockIds,
-      Arrays.asList(new PushBlockStream("app-id", "b0", 0),
-        new PushBlockStream("app-id", "b1", 1),
-        new PushBlockStream("app-id", "b2", 2)));
+      Arrays.asList(new PushBlockStream("app-id", 0, 0, 0, 0),
+        new PushBlockStream("app-id", 0, 1, 0, 1),
+        new PushBlockStream("app-id", 0, 2, 0, 2)));
 
-    verify(listener, times(1)).onBlockFetchSuccess(eq("b0"), any());
-    verify(listener, times(1)).onBlockFetchFailure(eq("b1"), any());
-    verify(listener, times(1)).onBlockFetchFailure(eq("b2"), any());
+    verify(listener, times(1)).onBlockFetchSuccess(eq("shufflePush_0_0_0"), any());
+    verify(listener, times(1)).onBlockFetchFailure(eq("shufflePush_0_1_0"), any());
+    verify(listener, times(1)).onBlockFetchFailure(eq("shufflePush_0_2_0"), any());
   }
 
   @Test
   public void testHandlingRetriableFailures() {
     LinkedHashMap<String, ManagedBuffer> blocks = Maps.newLinkedHashMap();
-    blocks.put("b0", new NioManagedBuffer(ByteBuffer.wrap(new byte[12])));
-    blocks.put("b1", null);
-    blocks.put("b2", new NioManagedBuffer(ByteBuffer.wrap(new byte[0])));
+    blocks.put("shufflePush_0_0_0", new NioManagedBuffer(ByteBuffer.wrap(new byte[12])));
+    blocks.put("shufflePush_0_1_0", null);
+    blocks.put("shufflePush_0_2_0", new NioManagedBuffer(ByteBuffer.wrap(new byte[0])));
     String[] blockIds = blocks.keySet().toArray(new String[blocks.size()]);
 
     BlockFetchingListener listener = pushBlocks(
       blocks,
       blockIds,
-      Arrays.asList(new PushBlockStream("app-id", "b0", 0),
-        new PushBlockStream("app-id", "b1", 1),
-        new PushBlockStream("app-id", "b2", 2)));
-
-    verify(listener, times(1)).onBlockFetchSuccess(eq("b0"), any());
-    verify(listener, times(0)).onBlockFetchSuccess(not(eq("b0")), any());
-    verify(listener, times(0)).onBlockFetchFailure(eq("b0"), any());
-    verify(listener, times(1)).onBlockFetchFailure(eq("b1"), any());
-    verify(listener, times(2)).onBlockFetchFailure(eq("b2"), any());
+      Arrays.asList(new PushBlockStream("app-id", 0, 0, 0, 0),
+        new PushBlockStream("app-id", 0, 1, 0, 1),
+        new PushBlockStream("app-id", 0, 2, 0, 2)));
+
+    verify(listener, times(1)).onBlockFetchSuccess(eq("shufflePush_0_0_0"), any());
+    verify(listener, times(0)).onBlockFetchSuccess(not(eq("shufflePush_0_0_0")), any());
+    verify(listener, times(0)).onBlockFetchFailure(eq("shufflePush_0_0_0"), any());
+    verify(listener, times(1)).onBlockFetchFailure(eq("shufflePush_0_1_0"), any());
+    verify(listener, times(2)).onBlockFetchFailure(eq("shufflePush_0_2_0"), any());
   }
 
   /**