[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,60 @@
+/*
+ * 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;
+
+import com.google.common.collect.Sets;
+
+import org.apache.uniffle.common.util.JavaUtils;
+
+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 RssStageResubmitManager() {
+    this.failuresShuffleServerIds = Sets.newConcurrentHashSet();
+    this.serverAssignedInfos = JavaUtils.newConcurrentMap();
+  }
+
+  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,72 @@
+/*
+ * 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 com.google.common.collect.Lists;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.uniffle.client.PartitionDataReplicaRequirementTracking;
+import org.apache.uniffle.common.ShuffleServerInfo;
+
+public class ChainShuffleHandleInfo extends ShuffleHandleInfoBase {
+  private static final Logger LOGGER = LoggerFactory.getLogger(MutableShuffleHandleInfo.class);
+
+  private ShuffleHandleInfo current;
+  private LinkedList<ShuffleHandleInfo> historyHandles;
+
+  public ChainShuffleHandleInfo(ShuffleHandleInfo shuffleServerInfo) {
+    super(0, null);
+    this.current = shuffleServerInfo;
+    this.historyHandles = Lists.newLinkedList();
+  }
+
+  @Override
+  public Set<ShuffleServerInfo> getServers() {
+    return current.getServers();
+  }
+
+  @Override
+  public Map<Integer, List<ShuffleServerInfo>> getAvailablePartitionServersForWriter() {
+    return current.getAvailablePartitionServersForWriter();
+  }
+
+  @Override
+  public Map<Integer, List<ShuffleServerInfo>> getAllPartitionServersForReader() {
+    return current.getAllPartitionServersForReader();
+  }
+
+  @Override
+  public PartitionDataReplicaRequirementTracking createPartitionReplicaTracking() {
+    return current.createPartitionReplicaTracking();
+  }
+
+  /**
+   * When a Stage retry occurs, replace the current shuffleHandleInfo and record the
+   * historical shuffleHandleInfo.
+   */
+  public void replaceCurrentShuffleHandleInfo(ShuffleHandleInfo shuffleHandleInfo) {
+    this.historyHandles.add(current);
+    this.current = shuffleHandleInfo;
+  }
+}