DAGScheduler的原理剖析以及源码解析

原理讲解

DAGScheduler的stage划分算法：会从触发的action操作的那个rdd开始往前倒推，首先会为最后一个rdd创建一个stage，然后往前倒推的时候，如果发现对某个rdd是宽依赖，那么就会将宽依赖的那个rdd创建一个新的stage，那个rdd就是对新的stage的最后一个rdd，然后依次类推，继续往前倒推，根据宽窄依赖，进行stage的划分，直到所有的rdd全部遍历完了为之。

源码讲解

在代码执行了算子之后，比如count(),代码依次如下

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

  def runJob[T, U: ClassTag](rdd: RDD[T], func: Iterator[T] => U): Array[U] = {
    runJob(rdd, func, 0 until rdd.partitions.length)
  }

  def runJob[T, U: ClassTag](
      rdd: RDD[T],
      func: Iterator[T] => U,
      partitions: Seq[Int]): Array[U] = {
    val cleanedFunc = clean(func)
    runJob(rdd, (ctx: TaskContext, it: Iterator[T]) => cleanedFunc(it), partitions)
  }

  def runJob[T, U: ClassTag](
      rdd: RDD[T],
      func: (TaskContext, Iterator[T]) => U,
      partitions: Seq[Int]): Array[U] = {
    val results = new Array[U](partitions.size)
    runJob[T, U](rdd, func, partitions, (index, res) => results(index) = res)
    results
  }

  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)
    }

    // 调用SparkContext,之前初始化创建的DAGScheduler的runJob()方法
    dagScheduler.runJob(rdd, cleanedFunc, partitions, callSite, resultHandler, localProperties.get)
    progressBar.foreach(_.finishAll())
    rdd.doCheckpoint()
  }

经过一系列的runJob调用，最后走到了具体功能实现的函数，
这个函数中最重要的就是dagScheduler.runJob()方法，接着进入DAGScheduler的runJob函数，然后会调用submitJob() 函数，进入submitJob函数，DAGSchedulerEventProcessLoop 会post JobSubmitted的消息。

 private def doOnReceive(event: DAGSchedulerEvent): Unit = event match {
    case JobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties) =>
      dagScheduler.handleJobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties)

最终会调用DAGScheduler的handleJobSubmitted函数，这个函数是DAGScheduler的job调度的核心入口，接下来说明一下这个函数具体的实现步骤：

    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.

      finalStage = newResultStage(finalRDD, func, partitions, jobId, callSite)
    } catch {
      case e: Exception =>
        logWarning("Creating new stage failed due to exception - job: " + jobId, e)
        listener.jobFailed(e)
        return
    }

第一步：使用触发job的最后一个rdd，创建finalStage , 并且将stage渐入DAGScheduler内部的内存缓存区

finalStage
    val job = new ActiveJob(jobId, finalStage, callSite, listener, properties)

第二步，用finalStage创建一个job, 就是说，这个job的最后一个stage，当然就是我们的finalStage

    jobIdToActiveJob(jobId) = job

第三部，将job加入内存缓存中

submitStage(finalStage)

第四部，使用submitStage方法提交finalStage，这个方法的调用，其实会导致第一个stage提交，并且导致其他所有的stage，都给放入waitingStage队列里。接下来我们看一下submitStage函数

  // 其实就是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)) {
        //调用getMissingParentStages方法，获取
        val missing = getMissingParentStages(stage).sortBy(_.id)

        logDebug("missing: " + missing)
        // 这里其实会反复调用
        // 直到最初的stage，它没有父stage了
        // 那么，此时，就会去首提交这个第一个stage
        // 其余的stage，此时全部都在waitingstage里面
        if (missing.isEmpty) {
          logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
          submitMissingTasks(stage, jobId.get)
        } else {
          // 递归调用submit方法，去提交父stage
          // 这里的递归，就是stage划分算法的推动者和精髓
          for (parent <- missing) {
            submitStage(parent)
          }
          // 并且将当前stage，放入waitingStage等待执行的stage的队列中
          waitingStages += stage
        }
      }
    } else {
      abortStage(stage, "No active job for stage " + stage.id, None)
    }
  }

其中比较重要的函数是getMissingParentStages，进入函数内部

// 获取某个stage的父stage
  // 对一个stage，如果它的最后一个rdd的所有依赖都是窄依赖，那么就不会创建任何新的stage
  // 但是，只要发现这个stage的rdd宽依赖了某个rdd，那么就用宽依赖的那个rdd，创建一个新的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[_]]
    def visit(rdd: RDD[_]) {
      if (!visited(rdd)) {
        visited += rdd
        val rddHasUncachedPartitions = getCacheLocs(rdd).contains(Nil)

        if (rddHasUncachedPartitions) {
          // 遍历rdd的依赖
          // 其实对于每一种有shuffle的操作，比如groupByKey、reduceByKey、countByKey
          // 底层对应于三个RDD： Map.PartitionRDD(会归入的新的stage),ShuffleRDD,MapPartitionRDD
          //

          for (dep <- rdd.dependencies) {
            dep match {
                // 如果是宽依赖的话，
              case shufDep: ShuffleDependency[_, _, _] =>
                // 那么使用宽依赖的那个rdd，创建一个stage，并且会将isShuffleMap设置为true
                // 默认最后一个stage，不是shuffle stage
                // 但是finalStage之前的所有stage，都是shuffle stage
                val mapStage = getShuffleMapStage(shufDep, stage.firstJobId)
                if (!mapStage.isAvailable) {
                  missing += mapStage
                }
                // 如果是窄依赖，那么将依赖的rdd放入栈
              case narrowDep: NarrowDependency[_] =>
                waitingForVisit.push(narrowDep.rdd)
            }
          }
        }
      }
    }
    // 首先往栈中推入了stage最后一个rdd
    waitingForVisit.push(stage.rdd)
    while (waitingForVisit.nonEmpty) {
      // 对stage最后一个rdd，调用visit方法
      visit(waitingForVisit.pop())
    }
    missing.toList
  }

stage划分算法很重要，因为对于spark高手，或者spark精通人员来说 , 必须对stage划分算法很清晰，直到你自己编写的spark application被划分为几个job , 每个job划分成了几个stage
每个stage包括哪些代码 , 只有知道了每个stage包括了你的哪些代码之后 . 在线上，如果你发现某个stage执行特别慢，或者某个stage一致报错,你才能针对那个stage对应的代码去排查问题，或者性能调优