spark内存管理源码分析系列之执行内存管理器TaskMemoryManager

Spark中Task的执行内存是通过TaskMemoryManger统一管理的，不论是ShuffleMapTask还是ResultTask，Spark都会生成一个专用的TaskMemoryManger对象，然后通过TaskContext将TaskMemoryManger对象共享给该task attempt的所有memory consumers。 TaskMemoryManger自建了一套内存页管理机制，并统一对ON_HEAP和OFF_HEAP内存进行编址，分配和释放。

概述

TaskMemoryManger负责管理单个任务的堆外执行内存和堆内执行内存，是Tungsten内存管理机制的核心实现类。对于堆外内存，可以内存地址直接使用64位长整型地址寻址；对于堆内内存，内存地址由一个obj对象和一个offset对象组合起来表示，主要有以下三方面作用：

1. 建立类似于操作系统内存页管理的机制，对ON_HEAP和OFF_HEAP内存统一编址和管理;

2. 通过调用MemoryManager和MemoryAllocator，将逻辑内存的申请&释放与物理内存的分配&释放结合起来;

3. 记录和管理task attempt的所有memory consumer。

那么页管理机制到底是如何设计的，堆内内存如何管理，如何避免堆内内存由于JVM的GC的存在引起的内存地址变化，数据在内存页中是如何寻址的？本文通过分析整体设计&实现细节来解决这些问题。

MemoryLocation

为了统一on-heap<JVM管理>和off-heap<自行管理>的执行内存，抽象出来一个MemoryLocation，包含了obj对象和offset属性，这个类可以用来内存寻址，具体的寻址如下:

1. Object obj处于堆内内存模式时，数据作为对象存储在JVM的堆上，此时的obj不为空；处于堆外内存模式时，数据存储在JVM的堆外内存[操作系统内存]中，因而不会在JVM中存在对象，所以obj为NULL;
2. long offsetoffset属性主要用来定位数据，处于堆内内存模式时，首先从堆内找到对象MemoryBlock，然后使用offset定位数据的具体位置；处于堆外内存模式时，则直接使用offset从堆外内存中定位。

MemoryBlock

MemoryBlock继承于MemoryLocation，它代表的是一个Page对象，表示从obj和offset定位的起始位置开始的固定长度[length]的连续内存块。Page代表了具体的内存区域以及内存里面具体的数据，Page中的数据可能是On-heap的数据，也可能是Off-heap中的数据。

另外提供了通过已经分配的array创建MemoryBlock和以指定的字节填充整个MemoryBlock的方法。

/**  创建一个指向由长整型数组使用的内存的MemoryBlock */
public static MemoryBlock fromLongArray(final long[] array) {
   
  return new MemoryBlock(array, Platform.LONG_ARRAY_OFFSET, array.length * 8L);
}

/** 以指定的字节填充整个MemoryBlock， 即将obj对象从offset开始，长度为length的堆内存替换为指定字节的值。
   Platform中封装了对sun.misc.Unsafe的API调用，Platform的setMemory方法实际调用了sun.misc.Unsafe的setMemory<在给定的内存块中设置值> */
public void fill(byte value) {
   
  Platform.setMemory(obj, offset, length, value);
}

1. Platform.setMemory(obj, offset, length, value);在给定的内存块中设置值，这里是指在将obj对象从offset开始，长度为length的堆内存替换为指定字节的值
2. Platform.LONG_ARRAY_OFFSET是long array数组类型中，数组第一个元素相对数组的偏移。

MemoryAllocator

有了内存页的抽象，就需要给它分配内存页，是通过MemoryAllocator实现的，该接口定义了allocate和free方法，等待具体类实现，目前有两个实现类HeapMemoryAllocator和UnsafeMemoryAllocator，分别负责堆内和堆外内存页分配。

HeapMemoryAllocator

HeapMemoryAllocator是Tungsten内存管理中在堆内存模式下使用的内存分配器，与onHeapExecutionMemoryPool配合使用，主要负责分配堆内内存，其主要分配long型数组，最大分配内存为16GB。实现了内存页分配以及回收的方法，另外维护了一个MemoryBlock的弱引用[只有弱引用时候，如果GC运行, 那么这个对象就会被回收，不论当前的内存空间是否足够，这个对象都会被回收]的缓冲池，用于Page页[即MemoryBlock]的快速分配。

@GuardedBy("this")
private final Map<Long, LinkedList<WeakReference<long[]>>> bufferPoolsBySize = new HashMap<>();

// 池化阈值，只有在池化的MemoryBlock大于该值时，才需要被池化: 1M
private static final int POOLING_THRESHOLD_BYTES = 1024 * 1024;

申请内存页

申请内存页的步骤如下:

1. 申请一个内存页时，首先需要进行内存对齐<8字节对齐>得到对齐后的申请的内存大小，然后根据是否满足池化条件<大于1M>，进行不同操作；
2. 如果满足池化条件，从缓存池中如果可以拿取到相同大小的内存，进行构建MemoryBlock；
3. 不满足池化条件或者缓存池中没有同样大小的array，则HeapMemoryAllocator利用long数组[new long[size]]向JVM堆申请内存，通过在MemoryBlock对象中维护对long数组的引用，来防止JVM将long数组所占内存垃圾回收掉。

@Override
public MemoryBlock allocate(long size) throws OutOfMemoryError {
   
  /**
     * MemoryBlock中以Long类型数组装载数据，所以需要对申请的大小进行转换，
     *  由于申请的是字节数，因此先为其多分配7个字节，避免最终分配的字节数不够，除以8是按照Long类型由8个字节组成来计算的。
     *  例如：申请字节数为50，理想情况应该分配56字节，即7个Long型数据。
     *  如果直接除以8，会得到6，即6个Long型数据，导致只会分配48个字节，
     *  但先加7后再除以8，即 (50 + 7) / 8 = 7个Long型数据，满足分配要求。
     */
  int numWords = (int) ((size + 7) / 8);
  long alignedSize = numWords * 8L; // 补齐后的字节数
  assert (alignedSize >= size);
  if (shouldPool(alignedSize)) {
    // 需要从池化中拿取
    synchronized (this) {
   
      final LinkedList<WeakReference<long[]>> pool = bufferPoolsBySize.get(alignedSize);
      if (pool != null) {
   
        while (!pool.isEmpty()) {
   
          // 取出池链头的MemoryBlock
          final WeakReference<long[]> arrayReference = pool.pop();
          final long[] array = arrayReference.get(); // 拿取array
          if (array != null) {
   
            // MemoryBlock的大小要比分配的大小大
            assert (array.length * 8L >= size);
            // 从MemoryBlock的缓存中获取指定大小的MemoryBlock并返回
            MemoryBlock memory = new MemoryBlock(array, Platform.LONG_ARRAY_OFFSET, size);
            if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
   
              memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_CLEAN_VALUE);
            }
            return memory;
          }
        }
        bufferPoolsBySize.remove(alignedSize);
      }
    }
  }
  /**
     *  走到此处，说明满足以下任意一点：
     *   1. 指定大小的MemoryBlock不需要采用池化机制。
     *   2. bufferPoolsBySize中没有指定大小的MemoryBlock。
     *
     */
  long[] array = new long[numWords];
  // 创建MemoryBlock并返回
  MemoryBlock memory