SparkCore — HashShuffle源码分析

HashShuffle源码分析

  之前分析了两种Shuffle的区别，现在我们通过源码来进行分析，首先看HashShuffle，回顾之前流程，Executor在接收到LaunchTask的消息后，调用executor的launchTask()方法，将Task封装为一个TaskRunner（线程），然后放入线程池中执行，在执行的时候最终会调用Task.run()方法，这里面调用了runTask()方法，在runTask里面就是真正执行task的地方了，前面也分析过了相应的源码。
  在runTask中首先获取ShuffleManager，它有两个子类HashShuffleManager和SortShuffleManager，我们先分析HashShuffleManager，它通过getWriter，获取一个HashShuffleWriter，接着调用它的write()方法，进行数据处理，和结果的文件写入，我们看一下HashShuffleWriter的write()方法。

override def write(records: Iterator[_ <: Product2[K, V]]): Unit = {
    // 首先判断是否需要在Map端本地聚合
    // 这里的话，如果是reduceByKey操作，那么dep.aggregator.isDefined就是true
    // dep.mapSideCombine也相应的是true
    val iter = if (dep.aggregator.isDefined) {
      if (dep.mapSideCombine) {
        // 这里就会执行本地聚合，比如本地有(hello, 1) (hello, 1) => (hello, 2)
        dep.aggregator.get.combineValuesByKey(records, context)
      } else {
        records
      }
    } else {
      require(!dep.mapSideCombine, "Map-side combine without Aggregator specified!")
      records
    }
    // 如果需要聚合，那么先进行本地聚合操作
    // 接着遍历数据，对每个数据调用partitioner，默认是HashPartitioner，生成bucketId。
    // 这就决定了，每一份数据要写入哪个bucket中，相同key一定写入同一个bucket中
    for (elem <- iter) {
      val bucketId = dep.partitioner.getPartition(elem._1)
      // 获取到bucketId之后，会调用shuffleBlockManager.forMapTask()方法，生成bucketId对应的writer
      // 然后调用writer将数据写入bucket
      shuffle.writers(bucketId).write(elem)
    }
  }

  这里比较重要的是，生成bucketId，对每个数据调用HashPartitioner，也就是对每个数据进行hash操作，那么对于相同key的数据就会分到同一个bucket（缓存）中，因为他们的bucketId是相同的。
  接着调用ShuffleBlockManager的forMapTask生成writer，然后调用它的write方法将数据写入磁盘文件中。其中这里的ShuffleBlockManager是一个trait，它的子类是FileShuffleBlockManager，我们去这里面看forMapTask()方法：

def forMapTask(shuffleId: Int, mapId: Int, numBuckets: Int, serializer: Serializer,
      writeMetrics: ShuffleWriteMetrics) = {
    new ShuffleWriterGroup {
      shuffleStates.putIfAbsent(shuffleId, new ShuffleState(numBuckets))
      private val shuffleState = shuffleStates(shuffleId)
      private var fileGroup: ShuffleFileGroup = null

      val openStartTime = System.nanoTime
      // 这里就很关键，对应我们之前说的，HashShuffle有两种模式，一种普通的，一种是优化后的，这里就会判断,
      // 如果开启了consolidation机制，即consolidateShuffleFile为true的话
      // 不会给每个bucket都获取一个独立的文件
      // 而是为这个bucket获取一个ShuffleGroup的Writer
      val writers: Array[BlockObjectWriter] = if (consolidateShuffleFiles) {
        fileGroup = getUnusedFileGroup()
        Array.tabulate[BlockObjectWriter](numBuckets) { bucketId =>
          // 首先，用shuffleId，mapId，bucketId(reduceId)生成一个唯一的ShuffleBlockId
          // 然后用bucketId，来调用ShuffleFileGroup的apply()函数，为bucket获取一个ShuffleFileGroup
          val blockId = ShuffleBlockId(shuffleId, mapId, bucketId)
          // 然后用BlockManager的getDiskWriter()方法，针对ShuffleFileGroup获取一个Writer
          // 这样的话，我们就清楚了，如果开启了consolidation机制，对于每一个bucket，都会获取一个针对ShuffleFileGroup的writer
          // 而不是一个独立的ShuffleBlockFile的writer，这样就实现了，多个ShuffleMapTask的输出数据的合并。
          blockManager.getDiskWriter(blockId, fileGroup(bucketId), serializer, bufferSize,
            writeMetrics)
        }
      } else {
        // 如果没有开启consolidation机制，也就是普通shuffle操作
        Array.tabulate[BlockObjectWriter](numBuckets) { bucketId =>
          // 同样生成一个ShuffleBlockId
          val blockId = ShuffleBlockId(shuffleId, mapId, bucketId)
          // 然后调用BlockManager的diskBlockManager，获取一个代表了要写入本地磁盘文件的blockFile
          val blockFile = blockManager.diskBlockManager.getFile(blockId)
          // Because of previous failures, the shuffle file may already exist on this machine.
          // If so, remove it.
          // 假如这个blockFile存在的话，就删除它 -- 因为一个bucket对应一个blockFile
          if (blockFile.exists) {
            if (blockFile.delete()) {
              logInfo(s"Removed existing shuffle file $blockFile")
            } else {
              logWarning(s"Failed to remove existing shuffle file $blockFile")
            }
          }
          // 然后调用blockManager的getDiskWriter()方法，针对那个blockFile生成writer
          blockManager.getDiskWriter(blockId, blockFile, serializer, bufferSize, writeMetrics)
        }

        // 使用普通的shuffle操作的话，对于每一个ShuffleMapTask输出的bucket，
        // 那么都会在本地获取一个单独的ShuffleBlockFile
      }
      // 省略一些代码
      ........
   }
 }

  这个方法主要就是给每个map task返回一个ShuffleWriterGroup，从这个方法里面我们就能清晰的看到开启Consolidation机制和未开启Consolidation机制的区别了。
  如果开启了Consolidation机制，首先会去获取一个filegroup，如果这个filegroup没有被创建，那么会新建，如果已经存在，那么就返回已经存在的filegroup，这就是复用第一个Task创建的filegroup（复用同一个文件）。然后利用shuffleId，mapId，bucketId创建一个唯一的ShuffleBlockID，然后使用BlockManager针对ShuffleGroupFile生成一个Writer，里面包含了blockId和filegroup，以及待写入的缓存bucket等。
  针对没有开启Consolidation机制而言，同样先生成一个ShuffleBlockId，接着会生成一个blockFile文件，假如这个文件已经存在，那么是之前某个task创建的，先删除再创建。然后同样获得一个writer。
  从上面的源码中，我们就能看出其中的区别了，开启了Consolidation机制会复用第一个Task创建的文件，把它封装为了一个FileGroup，而没有开启则每次写的时候都会创建一个新的文件，这就是他们的最大区别，从源码中也体现出来了。
  这个区别就导致Task创建文件数量的不同，Task map端产生的文件数量在很大程度上会影响Spark的性能，因此假如现在还在使用老版本中的HashShuffle，那么在实际生产环境中，强烈建议开启Consolidation机制（SparkConf设置spark.shuffle.consolidateFiles为true即可）。