前言
爲了解決Flink作業使用RocksDB狀態後端時的內存超用問題，Flink早在1.10版本就實現了RocksDB的託管內存（managed memory）機制。用戶只需啓用state.backend.rocksdb.memory.managed參數（默認即爲true），再設定合適的TaskManager託管內存比例taskmanager.memory.managed.fraction，即可滿足多數情況的需要。
關於RocksDB使用託管內存，Flink官方文檔給出了一段簡短的解釋：
Flink does not directly manage RocksDB’s native memory allocations, but configures RocksDB in a certain way to ensure it uses exactly as much memory as Flink has for its managed memory budget. This is done on a per-slot level (managed memory is accounted per slot).
To set the total memory usage of RocksDB instance(s), Flink leverages a shared cache and write buffer manager among all instances in a single slot. The shared cache will place an upper limit on the three components that use the majority of memory in RocksDB: block cache, index and bloom filters, and MemTables.
本文先簡單介紹一下RocksDB（版本5.17.2）內部的Cache和Write Buffer Manager這兩個組件，然後看一眼Flink是如何藉助它們來實現RocksDB內存託管的。
[LRU]Cache
Cache組件負責管理Block Cache，在RocksDB中的實現有兩種，分別對應兩種常用的緩存置換算法：LRUCache和ClockCache。由於ClockCache目前仍有bug，所以在生產環境總是使用默認的LRUCache。注意Cache有壓縮的和非壓縮的兩種，這裏只考慮默認的非壓縮Cache。
LRUCache最核心的四個參數列舉如下：
-
capacity：緩存的總大小。 -
num_shard_bits：按2num_shard_bits的規則確定整個緩存區域的分片（CacheShard）總數，也就是分片編號的比特數。每個CacheShard均分緩存容量，讀寫時，會根據key哈希值的高num_shard_bits位來確定路由。 -
strict_capacity_limit：是否嚴格控制單個緩存分片的容量限制，默認爲false。RocksDB的Iterator在遍歷數據時，會將它要讀取的一部分塊暫時固定在Cache內，稱爲Iterator-pinned blocks。如果Iterator-pinned blocks的大小超過了分片容量，再插入數據就有造成OOM的風險。開啓這個參數後，超額的緩存寫入就會直接失敗。 -
high_pri_pool_ratio：高優先級緩存區域佔整個Cache的比例。所謂高優先級緩存一般是指SST文件索引和布隆過濾器對應的塊，通過cache_index_and_filter_blocks和cache_index_and_filter_blocks_with_high_priority參數控制。
筆者之前講過實現LRU緩存的經典數據結構，即哈希表+雙鏈表（參見《手撕一個LRU Cache》）。RocksDB的LRUCache也是以這種思路爲基礎實現，簡單的示意圖如下。
每個緩存分片LRUCacheShard都有一套哈希表+循環雙鏈表的結構。哈希表稱爲LRUHandleTable，是RocksDB自己實現的鏈地址法分桶，且每個分片上都有互斥鎖，整體與JDK中的舊版ConcurrentHashMap非常相似。哈希桶的擴容和縮容也是按照2的冪次，並且會盡量保證扁平（即每個桶中儘量只有1個元素）。
一個低優先級指針（圖中Low-Pri）用於指示低優先級區域與高優先級區域的邊界。如果高優先級LRUHandle的量超過了high_pri_pool_ratio比例規定的量，就會將溢出的高優先級LRUHandle降格成低優先級。當然，淘汰LRUHandle時也是從低優先級區域開始淘汰。
LRUHandle是LRUCache的最小單元，其key是SST文件的ID加上塊在SST內的偏移量，value則是緩存的塊數據（代碼中爲void*類型），另外還有數據大小、指針域和引用計數域等。
爲什麼要有引用計數呢？因爲RocksDB的實現方法與傳統結構略有不同，鏈表中保存的並不是全部LRUHandle，而是可以被淘汰的那些LRUHandle，“可以被淘汰”的標準就是LRUHandle的引用計數爲1——只有哈希表中存在，而沒有外部引用者。也就是說，如果LRUHandle在鏈表中，那麼一定在哈希表中，反之則不成立。
Write Buffer Manager（WBM）
顧名思義，Write Buffer Manager（以下簡稱WBM）是用來管理寫緩存的組件。除了負責MemTable分配、Flush等細節，我們所關注的另一個作用則是追蹤和控制MemTable的內存用量，它可以以兩種形式生效：
- 傳入一個設定的閾值，WBM將多個列族或RocksDB實例的MemTable總大小限制在閾值內；
- 將WBM傳給Cache，可以使兩者共同控制RocksDB總內存佔用量的上限。
Flink也正是利用了上述特性來實現RocksDB託管內存的。那麼WBM與Cache如何協同工作？如下圖所示。
RocksDB Wiki中用了一句不符合英語語法的話來描述，即"Cost memory used in memtable to block cache"，此時Block Cache的內存配額就是RocksDB全部的內存配額。
MemTable的分配單元稱爲Arena Block，默認大小爲8MB。每分配一個Arena Block，WBM就會將它的內存消耗向LRUCache記賬——所謂“記賬”就是向Cache的低優先級區域內寫入Dummy LRUHandle。這些LRUHandle沒有value，只有key，但攜帶有Arena Block的內存消耗，且每個Dummy LRUHandle代表1MB的空間。也就是說它們僅佔用了邏輯配額，並未佔用物理空間，並且同樣受Cache的LRU規則的控制。由於MemTable本身既是讀緩存也是寫緩存，所以把它和Block Cache統一起來倒也合理。
WBM控制下的MemTable Flush策略也變得更加激進了一些：
- 當可變MemTable的大小超過WBM可用內存配額的7 / 8時，會觸發Flush；
- 當所有MemTable的大小超過內存配額，且可變MemTable的大小超過配額的一半時，也會觸發Flush。
下面來簡單看看Flink是如何利用WBM和Cache的。
To RocksDB Backend
直接上源碼，即org.apache.flink.contrib.streaming.state.RocksDBMemoryControllerUtils類。
public class RocksDBMemoryControllerUtils {
public static RocksDBSharedResources allocateRocksDBSharedResources(
long totalMemorySize, double writeBufferRatio, double highPriorityPoolRatio) {
long calculatedCacheCapacity =
RocksDBMemoryControllerUtils.calculateActualCacheCapacity(
totalMemorySize, writeBufferRatio);
final Cache cache =
RocksDBMemoryControllerUtils.createCache(
calculatedCacheCapacity, highPriorityPoolRatio);
long writeBufferManagerCapacity =
RocksDBMemoryControllerUtils.calculateWriteBufferManagerCapacity(
totalMemorySize, writeBufferRatio);
final WriteBufferManager wbm =
RocksDBMemoryControllerUtils.createWriteBufferManager(
writeBufferManagerCapacity, cache);
return new RocksDBSharedResources(cache, wbm, writeBufferManagerCapacity);
}
@VisibleForTesting
static long calculateActualCacheCapacity(long totalMemorySize, double writeBufferRatio) {
return (long) ((3 - writeBufferRatio) * totalMemorySize / 3);
}
@VisibleForTesting
static long calculateWriteBufferManagerCapacity(long totalMemorySize, double writeBufferRatio) {
return (long) (2 * totalMemorySize * writeBufferRatio / 3);
}
@VisibleForTesting
static Cache createCache(long cacheCapacity, double highPriorityPoolRatio) {
// TODO use strict capacity limit until FLINK-15532 resolved
return new LRUCache(cacheCapacity, -1, false, highPriorityPoolRatio);
}
@VisibleForTesting
static WriteBufferManager createWriteBufferManager(
long writeBufferManagerCapacity, Cache cache) {
return new WriteBufferManager(writeBufferManagerCapacity, cache);
}
static long calculateRocksDBDefaultArenaBlockSize(long writeBufferSize) {
long arenaBlockSize = writeBufferSize / 8;
// Align up to 4k
final long align = 4 * 1024;
return ((arenaBlockSize + align - 1) / align) * align;
}
static long calculateRocksDBMutableLimit(long bufferSize) {
return bufferSize * 7 / 8;
}
@VisibleForTesting
static boolean validateArenaBlockSize(long arenaBlockSize, long mutableLimit) {
return arenaBlockSize <= mutableLimit;
}
}
其中的writeBufferRatio就是state.backend.rocksdb.write-buffer-ratio參數，表示MemTable佔託管內存（即Block Cache）的比例，默認0.5。同理，highPriorityPoolRatio就是state.backend.memory.high-prio-pool-ratio參數，表示高優先級內存佔託管內存的比例，默認0.1。
託管內存在TaskManager的Slot之間平均分配，每個Slot都會有一組Cache和WBM。需要特別注意，實際的Cache和WBM配額是：
cache_capacity = (3 - write_buffer_ratio) * total_memory_size / 3
write_buffer_manager_capacity = 2 * total_memory_size * write_buffer_ratio / 3
也就是說，如果TM總的託管內存的大小是3GB，默認比例下的Block Cache大小其實是2.5GB，MemTable配額其實是1GB，都略偏小一些。這是因爲FLINK-15532尚未解決，strict_capacity_limit在Flink的場景下暫時不能生效，所以要留出一部分緩衝。推算的依據就是上一節提到的MemTable Flush策略，具體的關係如下：
write_buffer_manager_memory = 1.5 * write_buffer_manager_capacity
write_buffer_manager_memory = total_memory_size * write_buffer_ratio
write_buffer_manager_memory + other_part = total_memory_size
write_buffer_manager_capacity + other_part = cache_capacity
The End
寫的很潦草，就這樣吧。
下班跑路，民那晚安。