spark 原始碼分析之十五 -- Spark記憶體管理剖析

本篇文章主要剖析Spark的記憶體管理體系。

在上篇文章 spark 原始碼分析之十四 -- broadcast 是如何實現的？中對儲存相關的內容沒有做過多的剖析，下面計劃先剖析Spark的記憶體機制，進而進入記憶體儲存，最後再剖析磁碟儲存。本篇文章主要剖析記憶體管理機制。

整體介紹

Spark記憶體管理相關類都在 spark core 模組的 org.apache.spark.memory 包下。

文件對這個包的解釋和說明如下：

This package implements Spark's memory management system. This system consists of two main components, a JVM-wide memory manager and a per-task manager:

- org.apache.spark.memory.MemoryManager manages Spark's overall memory usage within a JVM. This component implements the policies for dividing the available memory across tasks and for allocating memory between storage (memory used caching and data transfer) and execution (memory used by computations, such as shuffles, joins, sorts, and aggregations).

- org.apache.spark.memory.TaskMemoryManager manages the memory allocated by individual tasks. Tasks interact with TaskMemoryManager and never directly interact with the JVM-wide MemoryManager. Internally, each of these components have additional abstractions for memory bookkeeping:

- org.apache.spark.memory.MemoryConsumers are clients of the TaskMemoryManager and correspond to individual operators and data structures within a task. The TaskMemoryManager receives memory allocation requests from MemoryConsumers and issues callbacks to consumers in order to trigger spilling when running low on memory.

- org.apache.spark.memory.MemoryPools are a bookkeeping abstraction used by the MemoryManager to track the division of memory between storage and execution. 

 

即記憶體管理主要涉及了兩個元件：JVM 範圍的記憶體管理和單個任務的記憶體管理。

1. MemoryManager管理Spark在JVM中的總體記憶體使用情況。該元件實現了跨任務劃分可用記憶體以及在儲存（記憶體使用快取和資料傳輸）和執行（計算使用的記憶體，如shuffle，連線，排序和聚合）之間分配記憶體的策略。
2. TaskMemoryManager管理由各個任務分配的記憶體。任務與TaskMemoryManager互動，永遠不會直接與JVM範圍的MemoryManager互動。

 

在TaskMemoryManager內部，每個元件都有額外的記憶簿來記錄記憶體使用情況：

 

• MemoryConsumers是TaskMemoryManager的客戶端，對應於任務中的各個運算子和資料結構。TaskMemoryManager接收來自MemoryConsumers的記憶體分配請求，並向消費者發出回撥，以便在記憶體不足時觸發溢位。
• MemoryPools是MemoryManager用來跟蹤儲存和執行之間記憶體劃分的薄記抽象。 

如圖：

 

MemoryManager的兩種實現： 

There are two implementations of org.apache.spark.memory.MemoryManager which vary in how they handle the sizing of their memory pools: 
- org.apache.spark.memory.UnifiedMemoryManager, the default in Spark 1.6+, enforces soft boundaries between storage and execution memory, allowing requests for memory in one region to be fulfilled by borrowing memory from the other. 
- org.apache.spark.memory.StaticMemoryManager enforces hard boundaries between storage and execution memory by statically partitioning Spark's memory and preventing storage and execution from borrowing memory from each other. This mode is retained only for legacy compatibility purposes.

 

org.apache.spark.memory.MemoryManager有兩種實現，它們在處理記憶體池大小方面有所不同：

• org.apache.spark.memory.UnifiedMemoryManager，Spark 1.6+中的預設值，強制儲存記憶體和執行記憶體之間的軟邊界，允許通過從另一個區域借用記憶體來滿足一個區域中的記憶體請求。
• org.apache.spark.memory.StaticMemoryManager 通過靜態分割槽Spark的記憶體，強制儲存記憶體和執行記憶體之間的硬邊界並防止儲存和執行從彼此借用記憶體。 僅為了傳統相容性目的而保留此模式。

先來一張自己畫的類圖，對涉及類之間的關係有一個比較直接的認識：

 

下面我們逐一對涉及的類做說明。

MemoryMode

記憶體模式：主要分堆內記憶體和堆外記憶體，MemoryMode是一個列舉類，從本質上來說，ON_HEAP和OFF_HEAP都是MemoryMode的子類。

MemoryPool

文件說明如下：

Manages bookkeeping for an adjustable-sized region of memory. This class is internal to the MemoryManager. 

 

即它負責管理可調大小的記憶體區域的簿記工作。可以這樣理解，記憶體就是一個金庫，它是一個負責記賬的管家，主要負責記錄記憶體的借出歸還。這個類專門為MempryManager而設計。

