Spark 的shuffle流程框架以及源码详解（匠心巨作）（3） 之UnsafeShuffleWriter

1. 简介

Spark 的shuffle流程框架以及源码详解（匠心巨作）（1） 和 Spark 的shuffle流程框架以及源码详解（匠心巨作）（2）这两篇博客介绍了SparkShuffle的发展，以及BypassShuffleWriter的源码分析。本篇博客主要解析SparkShuffleWriter三种方式中的UnsafeShuffleWriter，它还有另外的一个名字，Tungsten-ShuffleWriter，这种方式是根据java中的unsafe API 完成的，一般情况下，我们认为Java语言中是没有指针存在的，其实不然，这个unsafe类中就封装了一些像C/C++中的指针操作，可以直接访问系统的内存。既然能够像指针一样直接访问内存，那就会存在一定的风险，比如给指针赋值或者指针指向的地址加上偏移量后可能会指向系统程序的地址，从而覆盖原来地址中的内容，很有可能会直接导致系统的崩溃，所以说指针操作是不安全的，因此给它取名为UnsafeShuffleWriter。

2. UnsafeShuffleWriter详解

2.1 使用条件

• no Aggregator is specified； map端不会有聚合操作
• supportsRelocationOfSerializedObjects；序列化器支持对已经序列化的对象重定位
• the number of partitions is less than 16777216；分区数量小于16777216

UnsafeShuffleWriter是对普通sort的一种优化，排序的不是内容本身，而是内容序列化后字节数组的指针(元数据)，把数据的排序转变为了指针数组的排序，实现了直接对序列化后的二进制数据进行排序。由于直接基于二进制数据进行操作，所以在这里面没有序列化和反序列化的过程。内存的消耗降低，相应的也会减少gc的开销。

这些优化的实现导致引入了一个新的内存管理模型，即Page，Page其实也是一个MemoryBlock, 不过因为它是引入了指针Pointer的特性，所以又支持off-heap以及 on-heap 两种模式，可以直接访问系统的内存。处理模型如下图所示：


采用了用8位byte的数组指针，这个数组指针的前24位存储着数据记录的的PartitionId，中间13位存储着内存页的编号，其实也是内存块的编号，最后24位存储着数据记录在这个内存页中的地址偏移量。

2.2 代码流程

上图是整个UnsafeShuffleWriter的代码流程，基本上也是分为三个步骤，第一个步骤是ShuffleMapTask的数据输入到一个外部的排序器里面ShuffleExternalSorter，第二个步骤是把内存中放不下的文件溢写到磁盘，第三步是把溢写出来的文件最终合并为一个文件。

2.3 源码详解

根据上面的框架流程图，我们来详细看一下源码。首先进入UnsafeShuffleWriter 的write方法，基于spark-2.4.0版本。

public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOException {
    // Keep track of success so we know if we encountered an exception
    // We do this rather than a standard try/catch/re-throw to handle
    // generic throwables.
    boolean success = false;
    try {
    //判断shufflemaptask是否有数据输出，即records迭代器里面是否有未输出的数据
      while (records.hasNext()) {
      //把数据插入到sorter里面，即外部排序器
        insertRecordIntoSorter(records.next());
      }
      //合并溢写文件，输出最终的一个文件
      closeAndWriteOutput();
      success = true;
    } finally {
      if (sorter != null) {
        try {
        
          sorter.cleanupResources();
        } catch (Exception e) {
          // Only throw this error if we won't be masking another
          // error.
          if (success) {
            throw e;
          } else {
            logger.error("In addition to a failure during writing, we failed during " +
                         "cleanup.", e);
          }
        }
      }
    }
  }

这是UnsafeShuffleWriter 最顶层的调用，具体的整个过程的实现，在接下来要讲解的源码，现在进入到inserRecordIntoSorter里面。

@VisibleForTesting
  void insertRecordIntoSorter(Product2<K, V> record) throws IOException {
  //认定这个sorter 里面不是空的，
    assert(sorter != null);
    //选定数据记录的Key值
    final K key = record._1();
    //根据Key值对数据记录进行分区
    final int partitionId = partitioner.getPartition(key);
    //清空serBuffer的内存空间，这个内存空间用于存放字节输出流，
    serBuffer.reset();
    //把已经序列化的流对象写入到内存中，key值和value分别写入。
    serOutputStream.writeKey(key, OBJECT_CLASS_TAG);
    serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);
    serOutputStream.flush();

    final int serializedRecordSize = serBuffer.size();
    //如果serBuffer不为空，就说明里面是有shufflemaptask输出的序列化后的流数据，那么就会调用insertRecord函数，把序列化后的数据插入到sorter里面。
    assert (serializedRecordSize > 0);
//插入数据
    sorter.insertRecord(
      serBuffer.getBuf(), Platform.BYTE_ARRAY_OFFSET, serializedRecordSize, partitionId);
  }

在上面主要是数据记录首先放到serBuffer中，这里面存放的是序列化后的二进制数据，然后把serBuffer中的数据插入到sorter中，这里的serBuffer的默认大小为1M。接下来进入到insertRecord函数里面。

public void insertRecord(Object recordBase, long recordOffset, int length, int partitionId)
    throws IOException {

    // for tests
    //假设inMemSorter里面不是空的
    assert(inMemSorter != null);
    //如果不是空的，判断是否大于存储数据记录的上限值
    if (inMemSorter.numRecords() >= numElementsForSpillThreshold) {
    
      logger.info("Spilling data because number of spilledRecords crossed the threshold " $numElementsForSpillThreshold);
      //如果数据记录的个数大于能够存储的上限值，那么久把文件溢写出来。
      spill();
    }
//如果没有达到存储的上限值，就判断一下数组指针是否有空间存下一个数据，如果没有，则申请空间，申请不到就会发生溢写。
    growPointerArrayIfNecessary();
    final int uaoSize = UnsafeAlignedOffset.getUaoSize();
    // Need 4 or 8 bytes to store the record length.
    final int required = length + uaoSize;
    //判断当前的内存页是否有足够的空间，存放下一个数据，如果不够则申请一个新的内存页。
    acquireNewPageIfNecessary(required);
//如果当前页不是空的
    assert(currentPage != null);
    final Object base = currentPage.getBaseObject();
    //记录下数据的page的编号和地址偏移量
    final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor)
    UnsafeAlignedOffset.putSize(base, pageCursor, length);
    pageCursor += uaoSize;
    //把数据记录的长度和序列化后的字节序列写入
    Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
    pageCursor += length;
    //把这条记录转化为指针，传递到ShuffleinMemSorter中去。
    inMemSorter.insertRecord(recordAddress, partitionId);
  }

上面主要完成把数据记录插入到ShuffleExternalSorter中，期间会对内存的剩余空间进行判断，内存空间不够的时候会发生溢写。以上既是整个UnsafeShuffleWriter的第一个步骤，把ShuffleMapTask的输出插入到外部的排序器里面。（未完待续！！！）