Spark-Core源码精读(3)、Stage划分算法

Spark中对RDD的操作大体上可以分为transformation级别的操作和action级别的操作，transformation是lazy级别的操作，action操作(count、collect等)会触发具体job的执行，而每个job又会被划分成一个或者多个Stage，后面的Stage会依赖前面的Stage，而Stage划分的依据就是是否为宽依赖(Spark中RDD的依赖关系分成宽依赖和窄依赖)，所有的Stage会形成一个有向无环图(DAG)，最后依据Task的数据本地性将Task发送到指定的Executor上运行，下面我们就详细分析这一过程。

Stage的划分

首先从一个Action级别的操作开始，此处以count为例：

def count(): Long = sc.runJob(this, Utils.getIteratorSize _).sum

可以看到调用了runJob方法，我们点击runJob方法，中间执行的runJob都是sparkContext类中的，

  def runJob[T, U: ClassTag](
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U,
      partitions: Seq[Int],
      resultHandler: (Int, U) => Unit): Unit = {
    if (stopped.get()) {
      throw new IllegalStateException("SparkContext has been shutdown")
    }
    val callSite = getCallSite
    val cleanedFunc = clean(func)
    logInfo("Starting job: " + callSite.shortForm)
    if (conf.getBoolean("spark.logLineage", false)) {
      logInfo("RDD's recursive dependencies:\n" + rdd.toDebugString)
    }
    //进入spark最核心的DAG Schedular
    dagScheduler.runJob(rdd, cleanedFunc, partitions, callSite, resultHandler, localProperties.get)
    progressBar.foreach(_.finishAll())
    rdd.doCheckpoint()
  }