給記憶體記賬，其實從本質上來說，它不是Spark記憶體管理部分的核心功能，但是又很重要，它的核心方法都是被MemoryManager來呼叫的。

理解了這個類，其子類就比較好理解了。記賬的管家有兩種實現，分別是StorageMemoryPool和ExecutionMemoryPool。

StorageMemoryPool

文件解釋：

Performs bookkeeping for managing an adjustable-size pool of memory that is used for storage (caching).

 

說白了，它就是專門給負責儲存或快取的記憶體區域記賬的。

其類結構如下：

它有三種方法：

1. acquireMemory：獲取N個位元組的記憶體給指定的block，如果有必要，即記憶體不夠用了，可以將其他的從記憶體中驅除。原始碼如下：

圖中標記的邏輯，參照下文MemoryStore的剖析。

2. releaseMemory：釋放記憶體。原始碼如下：

很簡單，就只是在統計值_memoryUsed 上面做減法。

3. freeSpaceToShrinkPool：可用空間通過`spaceToFree`位元組縮小此儲存記憶體池的大小。原始碼如下：

 

簡單地可以看出，這個方法是在收縮儲存記憶體池之前呼叫的，因為這個方法返回值是要收縮的值。

收縮儲存記憶體池是為了擴大執行記憶體池，即這個方法是在收縮儲存記憶體，擴大執行記憶體時用的，這個方法只是為了縮小儲存記憶體池作準備的，並沒有真正的縮小儲存記憶體池。

實現思路，首先先計算需要驅逐的記憶體大小，如果需要驅逐記憶體，則跟 acquireMemory 方法類似，呼叫MemoryStore 的 evictBlocksToFreeSpace方法，否則直接返回。

總結：這個類是給儲存記憶體池記賬的，也負責不夠時或記憶體池不滿足縮小條件時，通知MemoryStore驅逐記憶體。

 

ExecutionMemoryPool

文件解釋：

Implements policies and bookkeeping for sharing an adjustable-sized pool of memory between tasks. 
Tries to ensure that each task gets a reasonable share of memory, 
instead of some task ramping up to a large amount first and then causing others to spill to disk repeatedly. 
If there are N tasks, it ensures that each task can acquire at least 1 / 2N of the memory before it has to spill, 
and at most 1 / N. Because N varies dynamically, we keep track of the set of active tasks and redo the calculations 
of 1 / 2N and 1 / N in waiting tasks whenever this set changes. This is all done by synchronizing access to mutable 
state and using wait() and notifyAll() to signal changes to callers. Prior to Spark 1.6, this arbitration of memory 
across tasks was performed by the ShuffleMemoryManager.

 

實現策略和簿記，以便在任務之間共享可調大小的記憶體池。 嘗試確保每個任務獲得合理的記憶體份額，而不是首先增加大量任務然後導致其他任務重複溢位到磁碟。

如果有N個任務，它確保每個任務在溢位之前至少可以獲取1 / 2N的記憶體，最多1 / N.

由於N動態變化，我們會跟蹤活動任務的集合並在每當任務集合改變時重做等待任務中的1 / 2N和1 / N的計算。

這一切都是通過同步對可變狀態的訪問並使用 wait() 和 notifyAll() 來通知對呼叫者的更改來完成的。 在Spark 1.6之前，跨任務的記憶體仲裁由ShuffleMemoryManager執行。 

 

類內部結構如下：

memoryForTask宣告如下：

@GuardedBy("lock")
private val memoryForTask = new mutable.HashMap[Long, Long]()

其中，key 指的是 taskAttemptId， value 是記憶體使用情況（以byte計算）。它用來記錄每一個任務記憶體使用情況。

它也有三類方法：

1. 獲取總的或每一個任務的記憶體使用大小，原始碼如下：

memoryForTask 記錄了每一個task使用的記憶體大小。

 

2. 給一個任務分配記憶體，原始碼如下：

numBytes表示申請的記憶體大小（in byte），taskAttemptId 表示申請記憶體的 task id，maybeGrowPool 表示一個可能會增加執行池大小的回撥。 它接受一個引數（Long），表示應該擴充套件此池的所需記憶體量。computeMaxPoolSize 表示在此給定時刻返回此池的最大允許大小的回撥。這不是欄位，因為在某些情況下最大池大小是可變的。 例如，在統一記憶體管理中，可以通過驅逐快取塊來擴充套件執行池，從而縮小儲存池。

如果之前該任務沒有申請過，則將(taskAttemptId <- 0) 放入到 memoryForTask map 中， 然後釋放鎖並喚醒lock鎖等待區的執行緒。

