Spark 源碼分析（七）: DAGScheduler 源碼分析1（stage 劃分算法）

前面幾篇文章已經說清楚了從 spark 任務提交到 driver 啓動，然後執行 main 方法，初始化 SparkContext 對象。

在初始化 SparkContext 對象的過程中創建了兩個重要組件：

一個是 TaskScheduler（實際上是他的實現類 TaskSchedulerImpl 對象），這個對象內部會持有一個 SchedulerBackend 對象，SchedulerBackend 內部會又會持有一個 DriverEndpoint 對象（實際上就是一個 RpcEndpoint）。這樣 TaskScheduler 就可以通過 SchedulerBackend 和集羣資源管理器或者 Executor 對應 worker 節點進行通信做一些事情。比如向 master 節點去註冊 application，master 在註冊 application 的過程中會分配 worker 去啓動 Executor，當 Executor 啓動後又會和 TaskScheduler 進行註冊。

另一個是 DAGScheduler，關於這個對象的創建過程前面沒有詳細講，主要是因爲 DAGScheduler 是在 SparkContext 初始化結束後，執行到 RDD 的 Action 操作的時候纔會開始工作，下面就從 RDD 的 action 操作說起，看看 DAGScheduler 是怎麼工作的。

還是以 wordcount 程序爲例：

		val conf = new SparkConf()
      .setAppName("WordCount")
      .setMaster("local")
    val sc = new SparkContext(conf)
    val lines = sc.textFile("./file/localfile")
    val words = lines.flatMap(line => line.split(" "))
    val wordPairs = words.map(word => (word, 1))
    val wordCounts = wordPairs.reduceByKey(_ + _)
    wordCounts.foreach(wordCount => println(wordCount._1 + "  " + wordCount._2))

當代碼執行到 wordCounts.foreach 時候會調用到 RDD 的 foreach 方法，RDD 的 foreach 方法會去調用 SparkContext 的 runjob 方法。

SparkContext 中會有多個 runjob 方法，最後都會走到一個 runjob 那裏去，這個 runjob 方法最終會調用 DAGScheduler 的 runJob 的方法，具體可以先看下這個 SparkContext 的 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 方法
    dagScheduler.runJob(rdd, cleanedFunc, partitions, callSite, resultHandler, localProperties.get)
    progressBar.foreach(_.finishAll())
    rdd.doCheckpoint()
  }

最主要的還是 DAGScheduler 中的 runjob 方法。

這個 runjob 方法內部實際上調用了 submitJob 方法，用於提交 job。該方法返回一個 JobWaiter，用於等待 DAGScheduler 任務的完成。

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
    // 調用 submitJob 方法
    val waiter = submitJob(rdd, func, partitions, callSite, resultHandler, properties)

submitJob 方法是調用 eventProcessLoop 的 post 方法將 JobSubmitted 事件添加到 DAGScheduler 的事件隊列中去。

eventProcessLoop.post(JobSubmitted(
      jobId, rdd, func2, partitions.toArray, callSite, waiter,
      SerializationUtils.clone(properties)))

這裏的 eventProcessLoop 是 DAGSchedulerEventProcessLoop 對象，在 DAGScheduler 的初始化代碼中可以看到。DAGSchedulerEventProcessLoop 實際上內部有一個線程，用來處理事件隊列。

事件隊列的處理最後會走到 DAGSchedulerEventProcessLoop 的 onReceive 的回調方法裏面去。

/**
   * The main event loop of the DAG scheduler.
   */
  override def onReceive(event: DAGSchedulerEvent): Unit = {
    val timerContext = timer.time()
    try {
      // 調用 doOnReceive 方法
      doOnReceive(event)
    } finally {
      timerContext.stop()
    }
  }

後面會去調用 doOnReceive 方法，根據 event 進行模式匹配，匹配到 JobSubmitted 的 event 後實際上是去調用 DAGScheduler 的 handleJobSubmitted 這個方法。

private def doOnReceive(event: DAGSchedulerEvent): Unit = event match {
  	// 模式匹配
    case JobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties) =>
      // 調用 handleJobSubmitted 方法
      dagScheduler.handleJobSubmitted(jobId, rdd, func, partitions, callSite, listener, properties)

下面來看 handleJobSubmitted 這個方法做了哪些操作：

1，使用觸發 job 的最後一個 rdd，創建 finalStage；

注： Stage 是一個抽象類，一共有兩個實現，一個是 ResultStage，是用 action 中的函數計算結果的 stage；另一個是 ShuffleMapStage，是爲 shuffle 準備數據的 stage。

2，構造一個 Job 對象，將上面創建的 finalStage 封裝進去，這個 Job 的最後一個 stage 也就是這個 finalStage；

3，將 Job 的相關信息保存到內存的數據結構中；

4，調用 submitStage 方法提交 finalStage。

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 {
      // 使用觸發 job 的最後一個 RDD 創建一個 ResultStage
      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
    }

  	// 使用前面創建好的 ResultStage 去創建一個 job
  	// 這個 job 的最後一個 stage 就是 finalStage
    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))

    // 將 job 的相關信息存儲到內存中
    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))
    // 提交 finalStage
    submitStage(finalStage)
  }

下面就會走進 submitStage 方法，這個方法是用來提交 stage 的，具體做了這些操作：

1，首先會驗證 stage 對應的 job id 進行校驗，存在纔會繼續執行；

