[apache#1464] improvement(client): Improve the error log message for …

…checkBlockSendResult

client-spark/spark2/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java

@@ -24,6 +24,7 @@
 import java.util.Set;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.LinkedBlockingDeque;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 
@@ -33,6 +34,7 @@
 import scala.collection.Seq;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.spark.ShuffleDependency;
@@ -658,6 +660,14 @@ public void addFailedBlockIds(String taskId, Set<Long> blockIds) {
     taskToFailedBlockIds.get(taskId).addAll(blockIds);
   }
 
+  @VisibleForTesting
+  public void addTaskToFailedBlockIdsAndServer(
+      String taskId, Long blockId, ShuffleServerInfo shuffleServerInfo) {
+    taskToFailedBlockIdsAndServer.putIfAbsent(taskId, Maps.newHashMap());
+    taskToFailedBlockIdsAndServer.get(taskId).putIfAbsent(blockId, new LinkedBlockingDeque<>());
+    taskToFailedBlockIdsAndServer.get(taskId).get(blockId).add(shuffleServerInfo);
+  }
+
   @VisibleForTesting
   public void addSuccessBlockIds(String taskId, Set<Long> blockIds) {
     if (taskToSuccessBlockIds.get(taskId) == null) {

client-spark/spark2/src/test/java/org/apache/spark/shuffle/writer/RssShuffleWriterTest.java

@@ -146,6 +146,8 @@ public void checkBlockSendResultTest() {
     // case 3: partial blocks are sent failed, Runtime exception will be thrown
     manager.addSuccessBlockIds(taskId, Sets.newHashSet(1L, 2L));
     manager.addFailedBlockIds(taskId, Sets.newHashSet(3L));
+    ShuffleServerInfo shuffleServerInfo = new ShuffleServerInfo("127.0.0.1", 20001);
+    manager.addTaskToFailedBlockIdsAndServer(taskId, 3L, shuffleServerInfo);
     Throwable e3 =
         assertThrows(
             RuntimeException.class,

client-spark/spark3/src/test/java/org/apache/spark/shuffle/writer/RssShuffleWriterTest.java

@@ -26,6 +26,7 @@
 import java.util.Set;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.LinkedBlockingQueue;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -85,10 +86,12 @@ public void checkBlockSendResultTest() {
         .set(RssSparkConfig.RSS_COORDINATOR_QUORUM.key(), "127.0.0.1:12345,127.0.0.1:12346");
     Map<String, Set<Long>> failBlocks = JavaUtils.newConcurrentMap();
     Map<String, Set<Long>> successBlocks = JavaUtils.newConcurrentMap();
+    Map<String, Map<Long, BlockingQueue<ShuffleServerInfo>>> taskToFailedBlockIdsAndServer =
+        JavaUtils.newConcurrentMap();
     Serializer kryoSerializer = new KryoSerializer(conf);
     RssShuffleManager manager =
         TestUtils.createShuffleManager(
-            conf, false, null, successBlocks, failBlocks, JavaUtils.newConcurrentMap());
+            conf, false, null, successBlocks, failBlocks, taskToFailedBlockIdsAndServer);
 
     ShuffleWriteClient mockShuffleWriteClient = mock(ShuffleWriteClient.class);
     Partitioner mockPartitioner = mock(Partitioner.class);
@@ -149,6 +152,13 @@ public void checkBlockSendResultTest() {
     // case 3: partial blocks are sent failed, Runtime exception will be thrown
     successBlocks.put("taskId", Sets.newHashSet(1L, 2L));
     failBlocks.put("taskId", Sets.newHashSet(3L));
+    Map<Long, BlockingQueue<ShuffleServerInfo>> blockIdToShuffleServerInfoMap =
+        JavaUtils.newConcurrentMap();
+    BlockingQueue blockingQueue = new LinkedBlockingQueue<>();
+    ShuffleServerInfo shuffleServerInfo = new ShuffleServerInfo("127.0.0.1", 20001);
+    blockingQueue.add(shuffleServerInfo);
+    blockIdToShuffleServerInfoMap.put(3L, blockingQueue);
+    taskToFailedBlockIdsAndServer.put("taskId", blockIdToShuffleServerInfoMap);
     Throwable e3 =
         assertThrows(
             RuntimeException.class,