Spark TaskMemoryManager如何为task分配执行内存?

本文讨论跟执行内存密切相关的一个组件:TaskMemoryManager(下文简称TMM)。TMM是tungsten内存管理机制的核心实现类(tungsten内存管理只作用于执行内存),它的功能包括:

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

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

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

1 TMM内部结构及功能

下图是TMM的内部结构,图中的各个模块是笔者抽象总结的,在TMM类中,这些模块各自对应一些成员变量和方法。

TMM内部结构

再来看看TMM的类图:


TMM类图

1.1 内存页管理模块

TMM实现了一套类似于操作系统内存页管理的机制:

    1. TMM用MemoryBlock表示内存页(下文称page),每个MemoryBlock对象就是一个page

    2. TMM维护了一个MemoryBlock数组用于存放该TMM分配得到的pages, 即TMM类图中的pageTable

    3. pageTable中每个page的每个字节都有一个地址,这个地址是long类型,共64位,由page number和页内offset组成, 高13位表示page number,低51位表示offset:

TaskMemoryManager.encodePageNumberAndOffset

这里的13和51就是TMM类图中的PAGE_NUMBER_BITS和OFFSET_BITS.

    4. TMM用一个bit set(即TMM类图中的allocatedPages)表示哪些page已经分配到内存,每一位对应一个page,所以这个bit set包含PAGE_TABLE_SIZE, 即2^13=8192个bits.

    5. 理论上讲,用于offset编址的位数为51,那么每个page最大容量为2^51(2+ PB)。但是,由于HeapMemoryAllocator中用于占内存的long数组的最大容量为(2^31-1)*8≈17GB,这里为了保持ON_HEAP和OFF_HEAP内存编址的一致性, 所以将单个page的最大容量限制为17GB :

TaskMemoryManager.MAXIMUM_PAGE_SIZE_BYTES

下文会详细讨论HeapMemoryAllocator(负责分配堆内存)和UnsafeMemoryAllocator(负责分配堆外内存)是如何进行内存分配的。

1.2 内存申请&释放模块

TMM对内存的申请和释放主要通过两组方法完成:

1. acquireExecutionMemory + releaseExecutionMemory

acquireExecutionMemory和releaseExecutionMemory调用MemoryManager(关于MemoryManger的实现,请参考漫谈Spark内存管理(二):spark自建的逻辑内存管理器是怎么申请和释放内存的?)进行逻辑内存的申请和释放,并且在逻辑空闲内存不足的情况下,会尝试从当前持有的memory consumers中spill内存到disk以释放空间。

2. allocatePage + freePage

    allocatePage和freePage包含了逻辑内存和物理内存的申请和释放。对于内存申请,allocatePage会先调用TMM的acquireExecutionMemory申请逻辑内存,然后根据申请到的内存大小调用memoryManager.tungstenMemoryAllocator分配物理内存(可能是堆内存,也可能是堆外内存,这个由参数spark.memory.offHeap.enabled决定);对于内存释放,顺序刚好相反,freePage会先调用memoryManager.tungstenMemoryAllocator释放物理内存,然后再调用TMM的releaseExecutionMemory释放逻辑内存。

这里有一点需要注意,allocatePage在调用acquireExecutionMemory获得逻辑内存后,可能在物理内存分配阶段失败(遇到OutOfMemoryError),这个时候,申请到的逻辑内存会被TMM添加到acquiredButNotUsed中,TMM类图中可以看到acquiredButNotUsed是long类型的。acquiredButNotUsed中记录的逻辑内存会在cleanUpAllAllocatedMemory方法中被释放,以避免发生已用逻辑内存量大于实际已用物理内存量的情况。

对于一些没有用到tungsten内存管理机制的memory consumer (比如都继承自Spillable抽象类的ExternalSorter和ExternalAppendOnlyMap),它们会通过MemoryConsumer的acquireMemory和freeMemory方法调用TMM的acquireExecutionMemory和releaseExecutionMemory进行逻辑内存的申请和释放,不会使用TMM的物理内存管理功能,也就是不会调用TMM的allocatePage和freePage方法;当然,也有很多memory consumer用到了TMM的物理内存管理功能,比如UnsafeShuffleWriter中的ShuffleExternalSorter.

1.3 内存消费者管理模块

每个task attempt都有对应的TMM对象,而一个task attempt可能有不止一个memory consumer需要消耗内存,所以TMM会维护当前task attempt的所有memory consumers,就是TMM类图中的私有成员变量consumers,可以看到它就是一个MemoryConsumer类型的HashSet.

1. consumers在TMM构造函数中被初始化为一个空的HashSet,并且在每次acquireExecutionMemory方法调用的最后都会将申请逻辑内存的memory consumer添加到TMM的consumers中。

2. 当某个memory consumer调用acquireExecutionMemory申请逻辑内存时遇到可用逻辑内存不足的情况,TMM首先会遍历consumers集合,生成一个以consumer占用的memory大小为key的TreeMap对象,也就是生成一颗以已用内存大小为key的红黑树。然后,在这个treeMap中查找占有目标内存大小的memory consumer,找到后依次调用这些consumer的spill方法,直到释放出足够的内存空间或treeMap遍历完毕。

2 MemoryAllocator的实现

2.1 Tungsten内存管理的Memory Mode

如上文所述,TMM中对物理内存的申请和释放最终交给memoryManager中的tungstenMemoryAllocator完成,tungstenMemoryAllocator是一个MemoryAllocator对象,而MemoryAllocator是一个接口,其实现类有两个:HeapMemoryAllocator(负责ON_HEAP内存的分配和释放)和UnsafeMemoryAllocator(负责OFF_HEAP内存的分配和释放).

