Flink堆状态后端核心：CopyOnWriteStateMap解析

CopyOnWriteStateMap

CopyOnWriteStateMap<K, N, S> 是 Flink 堆状态后端中 StateMap 的一个核心实现，它通过写时复制 (Copy-on-Write, COW) 机制来支持高效的异步快照和增量式哈希重组 (incremental rehashing)。它在性能和内存效率之间做了一些权衡，以换取这些高级特性。

下面我们来分析其实现方式和关键方法：

核心设计与数据结构

1. 继承 StateMap<K, N, S> : 实现了 StateMap 定义的接口，用于存储键(K)-命名空间(N)到状态值(S)的映射。

2. 哈希表结构:

   • 类似于 java.util.HashMap，它内部使用一个哈希表数组 （StateMapEntry<K, N, S>[]）来存储条目。哈希冲突通过链地址法解决，即每个数组桶位存储一个链表的头节点。
   • StateMapEntry<K, N, S>: 这是哈希表中的基本单元，包含了键、命名空间、状态值、指向链表下一个节点的指针 (next)，以及用于COW和版本控制的元数据。

3. 写时复制 (Copy-on-Write) 的核心字段:

   • private StateMapEntry<K, N, S>[] primaryTable;: 主要的哈希表数组。
   • private StateMapEntry<K, N, S>[] incrementalRehashTable;: 在进行增量哈希重组时使用的辅助哈希表。
   • private int stateMapVersion;: 当前 StateMap 的版本号。每次快照时会递增。
   • private int highestRequiredSnapshotVersion;: 当前所有未释放的快照中，最高的版本号。这个版本号用于判断一个条目或其状态是否需要被复制。
   • private final TreeSet<Integer> snapshotVersions;: 维护一个有序集合，存储所有当前正在使用（即尚未释放）的快照的版本号。当一个快照被释放后，其版本号会从这个集合中移除，并可能更新 highestRequiredSnapshotVersion。
   • StateMapEntry.entryVersion: 每个 StateMapEntry 实例的版本号。表示该条目结构（即其在链表中的位置或自身）最后一次被修改时的 stateMapVersion。
   • StateMapEntry.stateVersion: 每个 StateMapEntry 中状态值 (state) 的版本号。表示该状态值最后一次被修改时的 stateMapVersion。

4. 增量哈希重组 (Incremental Rehashing):

   • 当 StateMap 中的条目数量超过阈值 (threshold) 时，会触发哈希表的扩容（通常是容量翻倍）。
   • 为了避免在扩容时产生长时间的停顿（STW），CopyOnWriteStateMap 采用了增量哈希重组。
   • primaryTable 成为旧表，incrementalRehashTable 成为新表（扩容后的）。
   • 数据迁移不是一次性完成的，而是在后续的读写操作中逐步 将条目从 primaryTable 迁移到 incrementalRehashTable。
   • private int rehashIndex;: 指示在 primaryTable 中下一个需要进行迁移的桶位索引。
   • private int primaryTableSize; 和 private int incrementalRehashTableSize;: 分别记录两个表中的条目数量。

写时复制 (Copy-on-Write) 逻辑:

• 快照创建 (snapshotMapArrays() 和 snapshot()) :
  • 当需要创建快照时，stateMapVersion 会递增。
  • 新的 stateMapVersion 会被添加到 snapshotVersions 集合中。
  • highestRequiredSnapshotVersion 会根据 snapshotVersions 中的最大值更新。
  • 快照操作本身并不会立即复制所有数据，而是"记住"了当前 stateMapVersion。
• 修改操作 (如 put, remove, transform) :
  • 当修改一个 StateMapEntry 时（例如，更新其状态值或改变其在哈希链中的 next 指针）：
    • 检查条目版本 (e.entryVersion < highestRequiredSnapshotVersion) : 如果条目的版本低于当前最高要求的快照版本，意味着这个条目结构（或其在链中的前驱节点）可能被某个正在进行的快照所引用。此时，需要复制这个条目（或其前驱节点）以创建一个新版本的条目，并将修改应用到新条目上。旧版本的条目保持不变，供快照使用。这个过程通过 handleChainedEntryCopyOnWrite() 实现。
    • 检查状态版本 (e.stateVersion < highestRequiredSnapshotVersion) : 如果条目内的状态值的版本低于当前最高要求的快照版本，意味着这个状态值可能被某个快照引用。在修改状态值之前，需要使用 stateSerializer.copy(e.state) 复制一份状态值，然后对副本进行修改。
  • 修改完成后，被修改的条目（或其状态）的版本号会更新为当前的 stateMapVersion。
