[Bug] Long tail tasks in the Write Stage retry phase results in data loss.

Author: yl09099

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

@@ -45,6 +45,7 @@
 import org.apache.spark.MapOutputTracker;
 import org.apache.spark.MapOutputTrackerMaster;
 import org.apache.spark.SparkConf;
+import org.apache.spark.SparkContext;
 import org.apache.spark.SparkEnv;
 import org.apache.spark.SparkException;
 import org.apache.spark.shuffle.RssShuffleHandle;
@@ -1131,4 +1132,8 @@ public Map<String, String> sparkConfToMap(SparkConf sparkConf) {
   public ShuffleWriteClient getShuffleWriteClient() {
     return shuffleWriteClient;
   }
+
+  public void killAllTaskByStageId(int stageAttemptId, String cancelReason) {
+    SparkContext.getOrCreate().cancelStage(stageAttemptId, cancelReason);
+  }
 }

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

@@ -96,4 +96,12 @@ MutableShuffleHandleInfo reassignOnBlockSendFailure(
    * @return
    */
   ShuffleWriteClient getShuffleWriteClient();
+
+  /**
+   * To cancel all tasks under current stage by stageId.
+   *
+   * @param stageAttemptId
+   * @param cancelReason
+   */
+  void killAllTaskByStageId(int stageAttemptId, String cancelReason);
 }

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

@@ -115,9 +115,15 @@ public void reportShuffleWriteFailure(
                     shuffleManager.getMaxFetchFailures());
             if (!shuffleServerWriterFailureRecord.isClearedMapTrackerBlock()) {
               try {
+                // Cancel all tasks under the current StageId.
+                shuffleManager.killAllTaskByStageId(
+                    stageAttemptId, "Write failure triggers server reallocation.");
                 // Clear the metadata of the completed task, otherwise some of the stage's data will
                 // be lost.
                 shuffleManager.unregisterAllMapOutput(shuffleId);
+                // Need to clear the mapStatus twice to prevent partition data loss due to the
+                // long-tail task performed before the stage retry.
+                shuffleManager.unregisterAllMapOutput(shuffleId);
                 // Deregister the shuffleId corresponding to the Shuffle Server.
                 shuffleManager.getShuffleWriteClient().unregisterShuffle(appId, shuffleId);
                 shuffleServerWriterFailureRecord.setClearedMapTrackerBlock(true);

client-spark/common/src/test/java/org/apache/uniffle/shuffle/manager/DummyRssShuffleManager.java

@@ -90,4 +90,7 @@ public MutableShuffleHandleInfo reassignOnBlockSendFailure(
   public ShuffleWriteClient getShuffleWriteClient() {
     return null;
   }
+
+  @Override
+  public void killAllTaskByStageId(int stageAttemptId, String cancelReason) {}
 }