第32课：彻底解密Spark 2.1.X中Shuffle 下Task视角内存分配管理

本文根据家林大神系列课程编写 http://weibo.com/ilovepains

Spark 2.1.X内存管理包含2种类型：统一内存管理 UnifiedMemoryManager、静态内存StaticMemoryManager。这两种内存的管理方式最终要落实到Task的运行。我们先从源码角度对Spark内存管理进行回顾，从Spark Task的视角解析Task运行内存管理源码。在Spark 2.1.X中默认使用的UnifiedMemoryManager的统一内存管理的方式，从Task运行的角度来讲是Execution级别的，也就是UnifiedMemoryManager的核心是Execution Memory。Execution Memory主要做的事情是Shuffle、Sort、Aggregate等。我们为什么将Execution Memory的内存调大，理论上讲，Execution Memory的内存越大，IO就越来越少。Execution Memory相对于内存的占用比较强势。当Execution Memory空间不足而且Storage Memory空间也不足的时候，Storage Memory空间会被强制drop掉一部分数据来解决Execution 的空间不足问题。

Spark2.1.X内存管理UnifiedMemoryManager的管理方式从源码的角度看有2个核心方法：

• acquireExecutionMemory
• acquireStorageMemory

我们看一下acquireExecutionMemory的源代码：