• 读取操作 (get) :
  • 当读取一个状态值时，同样会检查其 stateVersion。如果低于 highestRequiredSnapshotVersion，并且该条目本身也需要COW（e.entryVersion < highestRequiredSnapshotVersion），则会先处理条目的COW，然后复制状态值返回给用户。如果仅仅是状态值版本低，则只复制状态值。这是为了确保用户获得的是一个"安全"的副本，对它的修改不会影响到正在进行的快照。

增量哈希重组

• doubleCapacity(): 当需要扩容时，会初始化 incrementalRehashTable，并将 primaryTable 设为旧表。
• computeHashForOperationAndDoIncrementalRehash(key, namespace): 在每次读写操作计算哈希值之前，会调用 incrementalRehash()。
• incrementalRehash(): 每次被调用时，会从 primaryTable 的 rehashIndex 位置开始，迁移一小部分条目（MIN_TRANSFERRED_PER_INCREMENTAL_REHASH 个）到 incrementalRehashTable 中。
  • 迁移过程中，如果被迁移的条目 e 的 entryVersion 低于 highestRequiredSnapshotVersion，则需要创建一个新的 StateMapEntry 放入新表，旧表中的条目保持不变（COW）。
• selectActiveTable(hash): 根据当前是否正在进行哈希重组以及给定哈希值，决定操作应该在 primaryTable 还是 incrementalRehashTable 上进行。如果一个条目已经被迁移到新表，后续操作就会在新表上进行。
• 当所有条目都从 primaryTable 迁移到 incrementalRehashTable 后，incrementalRehashTable 会成为新的 primaryTable，旧的 primaryTable 会被清空，增量哈希重组完成。

版本管理与快照释放 (releaseSnapshot(int snapshotVersion)):

• 当一个异步快照完成后，外部会调用 releaseSnapshot(snapshotVersion)。
• 该方法会将指定的 snapshotVersion 从 snapshotVersions 集合中移除。
• 然后，它会重新计算 highestRequiredSnapshotVersion（通常是 snapshotVersions 中剩余的最大值，如果没有则为0）。
• 降低 highestRequiredSnapshotVersion 可以使得后续的写操作不再需要为那些已经完成的旧快照版本进行复制，从而减少不必要的对象创建和内存开销。

关键方法

1. 构造函数 CopyOnWriteStateMap(int capacity, TypeSerializer<S> stateSerializer):

   • 初始化哈希表（primaryTable, incrementalRehashTable），版本信息（stateMapVersion, highestRequiredSnapshotVersion, snapshotVersions），以及其他控制参数。

2. 核心读写 API (实现 StateMap):

   • get(K key, N namespace): 获取状态值。包含COW逻辑。
   • put(K key, N namespace, S value): 设置状态值。包含COW逻辑。
   • putAndGetOld(K key, N namespace, S state): 设置新状态并返回旧状态。包含COW逻辑。
   • remove(K key, N namespace): 移除条目。包含COW逻辑。
   • removeAndGetOld(K key, N namespace): 移除条目并返回其状态。包含COW逻辑。
   • transform(K key, N namespace, T value, StateTransformationFunction<S, T> transformation): 原子性地转换状态。包含COW逻辑。

3. COW 核心逻辑:

   • private StateMapEntry<K, N, S> putEntry(K key, N namespace): put 和 transform 操作的基础，处理条目查找和COW。
   • private StateMapEntry<K, N, S> removeEntry(K key, N namespace): remove 操作的基础，处理条目查找和COW。
   • private StateMapEntry<K, N, S> handleChainedEntryCopyOnWrite(StateMapEntry<K, N, S>[] tab, int mapIdx, StateMapEntry<K, N, S> untilEntry): 处理哈希链中条目的写时复制。当链中的某个节点需要修改（例如 next 指针），并且该节点或其前驱节点的版本低于 highestRequiredSnapshotVersion 时，会从链头开始复制节点，直到目标节点，确保旧链结构不被破坏。

4. 增量哈希重组相关:

   • private void doubleCapacity(): 启动扩容和增量哈希重组过程。
   • private void incrementalRehash(): 执行一小步哈希重组，迁移部分数据。
   • private int computeHashForOperationAndDoIncrementalRehash(K key, N namespace): 在常规操作前触发一小步哈希重组。
   • private StateMapEntry<K, N, S>[] selectActiveTable(int hash): 根据哈希值和重组状态选择操作的目标表。