2，在提交這個 stage 之前會判斷當前 stage 的狀態。

如果是 running、waiting、failed 的話就不做任何操作。

如果不是這三個狀態則會根據當前 stage 去往前推前面的 stage，如果能找到前面的 stage 則繼續遞歸調用 submitStage 方法，直到當前 stage 找不到前面的 stage 爲止，這時候的 stage 就相當於當前 job 的第一個 stage，然後回去調用 submitMissingTasks 方法去分配 task。

private def submitStage(stage: Stage) {
    val jobId = activeJobForStage(stage)
    // 看看當前的 job 是否存在
    if (jobId.isDefined) {
      logDebug("submitStage(" + stage + ")")
       // 判斷當前 stage 的狀態
      if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {
        // 根據當前的 stage 去推倒前面的 stage
        val missing = getMissingParentStages(stage).sortBy(_.id)
        logDebug("missing: " + missing)
        // 如果前面已經沒有 stage 了，那麼久將當前 stage 去執行 submitMissingTasks 方法
        // 如果前面還有 stage 的話那麼遞歸調用 submitStage
        if (missing.isEmpty) {
          logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
          submitMissingTasks(stage, jobId.get)
        } else {
          for (parent <- missing) {
            submitStage(parent)
          }
          // 將當前 stage 加入等待隊列
          waitingStages += stage
        }
      }
    } else {
      // abortStage 終止提交當前 stage
      abortStage(stage, "No active job for stage " + stage.id, None)
    }
  }

上面最重要的一個地方就是使用當前 stage 向前推，找到前面的 stage，也是 stage 的劃分算法。下面就看看 getMissingParentStages 這個劃分算法做了哪些操作：

1，創建 missing 和 visited 兩個 HashSet，分別用來存儲根據當前 stage 向前找到的所有 stage 數據和已經調用過 visit 方法的 RDD；

2，創建一個存放 RDD 的棧，然後將傳進來的 stage 中的 rdd 也就是 finalStage 中的那個 job 觸發的最後一個 RDD 放入棧中；

3，然後將棧中的 RDD 拿出來調用 visit 方法，這個 visit 方法內部會根據當前 RDD 的依賴鏈逐個遍歷所有 RDD，並且會根據相鄰兩個 RDD 的依賴關係來決定下面的操作：

如果是寬依賴，即 ShuffleDependency ，那麼會調用 getOrCreateShuffleMapStage 創建一個新的 stage，默認每個 job 的最後一個 stage 是 ResultStage，剩餘的 job 中的其它 stage 均爲 ShuffleMapStage。然後會將創建的這個 stage 加入前面創建的 missing 的 HashSet 中；

如果是窄依賴，即 NarrowDependency，那麼會將該 RDD 加入到前面創建的 RDD 棧中，繼續遍歷調用 visit 方法。

直到所有的 RDD 都遍歷結束後返回前面創建的 missing 的集合。

private def getMissingParentStages(stage: Stage): List[Stage] = {
  	// 存放下面找到的所有 stage
    val missing = new HashSet[Stage]
    // 存放已經遍歷過的 rdd
    val visited = new HashSet[RDD[_]]
    // We are manually maintaining a stack here to prevent StackOverflowError
    // caused by recursively visiting
    // 創建一個維護 RDD 的棧
    val waitingForVisit = new Stack[RDD[_]]
    // visit 方法
    def visit(rdd: RDD[_]) {
      // 判斷當前 rdd 是否 visit 過
      if (!visited(rdd)) {
        visited += rdd
        val rddHasUncachedPartitions = getCacheLocs(rdd).contains(Nil)
        if (rddHasUncachedPartitions) {
          // 遍歷當前 RDD 的依賴鏈
          for (dep <- rdd.dependencies) {
            dep match {
              // 如果是寬依賴
              case shufDep: ShuffleDependency[_, _, _] =>
                // 創建 ShuffleMapStage 
                val mapStage = getOrCreateShuffleMapStage(shufDep, stage.firstJobId)
                if (!mapStage.isAvailable) {
                  // 加入 missing 集合
                  missing += mapStage
                }
              // 如果是窄依賴
              case narrowDep: NarrowDependency[_] =>
                // 加入等待 visit 的集合中，準備下一次遍歷
                waitingForVisit.push(narrowDep.rdd)
            }
          }
        }
      }
    }
  	// 將傳入的 stage 中的 rdd 拿出來壓入 waitingForVisit 的棧中
    waitingForVisit.push(stage.rdd)
    // 遍歷棧裏的所有 RDD 
    while (waitingForVisit.nonEmpty) {
      // 調用 visit 方法
      visit(waitingForVisit.pop())
    }
  	// 返回 missing 這個 stage 集合
    missing.toList
  }

至此，所有的 stage 都已經劃分結束了。可以看出每個 Spark Application 執行代碼的時候，每當碰到一個 Action 操作就會劃分出一個 Job，然後每個 Job 裏會根據寬窄依賴去劃分出多個 stage。

注：最後在補充一點，RDD 劃分寬窄依賴的方法是根據父子 RDD 的 Partition 對應關係來定的：窄依賴是父 RDD 的每一個分區最多被一個子 RDD 的分區所用，也就是說父 RDD 的一個或多個 Partition 對應一個子 RDD 的一個 Partition；寬依賴是指子 RDD 的分區依賴於父 RDD 的所有分區。