本文讨论跟执行内存密切相关的一个组件：TaskMemoryManager（下文简称TMM）。TMM是tungsten内存管理机制的核心实现类（tungsten内存管理只作用于执行内存），它的功能包括：
1. 建立类似于操作系统内存页管理的机制，对ON_HEAP和OFF_HEAP内存统一编址和管理。
2. 通过调用MemoryManager和MemoryAllocator，将逻辑内存的申请&释放与物理内存的分配&释放结合起来。
3. 记录和管理task attempt的所有memory consumer.
1 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：
这里的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 ：
下文会详细讨论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控制：
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也可以看出来：
allocatePage方法要求consumer的memory mode和tungstenMemoryMode是一致的。
2.2 MemoryBLock
如上文所述， TMM用MemoryBlock表示一个内存页，MemoryBlock继承自MemoryLocation，先来看看它们的类图：
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源码如下：
3. 释放一个内存页时，其实就是将MemoryBlock对象中的object引用置为null，这样的话MemoryBlock对象对应的long数组就没有被引用了，其所占内存空间在JVM下次GC时会被自动回收。
4. 释放内存页的时候会检查内存页大小，如果超过HeapMemoryAllocator.POOLING_THRESHOLD_BYTES=1024*1024 （1MB），则会将其加入#2中提到的HashMap中去，一共后续内存申请直接使用，这种做法可避免频繁地向操作系统申请和释放相同大小的内存空间，也就是一次申请，多次使用。HeapMemoryAllocator.free源码如下：
2.4 UnsafeMemoryAllocator
UnsafeMemoryAllocator用于对OFF_HEAP内存进行分配和释放，其直接调用sun.misc.Unsafe.allocateMemory和sun.misc.Unsafe.freeMemory方法进行内存分配和释放。
注意，OFF_HEAP的MemoryBlock对象的obj属性适中为null.
与HeapMemoryAllocator不同，UnsafeMemoryAllocator没有用HashMap来缓存申请得到的内存页，而是每次都向操作系统申请和释放。
3 TMM对象是如何被创建，被使用的？
下图描述了TMM对象被创建，被Task用来生成TaskContext的过程：
Task (ShuffleMapTask或ResultTask) 在运行时会生成TaskContext对象并赋值给Task.context属性，而spark运行过程中的各个memory consumer就是通过task的这个context属性来获取并使用TMM对象的。
举个例子，在shuffle过程中，ShuffleMapTask会通过SortShuffleManager.getWriter方法获取一个ShuffleWriter对象：
这里的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. 水平有限，如有错误，望读者指出