5. 快照与版本管理:

   • StateMapSnapshot<K, N, S> snapshot(): 创建当前 StateMap 的快照视图。内部会调用 snapshotMapArrays()。
   • StateMapEntry<K, N, S>[] snapshotMapArrays(): 核心快照逻辑，递增 stateMapVersion，注册快照版本，并返回一个包含当前状态数据的数组（可能是合并了 primaryTable 和 incrementalRehashTable 的部分数据）。
   • void releaseSnapshot(int snapshotVersion): 释放一个已完成的快照版本，更新 highestRequiredSnapshotVersion。

6. 迭代器与访问者:

   • iterator(): 返回一个迭代器，可以遍历 StateMap 中的所有条目。迭代器需要能正确处理增量哈希重组期间两个表的情况。
   • getStateIncrementalVisitor(...): 返回一个 StateIncrementalVisitor，用于增量快照，它能够高效地访问已更改或所有条目。其实现是 StateIncrementalVisitorImpl，内部使用 StateEntryChainIterator。

7. 内部辅助:

   • private static int compositeHash(Object key, Object namespace): 计算键和命名空间的组合哈希值。
   • private StateMapEntry<K, N, S>[] makeTable(int newCapacity): 创建一个新的哈希表数组。

通过这些机制，CopyOnWriteStateMap 能够在不阻塞正常数据处理的情况下进行异步快照，并且通过增量哈希重组避免了因哈希表扩容导致的长时间暂停，这对于流处理系统的低延迟和高吞吐至关重要。代价是可能需要复制更多的对象，以及更复杂的版本管理逻辑。

并发安全

CopyOnWriteStateMap (作为 StateMap 的一个实现) 的并发安全设计是针对 Flink 特定的执行模型和需求的，而不是一个通用的线程安全哈希映射（如 java.util.concurrent.ConcurrentHashMap）。

结合类注释和代码上下文，可以分析如下：

1. 写时复制 (Copy-on-Write) 的核心目的：

   • 注释明确指出 CopyOnWriteStateMap "sacrifices some peak performance and memory efficiency for features like incremental rehashing and asynchronous snapshots through copy-on-write."
   • 这意味着其主要并发处理机制是为了支持异步快照。当进行快照时，如果状态被修改，会创建数据的副本，从而保证快照读取的是一个一致性的版本，而在线处理可以继续进行。

2. 用户代码的执行上下文：

   • 注释中非常重要的一点是："IMPORTANT: the contracts for this class rely on the user not holding any references to objects returned by this map beyond the life cycle of per-element operations. Or phrased differently, all get-update-put operations on a mapping should be within one call of processElement."
   • 这强烈暗示了对于同一个键（key）的常规操作（get, put, update）通常期望在单个线程 的上下文中完成，即在 Flink 的 processElement 或类似的用户函数调用范围内。Flink 的 keyed state 本身保证了具有相同 key 的数据会被同一个 task (线程) 处理。

3. 并发修改检测：

   • 类中有一个 modCount 字段，注释为："Incremented by "structural modifications" to allow (best effort) detection of concurrent modification."
   • 这是 java.util.HashMap 等非线程安全集合中常见的机制，用于在迭代期间检测并发修改并抛出 ConcurrentModificationException。这进一步说明了该类不期望来自多个用户线程的对同一 StateMap 实例的任意并发修改。

4. 快照相关的同步：

   • snapshotVersions 字段是一个 TreeSet，用于维护快照版本。
   • 在 releaseSnapshot(int snapshotVersion) 和 snapshotMapArrays() 方法中，对 snapshotVersions 的访问被 synchronized (snapshotVersions) 代码块保护。这确保了快照版本管理的线程安全性。
   • snapshotMapArrays() 方法的注释提到："This method must be called by the same Thread that does modifications to the {@link CopyOnWriteStateMap}." 这说明快照的启动是由执行状态修改的同一个线程（通常是 Flink 的主任务线程）发起的。写时复制机制随后允许实际的快照数据读取（可能在另一个线程中进行，例如用于序列化和写入外部存储）与新的状态修改并发进行。

CopyOnWriteStateMap 不是为通用的多线程并发修改而设计的。它的并发控制主要集中在以下方面：

• 针对 Flink 执行模型的优化 ：Flink 的 keyed stream 处理模型天然地将对同一 key 的操作分配到单个线程执行。因此，对于单个 key 的 get/put/remove 操作，CopyOnWriteStateMap 并不需要像 ConcurrentHashMap 那样复杂的内部锁机制来处理来自多个用户线程的并发访问。
• 支持异步快照：这是其核心的并发特性。通过写时复制，它允许在创建快照（读取状态）的同时，主处理线程可以继续修改状态，而不会阻塞或产生数据不一致。

