概述
本文介绍DAGScheduler的实现原理。通过文章《Spark2原理分析-DAGScheduler(Stage调度器)的基本概念》我们学习了DAGScheduler的基本概念，并了解了它的功能。
这篇文章，介绍DAGScheduler的具体实现。为避免篇幅过长，本文先介绍DAGScheduler的一些关键框架，和成员变量和函数。并介绍一些关键函数的实现。具体的各个框架的详细实现，放到后面的文章进行讲解。
DAGScheduler功能概况
我们先通过查看该类的关键成员变量和成员函数的功能，可以大概了解一下该类实现了哪些功能。
DAGScheduler中的成员变量
| 变量名 | 类型 | 说明 |
|---|---|---|
| nextJobId | AtomicInteger | 用来生成Job的下一个id |
| numTotalJobs | Int | 获取目前的job总数：nextJobId.get() |
| jobIdToStageIds | HashMap[Int, HashSet[Int]] | 通过jobid获取对应的Stage的id集合 |
| stageIdToStage | HashMap[Int, Stage] | 通过StageId获取对应的stage集合 |
| shuffleIdToMapStage | HashMap[Int, ShuffleMapStage] | shuffle的依赖id和ShuffleMapStage的对应关系的集合 |
| jobIdToActiveJob | HashMap[Int, ActiveJob] | jobId对应的ActiveJob(正在运行的job)集合 |
| waitingStages | HashSet[Stage] | 需要运行的Stage的集合 |
| runningStages | HashSet[Stage] | 正在运行的Stage集合 |
| failedStages | HashSet[Stage] | 由于执行失败而必须要重新提交的Stage集合 |
| cacheLocs | HashMap[Int, IndexedSeq[Seq[TaskLocation]]] | 每个分区和被缓存的位置的映射集 |
| failedEpoch | HashMap[String, Long] | 为跟踪失败的node |
| outputCommitCoordinator | OutputCommitCoordinator | 对任务输出到HDFS的动作进行认证 |
| eventProcessLoop | DAGSchedulerEventProcessLoop | 事件处理框架对象 |
| DAGSchedulerEventProcessLoop | 处理框架类的定义 |
DAGScheduler中的成员函数
| 函数名 | 返回值类型 | 说明 |
|---|---|---|
| getCacheLocs | IndexedSeq[Seq[TaskLocation]] | 获取rdd分区缓存的位置 |
| getOrCreateShuffleMapStage | Int | 从shuffleIdToMapStage中获取Stage，若没有获取到，则创建一个 |
| createShuffleMapStage | ShuffleMapStage | 创建一个ShuffleMapStage，它用来产生所给shuffle依赖的分区 |
| createResultStage | ResultStage | 创建一个和所给jobId关联的ResultStage |
| getOrCreateParentStages | Int | 为给定的rdd获取或创建父stage的列表 |
| getMissingAncestorShuffleDependencies | ArrayStack[ShuffleDependency[_, _, _]] | 查找不在shuffleToMapStage结合中的祖先shuffle依赖 |
| getShuffleDependencies | HashSet[ShuffleDependency[_, _, _]] | 返回给定rdd的shuffle的依赖 |
| getMissingParentStages | List[Stage] | 返回缺失的Stage |
| submitJob | JobWaiter[U] | 向scheduler提交一个由rdd的action操作产生的job |
| runJob | Unit | 运行rdd的job，并把结果传递给resultHandler |
| submitMapStage | 提交并独立运行一个shuffle的map阶段，返回一个JobWaiter对象 | |
| handleJobSubmitted | 处理Job提交的事件 | |
| handleMapStageSubmitted | 处理map的stage提交事件 | |
| submitMissingTasks | 提交缺失的task，当父stage可用时该函数被调用 | |
| handleTaskCompletion | 处理task完成的事件 |
DAGScheduler实体的创建和初始化
DAGScheduler的创建
DAGScheduler只在Spark driver中运行，它作为SparkContext初始化过程的一部分。创建实体的步骤如下：
- 首先创建任务调度器(taskScheduler)和schedulerBackend
- 再创建DAGScheduler实体
- 启动任务调度器
通过下面的SparkContext的实现代码可以看到SparkContext是如何初始化DAGScheduler实例的：
class SparkContext(config: SparkConf) extends Logging {
...
// Create and start the scheduler
// 先创建任务调度器和schedulerBackend实体
val (sched, ts) = SparkContext.createTaskScheduler(this, master, deployMode)
_schedulerBackend = sched
_taskScheduler = ts
// 再创建DAGScheduler实体
_dagScheduler = new DAGScheduler(this)
_heartbeatReceiver.ask[Boolean](TaskSchedulerIsSet)
// start TaskScheduler after taskScheduler sets DAGScheduler reference in DAGScheduler's
// constructor
// 启动任务调度器（后面会重点讲到）
_taskScheduler.start()
...
}
DAGScheduler的初始化
DAGScheduler初始化的过程如下图所示：
可以看到，SparkContext会传递一个参数this给DAGScheduler构造函数：
_dagScheduler = new DAGScheduler(this)
我们进入到DAGScheduler内部查看一下具体的实现，DAGScheduler会调用以下构造函数进行构造：
def this(sc: SparkContext, taskScheduler: TaskScheduler) = {
this(
sc,
taskScheduler,
sc.listenerBus,
sc.env.mapOutputTracker.asInstanceOf[MapOutputTrackerMaster],
sc.env.blockManager.master,
sc.env)
}
理解SparkContext传递给DAGScheduler构造函数的参数很重要，因为这是实现DAGScheduler的基础。DAGScheduler从SparkContext中继承的参数说明如下：
| 参数名 | 类型 | 说明 |
|---|---|---|
| sparkContext | SparkContext上下文 | |
| taskScheduler | task调度器 | |
| sc.listenerBus | 事件监听总线 | |
| MapOutputTrackerMaster | map输出跟踪 | |
| BlockManagerMaster | 块管理器 | |
| SparkEnv | Spark环境信息 |
DAGScheduler的事件处理框架
DAGScheduler事件处理对象的创建和启动
在DAGScheduler类定义中，先构建事件处理框架对象。构建完成后，会调用start()函数来启动事件处理线程，代码实现如下：
class DAGScheduler(...) {
...
private[scheduler] val eventProcessLoop = new DAGSchedulerEventProcessLoop(this)
...
eventProcessLoop.start()
}
事件处理框架的实现
上图标识了DAGScheduler的事件处理框架。在发生事件时DAGScheduler会把事件保存到一个阻塞的事件队列中，此时会有一个处理线程在队列的另一端监听，若发现有事件放到队列，立即得到通知，并把该事件取出，并调用DAGScheduler对象中的事件处理函数进行处理。
DAGScheduler中能够处理的事件，和对应的处理函数如上图所示。
shuffleIdToMapStage成员变量
该成员变量是一个HashMap，它表示：从shuffle依赖项ID(shuffle dependency ID)到为该依赖项生成数据的ShuffleMapStage的映射。
shuffle dependency ID -> ShuffleMapStage
仅包括当前正在运行的作业的一部分的阶段（当shuffle阶段的作业完成时，映射将被删除，并且shuffle数据的唯一记录将在MapOutputTracker中）。
该成员变量的定义如下：
private[scheduler] val shuffleIdToMapStage = new HashMap[Int, ShuffleMapStage]
getOrCreateShuffleMapStage
这是DAGScheduler的一个成员函数，它的原型如下：
private def getOrCreateShuffleMapStage(
shuffleDep: ShuffleDependency[_, _, _],
firstJobId: Int): ShuffleMapStage
该函数的功能是：为ShuffleDependency查找或创建ShuffleMapStage。
该函数先在shuffleIdToMapStage中查找ShuffleMapStage，若找到了则返回，若找不到则创建一个。
getShuffleDependencies
该函数查找给定RDD的直接父shuffle依赖项。
但该函数不会返回更远的祖先。 例如，如果C对B有一个shuffle依赖，B对A有一个shuffle依赖，如下：
A<---B<---C
使用rdd C作为参数调用此函数，只会返回B <- C依赖项。
该函数的原型如下：
private[scheduler] def getShuffleDependencies(
rdd: RDD[_]): HashSet[ShuffleDependency[_, _, _]]
当DAGScheduler找到或创建缺少的直接父ShuffleMapStages（对于给定RDD的ShuffleDependencies）并且找到给定RDD的所有缺失的shuffle依赖性时，使用getShuffleDependencies。
getMissingParentStages函数
该函数的声明如下：
private def getMissingParentStages(stage: Stage): List[Stage]
主要功能是：在输入阶段的依赖关系图中找到缺少的父ShuffleMapStages（使用广度优先搜索算法），并按List返回。
该函数的实现流程如下：
- (1)从阶段的RDD开始，并向上遍历所有父RDD的树以查找未缓存的分区。
- (2)遍历RDD的父依赖项，并根据其类型进行操作，例如：ShuffleDependency或NarrowDependency。
- (3)对于每个NarrowDependency，getMissingParentStages只是标记要访问的相应RDD，然后转到RDD的下一个依赖项或在另一个未访问的父RDD上工作。
- (4)对于每个ShuffleDependency，getMissingParentStages都会找到ShuffleMapStage阶段。如果ShuffleMapStage不可用，则将其添加到缺失（映射）阶段的集合中。
submitJob函数：提交action任务
该函数向scheduler提交Action的job。具体来说，该函数会：创建一个JobWaiter实体，并发送JobSubmitted类型的事件。
该函数的实现逻辑如下：
(1) 获取参数rdd的分区长度，并检查这些分区是否选在参数partitions的集合中。
(2) 增加nextJobId内部工作计数器。
(3) 若分区集合参数partitions的长度为0，说明该job运行0个任务，直接返回总任务数为0的JobWaiter对象。
(4) 若分区集合不为0，获取每个分区的处理函数的实体
(5) 创建一个JobWaiter对象，并向DAGSchedulerEventProcessLoop发起一个JobSubmitted事件。
当SparkContext提交一个job，且DAGScheduler运行job时，才会调用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] = {...}
submitMissingTasks函数
该函数的主要功能是：在Spark Job中提交Stage缺失的任务。该函数的原型如下：
private def submitMissingTasks(stage: Stage, jobId: Int)
该函数的实现逻辑如下：
-
(1)查找stage的分区中缺失的分区的id
-
(2)标记stage的状态为running
-
(3)通知outputCommitCoordinator该stage启动了
-
(4)获取缺失分区的最合适的位置，若此时没有致命错误的异常，会尝试创建一个新的阶段。
-
(5)在LiveListenerBus上发布SparkListenerStageSubmitted消息。
-
(6) 根据阶段的类型和ShuffleDependency（对于ShuffleMapStage）或函数（对于ResultStage）对RDD进行序列化，用来向executors分发任务。
-
(7) 若发生异，中止该阶段（原因是“任务序列化失败”，然后是异常）并从阶段的内部runningStages集合中删除该阶段。submitMissingTasks退出。
-
(8) 分别为任务的每个缺失分区创建一个ShuffleMapTask或ResultTask，分别为ShuffleMapStage或ResultStage。submitMissingTasks使用每个分区的首选位置（前面已计算完成）。
-
(9) 记录stage的pendingPartitions属性中的（任务的）分区。
-
(10) 将任务提交给TaskScheduler执行（具有stage的id，尝试id，输入jobId以及带jobId的ActiveJob的属性）。
-
(11) 记录Stage的StageInfo中的提交时间并退出。
-
(12)最后，由于没有任务要提交执行，submitMissingTasks会提交等待子阶段以供执行和退出。
JobWaiter
该对象用来等待DAGScheduler的job完成。当任务完成时，该对象会把结果传输给给定的处理函数。
总结
本文介绍了DAGScheduler类的实现，包括对象初始化，事件处理框架，job的提交流程等。但为了不让整个篇幅过长，每个事件的具体处理的实现会放到后面的文章进行分析。