[FEATURE] Add an rpc method to obtain the uniffleId and delete the Write Stage for retry at the same time.

Author: yl09099

client-mr/core/src/test/java/org/apache/hadoop/mapred/SortWriteBufferManagerTest.java renamed to client-mr/core/src/test/resources/SortWriteBufferManagerTest.java

@@ -58,4 +58,30 @@ public int getOrCreateUniffleShuffleId(int shuffleId) {
   public void recordShuffleIdDeterminate(int shuffleId, boolean isDeterminate) {
     shuffleDeterminateMap.put(shuffleId, isDeterminate);
   }
+
+  public boolean getShuffleIdDeterminate(int shuffleId) {
+    return shuffleDeterminateMap.get(shuffleId);
+  }
+
+  public boolean hasUniffleShuffleIdByStageNumber(int shuffleId, int stageAttemptNumber) {
+    return shuffleIdMapping.get(shuffleId).get(stageAttemptNumber) != null;
+  }
+
+  public int getUniffleShuffleIdByStageNumber(int shuffleId, int stageAttemptNumber) {
+    return shuffleIdMapping.get(shuffleId).get(stageAttemptNumber);
+  }
+
+  public int getUniffleShuffleIdByStageNumber(int shuffleId) {
+    return shuffleIdMapping.get(shuffleId).entrySet().stream()
+        .sorted(Map.Entry.<Integer, Integer>comparingByKey().reversed())
+        .findFirst()
+        .get()
+        .getValue();
+  }
+
+  public int createUniffleShuffleId(int shuffleId, int stageAttemptNumber) {
+    int generatorShuffleId = shuffleIdGenerator.incrementAndGet();
+    shuffleIdMapping.get(shuffleId).put(stageAttemptNumber, generatorShuffleId);
+    return generatorShuffleId;
+  }
 }

client-spark/common/src/main/java/org/apache/ShuffleIdMappingManager.java

@@ -46,22 +46,6 @@ public class RssSparkConfig {
           .withDeprecatedKeys(RssClientConfig.RSS_RESUBMIT_STAGE)
           .withDescription("Whether to enable the resubmit stage for fetch/write failure");
 
-  public static final ConfigOption<Boolean> RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED =
-      ConfigOptions.key("rss.stageRetry.fetchFailureEnabled")
-          .booleanType()
-          .defaultValue(false)
-          .withFallbackKeys(RSS_RESUBMIT_STAGE_ENABLED.key(), RssClientConfig.RSS_RESUBMIT_STAGE)
-          .withDescription(
-              "If set to true, the stage retry mechanism will be enabled when a fetch failure occurs.");
-
-  public static final ConfigOption<Boolean> RSS_RESUBMIT_STAGE_WITH_WRITE_FAILURE_ENABLED =
-      ConfigOptions.key("rss.stageRetry.writeFailureEnabled")
-          .booleanType()
-          .defaultValue(false)
-          .withFallbackKeys(RSS_RESUBMIT_STAGE_ENABLED.key(), RssClientConfig.RSS_RESUBMIT_STAGE)
-          .withDescription(
-              "If set to true, the stage retry mechanism will be enabled when a write failure occurs.");
-
   public static final ConfigOption<Boolean> RSS_BLOCK_ID_SELF_MANAGEMENT_ENABLED =
       ConfigOptions.key("rss.blockId.selfManagementEnabled")
           .booleanType()

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

@@ -54,7 +54,7 @@
 import org.apache.uniffle.common.exception.RssFetchFailedException;
 import org.apache.uniffle.common.util.Constants;
 
-import static org.apache.spark.shuffle.RssSparkConfig.RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED;
+import static org.apache.spark.shuffle.RssSparkConfig.RSS_RESUBMIT_STAGE_ENABLED;
 
 public class RssSparkShuffleUtils {
 
@@ -369,7 +369,7 @@ public static RssException reportRssFetchFailedException(
       int stageAttemptId,
       Set<Integer> failedPartitions) {
     RssConf rssConf = RssSparkConfig.toRssConf(sparkConf);
-    if (rssConf.getBoolean(RSS_RESUBMIT_STAGE_WITH_FETCH_FAILURE_ENABLED)
+    if (rssConf.getBoolean(RSS_RESUBMIT_STAGE_ENABLED)
         && RssSparkShuffleUtils.isStageResubmitSupported()) {
       for (int partitionId : failedPartitions) {
         RssReportShuffleFetchFailureRequest req =

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

@@ -21,26 +21,26 @@
 import java.io.IOException;
 import java.util.Map;
 
-import org.apache.spark.shuffle.handle.ShuffleHandleInfo;
+import org.apache.spark.shuffle.handle.MutableShuffleHandleInfo;
 
 import org.apache.uniffle.common.util.JavaUtils;
 
 public class ShuffleHandleInfoManager implements Closeable {
-  private Map<Integer, ShuffleHandleInfo> shuffleIdToShuffleHandleInfo;
+  private Map<Integer, MutableShuffleHandleInfo> shuffleIdToShuffleHandleInfo;
 
   public ShuffleHandleInfoManager() {
     this.shuffleIdToShuffleHandleInfo = JavaUtils.newConcurrentMap();
   }
 
-  public ShuffleHandleInfo get(int shuffleId) {
+  public MutableShuffleHandleInfo get(int shuffleId) {
     return shuffleIdToShuffleHandleInfo.get(shuffleId);
   }
 
   public void remove(int shuffleId) {
     shuffleIdToShuffleHandleInfo.remove(shuffleId);
   }
 
-  public void register(int shuffleId, ShuffleHandleInfo handle) {
+  public void register(int shuffleId, MutableShuffleHandleInfo handle) {
     shuffleIdToShuffleHandleInfo.put(shuffleId, handle);
   }
 

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

@@ -77,31 +77,6 @@ public void reportShuffleResult(
     managerClientSupplier.get().reportShuffleResult(request);
   }
 
-  @Override
-  public void reportShuffleResult(
-      Map<ShuffleServerInfo, Map<Integer, Set<Long>>> serverToPartitionToBlockIds,
-      String appId,
-      int shuffleId,
-      long taskAttemptId,
-      int bitmapNum,
-      Set<ShuffleServerInfo> reportFailureServers,
-      boolean enableWriteFailureRetry) {
-    Map<Integer, List<Long>> partitionToBlockIds = new HashMap<>();
-    for (Map<Integer, Set<Long>> k : serverToPartitionToBlockIds.values()) {
-      for (Map.Entry<Integer, Set<Long>> entry : k.entrySet()) {
-        int partitionId = entry.getKey();
-        partitionToBlockIds
-            .computeIfAbsent(partitionId, x -> new ArrayList<>())
-            .addAll(entry.getValue());
-      }
-    }
-
-    RssReportShuffleResultRequest request =
-        new RssReportShuffleResultRequest(
-            appId, shuffleId, taskAttemptId, partitionToBlockIds, bitmapNum);
-    managerClientSupplier.get().reportShuffleResult(request);
-  }
-
   @Override
   public Roaring64NavigableMap getShuffleResult(
       String clientType,