这是最后一个runJob，最后会进入DagScheduler类中，让我们进到dagScheduler的runjob中，由于里面代码太多了，就粘贴一部分代码来看

  //DAG提交任务的入口 执行runjob
  def runJob[T, U](
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U,
      partitions: Seq[Int],
      callSite: CallSite,
      resultHandler: (Int, U) => Unit,
      properties: Properties): Unit = {
    val start = System.nanoTime
    //在这里回提交一个job任务，然后会返回一个阻塞的线程等待job任务完成

    val waiter = submitJob(rdd, func, partitions, callSite, resultHandler, properties)

可以看到这个runjob就是DAG提交任务的入口了，注意这里还没有开始stage的划分，这里提交了一个submitJob任务，我们进去看看

 //DAG 提交任务的第二部 进入submitjob方法
  def submitJob[T, U](
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U,
      partitions: Seq[Int],
      callSite: CallSite,
      resultHandler: (Int, U) => Unit,
      properties: Properties): JobWaiter[U] = {
    // Check to make sure we are not launching a task on a partition that does not exist.
    //检查分区是否存在保证Task正常运行
    val maxPartitions = rdd.partitions.length
    partitions.find(p => p >= maxPartitions || p < 0).foreach { p =>
      throw new IllegalArgumentException(
        "Attempting to access a non-existent partition: " + p + ". " +
          "Total number of partitions: " + maxPartitions)
    }
  //增加一个jobid当作job的标识符 +1
    val jobId = nextJobId.getAndIncrement()
    if (partitions.size == 0) {
      // Return immediately if the job is running 0 tasks
      return new JobWaiter[U](this, jobId, 0, resultHandler)
    }
  //为分区做个判断，确保分区大于0
    assert(partitions.size > 0)
    val func2 = func.asInstanceOf[(TaskContext, Iterator[_]) => _]
   //首先构造一个jobWaiter阻塞线程，等待job完成，然后把完成结果提交给resultHandler
    val waiter = new JobWaiter(this, jobId, partitions.size, resultHandler)
    //DAGSchduler 默认是FIFO先进先出的
    //这个对象很重要 eventProcessLoop
    /**
      * DAGSchedulerEventProcessLoop是一个DAGScheduler的一个内部类
      * 调用了post方法发送jobsubmited消息
      */
    eventProcessLoop.post(JobSubmitted(
      jobId, rdd, func2, partitions.toArray, callSite, waiter,
      SerializationUtils.clone(properties)))
    waiter
  }

这个方法前面一大堆代码判断，就是为stage划分阶段做一系列的准备工作，像：检查分区是否存在，如果分区都不存在，task肯定不能运行，这样就会出现异常，为每个job增加一个jobId，这个id后面会用到，最下面才是重要的方法，eventProcessLoop这个类调用了post方法，发送了一个提交信息，这个类他new了一个子类DAGSchedulerEventProcessLoop，用这个类来创建对象，相当于多态，
进入到这个子类里面

private[scheduler] class DAGSchedulerEventProcessLoop(dagScheduler: DAGScheduler)
  extends EventLoop[DAGSchedulerEvent]("dag-scheduler-event-loop") with Logging {

  private[this] val timer = dagScheduler.metricsSource.messageProcessingTimer

  /**
   * The main event loop of the DAG scheduler.
   */
  override def onReceive(event: DAGSchedulerEvent): Unit = {
    val timerContext = timer.time()
    try {
      doOnReceive(event)
    } finally {
      timerContext.stop()
    }
  }
//然后通过模式匹配哪个事件类型
  private def doOnReceive(event: DAGSchedulerEvent): Unit = event match {
    case JobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties) =>
      //开始stage划分的精髓所在 大家注意  HandleJobSubmit  重要
      dagScheduler.handleJobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties)

到现在可以看到。里面调用了doOnReveive方法，通过模式匹配来执行这个event事件，调用dagScheduler的handleJobSubmitted方法,让我们进到这个方法里面

 private[scheduler] def handleJobSubmitted(jobId: Int,
      finalRDD: RDD[_],
      func: (TaskContext, Iterator[_]) => _,
      partitions: Array[Int],
      callSite: CallSite,
      listener: JobListener,
      properties: Properties) {
    //第一步
    // 创建ResultStage   ，这里才是真正开始处理提交job划分stage的时候
    var finalStage: ResultStage = null
    try {
      // New stage creation may throw an exception if, for example, jobs are run on a
      // HadoopRDD whose underlying HDFS files have been deleted.
      //他会从后往前递归遍历他的每一个父RDD，从持久化中抽取反之重新计算
      //补充下 stage分为shuffleMapstage  跟ResultStage 俩种
      //每一个job都是有1个resultstage和0个shffleMapstage组成
      finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite)
    } catch {
      case e: Exception =>
        logWarning("Creating new stage failed due to exception - job: " + jobId, e)
        listener.jobFailed(e)
        return
    }
  //第二步  用finalStage创建一个job
    //就是说，这个job的最后一个stage，当然就是我们的最后一个stage了
    val job = new ActiveJob(jobId, finalStage, callSite, listener, properties)
    clearCacheLocs()
    logInfo("Got job %s (%s) with %d output partitions".format(
      job.jobId, callSite.shortForm, partitions.length))
    logInfo("Final stage: " + finalStage + " (" + finalStage.name + ")")
    logInfo("Parents of final stage: " + finalStage.parents)
    logInfo("Missing parents: " + getMissingParentStages(finalStage))

    val jobSubmissionTime = clock.getTimeMillis()
    jobIdToActiveJob(jobId) = job
    //第三部  将job加入到内存缓存中
    activeJobs += job
    finalStage.setActiveJob(job)
    val stageIds = jobIdToStageIds(jobId).toArray
    val stageInfos = stageIds.flatMap(id => stageIdToStage.get(id).map(_.latestInfo))
    listenerBus.post(
      SparkListenerJobStart(job.jobId, jobSubmissionTime, stageInfos, properties))
    //第四部，很关键 使用submitstage方法提交finalstage
    //这个方法会导致第一个stage提交，其他的stage放入waitingstages队列
    //使用递归优先提交父stage
    submitStage(finalStage)
  }

