Flink算子状态为何只能用ListState？

前言

Flink 将状态是否要按照 key 进行分类，将状态分为键值状态（Keyed State）和算子状态（Operator State）两种，两者除了状态本身的作用域不同外，其中算子状态的状态类型更是被 Flink 限制为 ListState，这是为什么呢？

使用算子状态

算子状态的作用域为当前 subTask，使用算子状态，Flink 算子的每个subTask只能访问当前subTask的数据，不能夸subTask访问。典型的应用场景就是 FlinkKafkaConsumer 使用算子状态保存 Kafka Topic 中的每个分区的消费偏移量。

在Flink中，要想使用算子状态，可以选择实现 CheckpointedFunction 接口

public interface CheckpointedFunction {
    void snapshotState(FunctionSnapshotContext var1) throws Exception;

    void initializeState(FunctionInitializationContext var1) throws Exception;
}

snapshotState Flink作业执行快照时调用该方法，开发者可以控制往ListState写入哪些数据
initializeState Flink作业启动或者异常容错从快照恢复时调用这个方法

Flink作业启动或异常恢复时会调用 CheckpointedFunction#initializeState，通过入参 FunctionInitializationContext 来获取算子状态。

要想获取算子状态，首先得先定义状态描述符，因为算子状态被强制限定为列表状态，所以只能用 ListStateDescriptor。然后通过入参 FunctionInitializationContext#getOperatorStateStore 对象来获取 ListState。

@Override
public void initializeState(FunctionInitializationContext functionInitializationContext) throws Exception {
    this.elementsState = functionInitializationContext.getOperatorStateStore().getListState(
            new ListStateDescriptor<>("elements", Integer.class)
    );
}

算子状态实战

算子状态在业务场景中并不常用，除了 FlinkKafkaConsumer 使用算子状态保存 Kafka Topic 中分区的消费偏移量外，Sink 算子使用算子状态作为写出数据的缓冲区也是一个较为常用的场景。

MySQL 是常用的关系型数据库，在流计算场景中，它也是一种常用的数据汇存储引擎，用来保存流计算的结果。但是MySQL的写入TPS通常不高，一般在几百甚至几千，上万已经是很夸张了。但是Flink作为一款高性能的流计算引擎，动辄十万百万的TPS数据流入，如果计算结果每次都写入MySQL，势必会压垮MySQL。此时可以在 Sink 算子上使用算子状态作为缓冲区，先缓存一部分数据，最后再一次性批量写MySQL，以此来减轻MySQL的压力。

举个例子，现在有一个数据源，会源源不断的产生一批数字，现在要开发一个 Flink 作业，计算这些数字的和，然后把结果写入到 MySQL，为了减轻MySQL的写入压力，要求 Sink 算子可以缓冲一部分数据再批量写。

如下代码所示，SumResultBufferingSink 实现了CheckpointedFunction 接口，元素到达时会先写入缓冲区 elements，缓冲区满才会累计求和后写入MySQL。同时，在执行快照时，也会把elements缓冲区的数据写入到elementsState，异常恢复时，再将elementsState数据恢复到缓冲区。

public class OperatorStateFuature {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
        environment.fromElements(1, 2, 3, 4, 5, 6)
                .keyBy(i -> "all")
                .sum(0)
                .addSink(new SumResultBufferingSink(3));
        environment.execute();
    }

    public static class SumResultBufferingSink implements SinkFunction<Integer>, CheckpointedFunction {

        private final int bufferSize;
        private final List<Integer> elements;
        private ListState<Integer> elementsState;

        public SumResultBufferingSink(int bufferSize) {
            this.bufferSize = bufferSize;
            this.elements = new ArrayList<>(bufferSize);
        }

        @Override
        public void invoke(Integer value, Context context) throws Exception {
            elements.add(value);
            if (elements.size() >= bufferSize) {
                int sum = elements.stream().mapToInt(Integer::intValue).sum();
                System.err.println("write to db : sum=" + sum);
                elements.clear();
            }
        }

        @Override
        public void snapshotState(FunctionSnapshotContext functionSnapshotContext) throws Exception {
            System.err.println("---snapshotState start---");
            elementsState.clear();
            elementsState.addAll(elements);
            System.err.println("---snapshotState end---");
        }

        @Override
        public void initializeState(FunctionInitializationContext functionInitializationContext) throws Exception {
            System.err.println("---initializeState start---");
            this.elementsState = functionInitializationContext.getOperatorStateStore().getListState(
                    new ListStateDescriptor<>("elements", Integer.class)
            );
            // 是否从故障中恢复
            if (functionInitializationContext.isRestored()) {
                Iterator<Integer> iterator = elementsState.get().iterator();
                while (iterator.hasNext()) {
                    elements.add(iterator.next());
                }
            }
            System.err.println("---initializeState end---");
        }
    }
}

使用 SumResultBufferingSink 后，缓冲区大小为3，六个元素只会写两次DB。

ListState和UnionListState

OperatorStateStore 提供了两个方法获取 ListState

public interface OperatorStateStore {
  <S> ListState<S> getListState(ListStateDescriptor<S> stateDescriptor) throws Exception;
  <S> ListState<S> getUnionListState(ListStateDescriptor<S> stateDescriptor) throws Exception;
}

ListState和UnionListState 有什么区别呢？

两者的区别在于，快照恢复或者算子并行度发生改变时，算子状态值的分配方式是不同的。

ListState 采用平均分割分配，状态重新分配时，所有subTask的ListState会先合并到一起，再采用 Round-Robin 策略将列表中的状态分配到各个subTask
UnionListState 采用合并分配，状态重新分配时，所有subTask的ListState合并到一起得到一个完整的列表，再将这个完整的列表发给每个subTask。

Tips：UnionListState要慎用，当列表中的元素非常多时，有内存溢出的风险。

算子状态为什么限制ListState

回到开篇提出的问题，为什么Flink要限制算子状态只能使用 ListState 类型？

本质上，是Flink异常恢复，或者算子并行度发生变化时，算子状态数据如何分配的问题。最简单公平的分配算法就是平均分配，那么除了 ListState 这种列表类型，其它如 ValueState，MapState 等数据结构实在是不方便数据划分啊，所以Flink才限制算子状态必须是 ListState 类型。当然，Flink 也给了开发者两种选择，一是 ListState 的平均分配，二是 UnionListState 给你全量的状态，程序自己来分配，更加灵活。