[#1177] improvement: Reduce the write time of tasks

client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java

@@ -24,10 +24,12 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
+import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 
@@ -93,6 +95,8 @@ public class RssShuffleWriter<K, V, C> extends ShuffleWriter<K, V> {
 
   protected final ShuffleWriteMetrics shuffleWriteMetrics;
 
+  private final BlockingQueue<Object> finishEventQueue = new LinkedBlockingQueue<>();
+
   // Only for tests
   @VisibleForTesting
   public RssShuffleWriter(
@@ -293,24 +297,52 @@ protected List<CompletableFuture<Long>> postBlockEvent(
       List<ShuffleBlockInfo> shuffleBlockInfoList) {
     List<CompletableFuture<Long>> futures = new ArrayList<>();
     for (AddBlockEvent event : bufferManager.buildBlockEvents(shuffleBlockInfoList)) {
+      event.addCallback(
+          () -> {
+            boolean ret = finishEventQueue.add(new Object());
+            if (!ret) {
+              LOG.error("Add event " + event + " to finishEventQueue fail");
+            }
+          });
       futures.add(shuffleManager.sendData(event));
     }
     return futures;
   }
 
   @VisibleForTesting
   protected void checkBlockSendResult(Set<Long> blockIds) {
-    long start = System.currentTimeMillis();
-    while (true) {
-      checkIfBlocksFailed();
-      Set<Long> successBlockIds = shuffleManager.getSuccessBlockIds(taskId);
-      blockIds.removeAll(successBlockIds);
-      if (blockIds.isEmpty()) {
-        break;
+    boolean interrupted = false;
+
+    try {
+      long remainingMs = sendCheckTimeout;
+      long end = System.currentTimeMillis() + remainingMs;
+
+      while (true) {
+        try {
+          finishEventQueue.clear();
+          checkIfBlocksFailed();
+          Set<Long> successBlockIds = shuffleManager.getSuccessBlockIds(taskId);
+          blockIds.removeAll(successBlockIds);
+          if (blockIds.isEmpty()) {
+            break;
+          }
+          LOG.info("Wait " + blockIds.size() + " blocks sent to shuffle server");
+          if (finishEventQueue.isEmpty()) {
+            Object event = finishEventQueue.poll(remainingMs, TimeUnit.MILLISECONDS);
+            if (event == null) {
+              break;
+            }
+          }
+        } catch (InterruptedException e) {
+          interrupted = true;
+        } finally {
+          remainingMs = end - System.currentTimeMillis();
+          if (remainingMs < 0) {
+            remainingMs = 0;
+          }
+        }
       }
-      LOG.info("Wait " + blockIds.size() + " blocks sent to shuffle server");
-      Uninterruptibles.sleepUninterruptibly(sendCheckInterval, TimeUnit.MILLISECONDS);
-      if (System.currentTimeMillis() - start > sendCheckTimeout) {
+      if (remainingMs <= 0) {
         String errorMsg =
             "Timeout: Task["
                 + taskId
@@ -322,6 +354,10 @@ protected void checkBlockSendResult(Set<Long> blockIds) {
         LOG.error(errorMsg);
         throw new RssException(errorMsg);
       }
+    } finally {
+      if (interrupted) {
+        Thread.currentThread().interrupt();
+      }
     }
   }