spark源码分析- shuffle read

Spark会将job划分为多个Stage，每个job会由多个ShuffleMapStage和一个ResultStage组成，然后每个Stage会由多个Task组成，Task数量和每个Stage的Partition的数量相同。每个Task任务由单独的线程执行，不同Stage的Task之间需要进行数据流动，并且下游Stage的Task会依赖上游Stage的多个Task，所以该过程需要将数据写入磁盘，并将属于下游Stage相同Task的数据汇总到一起，以供该Task进行拉取，该过程被称为Shuffle。 shuffle过程分为shuffle上游Stage的Task的写过程和shuffle下游Stage的Task的读过程。上次分析了shuffle write过程，接下来在看一下shuffle下游Stage如何拉取数据。
首先BlockStoreShuffleReader会从MapOutputTrackerMaster获取需要读取的shuffle数据的位置和大小。

override def read(): Iterator[Product2[K, C]] = {
    val wrappedStreams = new ShuffleBlockFetcherIterator(
      context,
      blockManager.shuffleClient,
      blockManager,
      // 通过shuffleId获取shuffle数据位置和大小
      mapOutputTracker.getMapSizesByExecutorId(handle.shuffleId, startPartition, endPartition),
      serializerManager.wrapStream,
      SparkEnv.get.conf.getSizeAsMb("spark.reducer.maxSizeInFlight", "48m") * 1024 * 1024,
      SparkEnv.get.conf.getInt("spark.reducer.maxReqsInFlight", Int.MaxValue),
      SparkEnv.get.conf.get(config.REDUCER_MAX_BLOCKS_IN_FLIGHT_PER_ADDRESS),
      SparkEnv.get.conf.get(config.MAX_REMOTE_BLOCK_SIZE_FETCH_TO_MEM),
      SparkEnv.get.conf.getBoolean("spark.shuffle.detectCorrupt", true))
    ...
  }

Executor端的MapOutputTrackerWorker服务会想MapOutputTrackerMaster发送GetMapOutputStatuses RPC调用获取shuffle数据位置和大小信息。

private def getStatuses(shuffleId: Int): Array[MapStatus] = {
    // 如果有task已经发送过GetMapOutputStatuses请求获取了shuffle相关信息，则其他task可以直接使用。不用再此获取
    val statuses = mapStatuses.get(shuffleId).orNull
    // 如果还没有本次shuffle相关的信息，则从MapOutputTrackerMaster获取
    if (statuses == null) {
      logInfo("Don't have map outputs for shuffle " + shuffleId + ", fetching them")
      val startTime = System.currentTimeMillis
      var fetchedStatuses: Array[MapStatus] = null
      // 如果有其他task正在获取相同shuffle的信息，直接等待该task获取完成，然后直接使用即可，不需再次拉取
      fetching.synchronized {
        while (fetching.contains(shuffleId)) {
          try {
            fetching.wait()
          } catch {
            case e: InterruptedException =>
          }
        }

        fetchedStatuses = mapStatuses.get(shuffleId).orNull
        if (fetchedStatuses == null) {
          fetching += shuffleId
        }
      }

      if (fetchedStatuses == null) {
        logInfo("Doing the fetch; tracker endpoint = " + trackerEndpoint)
        try {
        // 向MapOutputTrackerMaster发送GetMapOutputStatuses RPC请求获取shuffle信息
          val fetchedBytes = askTracker[Array[Byte]](GetMapOutputStatuses(shuffleId))
        // 反序列化数据
          fetchedStatuses = MapOutputTracker.deserializeMapStatuses(fetchedBytes)
          logInfo("Got the output locations")
          mapStatuses.put(shuffleId, fetchedStatuses)
        } finally {
          fetching.synchronized {
            fetching -= shuffleId
            fetching.notifyAll()
          }
        }
      }
      logDebug(s"Fetching map output statuses for shuffle $shuffleId took " +
        s"${System.currentTimeMillis - startTime} ms")

      if (fetchedStatuses != null) {
        fetchedStatuses
      } else {
        logError("Missing all output locations for shuffle " + shuffleId)
        throw new MetadataFetchFailedException(
          shuffleId, -1, "Missing all output locations for shuffle " + shuffleId)
      }
    } else {
      statuses
    }
  }

MapOutputTrackerMaster接收到GetMapOutputStatuses请求后，会将shuffle相关的MapStatus数据序列化，然后使用GZIP压缩算法进行压缩返回。压缩后的序列化数据如果大于spark.shuffle.mapOutput.minSizeForBroadcast配置项的设定值，则使用广播的方式进行shuffle数据的发送，该配置项默认值为512K。

// MapOutputTrackerMaster根据shuffleId查找到对应的ShuffleStatus，然后将shuffle数据序列化并压缩，然后返回给Executor端的MapOutputTrackerWorker
    val shuffleStatus = shuffleStatuses.get(shuffleId).head
    context.reply(
      shuffleStatus.serializedMapStatus(broadcastManager, isLocal, minSizeForBroadcast))
// 序列化并压缩ShuffleStatus对象持有的所有MapStatus，如果是首次获取该shuffleId的shuffle数据，则对所有MapStatus进行序列化并压缩，之后会将该shuffleId的序列化shuffle数据缓存起来，避免重复序列化，提高其他Executor获取shuffle数据的速度。  
def serializedMapStatus(
      broadcastManager: BroadcastManager,
      isLocal: Boolean,
      minBroadcastSize: Int): Array[Byte] = synchronized {
    // 如果是首次获取该shuffleId的MapStatus数据，则对shuffle数据进行序列化并压缩
    if (cachedSerializedMapStatus eq null) {
      val serResult = MapOutputTracker.serializeMapStatuses(
          mapStatuses, broadcastManager, isLocal, minBroadcastSize)
      cachedSerializedMapStatus = serResult._1
      cachedSerializedBroadcast = serResult._2
    }
    // 如果已有缓存数据，直接返回
    cachedSerializedMapStatus
  }

// 这里是真正序列化shuffle数据的方法，该方法会将整个MapStatus数组序列化，然后返回，Exeuctor端MapOutputTrackerWorker会把它重新反序列化。如果序列化并压缩后的MapStatus数组的大小大于spark.shuffle.mapOutput.minSizeForBroadcast配置项的设定值(默认为512K)，则使用广播的方式发送数据。Executor端MapOutputTrackerWorker通过判断返回数据的第一个byte值来区分数据的返回方式是直接返回还是使用广播的方式返回(0表示直接返回，1表示广播返回)
def serializeMapStatuses(statuses: Array[MapStatus], broadcastManager: BroadcastManager,
      isLocal: Boolean, minBroadcastSize: Int): (Array[Byte], Broadcast[Array[Byte]]) = {
    val out = new ByteArrayOutputStream
    // 首先尝试使用直接返回方式，所以返回的数据的第一个字节为0
    out.write(DIRECT)
    val objOut = new ObjectOutputStream(new GZIPOutputStream(out))
    Utils.tryWithSafeFinally {
      statuses.synchronized {
      // 序列化数据并压缩
        objOut.writeObject(statuses)
      }
    } {
      objOut.close()
    }
    val arr = out.toByteArray
    // 判断返回数据大小，大于spark.shuffle.mapOutput.minSizeForBroadcast设定值，则使用广播方式返回
    if (arr.length >= minBroadcastSize) {
    // 广播shuffle数据，并生成广播变量用于返回，并通过该对象获取广播的shuffle数据
      val bcast = broadcastManager.newBroadcast(arr, isLocal)
      out.reset()
      out.write(