Spark Shuffle源码分析系列之MapOutputTracker

MapOutputTracker是spark环境的主要组件之一，其功能有三方面，首先DAGScheduler使用MapOutputTrackerMaster来管理各个ShuffleMapTask的输出数据[MapStatus]，另外根据各ShuffleMapTask结果的统计信息来进行尽可能的本地化Scheduler；其次ShuffleMapStage使用MapOutputTrackerMaster来判断是否要进行ShuffleMapTask的计算；最后ShuffleReduce任务用来获取从哪些executor获取需要的map输出数据，读取数据进行处理。

架构概述

MapOutputTracker是Spark环境的主要组件之一，其功能是管理各个Shuffle输出的中间数据，主要有三个方面功能:

1. DAGScheduler使用MapOutputTrackerMaster来管理各个ShuffleMapTask的输出数据[MapStatus]，另外根据各ShuffleMapTask结果的统计信息来进行尽可能的本地化Scheduler。
   1. 创建ShuffleMapTask时候注册当前Shuffle，需要Shuffle的唯一标识ShuffleId和Map的个数，通过函数createShuffleMapStage调用registerShuffle来完成注册；
   2. ShuffleMapTask计算完成时候，计算成功则更新MapStatus信息，计算失败则删除Stage管理信息，通过handleTaskCompletion调用registerMapOutput来完成计算结果的注册；
   3. 在Executor/Worker删除时候，清理该Executor/Worker上面的所有的ShuffleId对应的Map结果信息，通过handleTaskCompletion调用removeOutputsOnHost或者removeOutputsOnExecutor进行；
2. ShuffleMapStage使用MapOutputTrackerMaster来判断是否要进行ShuffleMapTask的计算。
   1. ShuffleMapStage通过findMissingPartitions来查找还没计算好的map，进行计算;
3. ShuffleReduce任务使用MapOutputTrackerWorker用来获取从哪些executor获取需要的map输出数据，读取数据进行处理。
   1. ShuffleRDD在compute中，从MapOutputTrackerWorker中的getMapSizesByExecutorId函数获取当前Reduce需要的数据的位置信息，计算reduce计算。

所以，可以看出来要想完成上述功能，我们需要Driver端有一个MapOutputTrackerMaster，来管理所有的ShuffleMapTask的结果，Executor有一个MapOutputTrackerWorker，用来向Driver申请获取相应的ShuffleId对应的Map结果信息，所以还需要Driver端有一个常驻的RPCEndPoint MapOutputTrackerMasterEndpoint，用来接收Executor的获取ShuffleId对应Map结果信息，然后转交给MapOutputTrackerMaster进行数据获取然后进行结果数据获取，回传给Executor。整体架构图如下:

可以看出来，MapOutputTracker在executor和driver端都存在:

1. MapOutputTrackerMaster和MapOutputTrackerMasterEndpoint存在于driver端。MapOutputTrackerMasterEndpoint是MapOutputTrackerMaster的RPC Endpoint，处理请求；MapOutputTrackerMaster负责管理所有shuffleMapTask的输出数据，每个shuffleMapTask执行完后会把执行结果[MapStatus]注册到MapOutputTrackerMaster；另外MapOutputTrackerMaster还会处理executor发送的GetMapOutputStatuses请求，并返回serializedMapStatus给executor端。
2. MapOutputTrackerWorker存在于executor端，MapOutputTrackerWorker负责为Reduce任务提供需要的shuffleMapTask的输出数据信息[MapStatus]，如果MapOutputTrackerWorker在本地没有找到请求的Shuffle的mapStatuses，则会向MapOutputTrackerMasterEndpoint发送GetMapOutputStatuses请求获取对应的mapStatuses。

本文通过源码分析了各个组件的主要功能。

Shuffle结果

在Spark中，MapStatus用于表示ShuffleMapTask的运行结果，ShuffleStatus用于管理一个Shuffle对应的所有ShuffleMapTask的运行结果。

MapStatus

MapStatus用于表示ShuffleMapTask的运行结果, 包括map任务输出数据的location和size信息：

// org.apache.spark.scheduler.MapStatus
private[spark] sealed trait MapStatus {
   
  /**  用于返回ShuffleMapTask运行的位置，即所在节点的BlockManager的身份标识BlockManagerId **/
  def location: BlockManagerId

  /** 用于返回reduce任务需要拉取的Block的大小（单位为字节） */
  def getSizeForBlock(reduceId: Int): Long
}

根据ShuffleTask各个Partition的长度是否大于2000，分别创建HighlyCompressedMapStatus和CompressedMapStatus：

• 对于较大的数据量使用高度压缩的HighlyCompressedMapStatus。
• 一般的数据量则使用CompressedMapStatus。

这块我们只做简单的了解，后续文章我们会详细介绍。

ShuffleStatus

ShuffleStatus用于管理一个Shuffle对应的所有ShuffleMapTask的运行结果，ShuffleStatus对象只存在于MapOutputTrackerMaster中。ShuffleStatus中使用一个数组mapStatuses来保存所有的MapStatus，该数组的元素类型为MapStatus，数组下标是ShuffleMapTask的map id。由于executor获取时候，需要通过网络传输，使用的是序列化数据，为了加速访问，会对当前mapStatus的结果进行序列化缓存；另外有可能结果数据比较大<有参数spark.shuffle.mapOutput.minSizeForBroadcast控制，默认是512k>，就需要进行broadcast，使得executor可以从别的地方进行下载。

// org.apache.spark.ShuffleStatus
val mapStatuses = new Array[MapStatus](numPartitions)
private[this] var cachedSerializedMapStatus: Array[Byte] = _
private[this] var cachedSerializedBroadcast: Broadcast[Array[Byte]] = _

为了保证线程安全性，ShuffleStatus使用了读写锁来进行读取和写入的同步和互斥控制。另外ShuffleStatus还提供了用于添加，移除，序列化，缓存和广播mapStatus的方法。