Spark Core DAG生成源码分析

前提回顾

  之前分析了Spark Streaming关于DAG的生成，对于Spark Streaming而言，先是通过自己代码中的各种transform方法来构造各个DStream之间的关联关系，然后再通过最后调用action操作的算子处进行回溯，action算子操作的DStream作为outputStream存储到DStreamGraph中的数组中，回溯过程中会找到数据来源处的DStream，这个DStream作为元素存储到DStreamGraph中的inputStream数组中。通过以上两个inputStream和outputStream数组成功存储了源头和结尾的DStream，又依据之前构造的各个DStream之间的关联，从而成功的构造完毕了DAG关系图。
  回顾完Spark Streaming中DAG的生成流程，理论来说RDD的DAG生成逻辑应该和DStream是类似的，我们直接深入代码来验证我们的想法。

transform算子分析

  显然和DStream类似，RDD中肯定也是通过各种transform操作来提前构造好各个RDD之间关系的，为了验证这个想法，需要查看各个算子的代码，先看用的最多的map算子代码如下：

> rdd = rdd0.map(...)

--> RDD.scala
def map[U: ClassTag](f: T => U): RDD[U] = withScope {
    val cleanF = sc.clean(f)
    new MapPartitionsRDD[U, T](this, (context, pid, iter) => iter.map(cleanF))
  }

--> MapPartitionsRDD.scala
private[spark] class MapPartitionsRDD[U: ClassTag, T: ClassTag](
    var prev: RDD[T],
    f: (TaskContext, Int, Iterator[T]) => Iterator[U],  // (TaskContext, partition index, iterator)
    preservesPartitioning: Boolean = false,
    isOrderSensitive: Boolean = false)
  extends RDD[U](prev) {

  override val partitioner = if (preservesPartitioning) firstParent[T].partitioner else None

  override def getPartitions: Array[Partition] = firstParent[T].partitions

  override def compute(split: Partition, context: TaskContext): Iterator[U] =
    f(context, split.index, firstParent[T].iterator(split, context))

  override def clearDependencies() {
    super.clearDependencies()
    prev = null
  }

  override protected def getOutputDeterministicLevel = {
    if (isOrderSensitive && prev.outputDeterministicLevel == DeterministicLevel.UNORDERED) {
      DeterministicLevel.INDETERMINATE
    } else {
      super.getOutputDeterministicLevel
    }
  }
}

  以上代码分为三部分，实例用法、底层抽象类RDD.scala对map方法的实现以及map返回类MapPartitionsRDD的实现。可以很明显的看到在调用map方法后，生成新的RDD的时候会传入之前的rdd用于构造新的rdd，也就是说和dstream完全类似就是通过transform来关联不同rdd之间的关系。这样就可以很好的理解流程就是通过各种transform来关联rdd，然后再通过action来触发其它操作，这里说的其它操作应该包括有依据构造的rdd关系来构造dag、生成job、生成stage等等。

action算子分析

  为了验证我们的想法，那么我们就继续深入代码，我们挑选collect算子来作为入口，来通过collect算子入口看action算子内部主要做了哪些事情。

