update

Author: jiangxb1987

core/src/main/scala/org/apache/spark/BarrierCoordinator.scala

@@ -19,7 +19,7 @@ package org.apache.spark
 
 import java.util.{Timer, TimerTask}
 import java.util.concurrent.ConcurrentHashMap
-import java.util.function.Consumer
+import java.util.function.{Consumer, Function}
 
 import scala.collection.mutable.ArrayBuffer
 
@@ -93,8 +93,8 @@ private[spark] class BarrierCoordinator(
       val numTasks: Int) {
 
     // There may be multiple barrier() calls from a barrier stage attempt, `barrierEpoch` is used
-    // to identify each barrier() call. It shall get increased when a barrier() call succeed, or
-    // reset when a barrier() call fail due to timeout.
+    // to identify each barrier() call. It shall get increased when a barrier() call succeeds, or
+    // reset when a barrier() call fails due to timeout.
     private var barrierEpoch: Int = 0
 
     // An array of RPCCallContexts for barrier tasks that are waiting for reply of a barrier()
@@ -199,7 +199,10 @@ private[spark] class BarrierCoordinator(
     case request @ RequestToSync(numTasks, stageId, stageAttemptId, _, _) =>
       // Get or init the ContextBarrierState correspond to the stage attempt.
       val barrierId = ContextBarrierId(stageId, stageAttemptId)
-      states.putIfAbsent(barrierId, new ContextBarrierState(barrierId, numTasks))
+      states.computeIfAbsent(barrierId, new Function[ContextBarrierId, ContextBarrierState] {
+        override def apply(key: ContextBarrierId): ContextBarrierState =
+          new ContextBarrierState(key, numTasks)
+      })
       val barrierState = states.get(barrierId)
 
       barrierState.handleRequest(context, request)
@@ -220,7 +223,7 @@ private[spark] sealed trait BarrierCoordinatorMessage extends Serializable
  * @param stageId ID of current stage
  * @param stageAttemptId ID of current stage attempt
  * @param taskAttemptId Unique ID of current task
- * @param barrierEpoch ID of the `barrier()` call, a task may consists multiple `barrier()` calls.
+ * @param barrierEpoch ID of the `barrier()` call, a task may consist multiple `barrier()` calls.
  */
 private[spark] case class RequestToSync(
     numTasks: Int,

core/src/main/scala/org/apache/spark/BarrierTaskContext.scala

@@ -101,10 +101,9 @@ class BarrierTaskContext(
   @Experimental
   @Since("2.4.0")
   def barrier(): Unit = {
-    val callSite = Utils.getCallSite()
     logInfo(s"Task $taskAttemptId from Stage $stageId(Attempt $stageAttemptNumber) has entered " +
       s"the global sync, current barrier epoch is $barrierEpoch.")
-    logTrace(s"Current callSite: $callSite")
+    logTrace("Current callSite: " + Utils.getCallSite())
 
     val startTime = System.currentTimeMillis()
     val timerTask = new TimerTask {

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -575,6 +575,6 @@ package object config {
         "configed time, throw a SparkException to fail all the tasks. The default value is set " +
         "to 31536000(3600 * 24 * 365) so the barrier() call shall wait for one year.")
       .timeConf(TimeUnit.SECONDS)
-      .checkValue(v => v > 0, "The value should be a positive int value.")
+      .checkValue(v => v > 0, "The value should be a positive time value.")
       .createWithDefaultString("365d")
 }

core/src/main/scala/org/apache/spark/scheduler/TaskSchedulerImpl.scala

@@ -141,11 +141,15 @@ private[spark] class TaskSchedulerImpl(
 
   private lazy val barrierSyncTimeout = conf.get(config.BARRIER_SYNC_TIMEOUT)
 
-  private[scheduler] lazy val barrierCoordinator: RpcEndpoint = {
-    val coordinator = new BarrierCoordinator(barrierSyncTimeout, sc.listenerBus, sc.env.rpcEnv)
-    sc.env.rpcEnv.setupEndpoint("barrierSync", coordinator)
-    logInfo("Registered BarrierCoordinator endpoint")
-    coordinator
+  private[scheduler] var barrierCoordinator: RpcEndpoint = null
+
+  private def maybeInitBarrierCoordinator(): Unit = {
+    if (barrierCoordinator == null) {
+      barrierCoordinator = new BarrierCoordinator(barrierSyncTimeout, sc.listenerBus,
+        sc.env.rpcEnv)
+      sc.env.rpcEnv.setupEndpoint("barrierSync", barrierCoordinator)
+      logInfo("Registered BarrierCoordinator endpoint")
+    }
   }
 
   override def setDAGScheduler(dagScheduler: DAGScheduler) {
@@ -424,7 +428,7 @@ private[spark] class TaskSchedulerImpl(
               "been blacklisted or cannot fulfill task locality requirements.")
 
           // materialize the barrier coordinator.
-          barrierCoordinator
+          maybeInitBarrierCoordinator()
 
           // Update the taskInfos into all the barrier task properties.
           val addressesStr = addressesWithDescs