[CELEBORN-1776] Add UT when replica enable and workers randomly shutdown

tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/RetryReviveTest.scala

@@ -17,6 +17,8 @@
 
 package org.apache.celeborn.tests.spark
 
+import scala.util.Random
+
 import org.apache.spark.SparkConf
 import org.apache.spark.sql.SparkSession
 import org.scalatest.BeforeAndAfterEach
@@ -79,4 +81,28 @@ class RetryReviveTest extends AnyFunSuite
     assert(result.size == 1000)
     ss.stop()
   }
+
+  test("celeborn spark integration test - revive test replicate enabled when workers are randomly killed") {
+    setupMiniClusterWithRandomPorts()
+    ShuffleClient.reset()
+    val sparkConf = new SparkConf()
+      .set(s"spark.${CelebornConf.CLIENT_PUSH_REPLICATE_ENABLED.key}", "true")
+      .set(s"spark.${CelebornConf.CLIENT_STAGE_RERUN_ENABLED.key}", "false")
+      .setAppName("celeborn-demo").setMaster("local[2]")
+    val ss = SparkSession.builder()
+      .config(updateSparkConf(sparkConf, ShuffleMode.HASH))
+      .getOrCreate()
+
+    val startTime = System.currentTimeMillis()
+    val result = ss.sparkContext.parallelize(1 to 800, 100)
+      .flatMap(_ => (1 to 15000).iterator.map(num => num)).repartition(100).count()
+    val taskTime = System.currentTimeMillis() - startTime
+    val random = new Random()
+    val workerKillTime = random.nextInt(taskTime.toInt)
+    workerKiller(workerKillTime)
+    val result1 = ss.sparkContext.parallelize(1 to 800, 100)
+      .flatMap(_ => (1 to 15000).iterator.map(num => num)).repartition(100).count()
+    assert(result1 == result)
+    ss.stop()
+  }
 }

worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala

@@ -36,7 +36,9 @@ trait MiniClusterFeature extends Logging {
 
   var masterInfo: (Master, Thread) = _
   val workerInfos = new mutable.HashMap[Worker, Thread]()
+  val indToWorkerInfos = new mutable.HashMap[Int, (Worker, Thread)]()
   var workerConfForAdding: Map[String, String] = _
+  var testKillWorker: (Int, Thread) = (-1, null)
 
   val maxRetries = 4
   val masterWaitingTimeoutMs = TimeUnit.SECONDS.toMillis(60)
@@ -222,15 +224,17 @@ trait MiniClusterFeature extends Logging {
             worker.initialize()
           } catch {
             case ex: Exception =>
-              if (workers(i - 1) != null) {
-                workers(i - 1).shutdownGracefully()
-              }
-              workerStarted = false
-              workerStartRetry += 1
-              logError(s"cannot start worker $i, retrying: ", ex)
-              if (workerStartRetry == maxRetries) {
-                logError(s"cannot start worker $i, reached to max retrying", ex)
-                throw ex
+              if (testKillWorker._1 != (i - 1)) {
+                if (workers(i - 1) != null) {
+                  workers(i - 1).shutdownGracefully()
+                }
+                workerStarted = false
+                workerStartRetry += 1
+                logError(s"cannot start worker $i, retrying: ", ex)
+                if (workerStartRetry == maxRetries) {
+                  logError(s"cannot start worker $i, reached to max retrying", ex)
+                  throw ex
+                }
               }
           }
         }
@@ -250,6 +254,7 @@ trait MiniClusterFeature extends Logging {
               throw new IllegalStateException(s"worker $i hasn't been initialized")
             } else if (!workerInfos.contains(workers(i))) {
               workerInfos.put(workers(i), threads(i))
+              indToWorkerInfos.put(i, (workers(i), threads(i)))
             }
             if (!workers(i).registered.get()) {
               throw new IllegalStateException(s"worker $i hasn't been registered")
@@ -278,6 +283,30 @@ trait MiniClusterFeature extends Logging {
     (setUpMaster(masterConf), setUpWorkers(workerConf, workerNum))
   }
 
+  def workerKiller(sleepTime: Int, workerNum: Int = 3): Unit = {
+    var ind = 0
+    var workerInfo = indToWorkerInfos.get(ind)
+    while (workerInfo.isEmpty && (ind + 1 < workerNum)) {
+      ind += 1
+      workerInfo = indToWorkerInfos.get(ind)
+    }
+    if (workerInfo.nonEmpty && ind < workerNum) {
+      testKillWorker = (ind, null)
+    }
+    val killerThread = new RunnerWrap({
+      Thread.sleep(sleepTime)
+      if (testKillWorker._1 != -1) {
+        workerInfo.get._1.stop(CelebornExitKind.EXIT_IMMEDIATELY)
+        workerInfo.get._1.rpcEnv.shutdown()
+        workerInfo.get._2.interrupt()
+      }
+    })
+    killerThread.start()
+    if (testKillWorker._1 != -1) {
+      testKillWorker = (ind, killerThread)
+    }
+  }
+
   def shutdownMiniCluster(): Unit = {
     // shutdown workers
     workerInfos.foreach {
@@ -297,6 +326,10 @@ trait MiniClusterFeature extends Logging {
         worker.stop(CelebornExitKind.EXIT_IMMEDIATELY)
         thread.interrupt()
     }
+    if (testKillWorker._1 != -1) {
+      testKillWorker._2.interrupt()
+      testKillWorker = (-1, null)
+    }
     workerInfos.clear()
     masterInfo._2.interrupt()
     MemoryManager.reset()