fix

Author: summaryzb

client-spark/common/src/main/java/org/apache/spark/shuffle/writer/DataPusher.java

@@ -118,7 +118,7 @@ public Future<Long> send(AddBlockEvent event) {
                   .filter(x -> succeedBlockIds.contains(x.getBlockId()))
                   .map(x -> x.getFreeMemory())
                   .reduce((a, b) -> a + b)
-                  .get();
+                  .orElseGet(() -> 0L);
             });
     event.doPrepare(future);
     executorService.submit(future);

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

@@ -360,8 +360,7 @@ private void checkSentBlockCount() {
    *
    * @param shuffleBlockInfoList
    */
-  private List<Future<Long>> processShuffleBlockInfos(
-      List<ShuffleBlockInfo> shuffleBlockInfoList) {
+  private List<Future<Long>> processShuffleBlockInfos(List<ShuffleBlockInfo> shuffleBlockInfoList) {
     if (shuffleBlockInfoList != null && !shuffleBlockInfoList.isEmpty()) {
       shuffleBlockInfoList.stream()
           .forEach(
@@ -389,8 +388,7 @@ private List<Future<Long>> processShuffleBlockInfos(
 
   // don't send huge block to shuffle server, or there will be OOM if shuffle sever receives data
   // more than expected
-  protected List<Future<Long>> postBlockEvent(
-      List<ShuffleBlockInfo> shuffleBlockInfoList) {
+  protected List<Future<Long>> postBlockEvent(List<ShuffleBlockInfo> shuffleBlockInfoList) {
     List<Future<Long>> futures = new ArrayList<>();
     for (AddBlockEvent event : bufferManager.buildBlockEvents(shuffleBlockInfoList)) {
       futures.add(shuffleManager.sendData(event));

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

@@ -26,10 +26,13 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -474,6 +477,7 @@ protected void checkBlockSendResult(Set<Long> blockIds) {
 
       while (!interrupted) {
         try {
+          LOG.warn("checkBlockSendResult," + blockIds.size() + "," + inFlightEvent.size());
           checkDataIfAnyFailure();
           Set<Long> successBlockIds = shuffleManager.getSuccessBlockIds(taskId);
           blockIds.removeAll(successBlockIds);
@@ -492,9 +496,10 @@ protected void checkBlockSendResult(Set<Long> blockIds) {
             inFlightEvent.values().stream()
                 .filter(f -> !f.isDone())
                 .findAny()
-                .get()
+                .orElseGet(() -> CompletableFuture.completedFuture(0L))
                 .get(remainingMs, TimeUnit.MILLISECONDS);
           } else {
+            LOG.warn("blockSize:" + blockIds.size() + ",inflightEvent:" + inFlightEvent.size());
             // it seems never reach here, since `blockIds.isEmpty()` will break the loop first
             break;
           }
@@ -503,8 +508,9 @@ protected void checkBlockSendResult(Set<Long> blockIds) {
           interrupted = true;
           inFlightEvent.values().stream().forEach(f -> f.cancel(true));
           Thread.currentThread().interrupt();
-        } catch (Exception e) {
-          LOG.error("Exception happened when check block send result", e);
+        } catch (ExecutionException | TimeoutException e) {
+          LOG.error("check err", e);
+          break;
         }
       }
       if (!blockIds.isEmpty()) {

client/src/main/java/org/apache/uniffle/client/impl/ShuffleWriteClientImpl.java

@@ -249,6 +249,7 @@ private boolean sendShuffleDataAsync(
                 }
                 return true;
               });
+      dataTransferPool.submit(future);
       futures.add(future);
     }
 

client/src/main/java/org/apache/uniffle/client/util/ClientUtils.java

@@ -24,14 +24,19 @@
 import java.util.Set;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
 
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import org.apache.uniffle.client.api.ShuffleWriteClient;
 import org.apache.uniffle.common.ClientType;
 import org.apache.uniffle.common.RemoteStorageInfo;
 import org.apache.uniffle.storage.util.StorageType;
 
 public class ClientUtils {
+  private static final Logger LOG = LoggerFactory.getLogger(ClientUtils.class);
 
   public static RemoteStorageInfo fetchRemoteStorage(
       String appId,
@@ -65,11 +70,12 @@ public static boolean waitUntilDoneOrFail(List<Future<Boolean>> list, boolean al
       while (iterator.hasNext()) {
         Future<Boolean> future = iterator.next();
         if (future.isDone()) {
-          finished++;
           iterator.remove();
           try {
             if (!future.get()) {
               failed++;
+            } else {
+              finished++;
             }
           } catch (Exception e) {
             // cancel or execution exception
@@ -92,7 +98,10 @@ public static boolean waitUntilDoneOrFail(List<Future<Boolean>> list, boolean al
         futures.forEach(x -> x.cancel(true));
         Thread.currentThread().interrupt();
         return false;
+      } catch (TimeoutException e) {
+        // ignore
       } catch (Exception e) {
+        LOG.warn("Exception in waitUntilDoneOrFail", e);
         // ignore timeout or execution err
       }
     }