future true

Author: summaryzb

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

@@ -213,7 +213,7 @@ public WriteBufferManager(
       TaskMemoryManager taskMemoryManager,
       ShuffleWriteMetrics shuffleWriteMetrics,
       RssConf rssConf,
-      Function<List<ShuffleBlockInfo>, List<CompletableFuture<Long>>> spillFunc,
+      Function<List<ShuffleBlockInfo>, List<Future<Long>>> spillFunc,
       int stageAttemptNumber) {
     this(
         shuffleId,

client-spark/common/src/main/java/org/apache/uniffle/shuffle/manager/RssShuffleManagerBase.java

@@ -29,6 +29,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
@@ -1561,6 +1562,6 @@ public Future<Long> sendData(AddBlockEvent event) {
     if (dataPusher != null && event != null) {
       return dataPusher.send(event);
     }
-    return new Future<>();
+    return new CompletableFuture<>();
   }
 }

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

@@ -24,7 +24,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -361,8 +360,7 @@ private void checkSentBlockCount() {
    *
    * @param shuffleBlockInfoList
    */
-  private List<CompletableFuture<Long>> processShuffleBlockInfos(
-      List<ShuffleBlockInfo> shuffleBlockInfoList) {
+  private List<Future<Long>> processShuffleBlockInfos(List<ShuffleBlockInfo> shuffleBlockInfoList) {
     if (shuffleBlockInfoList != null && !shuffleBlockInfoList.isEmpty()) {
       shuffleBlockInfoList.stream()
           .forEach(
@@ -390,9 +388,8 @@ private List<CompletableFuture<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<CompletableFuture<Long>> postBlockEvent(
-      List<ShuffleBlockInfo> shuffleBlockInfoList) {
-    List<CompletableFuture<Long>> futures = new ArrayList<>();
+  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/spark2/src/test/java/org/apache/spark/shuffle/writer/RssShuffleWriterTest.java

@@ -22,6 +22,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
@@ -169,9 +170,9 @@ public void checkBlockSendResultTest() {
   }
 
   static class FakedDataPusher extends DataPusher {
-    private final Function<AddBlockEvent, CompletableFuture<Long>> sendFunc;
+    private final Function<AddBlockEvent, Future<Long>> sendFunc;
 
-    FakedDataPusher(Function<AddBlockEvent, CompletableFuture<Long>> sendFunc) {
+    FakedDataPusher(Function<AddBlockEvent, Future<Long>> sendFunc) {
       this(null, null, null, null, null, 1, 1, sendFunc);
     }
 
@@ -183,7 +184,7 @@ private FakedDataPusher(
         Set<String> failedTaskIds,
         int threadPoolSize,
         int threadKeepAliveTime,
-        Function<AddBlockEvent, CompletableFuture<Long>> sendFunc) {
+        Function<AddBlockEvent, Future<Long>> sendFunc) {
       super(
           shuffleWriteClient,
           taskToSuccessBlockIds,
@@ -195,7 +196,7 @@ private FakedDataPusher(
     }
 
     @Override
-    public CompletableFuture<Long> send(AddBlockEvent event) {
+    public Future<Long> send(AddBlockEvent event) {
       return sendFunc.apply(event);
     }
   }

client/src/test/java/org/apache/uniffle/client/ClientUtilsTest.java

@@ -19,9 +19,10 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.FutureTask;
 import java.util.concurrent.TimeUnit;
 
 import org.awaitility.Awaitility;
@@ -70,12 +71,12 @@ public void testGenerateTaskIdBitMap() {
     }
   }
 
-  private List<CompletableFuture<Boolean>> getFutures(boolean fail) {
-    List<CompletableFuture<Boolean>> futures = new ArrayList<>();
+  private List<Future<Boolean>> getFutures(boolean fail) {
+    List<Future<Boolean>> futures = new ArrayList<>();
     for (int i = 0; i < 3; i++) {
       final int index = i;
-      CompletableFuture<Boolean> future =
-          CompletableFuture.supplyAsync(
+      FutureTask<Boolean> future =
+          new FutureTask(
               () -> {
                 if (index == 2) {
                   try {
@@ -88,8 +89,8 @@ private List<CompletableFuture<Boolean>> getFutures(boolean fail) {
                   return true;
                 }
                 return !fail || index != 1;
-              },
-              executorService);
+              });
+      executorService.submit(future);
       futures.add(future);
     }
     return futures;
@@ -98,13 +99,13 @@ private List<CompletableFuture<Boolean>> getFutures(boolean fail) {
   @Test
   public void testWaitUntilDoneOrFail() {
     // case1: enable fail fast
-    List<CompletableFuture<Boolean>> futures1 = getFutures(true);
+    List<Future<Boolean>> futures1 = getFutures(true);
     Awaitility.await()
         .timeout(2, TimeUnit.SECONDS)
         .until(() -> !waitUntilDoneOrFail(futures1, true));
 
     // case2: disable fail fast
-    List<CompletableFuture<Boolean>> futures2 = getFutures(true);
+    List<Future<Boolean>> futures2 = getFutures(true);
     try {
       Awaitility.await()
           .timeout(2, TimeUnit.SECONDS)
@@ -115,7 +116,7 @@ public void testWaitUntilDoneOrFail() {
     }
 
     // case3: all succeed
-    List<CompletableFuture<Boolean>> futures3 = getFutures(false);
+    List<Future<Boolean>> futures3 = getFutures(false);
     Awaitility.await()
         .timeout(4, TimeUnit.SECONDS)
         .until(() -> waitUntilDoneOrFail(futures3, true));