被喚醒的因為synchronized實現的是一個互斥鎖，所以當前僅當只有一個執行緒執行while迴圈。

首先根據 （需要的記憶體大小 - 池總空閒記憶體大小）來確認是否需要擴大池，由於儲存池可能會偷執行池的記憶體，所以需要執行 maybeGrowPool 方法。

computeMaxPoolSize計算出此時該池允許的最大記憶體大小。然後分別算出每個任務最大分配記憶體和最小分配記憶體。進而計算出分配給該任務的最大分配大小（maxToGrant）和實際分配大小（toGrant）。

如果實際分配大小 小於需要分配的記憶體大小 並且 當前任務佔有記憶體 + 實際分配記憶體 < 每個任務最小分配記憶體，則該執行緒進入鎖wait區等待，等待記憶體可用時喚醒，否則將記憶體分配給任務。

可以看到這個方法中的wait和notify方法並不是成對的，因為新新增的taskAttemptId不能滿足記憶體可用的條件。因為這個鎖是從外部傳過來的，即MemoryManager也可能對其做了操作，使記憶體空餘下來，可供任務分配。

3. 釋放task記憶體，原始碼如下：

它有兩個方法，分別是釋放當前任務已經使用的所有記憶體空間 releaseAllMemoryForTask 和釋放當前任務的指定大小的記憶體空間 releaseMemory。

思路：

releaseAllMemoryForTask 先計算好當前任務使用的全部記憶體，然後呼叫 releaseMemory 方法釋放記憶體。

releaseMemory 方法則會比對當前使用記憶體和要釋放的記憶體，如果要釋放的記憶體大小小於 當前使用的 ，做減法即可。釋放之後的任務記憶體如果小於等於0，則移除task即可，最後通知lock鎖等待區的物件，讓其重新分配記憶體。

在這個記賬的實現裡，每一個來的task不一定是可以分配到記憶體的，所以，鎖在其中起了很大的資源協調的作用，也防止了記憶體的溢位。

 

MemoryManager

文件說明：

An abstract memory manager that enforces how memory is shared between execution and storage. In this context, execution memory refers to that used for computation in shuffles, joins, sorts and aggregations, while storage memory refers to that used for caching and propagating internal data across the cluster. There exists one MemoryManager per JVM.

一種抽象記憶體管理器，用於強制執行和儲存之間共享記憶體的方式。在這個上下文下，執行記憶體是指用於在shuffle，join，sort和aggregation中進行計算的記憶體，而儲存記憶體是指用於在群集中快取和傳播內部資料的記憶體。 每個JVM都有一個MemoryManager。

先來說一下其依賴的MemoryPool，原始碼如下：

MemoryPool中的lock物件就是MemoryManager物件

儲存記憶體池和執行記憶體池分別有兩個：堆內和堆外。

onHeapStorageMemory和onHeapExecutionMemory 是從構造方法傳過來的，先不予考慮。

maxOffHeapMemory 預設是 0， 可以根據 spark.memory.offHeap.size 引數設定，文件對這個引數的說明如下：

The absolute amount of memory in bytes which can be used for off-heap allocation. 
This setting has no impact on heap memory usage, so if your executors' total memory consumption must fit within some hard limit 
then be sure to shrink your JVM heap size  accordingly. This must be set to a positive value when spark.memory.offHeap.enabled=true.

 

儲存堆外記憶體 = 最大堆外記憶體（offHeapStorageMemory） X 堆外儲存記憶體佔比，這個佔比預設是0.5，可以根據 spark.memory.storageFraction 來調節

執行堆外記憶體 = 最大堆外記憶體 - 儲存堆外記憶體

還有跟 Tungsten 管理記憶體有關的常量：

這三個常量分別定義了tungsten的記憶體形式、記憶體頁大小和記憶體分配器。

 

其方法解釋如下：

1. 獲取儲存池最大使用記憶體，抽象方法，待子類實現。

 

2. 獲取已使用記憶體

3. 獲取記憶體，這也是抽象方法，待子類實現

 

4. 釋放記憶體

這些請求都委託給對應的MemoryPool來做了

1.6 之前 使用MemoryManager子類 StaticMemoryManager 來做記憶體管理。

StaticMemoryManager

這個靜態記憶體管理中的執行池和儲存池之間有嚴格的界限，兩個池的大小永不改變。

注意：如果想使用這個記憶體管理方式，設定 spark.memory.useLegacyMode 為 true即可（預設是false）

 

下面我們重點看1.6 之後的預設使用的MemoryManager子類 -- UnifiedMemoryManager

UnifiedMemoryManager

先來看文件說明：