override private[memory] def acquireExecutionMemory(
    numBytes: Long,
    taskAttemptId: Long,
    memoryMode: MemoryMode): Long = synchronized {
  assertInvariants()
  assert(numBytes >= 0)
  val (executionPool, storagePool, storageRegionSize, maxMemory) = memoryMode match {
    case MemoryMode.ON_HEAP => (
      onHeapExecutionMemoryPool,
      onHeapStorageMemoryPool,
      onHeapStorageRegionSize,
      maxHeapMemory)
    case MemoryMode.OFF_HEAP => (
      offHeapExecutionMemoryPool,
      offHeapStorageMemoryPool,
      offHeapStorageMemory,
      maxOffHeapMemory)
  }

acquireExecutionMemory根据ON_HEAP、OFF_HEAP两种不同的方式进行模式匹配，匹配到2种不同方式的时候，有自己的执行方式。内部onHeapExecutionMemoryPool, onHeapStorageMemoryPool无论是哪种方式都会调maybeGrowExecutionPool。

def maybeGrowExecutionPool(extraMemoryNeeded: Long): Unit = {
  if (extraMemoryNeeded > 0) {
    // There is not enough free memory in the execution pool, so try to reclaim memory from
    // storage. We can reclaim any free memory from the storage pool. If the storage pool
    // has grown to become larger than `storageRegionSize`, we can evict blocks and reclaim
    // the memory that storage has borrowed from execution.
    val memoryReclaimableFromStorage = math.max(
      storagePool.memoryFree,
      storagePool.poolSize - storageRegionSize)
    if (memoryReclaimableFromStorage > 0) {
      // Only reclaim as much space as is necessary and available:
      val spaceToReclaim = storagePool.freeSpaceToShrinkPool(
        math.min(extraMemoryNeeded, memoryReclaimableFromStorage))
      storagePool.decrementPoolSize(spaceToReclaim)
      executionPool.incrementPoolSize(spaceToReclaim)
    }
  }
}

每一个Task能够使用的内存理论上是多大呢？这是很重要的一件事情，因为能推测每一个Task能使用的大小，计算公式为poolSize/(2*numActiveTasks)到 maxPollSize/numActiveTasks。

在maybeGrowExecutionPool方法中通过比较 storagePool.memoryFree和storagePool.poolSize -storageRegionSize的大小，计算一个最大值：storage Memory剩余的内存和storagememory从Executionmemory借来的内存哪个大。Storage和Execution在适当的时候可以借用彼此的memory，Storage Memory在Execution执行的时候，Execution memory有空间，Storage Memory缓存的时候向Execution memory借用一些空间。如果Execution memory内存不够，这个时候进行一个比较，整个Execution memory能借到的最大内存是Storage Memory曾经找Execution memory借的内存+Storage Memory空闲内存；根据math.min(extraMemoryNeeded,memoryReclaimableFromStorage)如果Executionmemory需要的大小小于能借到的最大内存，那以实际需要的内存为准。

我们可以看一下freeSpaceToShrinkPool方法的实现：

def freeSpaceToShrinkPool(spaceToFree: Long): Long = lock.synchronized {
  val spaceFreedByReleasingUnusedMemory = math.min(spaceToFree, memoryFree)
  val remainingSpaceToFree = spaceToFree - spaceFreedByReleasingUnusedMemory
  if (remainingSpaceToFree > 0) {
    // If reclaiming free memory did not adequately shrink the pool, begin evicting blocks:
    val spaceFreedByEviction =
      memoryStore.evictBlocksToFreeSpace(None, remainingSpaceToFree, memoryMode)
    // When a block is released, BlockManager.dropFromMemory() calls releaseMemory(), so we do
    // not need to decrement _memoryUsed here. However, we do need to decrement the pool size.
    spaceFreedByReleasingUnusedMemory + spaceFreedByEviction
  } else {
    spaceFreedByReleasingUnusedMemory
  }
}

在acquireStorageMemory方法分别调用executionPool.decrementPoolSize方法减少内存，及storagePool.incrementPoolSize方法增加内容：

l executionPool.decrementPoolSize(memoryBorrowedFromExecution)

l storagePool.incrementPoolSize(memoryBorrowedFromExecution)

/**
 * Expands the pool by `delta` bytes.
 */
final def incrementPoolSize(delta: Long): Unit = lock.synchronized {
  require(delta >= 0)
  _poolSize += delta
}

/**
 * Shrinks the pool by `delta` bytes.
 */
final def decrementPoolSize(delta: Long): Unit = lock.synchronized {
  require(delta >= 0)
  require(delta <= _poolSize)
  require(_poolSize - delta >= memoryUsed)
  _poolSize -= delta
}

接下来我们阐述Storage Memory。Storage Memory只有一种情况：在Execution memory空闲的时候，StorageMemory能借走Execution memory的空闲内存。

override def acquireStorageMemory(
    blockId: BlockId,
    numBytes: Long,
    memoryMode: MemoryMode): Boolean = synchronized {
  assertInvariants()
  assert(numBytes >= 0)
  val (executionPool, storagePool, maxMemory) = memoryMode match {
    case MemoryMode.ON_HEAP => (
      onHeapExecutionMemoryPool,
      onHeapStorageMemoryPool,
      maxOnHeapStorageMemory)
    case MemoryMode.OFF_HEAP => (
      offHeapExecutionMemoryPool,
      offHeapStorageMemoryPool,
      maxOffHeapMemory)
  }
  if (numBytes > maxMemory) {
    // Fail fast if the block simply won't fit
    logInfo(s"Will not store $blockId as the required space ($numBytes bytes) exceeds our " +
      s"memory limit ($maxMemory bytes)")
    return false
  }
  if (numBytes > storagePool.memoryFree) {
    // There is not enough free memory in the storage pool, so try to borrow free memory from
    // the execution pool.
    val memoryBorrowedFromExecution = Math.min(executionPool.memoryFree, numBytes)
    executionPool.decrementPoolSize(memoryBorrowedFromExecution)
    storagePool.incrementPoolSize(memoryBorrowedFromExecution)
  }
  storagePool.acquireMemory(blockId, numBytes)
}

无论是storagePool还是executionPool，都要和memoryPoll打交道，我们看一下memoryPoll。

/**
 * Manages bookkeeping for an adjustable-sized region of memory. This class is internal to
 * the [[MemoryManager]]. See subclasses for more details.
 *
 * @param lock a [[MemoryManager]] instance, used for synchronization. We purposely erase the type
 *             to `Object` to avoid programming errors, since this object should only be used for
 *             synchronization purposes.
 */
private[memory] abstract class MemoryPool(lock: Object) {

  @GuardedBy("lock")
  private[this] var _poolSize: Long = 0

  /**
   * Returns the current size of the pool, in bytes.
   */
  final def poolSize: Long = lock.synchronized {
    _poolSize
  }

  /**
   * Returns the amount of free memory in the pool, in bytes.
   */
  final def memoryFree: Long = lock.synchronized {
    _poolSize - memoryUsed
  }

  /**
   * Expands the pool by `delta` bytes.
   */
  final def incrementPoolSize(delta: Long): Unit = lock.synchronized {
    require(delta >= 0)
    _poolSize += delta
  }

  /**
   * Shrinks the pool by `delta` bytes.
   */
  final def decrementPoolSize(delta: Long): Unit = lock.synchronized {
    require(delta >= 0)
    require(delta <= _poolSize)
    require(_poolSize - delta >= memoryUsed)
    _poolSize -= delta
  }

  /**
   * Returns the amount of used memory in this pool (in bytes).
   */
  def memoryUsed: Long
}

我们看一下StorageMemoryPool，其对于内存的记录和管理，一方面是内存使用的记录，一方面是可调整大小内存（要么借进新的内存，要么借出内存）的管理。StorageMemoryPool有acquireMemory，releaseMemory等方法。我们在shuffle的时候，如果我们使用了UnifiedMemoryManager的新型内存管理方式，同时开启了OFF_HEAP，对我们内存的使用有很大的影响，如果有offHeapExecutionMemoryPool，这个时候还存不存在一个概念，从StorageMemory中获取内存？实际上不需要找StorageMemory要内存。

每个Task能够分配的内存大小：poolSize/(2*numActiveTasks) 到 maxPollSize/numActiveTasks。

l val maxMemoryPerTask = maxPoolSize / numActiveTasks

l val minMemoryPerTask = poolSize / (2 * numActiveTasks)

如果ShuffleTask计算比较复杂的情况，业务逻辑比较复杂，使用新型的UnifiedMemoryManager内存管理方式能取得较好的效果；但是如果计算的业务逻辑不复杂，shuffle计算比较简单，这个时候可以回退到旧的内存管理方式，使用StaticMemory Management效果会更好，因为缓存空间更大。

在这个基础之上，从Task的角度来看，Task有个TaskMemoryManager。TaskMemoryManager管理单个任务分配的内存。 TaskMemoryManager类中的大多数复杂情况涉及到64位长的off-heap 非堆地址编码。在off-heap 非堆模式下，内存可以直接寻址64位。在堆模式下，内存是对象内的基本对象引用和64位偏移的组合寻址。这是一个问题，当我们想存储指针的数据结构内的其他结构，如在hashmaps或排序缓冲区的记录内的指针。即使我们决定使用128位来解决内存问题，我们不能只存储基本对象的地址，由于堆内存存在GC，因此它不能保证保持稳定。相反，我们使用64位地址方式来编码记录指针：off-heap 非堆模式，只需存储原始地址，并在堆模式上使用的地址的上13位存储“页码”和较低的51位来存储此页内的偏移量。这些页数字是用来索引为“页表”阵内以MemoryManager来检索基本对象。 这使我们能够解决8192页。在堆模式下，最大页大小受最大长度的数组限制，我们能够解决8192 2^32 字节的地址，大约35百万兆字节内存。

TaskMemoryManager 肯定要申请内存，管理内存以及释放内存，需要一个MemoryBlock[]类型的内存的pageTable，对堆的内存的分配、释放，这个都是JVM进行管理的，我们new出一个MemoryBlock，只是对象的引用，不是具体内存空间的地址。堆内存有GC的问题，是JVM的死穴，那我们做堆外内存，JAVA提供了ByteBufferBlock的工具类。

接下来我们看一下MemoryLocation，MemoryLocation是内存位置：跟踪内存地址（off-heap 非堆地址分配）；或者JVM对象的偏移量（in-heap堆分配）。

public class MemoryLocation {

  @Nullable
  Object obj;

  long offset;

  public MemoryLocation(@Nullable Object obj, long offset) {
    this.obj = obj;
    this.offset = offset;
  }

  public MemoryLocation() {
    this(null, 0);
  }

  public void setObjAndOffset(Object newObj, long newOffset) {
    this.obj = newObj;
    this.offset = newOffset;
  }

  public final Object getBaseObject() {
    return obj;
  }

  public final long getBaseOffset() {
    return offset;
  }
}

我们看一下MemoryBlock，MemoryBlock是一个连续的内存块，从{@link MemoryLocation} 开始具有固定大小的存储单元。pageNumber变量是可选的页码；当TaskMemoryManager分配内存使用MemoryBlock来表示页，pageNumber是公共变量，在不同的包里面可以通过TaskMemoryManager修改。

ublic class MemoryBlock extends MemoryLocation {

  private final long length;

  /**
   * Optional page number; used when this MemoryBlock represents a page allocated by a
   * TaskMemoryManager. This field is public so that it can be modified by the TaskMemoryManager,
   * which lives in a different package.
   */
  public int pageNumber = -1;

  public MemoryBlock(@Nullable Object obj, long offset, long length) {
    super(obj, offset);
    this.length = length;
  }

  /**
   * Returns the size of the memory block.
   */
  public long size() {
    return length;
  }

  /**
   * Creates a memory block pointing to the memory used by the long array.
   */
  public static MemoryBlock fromLongArray(final long[] array) {
    return new MemoryBlock(array, Platform.LONG_ARRAY_OFFSET, array.length * 8L);
  }

  /**
   * Fills the memory block with the specified byte value.
   */
  public void fill(byte value) {
    Platform.setMemory(obj, offset, length, value);
  }
}

接下来我们我们看一下UnsafeMemoryAllocator，一个简单的 {@linkMemoryAllocator} 使用 {@code Unsafe} 分配off-heap非堆内存分配。

public class UnsafeMemoryAllocator implements MemoryAllocator {

  @Override
  public MemoryBlock allocate(long size) throws OutOfMemoryError {
    long address = Platform.allocateMemory(size);
    MemoryBlock memory = new MemoryBlock(null, address, size);
    if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
      memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_CLEAN_VALUE);
    }
    return memory;
  }

  @Override
  public void free(MemoryBlock memory) {
    assert (memory.obj == null) :
      "baseObject not null; are you trying to use the off-heap allocator to free on-heap memory";
    if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
      memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_FREED_VALUE);
    }
    Platform.freeMemory(memory.offset);
  }
}

