/**

 * A stage is a set of independent tasks all computing the same function that need to run as part

 * of a Spark job, where all the tasks have the same shuffle dependencies. Each DAG of tasks run

 * by the scheduler is split up into stages at the boundaries where shuffle occurs, and then the

 * DAGScheduler runs these stages in topological order.

 *

 * Each Stage can either be a shuffle map stage, in which case its tasks' results are input for

 * another stage, or a result stage, in which case its tasks directly compute the action that

 * initiated a job (e.g. count(), save(), etc). For shuffle map stages, we also track the nodes

 * that each output partition is on.

 *

 * Each Stage also has a jobId, identifying the job that first submitted the stage.  When FIFO

 * scheduling is used, this allows Stages from earlier jobs to be computed first or recovered

 * faster on failure.

 */

private[spark] class Stage(

    val id: Int,

    val rdd: RDD[_], // final RDD

    val shuffleDep: Option[ShuffleDependency[_,_]],  // Output shuffle if stage is a map stage

    val parents: List[Stage], // 父stage

    val jobId: Int,

    callSite: Option[String])

  extends Logging {

  val isShuffleMap = shuffleDep != None  // 是否是shuffle map stage, 取决于是否有shuffleDep

  val numPartitions = rdd.partitions.size

  val outputLocs = Array.fill[List[MapStatus]](numPartitions)(Nil) // 用于buffer每个shuffle中每个maptask的MapStatus

  var numAvailableOutputs = 0

  private var nextAttemptId = 0

  def isAvailable: Boolean = {

    if (!isShuffleMap) {

      true

    } else {

      numAvailableOutputs == numPartitions

    }

  }

}

  /**

   * Create a Stage for the given RDD, either as a shuffle map stage (for a ShuffleDependency) or

   * as a result stage for the final RDD used directly in an action. The stage will also be

   * associated with the provided jobId.

   */

  private def newStage(

      rdd: RDD[_],

      shuffleDep: Option[ShuffleDependency[_,_]],

      jobId: Int,

      callSite: Option[String] = None)

    : Stage =

  {

    if (shuffleDep != None) {

      // Kind of ugly: need to register RDDs with the cache and map output tracker here

      // since we can't do it in the RDD constructor because # of partitions is unknown

      logInfo("Registering RDD " + rdd.id + " (" + rdd.origin + ")")

      mapOutputTracker.registerShuffle(shuffleDep.get.shuffleId, rdd.partitions.size)

    }

    val id = nextStageId.getAndIncrement()

    val stage = new Stage(id, rdd, shuffleDep, getParentStages(rdd, jobId), jobId, callSite)

    stageIdToStage(id) = stage

    stageToInfos(stage) = StageInfo(stage)

    stage

  }

  private def getMissingParentStages(stage: Stage): List[Stage] = {

    val missing = new HashSet[Stage]

    val visited = new HashSet[RDD[_]]

    def visit(rdd: RDD[_]) {

      if (!visited(rdd)) {

        visited += rdd

        if (getCacheLocs(rdd).contains(Nil)) {

          for (dep <- rdd.dependencies) {

            dep match {

              case shufDep: ShuffleDependency[_,_] => // 如果发现ShuffleDependency, 说明遇到新的stage

                val mapStage = getShuffleMapStage(shufDep, stage.jobId) // check shuffleToMapStage, 如果该stage已经被创建则直接返回, 否则newStage

                if (!mapStage.isAvailable) {

                  missing += mapStage

                }

              case narrowDep: NarrowDependency[_] => // 对于NarrowDependency, 说明仍然在这个stage中

                visit(narrowDep.rdd)

            }

          }

        }

      }

    }

    visit(stage.rdd)

    missing.toList

  }