本文仅为笔者平日学习记录之用，侵删
原文：https://mp.weixin.qq.com/s/eaALnpd_qHQg6fxI12fQjg
本文会详细分析 TM 端恢复及创建 KeyedState 的流程，恢复过程会分析 RocksDB 和 Fs 两种 StateBackend 的恢复流程，创建流程会介绍 Checkpoint 处恢复的 State 如何与代码中创建的 State 关联起来。
一、 RocksDBKeyedStateBackend 创建流程
从 RocksDBStateBackend 类的 createKeyedStateBackend 方法开始，createKeyedStateBackend 方法源码主要加载一些配置和创建 RocksDBKeyedStateBackend，就不贴出来了。简单介绍一下 createKeyedStateBackend 方法的功能：
- 加载 RocksDB JNI library
- 初始化 RocksDB 的本地数据目录，对应的是 RocksDB
state.backend.rocksdb.localdir参数配置的目录 - new RocksDBKeyedStateBackendBuilder 的构造器，所有状态的恢复及初始化都封装在 build 方法中。
RocksDBKeyedStateBackendBuilder 的 build 方法删减后源码如下所示：
@Override
public RocksDBKeyedStateBackend build() throws BackendBuildingException {
// 维护 StateName 与 StateInfo 的映射关系
LinkedHashMap
kvStateInformation = new LinkedHashMap<>();
RocksDB db = null;
AbstractRocksDBRestoreOperation restoreOperation = null;
SnapshotStrategy snapshotStrategy;
try {
// 保存 CheckpointID 与 sst 的映射
SortedMap> materializedSstFiles = new TreeMap<>();
long lastCompletedCheckpointId = -1L;
// 准备 instanceBasePath 目录，用于本地状态存储
prepareDirectories();
// 根据 restoreStateHandles 决定三种恢复模式：
// 1、 RocksDB 无需恢复，直接启动
// 2、 RocksDB 增量 Checkpoint 的状态恢复
// 3、 RocksDB 全量 Checkpoint 的状态恢复
restoreOperation = getRocksDBRestoreOperation(XXX);
// 恢复状态，并打开 db，具体 执行三种不同恢复流程
RocksDBRestoreResult restoreResult = restoreOperation.restore();
db = restoreResult.getDb();
// RocksDB 的增量 Checkpoint 模式，则获取 sst 文件，
if (restoreOperation instanceof RocksDBIncrementalRestoreOperation) {
backendUID = restoreResult.getBackendUID();
materializedSstFiles = restoreResult.getRestoredSstFiles();
lastCompletedCheckpointId = restoreResult.getLastCompletedCheckpointId();
}
// 初始化 Savepoint 和 Checkpoint 的 snapshot 策略
snapshotStrategy = initializeSavepointAndCheckpointStrategies(XXX);
} catch (Throwable e) {
// State 初始化异常，做一些清理操作
}
InternalKeyContext keyContext = new InternalKeyContextImpl<>(
keyGroupRange,
numberOfKeyGroups
);
// kvStateInformation 会传递给 RocksDBKeyedStateBackend
return new RocksDBKeyedStateBackend<>(XXX);
}
build 方法中同样定义了一个 Map kvStateInformation 用于维护 StateName 与 StateInfo 的映射关系。之后会准备本地目录，用于本地状态存储。
getRocksDBRestoreOperation 方法会创建 RestoreOperation，源码如下所示：
// 根据 restoreStateHandles 决定三种恢复模式：
// 1、 RocksDB 直接启动无需恢复
// 2、 增量 Checkpoint 的 RocksDB 恢复
// 3、 全量 Checkpoint 的 RocksDB 恢复
private AbstractRocksDBRestoreOperation getRocksDBRestoreOperation() {
// restoreStateHandles 为空表示没有 State 需要恢复，构造 NoneRestoreOperation
if (restoreStateHandles.isEmpty()) {
return new RocksDBNoneRestoreOperation<>(X);
}
KeyedStateHandle firstStateHandle = restoreStateHandles.iterator().next();
// StateHandle 是 Increment 模式，则构造 IncrementalRestoreOperation
// 否则构造 FullRestoreOperation
if (firstStateHandle instanceof IncrementalKeyedStateHandle) {
return new RocksDBIncrementalRestoreOperation<>(XXX);
} else {
return new RocksDBFullRestoreOperation<>(XXX);
}
}
getRocksDBRestoreOperation 方法会根据 restoreStateHandles 决定三种恢复模式：
restoreStateHandles 为空表示没有 State 需要恢复，构造 NoneRestoreOperation，即：不恢复 State 的方式，直接启动
restoreStateHandles 不为空的情况下，判断 StateHandle 的类型， StateHandle 是 Increment 模式，则构造 IncrementalRestoreOperation；否则构造 FullRestoreOperation
build 方法下一步就会执行具体 RocksDBRestoreOperation 的 restore 方法了，三种 RocksDBRestoreOperation 的 restore 流程完全不同，且比较复杂，后续单独介绍。后续会执行 initializeSavepointAndCheckpointStrategies 方法初始化 Savepoint 和 Checkpoint 的 snapshot 策略，方法的返回值是 SnapshotStrategy 类型，SnapshotStrategy 封装了 Checkpoint 和 Savepoint 两个策略。如果用户触发 Checkpoint，则执行 Checkpoint 策略，触发 Savepoint，则执行 Checkpoint 策略。
initializeSavepointAndCheckpointStrategies 方法源码如下所示：
// SnapshotStrategy 封装了 Checkpoint 和 Savepoint 两个策略。
class SnapshotStrategy {
final RocksDBSnapshotStrategyBase checkpointSnapshotStrategy;
final RocksDBSnapshotStrategyBase savepointSnapshotStrategy;
}
private SnapshotStrategy initializeSavepointAndCheckpointStrategies(XXX) {
// 创建 Savepoint 的 snapshot 类为 Full Snapshot 策略
RocksDBSnapshotStrategyBase savepointSnapshotStrategy =
new RocksFullSnapshotStrategy<>(XXX);
RocksDBSnapshotStrategyBase checkpointSnapshotStrategy;
// 如果开启了增量 Checkpoint，
// 则 Checkpoint 的 snapshot 类为 Increment Snapshot 策略
if (enableIncrementalCheckpointing) {
checkpointSnapshotStrategy = new RocksIncrementalSnapshotStrategy<>(XXX);
} else {
// 未开始增量 Checkpoint，
// 则 Checkpoint 的 snapshot 为 Savepoint 的 snapshot 策略
checkpointSnapshotStrategy = savepointSnapshotStrategy;
}
// 封装两个策略到 SnapshotStrategy 中
return new SnapshotStrategy<>(checkpointSnapshotStrategy, savepointSnapshotStrategy);
}
可以看到 RocksDB Savepoint 的 snapshot 类永远为 Full Snapshot 策略，如果开启增量 Checkpoint，则 Checkpoint 的 snapshot 类为 Increment Snapshot 策略。未开始增量 Checkpoint，则 Checkpoint 的 snapshot 为 Savepoint 的 snapshot 策略，最后封装两个策略到 SnapshotStrategy 中。
这块可以得出一个结论：使用 RocksDBStateBackend 时，如果不开启增量 Checkpoint，那么触发 Savepoint 和 Checkpoint 都是相同的策略，即：都是 Full Snapshot 模式。
build 方法中上述流程如果任何阶段抛出异常，都认为 State 初始化异常，会做一些清理操作，并认为本次任务恢复失败。如果一切都成功，最后会构建出 RocksDBKeyedStateBackend。
build 方法结束！下面重点关注三种不同的 RestoreOperation 具体是怎么 restore 的。
二、 RocksDB 的 NoneRestoreOperation 恢复流程
NoneRestoreOperation 表示没有 State 需要恢复，直接启动，所以该模式下 restore 流程特别简单。
restore 方法源码如下所示：
@Override
public RocksDBRestoreResult restore() throws Exception {
openDB();
return new RocksDBRestoreResult(this.db, defaultColumnFamilyHandle,
nativeMetricMonitor, -1, null, null);
}
直接打开一个空的 RocksDB 就恢复完成返回结果。
三、 RocksDB 的 IncrementalRestoreOperation 恢复流程
RocksDBIncrementalRestoreOperation 类的 restore 方法源码如下所示：
@Override
public RocksDBRestoreResult restore() throws Exception {
if (restoreStateHandles == null || restoreStateHandles.isEmpty()) {
return null;
}
final KeyedStateHandle theFirstStateHandle = restoreStateHandles.iterator().next();
// restoreStateHandles 数量大于 1，
// 或者 恢复的 keyGroupRange 与当前负责的 keyGroupRange 不同，
// 则使用 Rescaling 模式。如果没有改并发，则关闭 Rescaling 模式
boolean isRescaling = (restoreStateHandles.size() > 1 ||
!Objects.equals(theFirstStateHandle.getKeyGroupRange(), keyGroupRange));
if (isRescaling) {
// Rescaling 开启的恢复模式，相当于改并发恢复，需要依赖 KeyGroup 恢复
restoreWithRescaling(restoreStateHandles);
} else {
// Rescaling 关闭的恢复模式，相当于没有改变并发，直接恢复 sst 即可
// 没有改并发就只有一个 StateHandle，所以这里只需要将 firstStateHandle 当做参数传递即可
restoreWithoutRescaling(theFirstStateHandle);
}
return new RocksDBRestoreResult(this.db, defaultColumnFamilyHandle,
nativeMetricMonitor, lastCompletedCheckpointId, backendUID, restoredSstFiles);
}
如果 restoreStateHandles 为 null 或者集合为空，直接返回 null。否则开始后面的恢复流程。
恢复流程第一步检测是否是 rescale 模式，换言之：检测是否新旧 Job 之间修改并发了。isRescaling 为 true 表示修改并发了，isRescaling 为 false 表示没有修改并发。
判断是否修改并发的逻辑：
代码中判断逻辑：如果 restoreStateHandles 集合中元素数量大于 1 或者恢复的 keyGroupRange 与当前负责的 keyGroupRange 不同，则开启 Rescaling 模式。否则关闭 Rescaling 模式。
分析一波：如果不修改并发，那么新 Job 的 subtask 与旧 Job 的 subtask 是一对一的关系，每个 subtask 只会恢复旧 Job 对应的那一个 subtask 的 StateHandle，且新旧 subtask 负责的 KeyGroupRange 是相同的。
代码中 restoreStateHandles 集合中元素数量表示要恢复的 KeyedStateHandle 的数据，数量大于 1 表示当前 subtask 要恢复旧 Job 的多个 subtask 的 KeyedStateHandle。所以 restoreStateHandles.size() > 1 必然修改了并发。
代码中拿要恢复的 StateHandle 的 KeyGroupRange 与当前 subtask 负责的 KeyGroupRange 进行比较，两者不同则表示新旧 subtask 负责的 KeyGroupRange 不是完全相同的，也可以推断出一定修改并发了。
例如：旧的 subtask 0 负责的 KeyGroupRange(0,9)，新的 subtask 0 负责的 KeyGroupRange(0,6)，虽然 restoreStateHandles 集合中只有一个 StateHandle，但是 KeyGroupRange 变了，也可以推断出并发改变了。
判断完是否修改并发，就会按照是否修改并发，进行两种不同的模式开始恢复流程。如源码所示，restoreWithoutRescaling 方法表示为修改并发的恢复模式，这里有个小细节，restoreWithoutRescaling 方法的参数是 KeyedStateHandle 类型，而不用传整个集合，因为不修改并发只会恢复一个 KeyedStateHandle。而 restoreWithRescaling 方法的参数就是 KeyedStateHandle 的集合类型。
下面详细介绍两种恢复模式。
未修改并发的恢复流程
restoreWithoutRescaling 方法源码如下所示：
private void restoreWithoutRescaling(KeyedStateHandle keyedStateHandle) throws Exception {
if (keyedStateHandle instanceof IncrementalRemoteKeyedStateHandle) {
// 远程的 KeyedStateHandle 恢复
IncrementalRemoteKeyedStateHandle incrementalRemoteKeyedStateHandle =
(IncrementalRemoteKeyedStateHandle) keyedStateHandle;
// 保存 StateHandle 对应的那一次 Checkpoint 对应的文件状态，
// 包括 lastCompletedCheckpointId 和 chk id 与 sst 映射关系
restorePreviousIncrementalFilesStatus(incrementalRemoteKeyedStateHandle);
/**
* 从远程恢复 State 的过程：
* 1、 本地创建 tmp 目录
* 2、 从远程拉取 sst 文件到本地，将 远程的 StateHandle 转换为 本地 的 StateHandle
* 3、 调用 restoreFromLocalState 方法，从 local 恢复 State
* 4、 清理 tmp 文件
*/
restoreFromRemoteState(incrementalRemoteKeyedStateHandle);
} else if (keyedStateHandle instanceof IncrementalLocalKeyedStateHandle) {
// Local 的 KeyedStateHandle 恢复
IncrementalLocalKeyedStateHandle incrementalLocalKeyedStateHandle =
(IncrementalLocalKeyedStateHandle) keyedStateHandle;
// 保存 StateHandle 对应的那一次 Checkpoint 对应的文件状态，
// 包括 lastCompletedCheckpointId 和 chk id 与 sst 映射关系
restorePreviousIncrementalFilesStatus(incrementalLocalKeyedStateHandle);
// 从 local 恢复 State 的方法
restoreFromLocalState(incrementalLocalKeyedStateHandle);
} else {
throw new BackendBuildingException("XXX");
}
}
restoreWithoutRescaling 方法中首先会区分 KeyedStateHandle 到底是 IncrementalRemoteKeyedStateHandle 类型还是 IncrementalLocalKeyedStateHandle，区别在于一个是 Remote 一个是 Local。正常从 dfs 上恢复就属于 Remote 模式，但是 RocksDB 增量 Checkpoint 有个 local-recovery 的优化。
local-recovery 是指：Checkpoint 时会在本地目录保留一份状态快照，当任务重启时避免了从 dfs 上拉取状态文件的过程，加速任务恢复。
Local 模式的恢复流程：
restorePreviousIncrementalFilesStatus 方法保存 StateHandle 对应的那一次 Checkpoint 对应的文件状态，包括 lastCompletedCheckpointId 和 chk id 与 sst 映射关系。
然后调用 restoreFromLocalState 方法从 Local 恢复 State 即可。
Remote 模式的恢复流程：
同样也是调用 restorePreviousIncrementalFilesStatus 方法保存 StateHandle 对应的那一次 Checkpoint 对应的文件状态，包括 lastCompletedCheckpointId 和 chk id 与 sst 映射关系。
区别在于调用 restoreFromRemoteState 方法从 Remote 恢复 State，但其实 restoreFromRemoteState 最后还是调用的 restoreFromLocalState。restoreFromRemoteState 方法源码如下所示：
private void restoreFromRemoteState(IncrementalRemoteKeyedStateHandle stateHandle) {
// 创建临时目录
final Path tmpRestoreInstancePath = new Path(
instanceBasePath.getAbsolutePath(),
UUID.randomUUID().toString());
try {
// 从本地恢复
restoreFromLocalState(
// 从远程拉取 State 文件到本地
transferRemoteStateToLocalDirectory(tmpRestoreInstancePath, stateHandle));
} finally {
cleanUpPathQuietly(tmpRestoreInstancePath);
}
}
远程恢复 State 的过程如下所示：
- 本地创建 tmp 目录
- 从远程拉取 sst 文件到本地，将 远程的 StateHandle 转换为 本地 的 StateHandle
- 调用 restoreFromLocalState 方法，从 local 恢复 State
- 清理 tmp 文件
所以得出的结论是：Remote 模式相比 Local 模式而言，只是多了一个从 dfs 上下载文件到本地的过程，下载到本地后就转换成了 Local 模式进行恢复。所以先看一下 transferRemoteStateToLocalDirectory 方法是如何下载文件的，之后重点关注 restoreFromLocalState 即可。
下载状态文件到本地流程：
transferRemoteStateToLocalDirectory 方法源码如下所示：
private IncrementalLocalKeyedStateHandle transferRemoteStateToLocalDirectory(
Path temporaryRestoreInstancePath,
IncrementalRemoteKeyedStateHandle restoreStateHandle) throws Exception {
// try with resource 的方式创建 RocksDBStateDownloader
try (RocksDBStateDownloader rocksDBStateDownloader =
new RocksDBStateDownloader(numberOfTransferringThreads)) {
// 具体的从 dfs 上 Download 数据到本地
// 使用线程池，多线程拉取 所有的 sst 文件和 RocksDB 数据库的元数据
rocksDBStateDownloader.transferAllStateDataToDirectory(
restoreStateHandle,
temporaryRestoreInstancePath,
cancelStreamRegistry);
}
// 将 Remote 的 StateHandle 重新构建成 Local 的 StateHandle
return new IncrementalLocalKeyedStateHandle(
restoreStateHandle.getBackendIdentifier(),
restoreStateHandle.getCheckpointId(),
// 使用 DirectoryStateHandle，目录就是之前创建的临时数据目录
new DirectoryStateHandle(temporaryRestoreInstancePath),
restoreStateHandle.getKeyGroupRange(),
restoreStateHandle.getMetaStateHandle(),
restoreStateHandle.getSharedState().keySet());
}
创建 RocksDBStateDownloader 类，见名之意，用于下载 RocksDB 状态文件的类。RocksDBStateDownloader 的构造参数是拉取文件的线程数，具体可以进行配置的。然后使用 RocksDBStateDownloader 去 dfs 上多线程 Download 所有的 sst 文件和 RocksDB 数据库的元数据数据到本地的 temporaryRestoreInstancePath 临时目录下。
拉取完成后，将 Remote 的 StateHandle 重新构建成 Local 的 StateHandle，并且使用 DirectoryStateHandle，这里的目录就是之前创建的临时数据目录，刚才下载的数据也在该目录下。然后就开始
Increment 模式不修改并发，从 Local 恢复 State 流程
restoreFromLocalState 方法源码如下所示：
// 从本地 State 文件中恢复状态
private void restoreFromLocalState(IncrementalLocalKeyedStateHandle localKeyedStateHandle) throws Exception {
// 从 State 中获取元数据
KeyedBackendSerializationProxy serializationProxy = readMetaData(
localKeyedStateHandle.getMetaDataState());
List stateMetaInfoSnapshots = serializationProxy
.getStateMetaInfoSnapshots();
// 根据 State 的元数据，创建或注册 CF 描述符
columnFamilyDescriptors =
createAndRegisterColumnFamilyDescriptors(stateMetaInfoSnapshots, true);
Path restoreSourcePath = localKeyedStateHandle.getDirectoryStateHandle().getDirectory();
// 准备数据目录
restoreInstanceDirectoryFromPath(restoreSourcePath, dbPath);
// 打开 DB，打开 DB 时会将 columnFamilyHandles 填满，
// 即获取到 columnFamilyHandles 集合
columnFamilyHandles = new ArrayList<>(columnFamilyDescriptors.size() + 1);
openDB();
// 将 StateName 和 State 在 RocksDB 的 CF 句柄等元数据 映射信息保存到 kvStateInformation 中
registerColumnFamilyHandles(stateMetaInfoSnapshots);
}
首先从 State 中获取元数据，根据 State 的元数据，创建或注册 CF 描述符，然后 restoreInstanceDirectoryFromPath 方法用于准备数据目录，准备数据目录就是将那些从 Checkpoint 处拉取的文件从临时目录拷贝到真正的 DB 目录。restoreInstanceDirectoryFromPath 方法在拷贝数据时做了一个优化，以 .sst 结尾的文件不是真正拷贝，而是做了一个 link，其他文件真正的拷贝过去。RocksDB 真正存储数据的是 .sst，这些才是占用空间较大的文件，向其他的元数据占用空间较小，所以直接拷贝。
最后 registerColumnFamilyHandles 方法，将 StateName 和 State 在 RocksDB 的 CF 句柄等元数据映射信息保存到 kvStateInformation 中。kvStateInformation 是一个 Map，key 为 StateName，value 为 RocksDbKvStateInfo 类型。
RocksDbKvStateInfo 相关类源码如下所示：
public static class RocksDbKvStateInfo implements AutoCloseable{
public final ColumnFamilyHandle columnFamilyHandle;
public final RegisteredStateMetaInfoBase metaInfo;
}
RocksDbKvStateInfo 封装了 ColumnFamilyHandle 和 RegisteredStateMetaInfoBase，ColumnFamilyHandle 表示 RocksDB CF 句柄，RegisteredStateMetaInfoBase 中维护了 State 的 name、类型、序列化信息等。有了 kvStateInformation，就可以根据 StateName 拿到当前 State 对应的 CF 句柄读到数据，并拿到其对应的序列化规则对读到的数据进行反序列化。
到这里 RocksDB 的 Increment 模式在不改变并发的情况下，无论是 Remote 还是 Local，数据都正常恢复了（数据在本地的 RocksDB 实例中，可以根据 kvStateInformation 中维护的信息从 RocksDB 中读取到数据）。
修改并发的恢复流程
修改并发的情况，新的 subtask 可能要恢复多个 StateHandle 的数据，也就是多个 RocksDB 实例的数据。最笨的方法是将多个 RocksDB 的数据全拉取的本地，建立多个 RocksDB 实例，从中读取出当前 subtask 对应 KeyGroup 的数据，写入到一个新的 RocksDB 实例中。这样存在一个问题，所有数据都是一条条从旧的 RocksDB get 出来，再一条条 put 到一个新的 RocksDB 中 。
恢复 RocksDB 优化
FLINK-8790 基于上述方案做了一个优化：首先从多个 RocksDB 实例中选取一个最优的 RocksDB 实例，最优的标准是：RocksDB 负责 KeyGroupRange 与当前 subtask 负责的 KeyGroupRange 的交集占 RocksDB 负责的 KeyGroupRange 的百分比最高。
举个例子：RocksDB a 存储的 KeyGroupRange(0,9) 的数据，当前 subtask 负责的 KeyGroupRange(0,7) 的数据，那么交集就是 KeyGroupRange(0,7)。KeyGroupRange(0,7) 中包含 8 个 KeyGroup，KeyGroupRange(0,9) 包含 10 个 KeyGroup，8 /10 为 80%。
此时就认为 RocksDB a 上有 80% 的数据是有效的，所有要恢复的 RocksDB 都做一次上述运算，挑选出分数最高的 RocksDB 实例。源码中还加了一层限制，重叠率低于 75% 的 RocksDB 会直接被过滤掉。通过上述筛选，可能会得到一个相对来讲最优的 RocksDB 做为最终的 RocksDB，但是要对其进行裁剪。就拿上述例子来讲，RocksDB 中存储的 KeyGroupRange(0,9) 的数据，但只需要 KeyGroupRange(0,7) 的数据，所以会将 KeyGroupRange(8,9) 的数据裁掉。当然裁剪效率相对较高，RocksDB 中 key 的设计都是以 KeyGroup 开头的，LSM Tree 的底层存储都是按照 key 有序存储，所以直接按照前缀即可高效裁剪。
筛选最优 StateHandle 的代码，参考 RocksDBIncrementalCheckpointUtils 的 chooseTheBestStateHandleForInitial 方法和 STATE_HANDLE_EVALUATOR 函数式接口，源码如下所示：
public static KeyedStateHandle chooseTheBestStateHandleForInitial(
@Nonnull Collection restoreStateHandles,
@Nonnull KeyGroupRange targetKeyGroupRange) {
KeyedStateHandle bestStateHandle = null;
double bestScore = 0;
// 遍历所有 KeyedStateHandle
for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
// 计算分数
double handleScore = STATE_HANDLE_EVALUATOR.apply(rawStateHandle, targetKeyGroupRange);
if (handleScore > bestScore) {
// 保存最高分 及 对应的 KeyedStateHandle
bestStateHandle = rawStateHandle;
bestScore = handleScore;
}
}
return bestStateHandle;
}
private static final BiFunction
STATE_HANDLE_EVALUATOR = (stateHandle, targetKeyGroupRange) -> {
final KeyGroupRange handleKeyGroupRange = stateHandle.getKeyGroupRange();
// 计算当前 StateHandle 与 目标 Handle 在 KeyGroup 上的交集
final KeyGroupRange intersectGroup = handleKeyGroupRange.
getIntersection(targetKeyGroupRange);
// 计算 当前 StateHandle 对应的状态文件上 有 百分之多少的数据应该在当前 subtask 上
final double overlapFraction = (double) intersectGroup.getNumberOfKeyGroups() /
handleKeyGroupRange.getNumberOfKeyGroups();
// 概率小于 0.75 返回 -1，意味着该 StateHandle 是肯定会被丢弃的
if (overlapFraction < OVERLAP_FRACTION_THRESHOLD) {
return -1.0;
}
return intersectGroup.getNumberOfKeyGroups()
* overlapFraction * overlapFraction;
};
修改并发的恢复主流程
restoreWithRescaling 源码如下所示：
private void restoreWithRescaling(Collection restoreStateHandles) {
// 选取一个最好的 StateHandle 用于数据初始化，会有一个选择标准打分，分数最高，则被选中
// 分数的计算规则：主要依赖 StateHandle 的 KeyGroup 与 当前 subtask 处理的 KeyGroup 求一个交集，看重叠率
KeyedStateHandle initialHandle = RocksDBIncrementalCheckpointUtils.
chooseTheBestStateHandleForInitial(restoreStateHandles, keyGroupRange);
// Init base DB instance
if (initialHandle != null) {
// 打开选取的认为最好的 StateHandle 对应的 db，并对其多余的 KeyGroup 进行裁剪
restoreStateHandles.remove(initialHandle);
initDBWithRescaling(initialHandle);
} else {
// open 一个空的 db
openDB();
}
// 将 target 的 startKey 和 endKey 转换成 byte 形式
byte[] startKeyGroupPrefixBytes = new byte[keyGroupPrefixBytes];
RocksDBKeySerializationUtils.serializeKeyGroup(keyGroupRange.getStartKeyGroup(),
startKeyGroupPrefixBytes);
byte[] stopKeyGroupPrefixBytes = new byte[keyGroupPrefixBytes];
RocksDBKeySerializationUtils.serializeKeyGroup(keyGroupRange.getEndKeyGroup() + 1,
stopKeyGroupPrefixBytes);
// 将所有要恢复的 StateHandle 中对应的 RocksDB 恢复，
// 并将 target 的 startKey 和 endKey 之间的数据 put 到目标 db
// 同时还需要将 StateName 和 State 在 RocksDB 的 CF 句柄等元数据 映射信息保存到 kvStateInformation 中
for (KeyedStateHandle rawStateHandle : restoreStateHandles) {
XXX
}
}
首先根据分数，挑选一个最优的 StateHandle 作为 RocksDB 初始 DB，initDBWithRescaling 方法会对多余的 KeyGroup 进行裁剪。如果没有挑选出来说明都不太优，会直接创建一个空的 RocksDB 作为初始 DB。
根据当前 subtask 负责的 keyGroupRange 计算出 RocksDB 的 startKey 和 endKey，把其他剩余的所有 StateHandle 对应 RocksDB 数据库一一恢复，从中读取出 key 位于 startKey 和 endKey 之间的数据插入到初始 DB 中。从 RocksDB 读取数据时可以直接通过 startKey seek 到指定位置，因为是全局有序的，所以遍历过程中一旦读到 endKey 以外的数据，就认为遍历结束了，直接退出循环。
同时在恢复过程中，需要将 StateName 和 State 在 RocksDB 的 CF 句柄等元数据 映射信息保存到 kvStateInformation 中。
小优化思考
后续多个 RocksDB 实例恢复时的流程基本是串行操作，即：从 dfs 上拉取第一个 RocksDB 数据文件、本地构建 RocksDB 数据库，依次读出 startKey 和 endKey 之间的数据插入到新的 RocksDB。再拉取第二个、构建、读数据、写数据。再拉取第三个。。。
思考：所有从 dfs 上拉取 RocksDB 数据文件的过程，能不能完全异步化，即：读写第一个 RocksDB 的过程中，就开始拉取第二个、第三个等。
四、 RocksDB 的 FullRestoreOperation 恢复流程
RocksDBFullRestoreOperation 类的 restore 方法的源码如下所示：
@Override
public RocksDBRestoreResult restore()
throws IOException, StateMigrationException, RocksDBException {
// 打开空的 DB
openDB();
// 遍历所有的 restoreStateHandles
for (KeyedStateHandle keyedStateHandle : restoreStateHandles) {
if (keyedStateHandle != null) {
// RocksDB 的 Full 模式与 Savepoint 模式保存的状态文件都是 Flink 自己序列化好的问题，
// 其对应的 KeyedStateHandle 必然是 KeyGroupsStateHandle。
if (!(keyedStateHandle instanceof KeyGroupsStateHandle)) {
throw new IllegalStateException("Unexpected state handle type, " +
"expected: " + KeyGroupsStateHandle.class +
", but found: " + keyedStateHandle.getClass());
}
this.currentKeyGroupsStateHandle = (KeyGroupsStateHandle) keyedStateHandle;
// 根据 StateHandle 恢复
restoreKeyGroupsInStateHandle();
}
}
return new RocksDBRestoreResult(this.db, defaultColumnFamilyHandle,
nativeMetricMonitor, -1, null, null);
}
restore 方法会打开一个空的 RocksDB，遍历所有的 restoreStateHandles，之前强调过 Full 模式的 KeyStateHandle 对应的是 KeyGroupsStateHandle 类型。所以这里进行了判断，如果不是 KeyGroupsStateHandle 类型，直接抛出异常，恢复失败。然后 restoreKeyGroupsInStateHandle 方法用于恢复当前 keyedStateHandle 对应的数据。
restoreKeyGroupsInStateHandle 方法源码如下所示：
private void restoreKeyGroupsInStateHandle()
throws IOException, StateMigrationException, RocksDBException {
try {
// KeyGroupsStateHandle 的场景，并不会直接拉回文件，而是建立一个远程的输入流
currentStateHandleInStream = currentKeyGroupsStateHandle.openInputStream();
cancelStreamRegistry.registerCloseable(currentStateHandleInStream);
currentStateHandleInView = new DataInputViewStreamWrapper(currentStateHandleInStream);
// 注册 StateName 和 State 在 RocksDB 的 CF 句柄等元数据
// 映射信息保存到 kvStateInformation 中
restoreKVStateMetaData();
// 将当前 StateHandle 中属于当前 KeyGroupRange 的数据 put 到 db 中
restoreKVStateData();
} finally {
if (cancelStreamRegistry.unregisterCloseable(currentStateHandleInStream)) {
IOUtils.closeQuietly(currentStateHandleInStream);
}
}
}
restoreKeyGroupsInStateHandle 依然会将 StateName 和 State 在 RocksDB 的 CF 句柄等元数据映射信息保存到 kvStateInformation 中，并将 StateHandle 中属于当前 KeyGroupRange 的数据 put 到 db 中。KeyGroupsStateHandle 中能拿到状态文件的输入流，且有保存的每个 KeyGroup 在文件中的 offset，所以可以直接读取到数据并 put 到刚创建的 RocksDB 中。
小结
到这里 RocksDB 的三种模式都恢复完成，RocksDB 的三种恢复模式下，都会将 StateName 与具体 State 的信息维护在 kvStateInformation 中。后期在创建 State 的过程中，也会通过 kvStateInformation 将创建的 State 与 Checkpoint 中恢复的 State 进行关联。
下面分析 FsStateBackend 模式下的恢复流程。
五、 HeapKeyedStateBackend 创建恢复流程
FsStateBackend 模式下，createKeyedStateBackend 方法创建的是 HeapKeyedStateBackend，最后调用的 HeapKeyedStateBackendBuilder 的 build 方法。
build 方法源码如下所示：
public HeapKeyedStateBackend build() throws BackendBuildingException {
// Map of registered Key/Value states
Map> registeredKVStates = new HashMap<>();
// Map of registered priority queue set states
Map
registeredPQStates = new HashMap<>();
HeapSnapshotStrategy snapshotStrategy = initSnapshotStrategy(XXX);
InternalKeyContext keyContext = new InternalKeyContextImpl<>(
keyGroupRange,
numberOfKeyGroups
);
HeapRestoreOperation restoreOperation = new HeapRestoreOperation<>(XXX);
try {
// 恢复流程
restoreOperation.restore();
} catch (Exception e) {
throw new BackendBuildingException("XXX", e);
}
// 构建 HeapKeyedStateBackend
return new HeapKeyedStateBackend<>(XXX);
}
build 方法中首先创建出 Map 类型的 registeredKVStates，用于保存 StateName 及对应的 StateTable，每个 State 对应一个 StateTable 存储状态数据。将 registeredKVStates 传递给 HeapRestoreOperation 用于恢复，最后再传递给 HeapKeyedStateBackend 用于后续使用。
HeapRestoreOperation 类的 restore 方法会遍历所有的 StateHandle 恢复 State 信息，维护映射关系到 registeredKVStates 中，并恢复 State 信息到具体的 StateTable 中。StateTable 是 Heap 模式真正存储 State 的集合。
小结
HeapKeyedStateBackend 会将 StateName 与具体 State 的信息维护在 registeredKVStates 中。后期在创建 State 的过程中，也会通过 registeredKVStates 将创建的 State 与 Checkpoint 中恢复的 State 进行关联。
下面详细分析 KeyedState 的创建流程。
六、 用户定义的 KeyedState 创建流程
可以拿 IntervalJoinOperator 的例子来分析 KeyedState 的创建流程，IntervalJoin 用于对两个输入流的数据进行关联，两个流先到的数据会放到 buffer 中，左右两个流分别有各自的 buffer，使用 Flink 的 MapState 充当 buffer。
IntervalJoinOperator 类的 initializeState 方法源码如下所示：
public void initializeState(StateInitializationContext context) throws Exception {
super.initializeState(context);
// 左流的 buffer
this.leftBuffer = context.getKeyedStateStore().getMapState(new MapStateDescriptor<>(
LEFT_BUFFER,
LongSerializer.INSTANCE,
new ListSerializer<>(new BufferEntrySerializer<>(leftTypeSerializer))
));
// 右流的 buffer
this.rightBuffer = context.getKeyedStateStore().getMapState(new MapStateDescriptor<>(
RIGHT_BUFFER,
LongSerializer.INSTANCE,
new ListSerializer<>(new BufferEntrySerializer<>(rightTypeSerializer))
));
}
initializeState 方法会创建两个 State，即：LEFT_BUFFER 和 RIGHT_BUFFER。
context.getKeyedStateStore().getMapState 实际调用 DefaultKeyedStateStore 类的 getMapState 方法，整个创建 State 第一阶段的方法调用时序图如下所示：
调用关系从 DefaultKeyedStateStore 类的 getMapState 方法到 KeyedStateFactory 的 createInternalState 方法。
如下图所示，KeyedStateFactory 有两个子类，即：HeapKeyedStateBackend 和 RocksDBKeyedStateBackend。
如下图所示，HeapKeyedStateBackend 会对应 MemoryStateBackend 和 FsStateBackend，RocksDBKeyedStateBackend 对应 RocksDBStateBackend。
下面详细介绍 HeapKeyedStateBackend 和 RocksDBKeyedStateBackend 的 createInternalState 方法是如何创建 State 的。
HeapKeyedStateBackend 创建 State 流程
HeapKeyedStateBackend 类的 createInternalState 方法源码如下所示：
public IS createInternalState(XXX) {
StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
// 注册和恢复 StateTable
StateTable stateTable = tryRegisterStateTable(
namespaceSerializer, stateDesc,
getStateSnapshotTransformFactory(stateDesc, snapshotTransformFactory));
// 根据 stateDesc、StateTable 和 序列化信息，创建具体的 State
return stateFactory.createState(stateDesc, stateTable, getKeySerializer());
}
createInternalState 方法中首先会执行 tryRegisterStateTable 方法「注册和恢复」 StateTable，然后根据 stateDesc、StateTable 和 序列化信息，创建具体的 State。
重点的恢复逻辑就在 tryRegisterStateTable 方法中，tryRegisterStateTable 方法源码如下所示：
private StateTable tryRegisterStateTable(XXX) {
// 根据 StateName 从 registeredKVStates 中获取 StateTable
StateTable stateTable = (StateTable)
registeredKVStates.get(stateDesc.getName());
// stateTable 不为空，表示从 Checkpoint 中恢复了当前 State
if (stateTable != null) {
RegisteredKeyValueStateBackendMetaInfo restoredKvMetaInfo =
stateTable.getMetaInfo();
// 主要对 State 的兼容性进行校验，校验包括：StateName、状态类型、序列化校验
// 如果创建的 State 与 Checkpoint 恢复的 State 不匹配，
// 则抛出异常，不能成功恢复
restoredKvMetaInfo.updateSnapshotTransformFactory(snapshotTransformFactory);
TypeSerializerSchemaCompatibility namespaceCompatibility =
restoredKvMetaInfo.updateNamespaceSerializer(namespaceSerializer);
// 检查 State 的 name 和 Type 是否可以匹配
restoredKvMetaInfo.checkStateMetaInfo(stateDesc);
TypeSerializerSchemaCompatibility stateCompatibility =
restoredKvMetaInfo.updateStateSerializer(newStateSerializer);
if (stateCompatibility.isIncompatible()) {
throw new StateMigrationException("XXX");
}
stateTable.setMetaInfo(restoredKvMetaInfo);
} else {
// 没有从 Checkpoint 恢复，则创建 StateTable，存放到 registeredKVStates 中
RegisteredKeyValueStateBackendMetaInfo newMetaInfo = new
RegisteredKeyValueStateBackendMetaInfo<>(XXX);
stateTable = snapshotStrategy.newStateTable(keyContext, newMetaInfo, keySerializer);
registeredKVStates.put(stateDesc.getName(), stateTable);
}
return stateTable;
}
tryRegisterStateTable 方法首先会根据 StateName 从 registeredKVStates 中获取 StateTable 保存到 stateTable 中。
如果 stateTable 不为空，表示 Checkpoint 中有当前 StateName 对应的状态，应该恢复，此时会对新旧 Job 的 State 匹配性进行检测，校验项包括：StateName、状态类型、序列化校验。
否则 stateTable 为空，表示当前 StateName 不需要从 Checkpoint 恢复，直接创建一个新的 StateTable，存放到 registeredKVStates 中。
RocksDBKeyedStateBackend 创建 State 流程
RocksDBKeyedStateBackend 类的 createInternalState 方法源码如下所示：
public IS createInternalState(XXX) {
StateFactory stateFactory = STATE_FACTORIES.get(stateDesc.getClass());
// 注册和恢复 State 元信息
Tuple2>
registerResult = tryRegisterKvStateInformation(
stateDesc, namespaceSerializer, snapshotTransformFactory);
// 根据 stateDesc、State 元信息 和 RocksDBKeyedStateBackend，创建具体的 State
return stateFactory.createState(stateDesc, registerResult,
RocksDBKeyedStateBackend.this);
}
createInternalState 方法中首先会执行 tryRegisterKvStateInformation 方法「注册和恢复」 State 元信息，然后根据 stateDesc、State 元信息 和 RocksDBKeyedStateBackend，创建具体的 State。
tryRegisterKvStateInformation 与上述 HeapKeyedStateBackend 类的恢复逻辑类似，所以不贴代码了，tryRegisterKvStateInformation 方法的整体逻辑就是：
首先会根据 StateName 从 kvStateInformation 中获取 State 的元信息保存到 oldStateInfo 中。
如果 stateTable 不为空，表示 Checkpoint 中有当前 StateName 对应的状态，应该恢复，此时会对新旧 Job 的 State 匹配性进行校验。
否则 stateTable 为空，表示当前 StateName 不需要从 Checkpoint 恢复，直接在 RocksDB 中创建一个新的 ColumnFamily 存储当前 State 的数据。
小结
在恢复过程中主要依赖之前创建的 Map，Map 中保存的从 Checkpoint 中恢复出来的状态数据。如果 Map 中有对应 StateName 的数据，则对其进行校验并恢复；如果 Map 中找不到，则创建新的。
七、 总结
本文首先介绍了 RocksDBKeyedStateBackend 创建流程，并分别介绍 RocksDB 三种模式下的 State 恢复流程，分别是：NoneRestoreOperation、IncrementalRestoreOperation、FullRestoreOperation 三种模式。之后介绍 HeapKeyedStateBackend 恢复流程。最后介绍了用户定义的 KeyedState 创建流程，创建流程会介绍 Checkpoint 处恢复的 State 如何与代码中创建的 State 关联起来。