Flink 中的 CopyOnWriteStateTable

现如今想阅读 HashMap 源码实际上比较简单，因为网上一大堆博客去分析 HashMap 和 ConcurrentHashMap。而本文是全网首篇详细分析 CopyOnWriteStateTable 源码的博客，阅读复杂集合类源码的过程是相当有挑战的，笔者在刚开始阅读也遇到很多疑问，最后一一解决了。本文有一万两千多字加不少的配图，实属不易。详细阅读完本文，无论是针对面试还是开阔视野一定会对大家有帮助的。感觉有帮助的同学文末点个在看呗，如果能转发那更好了。

❝

声明：笔者的源码分析都是基于 flink-1.9.0 release 分支，其实阅读源码不用非常在意版本的问题，各版本的主要流程基本都是类似的。如果熟悉了某个版本的源码，之后新版本有变化，我们重点看一下变化之处即可。

笔者阅读源码中会加很多中文注释，对源码感兴趣且有需要的同学可以关注一下笔者的 github 仓库：https://github.com/1996fanrui/flink/tree/feature/source-code-read-1-9-0

注释都在 feature/source-code-read-1-9-0 分支，之后也会持续更新

❞

本文主要讲述 Flink 中 CopyOnWriteStateTable 相关的知识，当使用 MemoryStateBackend 和 FsStateBackend 时，默认情况下会将状态数据保存到 CopyOnWriteStateTable 中。CopyOnWriteStateTable 中保存多个 KeyGroup 的状态，每个 KeyGroup 对应一个 CopyOnWriteStateMap。

CopyOnWriteStateMap 是一个类似于 HashMap 的结构，但支持了两个非常有意思的功能：

• 1、hash 结构为了保证读写数据的高性能，都需要有扩容策略，CopyOnWriteStateMap 的扩容策略是一个渐进式 rehash 的策略，即：不是一下子将数据全迁移的新的 hash 表，而是慢慢去迁移数据到新的 hash 表中

• 2、Checkpoint 时 CopyOnWriteStateMap 支持异步快照，即：Checkpoint 时可以在做快照的同时，仍然对 CopyOnWriteStateMap 中数据进行修改。问题来了：数据修改了，怎么保证快照数据的准确性呢？

了解 Redis 的同学应该知道 Redis 也是一个大的 hash 结构，扩容策略也是渐进式 rehash。Redis 的 RDB 在持久化数据的过程中同时也是对外服务的，对外服务意味着数据可能被修改，那么 RDB 如何保证持久化好的数据一定是正确的呢？

举个例子：17 点00分00秒 RDB 开始持久化数据，过了 1 秒 Redis 中某条数据被修改了，过了一分钟 RDB 才持久化结束。RDB 预期的持久化结果应该是 17 点00分00秒那一刻 Redis 的完整快照，请问持久化过程中那些修改操作是否会影响 Redis 的快照。答：当然可以做到不影响。

Flink 在 Checkpoint 时的快照与 Redis 类似，都是想在快照时依然对外提供服务，减少服务停顿时间。Flink 具体如何实现上述功能的呢？带着问题详细阅读下文。

1、 StateTable 简介

MemoryStateBackend 和 FsStateBackend 的 KeyedStateBackend 都使用 HeapKeyedStateBackend 存储数据，HeapKeyedStateBackend 持有 Map<String, StateTable<K, ?, ?>> registeredKVStates 来存储 StateName 与具体 State 的映射关系。registeredKVStates 的 key 就是 StateName，value 为具体的 State 数据。具体 State 的数据存储在 StateTable 中。

StateTable 有两个实现：CopyOnWriteStateTable 和 NestedMapsStateTable。

• CopyOnWriteStateTable 属于 Flink 自己定制化的数据结构，Checkpoint 时支持异步 Snapshot。

• NestedMapsStateTable 直接嵌套 Java 的两层 HashMap 来存储数据，Checkpoint 时需要同步快照。

下面详细介绍 CopyOnWriteStateTable。

2、 CopyOnWriteStateTable

StateTable 中持有 StateMap[] keyGroupedStateMaps 真正的存储数据。StateTable 会为每个 KeyGroup 的数据初始化一个 StateMap 来对 KeyGroup 做数据隔离。对状态进行操作时，StateTable 会先根据 key 计算对应的 KeyGroup，拿到相应的 StateMap，才能对状态进行操作。

CopyOnWriteStateTable 中使用 CopyOnWriteStateMap 存储数据，这里主要介绍 CopyOnWriteStateMap 的实现。CopyOnWriteStateMap 中就是一个数组 + 链表构成的 hash 表。

CopyOnWriteStateMap 中元素类型都是是：StateMapEntry。hash 表的第一层先是一个 StateMapEntry 类型的数组，即：StateMapEntry[]。在 StateMapEntry 类中有个 StateMapEntry next 指针构成链表。

CopyOnWriteStateMap 相比普通的 hash 表，有以下几点需要重点关注：

• CopyOnWriteStateMap 的扩容策略是渐进式 rehash，而不是一下子扩容完

• 为了支持异步的 Snapshot，需要将 Snapshot 时 StateMap 的快照保存下来，具体的保存策略怎么实现的？

• 为了支持 CopyOnWrite 功能，所以在修改数据时，要进行一系列 copy 的操作，不能修改原始数据，否则会影响 Snapshot

• Snapshot 异步快照流程及 Snapshot 完成时，如何 release 掉旧版本数据？

3、 CopyOnWriteStateMap 的渐进式 rehash 策略

渐进式 rehash 策略表示 CopyOnWriteStateMap 中当前有一个 hash 表对外服务，但是当前 hash 表中元素太多需要扩容了，需要将数据迁移到一个容量更大的 hash 表中。

