fix(spark): add all faulty servers into blacklist while writing for stage recomputing

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

@@ -19,7 +19,6 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -350,17 +349,16 @@ public boolean incPartitionWriteFailure(
                 });
             List<Map.Entry<String, AtomicInteger>> list =
                 new ArrayList(shuffleServerFailureRecordCount.entrySet());
-            if (!list.isEmpty()) {
-              Collections.sort(list, (o1, o2) -> (o1.getValue().get() - o2.getValue().get()));
-              Map.Entry<String, AtomicInteger> shuffleServerInfoIntegerEntry = list.get(0);
+            boolean retry = false;
+            for (Map.Entry<String, AtomicInteger> shuffleServerInfoIntegerEntry : list) {
               if (shuffleServerInfoIntegerEntry.getValue().get()
                   > shuffleManager.getMaxFetchFailures()) {
                 shuffleManager.addFailuresShuffleServerInfos(
                     shuffleServerInfoIntegerEntry.getKey());
-                return true;
+                retry = true;
               }
             }
-            return false;
+            return retry;
           });
     }
   }

integration-test/common/src/test/java/org/apache/uniffle/test/IntegrationTestBase.java

@@ -183,18 +183,21 @@ protected static void createShuffleServer(ShuffleServerConf serverConf) throws E
     }
   }
 
-  protected static void createMockedShuffleServer(ShuffleServerConf serverConf) throws Exception {
+  protected static MockedShuffleServer createMockedShuffleServer(ShuffleServerConf serverConf)
+      throws Exception {
     ServerType serverType = serverConf.get(ShuffleServerConf.RPC_SERVER_TYPE);
+    MockedShuffleServer server = new MockedShuffleServer(serverConf);
     switch (serverType) {
       case GRPC:
-        grpcShuffleServers.add(new MockedShuffleServer(serverConf));
+        grpcShuffleServers.add(server);
         break;
       case GRPC_NETTY:
-        nettyShuffleServers.add(new MockedShuffleServer(serverConf));
+        nettyShuffleServers.add(server);
         break;
       default:
         throw new UnsupportedOperationException("Unsupported server type " + serverType);
     }
+    return server;
   }
 
   protected static void createAndStartServers(

integration-test/spark-common/src/test/java/org/apache/uniffle/test/RSSStageDynamicServerReWriteTest.java

@@ -26,6 +26,7 @@
 import org.apache.spark.shuffle.RssSparkConfig;
 import org.apache.spark.sql.Row;
 import org.apache.spark.sql.SparkSession;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -36,6 +37,7 @@
 import org.apache.uniffle.common.rpc.ServerType;
 import org.apache.uniffle.coordinator.CoordinatorConf;
 import org.apache.uniffle.server.MockedGrpcServer;
+import org.apache.uniffle.server.MockedShuffleServer;
 import org.apache.uniffle.server.ShuffleServerConf;
 import org.apache.uniffle.storage.util.StorageType;
 
@@ -52,14 +54,19 @@ public static void setupServers(@TempDir File tmpDir) throws Exception {
     addDynamicConf(coordinatorConf, dynamicConf);
     createCoordinatorServer(coordinatorConf);
     createServer(0, tmpDir, true, ServerType.GRPC);
-    createServer(1, tmpDir, false, ServerType.GRPC);
+    createServer(1, tmpDir, true, ServerType.GRPC);
     createServer(2, tmpDir, false, ServerType.GRPC);
     createServer(3, tmpDir, true, ServerType.GRPC_NETTY);
-    createServer(4, tmpDir, false, ServerType.GRPC_NETTY);
+    createServer(4, tmpDir, true, ServerType.GRPC_NETTY);
     createServer(5, tmpDir, false, ServerType.GRPC_NETTY);
     startServers();
   }
 
+  @AfterAll
+  public static void shutdownAll() throws Exception {
+    IntegrationTestBase.shutdownServers();
+  }
+
   public static void createServer(int id, File tmpDir, boolean abnormalFlag, ServerType serverType)
       throws Exception {
     ShuffleServerConf shuffleServerConf = getShuffleServerConf(serverType);
@@ -74,32 +81,25 @@ public static void createServer(int id, File tmpDir, boolean abnormalFlag, Serve
         shuffleServerConf.getInteger(ShuffleServerConf.RPC_SERVER_PORT) + id);
     shuffleServerConf.setInteger("rss.jetty.http.port", 19081 + id * 100);
     shuffleServerConf.setString("rss.storage.basePath", basePath);
+
+    MockedShuffleServer server = createMockedShuffleServer(shuffleServerConf);
     if (abnormalFlag) {
-      createMockedShuffleServer(shuffleServerConf);
       // Set the sending block data timeout for the first shuffleServer
       switch (serverType) {
         case GRPC:
-          ((MockedGrpcServer) grpcShuffleServers.get(0).getServer())
-              .getService()
-              .enableMockSendDataFailed(true);
-          break;
         case GRPC_NETTY:
-          ((MockedGrpcServer) nettyShuffleServers.get(0).getServer())
-              .getService()
-              .enableMockSendDataFailed(true);
+          ((MockedGrpcServer) server.getServer()).getService().enableMockSendDataFailed(true);
           break;
         default:
           throw new UnsupportedOperationException("Unsupported server type " + serverType);
       }
-    } else {
-      createShuffleServer(shuffleServerConf);
     }
   }
 
   @Override
   public Map runTest(SparkSession spark, String fileName) throws Exception {
     List<Row> rows =
-        spark.range(0, 1000, 1, 4).repartition(2).groupBy("id").count().collectAsList();
+        spark.range(0, 1000, 1, 4).repartition(4).groupBy("id").count().collectAsList();
     Map<String, Long> result = Maps.newHashMap();
     for (Row row : rows) {
       result.put(row.get(0).toString(), row.getLong(1));