Eagerly cancel rpc request

Author: summaryzb

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

@@ -200,7 +200,8 @@ private boolean sendShuffleDataAsync(
                               stageAttemptNumber,
                               retryMax,
                               retryIntervalMax,
-                              shuffleIdToBlocks);
+                              shuffleIdToBlocks,
+                              needCancelRequest);
                       long s = System.currentTimeMillis();
                       RssSendShuffleDataResponse response =
                           getShuffleServerClient(ssi).sendShuffleData(request);

internal-client/src/main/java/org/apache/uniffle/client/impl/grpc/ShuffleServerGrpcClient.java

@@ -647,7 +647,7 @@ public RssSendShuffleDataResponse sendShuffleData(RssSendShuffleDataRequest requ
             null,
             request.getRetryIntervalMax(),
             maxRetryAttempts,
-            t -> !(t instanceof OutOfMemoryError) && !(t instanceof NotRetryException));
+            t -> !request.needCancel() && !(t instanceof OutOfMemoryError) && !(t instanceof NotRetryException));
       } catch (Throwable throwable) {
         LOG.warn("Failed to send shuffle data due to ", throwable);
         isSuccessful = false;

internal-client/src/main/java/org/apache/uniffle/client/impl/grpc/ShuffleServerGrpcNettyClient.java

@@ -235,7 +235,7 @@ public RssSendShuffleDataResponse sendShuffleData(RssSendShuffleDataRequest requ
             null,
             request.getRetryIntervalMax(),
             maxRetryAttempts,
-            t -> !(t instanceof OutOfMemoryError) && !(t instanceof NotRetryException));
+            t -> !request.needCancel() && !(t instanceof OutOfMemoryError) && !(t instanceof NotRetryException));
       } catch (Throwable throwable) {
         LOG.warn("Failed to send shuffle data due to ", throwable);
         isSuccessful = false;

internal-client/src/main/java/org/apache/uniffle/client/request/RssSendShuffleDataRequest.java

@@ -19,6 +19,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 
 import org.apache.uniffle.common.ShuffleBlockInfo;
 
@@ -29,26 +30,29 @@ public class RssSendShuffleDataRequest {
   private int retryMax;
   private long retryIntervalMax;
   private Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleIdToBlocks;
+  private Supplier<Boolean> needCancel;
 
   public RssSendShuffleDataRequest(
       String appId,
       int retryMax,
       long retryIntervalMax,
       Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleIdToBlocks) {
-    this(appId, 0, retryMax, retryIntervalMax, shuffleIdToBlocks);
+    this(appId, 0, retryMax, retryIntervalMax, shuffleIdToBlocks, () -> false);
   }
 
   public RssSendShuffleDataRequest(
       String appId,
       int stageAttemptNumber,
       int retryMax,
       long retryIntervalMax,
-      Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleIdToBlocks) {
+      Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> shuffleIdToBlocks,
+      Supplier<Boolean> needCancel) {
     this.appId = appId;
     this.retryMax = retryMax;
     this.retryIntervalMax = retryIntervalMax;
     this.shuffleIdToBlocks = shuffleIdToBlocks;
     this.stageAttemptNumber = stageAttemptNumber;
+    this.needCancel = needCancel;
   }
 
   public String getAppId() {
@@ -70,4 +74,8 @@ public int getStageAttemptNumber() {
   public Map<Integer, Map<Integer, List<ShuffleBlockInfo>>> getShuffleIdToBlocks() {
     return shuffleIdToBlocks;
   }
+
+  public Boolean needCancel() {
+    return needCancel.get();
+  }
 }