[apache#1579][part-1] fix(spark): Adjust reassigned time to ensure th…

…at all previous data is cleared for stage retry

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

@@ -517,7 +517,8 @@ public void registerShuffle(
         List<PartitionRange> partitionRanges,
         RemoteStorageInfo remoteStorage,
         ShuffleDataDistributionType distributionType,
-        int maxConcurrencyPerPartitionToWrite) {}
+        int maxConcurrencyPerPartitionToWrite,
+        int stageAttemptNumber) {}
 
     @Override
     public boolean sendCommit(

client-mr/core/src/test/java/org/apache/hadoop/mapreduce/task/reduce/FetcherTest.java

@@ -504,7 +504,8 @@ public void registerShuffle(
         List<PartitionRange> partitionRanges,
         RemoteStorageInfo storageType,
         ShuffleDataDistributionType distributionType,
-        int maxConcurrencyPerPartitionToWrite) {}
+        int maxConcurrencyPerPartitionToWrite,
+        int stageAttemptNumber) {}
 
     @Override
     public boolean sendCommit(

client-spark/common/pom.xml

@@ -89,6 +89,12 @@
             <groupId>net.jpountz.lz4</groupId>
             <artifactId>lz4</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.apache.spark</groupId>
+            <artifactId>spark-catalyst_${scala.binary.version}</artifactId>
+            <version>${spark.version}</version>
+            <scope>provided</scope>
+        </dependency>
     </dependencies>
 
 </project>

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

@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle;
+
+import java.util.Map;
+import java.util.Set;
+
+public class RssStageResubmitManager {
+  /** A list of shuffleServer for Write failures */
+  private Set<String> failuresShuffleServerIds;
+  /**
+   * Prevent multiple tasks from reporting FetchFailed, resulting in multiple ShuffleServer
+   * assignments, stageID, Attemptnumber Whether to reassign the combination flag;
+   */
+  private Map<String, Boolean> serverAssignedInfos;
+
+  public Set<String> getFailuresShuffleServerIds() {
+    return failuresShuffleServerIds;
+  }
+
+  public void setFailuresShuffleServerIds(Set<String> failuresShuffleServerIds) {
+    this.failuresShuffleServerIds = failuresShuffleServerIds;
+  }
+
+  public void recordFailuresShuffleServer(String shuffleServerId) {
+    failuresShuffleServerIds.add(shuffleServerId);
+  }
+
+  public boolean recordAndGetServerAssignedInfo(String stageIdAndAttempt) {
+    return serverAssignedInfos.computeIfAbsent(stageIdAndAttempt, id -> false);
+  }
+
+  public void recordAndGetServerAssignedInfo(String stageIdAndAttempt, boolean isRetried) {
+    serverAssignedInfos.put(stageIdAndAttempt, isRetried);
+  }
+}

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

@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.handle;
+
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import com.google.common.collect.Lists;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.PartitionDataReplicaRequirementTracking;
+import org.apache.uniffle.common.RemoteStorageInfo;
+import org.apache.uniffle.common.ShuffleServerInfo;
+
+public class ChainShuffleHandleInfo extends ShuffleHandleInfoBase {
+  private static final Logger LOGGER = LoggerFactory.getLogger(MutableShuffleHandleInfo.class);
+
+  private Map<Integer, List<ShuffleServerInfo>> currentPartitionToServers;
+
+  private LinkedList<Map<Integer, List<ShuffleServerInfo>>> historyPartitionToServers;
+
+  public ChainShuffleHandleInfo(int shuffleId, RemoteStorageInfo remoteStorage) {
+    super(shuffleId, remoteStorage);
+  }
+
+  public ChainShuffleHandleInfo(
+      int shuffleId,
+      RemoteStorageInfo remoteStorage,
+      Map<Integer, List<ShuffleServerInfo>> partitionToServers) {
+    super(shuffleId, remoteStorage);
+    this.currentPartitionToServers = partitionToServers;
+    this.historyPartitionToServers = Lists.newLinkedList();
+  }
+
+  @Override
+  public Set<ShuffleServerInfo> getServers() {
+    return currentPartitionToServers.values().stream()
+        .flatMap(x -> x.stream())
+        .collect(Collectors.toSet());
+  }
+
+  @Override
+  public Map<Integer, List<ShuffleServerInfo>> getAvailablePartitionServersForWriter() {
+    return currentPartitionToServers;
+  }
+
+  @Override
+  public Map<Integer, List<ShuffleServerInfo>> getAllPartitionServersForReader() {
+    return currentPartitionToServers;
+  }
+
+  @Override
+  public PartitionDataReplicaRequirementTracking createPartitionReplicaTracking() {
+    return new PartitionDataReplicaRequirementTracking(currentPartitionToServers, shuffleId);
+  }
+
+  @Override
+  public void replaceCurrentShuffleHandleInfo(
+      Map<Integer, List<ShuffleServerInfo>> partitionToServers) {
+    this.historyPartitionToServers.add(currentPartitionToServers);
+    this.currentPartitionToServers = partitionToServers;
+  }
+}

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

@@ -29,6 +29,13 @@ public interface ShuffleHandleInfo {
   /** Get all the assigned servers including the excluded servers. */
   Set<ShuffleServerInfo> getServers();
 
+  /**
+   * When a Stage retry occurs, replace the current PartitionToShuffleServer and record the
+   * historical PartitionToShuffleServe.
+   */
+  default void replaceCurrentShuffleHandleInfo(
+      Map<Integer, List<ShuffleServerInfo>> partitionToServers) {}
+
   /**
    * Get the assignment of available servers for writer to write partitioned blocks to corresponding
    * shuffleServers. Implementations might return dynamic, up-to-date information here.