這個MemoryManager保證了儲存池和執行池之間的軟邊界，即可以互相借用記憶體來滿足彼此動態的記憶體需求變化。執行和儲存的佔比由 spark.memory.storageFraction 配置，預設是0.6，即偏向於儲存池。其中儲存池的預設佔比是由 spark.memory.storageFraction 引數決定，預設是 0.5 ，即 儲存池預設佔比 = 0.6 * 0.5 = 0.3 ，即儲存池預設佔比為0.3。儲存池可以儘可能多的向執行池借用空閒記憶體。但是當執行池需要它的記憶體的時候，會把一部分記憶體池的記憶體物件從記憶體中驅逐出，直到滿足執行池的記憶體需求。類似地，執行池也可以儘可能地借用儲存池中的空閒記憶體，不同的是，執行記憶體不會被儲存池驅逐出記憶體，也就是說，快取block時可能會因為執行池佔用了大量的記憶體池不能釋放導致快取block失敗，在這種情況下，新的block會根據StorageLevel做相應處理。

 

我們主要來看其實現的父類MemoryManager 的方法：

1. 獲取儲存池最大使用記憶體：

其中，maxHeapMemory 是從構造方法傳進來的成員變數，maxOffHeapMemory 是根據引數 spark.memory.offHeap.size 配置生成的。

可以看出，儲存池的允許的最大使用記憶體是實時變化的，因為總記憶體不變，執行池記憶體使用情況隨任務執行情況變化而變化。

 

2. 獲取記憶體，逐一來看：

實現思路：先根據儲存方式（堆內還是堆外）確定儲存池，執行池，儲存區域記憶體大小和最大總記憶體。

然後呼叫執行池的 acquireMemory 方法申請記憶體，computeMaxExecutionPoolSize是隨儲存的實時變化而變化的，增大ExecutionPool的回撥也被呼叫來確保有足夠空間可供執行池分配。

acquireUnrollMemory 直接呼叫 acquireStorageMemory 方法。

acquireStorageMemory實現思路：先根據儲存方式（堆內還是堆外）確定儲存池，執行池，儲存區域記憶體大小和最大總記憶體。

儲存記憶體如果大於最大記憶體，直接儲存失敗，否則，繼續檢視所需記憶體大小是否大於記憶體池最大空閒記憶體，如果大於，則從執行池中申請足夠的空閒空間，注意，真正申請的空間大小在0 和numBytes - storagePool.memoryFree 之間，繼續呼叫storagePool的acquireMemory 方法去申請記憶體，如果不夠申請，則會驅逐出舊或空的block塊。

最後，我們來看一下其伴生物件：

首先 apply 方法就類似於工廠方法的創造方法。我們對比下面的一張圖，來說明一下Spark記憶體結構：

系統記憶體：可以根據 spark.testing.memory 引數來配置（主要用於測試），預設是JVM 的可以使用的最大記憶體。

保留記憶體：可以根據 spark.testing.reservedMemory 引數來配置（主要用於測試）， 預設是 300M

最小系統記憶體：保留記憶體 * 1.5 後，再向下取整

系統記憶體的約束：系統記憶體必須大於最小保留記憶體，即 系統可用記憶體必須大於 450M， 可以通過 --driver-memory 或  spark.driver.memory 或 --executor-memory 或spark.executor.memory 來調節

可用記憶體 = 系統記憶體 - 保留記憶體

堆內記憶體佔比預設是0.6， 可以根據 spark.memory.fraction 引數來調節

最大堆內記憶體 = 堆內可用記憶體 * 堆內記憶體佔比

堆內記憶體儲存池佔比預設是 0.5 ，可以根據spark.memory.storageFraction 來調節。

預設堆記憶體儲記憶體大小 = 最大堆內記憶體 * 堆內記憶體儲存池佔比。即堆記憶體儲池記憶體大小預設是 （系統JVM最大可用記憶體 -  300M）* 0.6 * 0.5， 即約等於JVM最大可用記憶體的三分之一。

注意： 下圖中的spark.memory.fraction是0.75，是Spark 1.6 的預設配置。在Spark 2.4.3 中預設是0.6。

 圖片來源：https://0x0fff.com/spark-memory-management/

至此，Saprk 的記憶體管理模組基本上剖析完畢。

總結：先介紹了記憶體的管理池，即MemoryPool的實現，然後重點分析了Spark 1.6 以後的記憶體管理機制，著重說明Spark內部的記憶體是如何劃分以及如何動態調整記憶體的。

 

注，關於堆內記憶體和堆外記憶體的介紹，可參照：https://www.jianshu.com/p/50be08b54bee