又見Rescale
筆者在很久之前的一篇文章（傳送門）中講解過Flink的狀態縮放（Rescale）和鍵組（Key Group）設計，相信各位看官對下面這張圖已經很熟悉了。
簡言之，Flink通過引入Key Group，將狀態Rescale時從遠端DFS恢復數據的操作從隨機讀盡量優化成順序讀，I/O瓶頸大大減輕。
但是，當Flink應用的Sub-task和狀態Key非常多時，改變有狀態算子的并行度仍然可能要花費較長時間恢復。例如我們負責的比較大的一個Flink任務（版本1.14），狀態Key總數接近10億，并行度從240擴容到480，在TaskManager資源充足且HDFS吞吐無瓶頸的情況下，狀態數據完全恢復也需要20+分鐘時間。除了數據量大之外，還有一個關鍵原因是：并行度增加后，新的Sub-task拿到原來的狀態數據，需要將不屬于自己的Key Group裁剪掉，否則會與其他Sub-task沖突。以上圖為例，擴容后的Sub-task 1需要將擴容前Sub-task 0和1的狀態文件都恢復到RocksDB，并且裁剪掉KG-1、KG-2、KG-5、KG-6的數據，只保留KG-3和KG-4。而我們知道，RocksDB的刪除操作是產生Tombstone記錄，本質上與寫入無異，所以這種情況下TaskManager本地磁盤仍然有較大的I/O壓力。
不過，上述問題在Flink 1.16有了一定改善，因為RocksDB狀態后端運用了RocksDB原生的DeleteRange API來快速刪除指定區間內的Keys，在我們的實測中，大狀態任務恢復速度最多可以提升60%。下面討論DeleteRange的實現原理。
RocksDB原生DeleteRange原理
在沒有DeleteRange API的時候，區間刪除只能采用傳統的迭代器遍歷操作：
Slice start, end;
auto it = db->NewIterator(ReadOptions());
for (it->Seek(start); cmp->Compare(it->key(), end) < 0; it->Next()) {
db->Delete(WriteOptions(), it->key());
}
有了DeleteRange API，就簡單很多：
Slice start, end;
db->DeleteRange(WriteOptions(), start, end);
為了實現區間刪除，RocksDB在原始的MemTable（稱為Point-key MemTable）之外，又新增了Range Tombstone MemTable，專門緩存區間刪除的數據。同理，在SST文件中也對應新增了包含區間刪除信息的元數據塊Range Tombstone Block（Seqnum為寫入序列號）。如下兩圖所示。
由此可見，如果我們要刪除包含10000個連續Key的集合，傳統方式會產生10000個Tombstone，而DeleteRange方式只會產生1個Range Tombstone，能夠有效降低讀寫放大。
在寫入過程中，Range Tombstone也需要參與Compaction流程，以及時刪除無效Tombstone。此處細節很多，簡單概括來講：
- 在Compaction開始時，收集所有源SST文件的Range Tombstone區間，形成一個包含所有區間刪除Key的最小堆。
- 對于每個輸入Key，判斷它是Merge類型還是Put類型：Merge操作則將該Key的所有歷史版本合并，如果歷史版本沒有被Snapshot引用，則可以刪除對應的Tombstone；Put操作說明該Key是新寫入數據，所有Tombstone都可以被清理掉。
- 清理完成后，將剩余的有效Tombstone重新寫回新SST文件的Range Tombstone Block。
引入Range Tombstone后，RocksDB讀取操作面臨一個新問題：如何快速判斷要讀取的Key是否位于某個已經標記刪除的區間中？答案是分段（RocksDB內部稱為"Fragmentation"），本質上與天際線問題（The Skyline Problem）的解法相同，見Leetcode 218。
如上圖所示，在RocksDB的語義下，X軸表示Key，Y軸表示Seqnum。在打開一個SST文件時，RocksDB會掃描該文件中所有的Range Tombstone區間（圖A中不同顏色的色塊），并將它們整合成互不重疊的子區間。將這些子區間按照左值升序排序并緩存下來（圖B），就可以根據Key進行高效的二分查找了。
Flink對DeleteRange的運用
回到Flink，以1.16版本為例，當任務發生Rescale并從狀態恢復到RocksDB時，實際上是調用RocksDBIncrementalRestoreOperation#restoreWithRescaling()方法：
private void restoreWithRescaling(Collection<KeyedStateHandle> restoreStateHandles)
throws Exception {
// Prepare for restore with rescaling
KeyedStateHandle initialHandle =
RocksDBIncrementalCheckpointUtils.chooseTheBestStateHandleForInitial(
restoreStateHandles, keyGroupRange, overlapFractionThreshold);
// Init base DB instance
if (initialHandle != null) {
restoreStateHandles.remove(initialHandle);
initDBWithRescaling(initialHandle);
} else {
this.rocksHandle.openDB();
}
// ..................................
}
其中，RocksDBIncrementalCheckpointUtils#chooseTheBestStateHandleForInitial()方法負責從所有要恢復的狀態句柄中，盡量選擇出與當前Sub-task負責的KeyGroupRange重合比例最高的一個，用來初始化本地RocksDB實例，以盡量降低后續裁剪的壓力。
接下來通過initDBWithRescaling()方法調用RocksDBIncrementalCheckpointUtils#clipDBWithKeyGroupRange()方法，按照KeyGroupRange的范圍進行裁剪。從文章開頭的圖示可知，由于Key Group已經是有序的，因此在擴容的情況下，新Sub-task不再負責的Key Group一定位于頭尾，因此只需要比較兩者的startKeyGroup和endKeyGroup即可。
public static void clipDBWithKeyGroupRange(
@Nonnull RocksDB db,
@Nonnull List<ColumnFamilyHandle> columnFamilyHandles,
@Nonnull KeyGroupRange targetKeyGroupRange,
@Nonnull KeyGroupRange currentKeyGroupRange,
@Nonnegative int keyGroupPrefixBytes)
throws RocksDBException {
final byte[] beginKeyGroupBytes = new byte[keyGroupPrefixBytes];
final byte[] endKeyGroupBytes = new byte[keyGroupPrefixBytes];
if (currentKeyGroupRange.getStartKeyGroup() < targetKeyGroupRange.getStartKeyGroup()) {
CompositeKeySerializationUtils.serializeKeyGroup(
currentKeyGroupRange.getStartKeyGroup(), beginKeyGroupBytes);
CompositeKeySerializationUtils.serializeKeyGroup(
targetKeyGroupRange.getStartKeyGroup(), endKeyGroupBytes);
deleteRange(db, columnFamilyHandles, beginKeyGroupBytes, endKeyGroupBytes);
}
if (currentKeyGroupRange.getEndKeyGroup() > targetKeyGroupRange.getEndKeyGroup()) {
CompositeKeySerializationUtils.serializeKeyGroup(
targetKeyGroupRange.getEndKeyGroup() + 1, beginKeyGroupBytes);
CompositeKeySerializationUtils.serializeKeyGroup(
currentKeyGroupRange.getEndKeyGroup() + 1, endKeyGroupBytes);
deleteRange(db, columnFamilyHandles, beginKeyGroupBytes, endKeyGroupBytes);
}
}
而deleteRange()方法就是代理了RocksDB JNI的同名方法，進行高效的區間刪除。
private static void deleteRange(
RocksDB db,
List<ColumnFamilyHandle> columnFamilyHandles,
byte[] beginKeyBytes,
byte[] endKeyBytes)
throws RocksDBException {
for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandles) {
// Using RocksDB's deleteRange will take advantage of delete
// tombstones, which mark the range as deleted.
//
// https://github.com/ververica/frocksdb/blob/FRocksDB-6.20.3/include/rocksdb/db.h#L363-L377
db.deleteRange(columnFamilyHandle, beginKeyBytes, endKeyBytes);
}
}
讀者也可以手動翻一下Flink 1.14或更早版本的源碼， deleteRange()方法是通過遍歷Key進行前綴比較，并執行WriteBatch操作批量刪除不符合條件的Key，相比原生DeleteRange的效率要低很多。