再回到TaskMemoryManager，关注一下关键点：

1，pageTable：与操作系统的页表类似，此数组将页编号映射为基准对象指针，允许我们把哈希表的内部64位地址表示、以及baseObject+offset偏移表示，这样可以支持

堆地址和非堆地址。当使用“非堆地址分配”时，此映射中的每个条目将为“null”。使用堆分配器时，此映射中的项将指向页的基准对象。当新数据页被分配时，条目将被添加到该map。

2，PAGE_TABLE_SIZE是页表中的条目数。

3，PAGE_NUMBER_BITS：用于处理页表的位数。

OFFSET_BITS：用于在数据页中编码偏移量的位数。实际为51

public class TaskMemoryManager {

  private static final Logger logger = LoggerFactory.getLogger(TaskMemoryManager.class);

  /** The number of bits used to address the page table. */
  private static final int PAGE_NUMBER_BITS = 13;

  /** The number of bits used to encode offsets in data pages. */
  @VisibleForTesting
  static final int OFFSET_BITS = 64 - PAGE_NUMBER_BITS;  // 51

  /** The number of entries in the page table. */
  private static final int PAGE_TABLE_SIZE = 1 << PAGE_NUMBER_BITS;

  /**
   * Maximum supported data page size (in bytes). In principle, the maximum addressable page size is
   * (1L &lt;&lt; OFFSET_BITS) bytes, which is 2+ petabytes. However, the on-heap allocator's
   * maximum page size is limited by the maximum amount of data that can be stored in a long[]
   * array, which is (2^32 - 1) * 8 bytes (or 16 gigabytes). Therefore, we cap this at 16 gigabytes.
   */
  public static final long MAXIMUM_PAGE_SIZE_BYTES = ((1L << 31) - 1) * 8L;

  /** Bit mask for the lower 51 bits of a long. */
  private static final long MASK_LONG_LOWER_51_BITS = 0x7FFFFFFFFFFFFL;