到这里，程序已经到了stage划分算法的门口了，代码开始创建了resultStage，也就是最后一个stage，并用最后这个stage创建了一个job，将job加入缓存中，最最重要的一步来了，开始提交这个finalstage，进入stage划分算法

  //但是stage划分算法是有submitStage和getMissingParentStages方法共同组成的。
  private def submitStage(stage: Stage) {
    val jobId = activeJobForStage(stage)
    if (jobId.isDefined) {
      logDebug("submitStage(" + stage + ")")
      if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {
        很关键的一行，调用getMissingParentStage方法去获取这个stage的父stage
        //这里会循环递归调用，直到最初的stage没有父stage，其余的stage放在
        //waitingMissingStages
        val missing = getMissingParentStages(stage).sortBy(_.id)
        logDebug("missing: " + missing)
        if (missing.isEmpty) {
          logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
          submitMissingTasks(stage, jobId.get)
        } else {
          //如果不为空，就是有父Stage，递归调用submitStage方法去提交父Stage，这里是stage划分算法的精髓。
          for (parent <- missing) {
            submitStage(parent)
          }
          //并且将当前stage放入等待的stage队列中
          waitingStages += stage
        }
      }
    }

stage划分算法里面有2个方法构成，也就是这2个方法实现的划分，submitStage跟getMissingParentStages，顾名思义，一个是提交stage，另一个是获取父stage，
看上面的代码可知，程序先把这个stage传入到了getMissingParentStages中，先去判断是否有父stage，如果没有的话，就把当前stage提交给了tasks，如果有的话，则递归调用submitStage方法去提交父stage，也就是把父stage再次传入到这个方法里面，去做相同的操作，把当前stage放入到等待队列里面，

getMissingParentStages方法(stage划分算法的精髓)

 //获取某个stage的父stage
  private def getMissingParentStages(stage: Stage): List[Stage] = {
    val missing = new HashSet[Stage]
    val visited = new HashSet[RDD[_]]
    // We are manually maintaining a stack here to prevent StackOverflowError
    // caused by recursively visiting
    val waitingForVisit = new Stack[RDD[_]]
    //重要
    //visit实现stage划分  供stage的rdd使用
    def visit(rdd: RDD[_]) {
      if (!visited(rdd)) {
        visited += rdd
        val rddHasUncachedPartitions = getCacheLocs(rdd).contains(Nil)
        if (rddHasUncachedPartitions) {
          //遍历rdd的依赖
          for (dep <- rdd.dependencies) {
            //用到了模式匹配
            dep match {
              case shufDep: ShuffleDependency[_, _, _] =>
                //如果是宽依赖，使用宽依赖的RDD创建一个新的stage，并且会把isShuffleMap变量设置为true
                //默认最后一个stage不是shufflemapstage
                val mapStage = getOrCreateShuffleMapStage(shufDep, stage.firstJobId)
                if (!mapStage.isAvailable) {
                  //将stage放到缓存中
                  missing += mapStage
                }
              case narrowDep: NarrowDependency[_] =>
                //继续，把这个rdd推入栈当中
                waitingForVisit.push(narrowDep.rdd)
            }
          }
        }
      }
    }
    //首先往栈中，推入了stage中最后一个rdd
    然后进行while循环
    waitingForVisit.push(stage.rdd)
    while (waitingForVisit.nonEmpty) {
      对stage的最后一个rdd，调用自己内部的visit方法
      visit(waitingForVisit.pop())
    }
    missing.toList
  }

里面有一个visit方法，默认是不调用的，先把当前stage的rdd压入栈中，然后开始遍历，对stage刚才压入的最后一个rdd调用内部的visit方法，方法刚开始声明了2个数字，一个是存放shufflestage的数组，一个是存放rdd的数组，首先刚进入visit方法中，判断当前rdd是否在数组中，如果没有，将其加入到数组中，然后才开始进行操作，遍历这个rdd的依赖，如果是窄依赖，则把窄依赖的rdd放入到栈中，继续while循环，往前推rdd，如果发现是宽依赖，会创建一个mapstage，加入到数组中，将这个数组转化成list返回，

让我们看看上一个方法submitStage

当返回过来之后会进行判断

 if (missing.isEmpty) {
          logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
          submitMissingTasks(stage, jobId.get)
        } else {
          //如果不为空，就是有父Stage，递归调用submitStage方法去提交父Stage，这里是stage划分算法的精髓。
          for (parent <- missing) {
            submitStage(parent)
          }
          //并且将当前stage放入等待的stage队列中
          waitingStages += stage

如果获取到shufflestage的话missing数组就不是空的，不为空说明就有父stage，然后，然后开始循环这个missing里面的stage，再次调用submitstage方法，并把当前stage放入等待队列中，这样一直递归调用submitstage，如果到第一个stage的时候，再往前没有stage了，这missing就是空的，开始提交到TaskScheduler，关于这DAGScheduler的理解就到这里，TaskScheduler下次有机会了在跟大家分享。