memoryManager.tungstenMemoryAllocator中使用哪种MemoryAllocator由memoryManager.tungstenMemoryMode决定,而tungstenMemoryMode则由参数spark.memory.offHeap.enabled控制:

MemoryManager.tungstenMemoryAllocator
MemoryManager.tungstenMemoryMode

spark.memory.offHeap.enabled默认为false,即memoryManager.tungstenMemoryAllocator默认为HeapMemoryAllocator.

从这里可以发现,在目前的spark实现中(spark 2.4.0),tungsten内存管理要么使用堆内存,要么使用堆外内存,不支持两种内存的混合模式。

也就是说各个memory consumer通过TMM的allocatePage方法分配得到的内存的memory mode是一致的,要么全都是ON_HEAP,要么全都是OFF_HEAP.

这一点从TMM的allocatePage方法中的assert也可以看出来:

TaskMemoryManager.allocatePage

allocatePage方法要求consumer的memory mode和tungstenMemoryMode是一致的。

2.2 MemoryBLock

如上文所述, TMM用MemoryBlock表示一个内存页,MemoryBlock继承自MemoryLocation,先来看看它们的类图:

MemoryLocation和MemoryBlock的类图

MemoryLocation有两个属性,一个是object,一个是offset.

对于ON_HEAP内存页,object是一个long类型数组;对于OFF_HEAP内存页,object为null.

MemoryBlock在MemoryLocation的基础上添加了属性length,用于表示内存页中数据的长度(字节数)。

2.3 HeapMemoryAllocator

下面说一下笔者对HeapMemoryAllocator工作原理的理解:

1. 申请一个内存页时,HeapMemoryAllocator利用long数组向JVM堆申请内存,通过在MemoryBlock对象中维护对long数组的引用来防止JVM将long数组所占内存垃圾回收掉。

2. HeapMemoryAllocator还会维护一个从内存大小到long数组引用列表的HashMap,

    比如,当需要申请10MB大小的MemoryBlock时,HeapMemoryAllocator会先以10*1024*1024=10485760为key到这个HashMap中去查找,如果找到了则直接从HashMap中拿一个long数组并封装为MemoryBlock对象返回;如果在HashMap中没找到指定大小的long数组,则新建一个目标大小的long数组并封装为MemoryBlock对象返回。HeapMemoryAllocator.allocate源码如下:


HeapMemoryAllocator.allocate

3. 释放一个内存页时,其实就是将MemoryBlock对象中的object引用置为null,这样的话MemoryBlock对象对应的long数组就没有被引用了,其所占内存空间在JVM下次GC时会被自动回收。

4. 释放内存页的时候会检查内存页大小,如果超过HeapMemoryAllocator.POOLING_THRESHOLD_BYTES=1024*1024 (1MB),则会将其加入#2中提到的HashMap中去,一共后续内存申请直接使用,这种做法可避免频繁地向操作系统申请和释放相同大小的内存空间,也就是一次申请,多次使用。HeapMemoryAllocator.free源码如下:

HeapMemoryAllocator.free

2.4 UnsafeMemoryAllocator

UnsafeMemoryAllocator用于对OFF_HEAP内存进行分配和释放,其直接调用sun.misc.Unsafe.allocateMemory和sun.misc.Unsafe.freeMemory方法进行内存分配和释放。

UnsafeMemoryAllocator.allocate

注意,OFF_HEAP的MemoryBlock对象的obj属性适中为null.

UnsafeMemoryAllocator.free

与HeapMemoryAllocator不同,UnsafeMemoryAllocator没有用HashMap来缓存申请得到的内存页,而是每次都向操作系统申请和释放。

3 TMM对象是如何被创建,被使用的?

下图描述了TMM对象被创建,被Task用来生成TaskContext的过程:

TMM被创建和使用的过程

Task (ShuffleMapTask或ResultTask) 在运行时会生成TaskContext对象并赋值给Task.context属性,而spark运行过程中的各个memory consumer就是通过task的这个context属性来获取并使用TMM对象的。

举个例子,在shuffle过程中,ShuffleMapTask会通过SortShuffleManager.getWriter方法获取一个ShuffleWriter对象:

SortShuffleManager.getWriter

这里的context就是Task对象中的context属性。不管是SortShuffleWriter中的externalSorter对象,还是UnsafeShuffleWriter中的shuffleExternalSorter,都是通过这个context来获取并使用TMM对象的。

4 TMM如何为task分配执行内存?

回到文章标题,总结一下本文内容:

1. TMM是spark tungsten内存管理机制的核心实现类,用于管理spark任务使用的执行内存。

2. Spark中的每个task attempt,无论是ShuffleMapTask还是ResultTask,都会生成一个专用的TMM对象,然后通过TaskContext将TMM对象共享给该task attempt的所有memory consumers.

3. TMM自建了一套内存页管理机制,并统一对ON_HEAP和OFF_HEAP内存进行编址,分配和释放。

4. TMM结合了逻辑内存和物理内存管理。

5. 部分memory consumer(比如Spillble的两个实现类)只用到了TMM的逻辑内存管理功能,物理内存的分配和释放还是依靠JVM完成的;

6. 部分memory consumer(比如ShuffleExternalSorter)则使用了TMM的逻辑和物理内存管理功能:

    a. 对于OFF_HEAP内存的使用绕开了JVM,避免了GC和java对象额外内存负载带来的性能损耗。

    b. 对于ON_HEAP内存也通过TMM的内存页管理和内存页缓存机制提高了spark对内存的使用效率。

5 说明

1. 源码版本:2.4.0

2. 水平有限,如有错误,望读者指出

推荐阅读更多精彩内容