Java 的 HashMap 在扩容时会一下子将旧 hash 表中所有数据都移动到大 hash 表中，这样的策略存在的问题是如果 HashMap 当前存储了 1 G 的数据，那么瞬间需要将 1 G 的数据迁移完，可能会比较耗时。而 CopyOnWriteStateMap 在扩容时，不会一下子将数据全部迁移完，而是在每次操作 CopyOnWriteStateMap 时，慢慢去迁移数据到大的 hash 表中。

例如：可以在每次 get、put 操作时，迁移 4 条数据到大 hash 表中，这样经过一段时间的 get 和 put 操作，所有的数据就能迁移完成。所以渐进式 rehash 策略，会分很多次将所有的数据迁移到新的 hash 表中。

3.1 扩容简述

在内存中有两个 hash 表，一个是 primaryTable 作为主桶，一个是 rehashTable 作为扩容期间用的桶。初始阶段只有 primaryTable，当 primaryTable 中元素个数大于设定的阈值时，就要开始扩容。

扩容过程：申请一个相比 primaryTable 容量大一倍的 hash 表保存到 rehashTable 中，慢慢地将 primaryTable 中的元素迁移到 rehashTable 中。对应到源码中：putEntry 方法中判断 size() > threshold 时，会调用 doubleCapacity 方法申请新的 hash 表赋值给 rehashTable。

如下图所示 primaryTable 中桶的个数为 4，rehashTable 中桶的个数为 8。

扩容时 primaryTable 中 0 位置上的元素会迁移到 rehashTable 的 0 和 4 位置上，同理 primaryTable 中 1 位置上的元素会迁移到 rehashTable 的 1 和 5 位置上。

3.2 选择 Table 的策略

假设 primaryTable 中 0 桶的数据已经迁移到 rehashTable 桶了，那么之后无论是 put 还是 get 操作 0 桶的数据，那么都会去操作 rehashTable。而 1、2、3 桶还未迁移，所以 1、2、3 桶还需要操作 primaryTable 桶。对应到源码中会有一个选桶的操作，选择到底使用 primaryTable 还是 rehashTable。

源码实现如下所示：

// 选择当前元素到底使用 primaryTable 还是 incrementalRehashTable
private StateMapEntry<K, N, S>[] selectActiveTable(int hashCode) {
  // 计算 hashCode 应该被分到 primaryTable 的哪个桶中
  int curIndex = hashCode & (primaryTable.length - 1);
  // 大于等于 rehashIndex 的桶还未迁移，应该去 primaryTable 中去查找。
  // 小于 rehashIndex 的桶已经迁移完成，应该去 incrementalRehashTable 中去查找。
 return curIndex >= rehashIndex ? primaryTable : incrementalRehashTable;
}

首先通过 int curIndex = hashCode & (primaryTable.length - 1); 计算当前 hashCode 应该分到 primaryTable 的哪个桶中。

rehashIndex 用来标记当前 rehash 迁移的进度，即：rehashIndex 之前的数据已经从 primaryTable 迁移到 rehashTable 桶中。假设 rehashIndex = 1，表示 primaryTable 1 桶之前的数据全部迁移完成了，即：0 桶数据全部迁移完了。

策略：大于等于 rehashIndex 的桶还未迁移，应该去 primaryTable 中去查找。小于 rehashIndex 的桶已经迁移完成，应该去 incrementalRehashTable 中去查找。

3.3 迁移过程

每次有 get、put、containsKey、remove 操作时，都会调用 computeHashForOperationAndDoIncrementalRehash 方法触发迁移操作。

computeHashForOperationAndDoIncrementalRehash 方法作用：

• 检测是否处于 rehash 中，如果正在 rehash 就会调用 incrementalRehash 迁移一波数据

• 计算 key 和 namespace 对应的 hashCode

重点关注 incrementalRehash 方法实现：

private void incrementalRehash() {

 StateMapEntry<K, N, S>[] oldMap = primaryTable;
 StateMapEntry<K, N, S>[] newMap = incrementalRehashTable;

 int oldCapacity = oldMap.length;
 int newMask = newMap.length - 1;
 int requiredVersion = highestRequiredSnapshotVersion;
 int rhIdx = rehashIndex;
  // 记录本次迁移了几个元素
 int transferred = 0;

 // 每次至少迁移 MIN_TRANSFERRED_PER_INCREMENTAL_REHASH 个元素到新桶、
 // MIN_TRANSFERRED_PER_INCREMENTAL_REHASH 默认为 4
 while (transferred < MIN_TRANSFERRED_PER_INCREMENTAL_REHASH) {

    // 遍历 oldMap 的第 rhIdx 个桶
  StateMapEntry<K, N, S> e = oldMap[rhIdx];

    // 每次 e 都指向 e.next，e 不为空，表示当前桶中还有元素未遍历，需要继续遍历
  // 每次迁移必须保证，整个桶被迁移完，不能是某个桶迁移到一半
  while (e != null) {
   // 遇到版本比 highestRequiredSnapshotVersion 小的元素，则 copy 一份
   if (e.entryVersion < requiredVersion) {
    e = new StateMapEntry<>(e, stateMapVersion);
   }
   // 保存下一个要迁移的节点节点到 n
   StateMapEntry<K, N, S> n = e.next;

   // 迁移当前元素 e 到新的 table 中，插入到链表头部
   int pos = e.hash & newMask;
   e.next = newMap[pos];
   newMap[pos] = e;

   // e 指向下一个要迁移的节点
   e = n;
   // 迁移元素数 +1
   &