如果用户试图在 Flink 的正常处理流程之外，从多个自定义线程并发地修改同一个 CopyOnWriteStateMap 实例中的相同条目，那么很可能会遇到并发问题。注释中的警告（用户不应在用户函数调用范围之外持有状态对象的引用，并且所有 get-update-put 操作应在单个 processElement 调用内）正是为了防止这类不当使用。

深入理解快照机制

一个简单的比喻:

想象你在编辑一份文档（CopyOnWriteStateMap）。

1. 创建快照 : 你想保存当前版本（比如 V1）。你按下了"另存为 V1快照"的按钮。
   • 系统内部记录："V1快照需要的是文档在按下按钮这一刻的内容"。
   • 同时，系统标记文档当前活跃的最高快照需求是 V1。
2. 后续编辑 : 你继续编辑文档。
   • 当你修改文档中某个段落时，如果这个段落是 V1 快照创建时就存在的，系统会悄悄地把这个段落复制一份，你的修改会发生在这个副本上。V1 快照仍然指向那个未修改的原始段落。
   • 如果你添加了全新的段落，它自然不属于 V1 快照。
3. 读取快照: 当外部程序要读取 V1 快照时，它看到的就是你按下"另存为 V1快照"那一刻的完整内容，不受你后续编辑的影响。
4. 释放快照: 当 V1 快照不再需要时，你告诉系统可以删除它了。系统会更新记录，可能就不再需要特意保护 V1 版本的数据了（除非还有其他快照也需要它）。

highestRequiredSnapshotVersion 就像一个"保护线"，版本低于这条线的数据在被修改前都需要先复制一份，解除对原来数据的引用，因为旧版本会引用这些条目，因此内存不会被释放。

highestRequiredSnapshotVersion 字段分析

highestRequiredSnapshotVersion 字段代表的是当前所有尚未被释放的快照所依赖的最高的 StateMap 版本号。

在 CopyOnWriteStateMap 中，为了支持高效的异步快照和增量重哈希 (incremental rehashing)，采用了写时复制 (copy-on-write) 的策略。这意味着当状态或映射结构需要被修改时，如果这部分数据仍然被某个活动的（未释放的）快照所引用，那么在修改之前会先复制一份。

• stateMapVersion：这是 CopyOnWriteStateMap 自身的当前版本号。每次创建快照时，这个版本号会递增。
• snapshotVersions：这是一个 TreeSet，用于存储所有当前活动的、尚未释放的快照所对应的 stateMapVersion。
• highestRequiredSnapshotVersion：这个字段的值等于 snapshotVersions 集合中最大的版本号（如果集合不为空的话）。它的核心作用是作为一个阈值，用于判断某个 StateMapEntry（状态条目）或其内部状态 (state) 是否需要进行写时复制。
  • 如果一个条目的版本 (entryVersion) 或其状态的版本 (stateVersion) 小于 highestRequiredSnapshotVersion，就意味着这个条目或状态的当前副本仍然被至少一个快照所需要。因此，在对其进行修改（如更新、删除）之前，必须先创建一个新的副本，并将新副本的版本更新为当前的 stateMapVersion。旧的副本则保留给那些较旧的快照使用。
  • 如果版本号不小于 highestRequiredSnapshotVersion，则可以直接修改，因为没有活动的快照依赖这个特定版本的数据。

简单来说，highestRequiredSnapshotVersion 帮助 CopyOnWriteStateMap 决定何时需要为旧快照保留数据的旧版本，何时可以直接修改数据。

更新时机

highestRequiredSnapshotVersion 字段主要在以下几种情况下会被更新：