def collect(): Array[T] = withScope {
    val results = sc.runJob(this, (iter: Iterator[T]) => iter.toArray)
    Array.concat(results: _*)
  }

  如上为collect的实现，可以看到主要就是调用了SparkContext的runJob方法，那么很显然了，这里就是深入源码内部的入口。我们来继续看runJob方法的实现：

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)
    }
    dagScheduler.runJob(rdd, cleanedFunc, partitions, callSite, resultHandler, localProperties.get)
    progressBar.foreach(_.finishAll())
    rdd.doCheckpoint()
  }

  以上代码省略了各种重载方法的调用链直接到最后的实现方法处，这里实现也很简单，主要就是去调用dagScheduler的runJob方法，我们来继续查看dagScheduler中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
    val waiter = submitJob(rdd, func, partitions, callSite, resultHandler, properties)
    ThreadUtils.awaitReady(waiter.completionFuture, Duration.Inf)
    waiter.completionFuture.value.get match {
      case scala.util.Success(_) =>
        logInfo("Job %d finished: %s, took %f s".format
          (waiter.jobId, callSite.shortForm, (System.nanoTime - start) / 1e9))
      case scala.util.Failure(exception) =>
        logInfo("Job %d failed: %s, took %f s".format
          (waiter.jobId, callSite.shortForm, (System.nanoTime - start) / 1e9))
        // SPARK-8644: Include user stack trace in exceptions coming from DAGScheduler.
        val callerStackTrace = Thread.currentThread().getStackTrace.tail
        exception.setStackTrace(exception.getStackTrace ++ callerStackTrace)
        throw exception
    }
  }

  这里主要也是为了调用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] = {

    val jobId = nextJobId.getAndIncrement()
    if (partitions.size == 0) {
      // 任务分区数为0则直接返回。
      return new JobWaiter[U](this, jobId, 0, resultHandler)
    }

    val func2 = func.asInstanceOf[(TaskContext, Iterator[_]) => _]
     // job任务监听器，任务完成后的数据交给resultHandle来处理
    val waiter = new JobWaiter(this, jobId, partitions.size, resultHandler)
    // 发布消息
    eventProcessLoop.post(JobSubmitted(
      jobId, rdd, func2, partitions.toArray, callSite, waiter,
      SerializationUtils.clone(properties)))
    waiter
  }

  如上submitJob主要就是创建一个job监听器，然后和任务一起发布到eventProcessLoop中，然后会有对应的事件响应处理，对于任务提交JobSubmitted类型事件的处理方法为dagScheduler.handleJobSubmitted方法，具体实现如下：

private[scheduler] def handleJobSubmitted(jobId: Int,
      finalRDD: RDD[_],
      func: (TaskContext, Iterator[_]) => _,
      partitions: Array[Int],
      callSite: CallSite,
      listener: JobListener,
      properties: Properties) {
    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 = 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
    }

    // 获取这个stage触发的job
    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
    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)
  }

  这里可以发现此处会有一个变量叫finalStage，这个可以理解就是最后一个stage，我们再来思考我们是如何一路走到这里的，就是从action算子触发一路走到这里，所以很显然可以知道这个finalStage就是DAG中进行stage划分时最后的一个stage。因为action算子就已经是触发计算任务的地方，确实就应该是一个DAG的末尾了。
  以上分析，可以看出这里应该也是和dstream类似，是通过rdd的action算子来进行回溯的方式找到源头rdd，中间会触发一些事件暂时用于他处这里不进行分析，主要围绕stage的划分和启动流程，那么对应的就直接看到最后的submitStage方法，代码如下：

private def submitStage(stage: Stage) {
    val jobId = activeJobForStage(stage)
    if (jobId.isDefined) { // 之前已经定义过了jobId
      logDebug("submitStage(" + stage + ")")
      // 本次提交的stage状态不是等待、运行、失败状态，说明应该是还未正式提交启动运行的
      if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {
        // 获取当前stage所依赖的且一样没有提交的stage，这个过程就是一个回溯的过程
        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 {
          for (parent <- missing) {
            submitStage(parent)
          }
          waitingStages += stage
        }
      }
    } else {
      abortStage(stage, "No active job for stage " + stage.id, None)
    }
  }

  到上面这一步的时候，注意最开始调用submitStage这个方法的时候，传过来的stage是由action触发得到的，是finalStage。这里需要知道的一点就是，每个stage刚开始创建出来的时候都只是对应一个rdd，然后在进行回溯的时候，其实是从最开始创建好的DAG图谱中最后的finalStage只是一个调用action操作的rdd，通过这个rdd来回溯它所依赖的父类rdd，然后判断和所依赖的rdd的依赖关系是宽依赖还是窄依赖，如果是窄依赖（即分区之间一对一的依赖）关系，则把依赖的rdd直接吸收到本次的stage中；反之如果是宽依赖，则会依据所依赖的rdd生成一个新的stage，然后这两个stage之间建立联系。如此往复一直沿着之前构建的rdd关系链寻找所依赖的rdd和被依赖rdd之间的关系来创建stage，最终会把所有的rdd划分为若干个stage。

以上介绍的是stage划分的逻辑，还可以补充的一点是关于task以及job的产生，对于job而言，对应的就是一个action算子，每一个action操作都会触发任务执行得到结果，对应会有一个小的DAG关系图，这个action所触发计算所涉及的计算任务就是一个job。而具体到划分完stage之后，每个stage所关联的rdd中会有若干个partition，每个partition计算都会产生一个task，以上所有的stage、job最终都是为了最终的计算，而最终的计算单位就是以一个一个的partition作为单位得到一个交给Executor执行的task。