本文討論跟執行內存密切相關的一個組件：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還是ResultTas，都會生成一個專用的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. 水平有限，如有錯誤，望讀者指出