1. 初始化时 ： 在 CopyOnWriteStateMap 的构造函数中，它被初始化为 0。

   // ... existing code ...
   @SuppressWarnings("unchecked")
   private CopyOnWriteStateMap(int capacity, TypeSerializer<S> stateSerializer) {
       this.stateSerializer = Preconditions.checkNotNull(stateSerializer);
   
       // initialized maps to EMPTY_TABLE.
       this.primaryTable = (StateMapEntry<K, N, S>[]) EMPTY_TABLE;
       this.incrementalRehashTable = (StateMapEntry<K, N, S>[]) EMPTY_TABLE;
   
       // initialize sizes to 0.
       this.primaryTableSize = 0;
       this.incrementalRehashTableSize = 0;
   
       this.rehashIndex = 0;
       this.stateMapVersion = 0;
       this.highestRequiredSnapshotVersion = 0; // 初始化
       this.snapshotVersions = new TreeSet<>();
   
       if (capacity < 0) {
           throw new IllegalArgumentException("Capacity: " + capacity);
       }
   
       if (capacity == 0) {
           threshold = -1;
           return;
       }
   
       if (capacity < MINIMUM_CAPACITY) {
           capacity = MINIMUM_CAPACITY;
       } else if (capacity > MAXIMUM_CAPACITY) {
           capacity = MAXIMUM_CAPACITY;
       } else {
           capacity = MathUtils.roundUpToPowerOfTwo(capacity);
       }
       primaryTable = makeTable(capacity);
   }
   // ... existing code ...

2. 创建快照时 (snapshotMapArrays() 方法) ： 当系统需要为当前的 CopyOnWriteStateMap 创建一个快照时（通常在 checkpoint 过程中），会调用 snapshotMapArrays() 方法。在此方法内部：

   • stateMapVersion 会首先自增，代表 StateMap 进入了一个新的版本。
   • 然后，highestRequiredSnapshotVersion 会被更新为这个新的 stateMapVersion。
   • 同时，这个新的 stateMapVersion 也会被添加到 snapshotVersions 集合中，表示有一个新的活动快照依赖于这个版本。 这个操作在 synchronized (snapshotVersions) 块中完成，以确保线程安全。

   // ... existing code ...
   @VisibleForTesting
   @SuppressWarnings("unchecked")
   StateMapEntry<K, N, S>[] snapshotMapArrays() {
   
       // we guard against concurrent modifications of highestRequiredSnapshotVersion between
       // snapshot and release.
       // Only stale reads of from the result of #releaseSnapshot calls are ok. This is why we must
       // call this method
       // from the same thread that does all the modifications to the map.
       synchronized (snapshotVersions) {
   
           // increase the map version for copy-on-write and register the snapshot
           if (++stateMapVersion < 0) {
               // this is just a safety net against overflows, but should never happen in practice
               // (i.e., only after 2^31 snapshots)
               throw new IllegalStateException(
                       "Version count overflow in CopyOnWriteStateMap. Enforcing restart.");
           }
   
           highestRequiredSnapshotVersion = stateMapVersion; // 更新为新的 stateMapVersion
           snapshotVersions.add(highestRequiredSnapshotVersion);
       }
   
       StateMapEntry<K, N, S>[] table = primaryTable;
   // ... existing code ...

3. 释放快照时 (releaseSnapshot(int snapshotVersion) 方法) ： 当一个快照不再需要（例如，checkpoint 完成并被后续的 checkpoint 取代后，旧的 checkpoint 相关资源可以被清理）时，会调用 releaseSnapshot(int snapshotVersion) 方法（或者其公共包装方法 releaseSnapshot(StateMapSnapshot<K, N, S, ? extends StateMap<K, N, S>> snapshotToRelease)）。在此方法内部：

   • 指定的 snapshotVersion 会从 snapshotVersions 集合中移除。
   • 然后，highestRequiredSnapshotVersion 会被更新为 snapshotVersions 集合中当前剩余的最大版本号。
   • 如果 snapshotVersions 集合变为空（即没有活动的快照了），highestRequiredSnapshotVersion 会被重置为 0。 这个操作同样在 synchronized (snapshotVersions) 块中完成。

   // ... existing code ...
   @VisibleForTesting
   void releaseSnapshot(int snapshotVersion) {
       // we guard against concurrent modifications of highestRequiredSnapshotVersion between
       // snapshot and release.
       // Only stale reads of from the result of #releaseSnapshot calls are ok.
       synchronized (snapshotVersions) {
           Preconditions.checkState(
                   snapshotVersions.remove(snapshotVersion),
                   "Attempt to release unknown snapshot version");
           highestRequiredSnapshotVersion = // 更新为 snapshotVersions 中最大的，或者 0
                   snapshotVersions.isEmpty() ? 0 : snapshotVersions.last();
       }
   }
   // ... existing code ...

通过这种方式，highestRequiredSnapshotVersion 动态地反映了为了维护所有活动快照的一致性视图，数据修改操作需要考虑到的最旧版本界限。

CopyOnWriteSkipListStateMap

使用跳表实现的，参考：深入解析ConcurrentSkipListMap源码-CSDN博客

不过这里是用版本判断，和CopyOnWriteStateMap类似，一个旧版本，就用新的替换。