前言

前文砥砺前行-初学Flink的我如何快速定位并解决数据同步问题解决了数据同步的执行顺序问题后，接着去定位TIDB的问题，但是查了一阵之后发现Flink这块仍然存在问题，很明显的数据倾斜问题，十分影响整体效率。把这个问题抛给同事去研究，我继续去定位TIDB问题，但是最后没有收到满意的结果，于是我再次上手操作。本次问题也是一个思路的问题，解决起来并不复杂，涉及到观测和逆向思维，我觉得还是蛮有意思的，就拎出来写一篇文章记录一下。TIDB是下一篇，不要急啦，再下一周就OK。

问题现状

这是Flink WebUI上输出流的状态截图，我设置了30并行度，这是其中一部分subTask的截图。因为对表名KeyBy的原因，在未变更并行度的时候，每个表都会有固定的subTask消费，这保证了执行的顺序(前文有解释)，但是带来了新的问题，数据倾斜。如上图，这是几个小时的消费情况，有的subTask都消费了几十万了，有的才几十，这是很明显的数据倾斜问题，那么下一步就是思考如何解决这个问题。

Flink的KeyBy源码分析

当然咱们不能盲目修改，首先要看下KeyBy的源码，了解他的原始逻辑。我自己看源码的时候，不喜欢直接生啃，还是先找找资料比较友善，最后自己对比源码看下资料是否正确即可。Flink key 应该分配到哪个 KeyGroup 以及 KeyGroup 分配在哪个subtask，这一篇文章我感觉可以参考，看了下源码，这块和我现在使用的1.15.2的Flink源码一致，所以我就看这个理解的。还有点小庆幸，因为这个KeyBy的算法在Flink的另一个依赖里面，直接代码用IDEA一路点进去看源码可能得找好一会儿，这下省时间了，真不错啊。

简单总结一下，根据key计算其对应的subTaskIndex，即应该分配给哪个subTask运行，计算过程包括以下两步，源码都在相应的org.apache.flink.runtime.state.KeyGroupRangeAssignment类中

第一步: 根据 key 计算其对应哪个 KeyGroup

第二步: 计算 KeyGroup 属于哪个并行度

这个Flink默认算法是很简单的，逻辑就两步，会根据并行度以及算出来的最大并行度动态调整槽位。

逆向思维

我一般是不喜欢动源码的，特别是这种成熟中间件，在非必要情况下还是尽量别动。因此在观察了源码之后，我得出了一个结论，在保持并行度不变的情况下，相同的Key值总是能计算出固定的KeyGroup，而KeyGroup也能求出固定的subTaskIndex。那么反向思考，选择修改我传给他的Key，而不是费劲修改KeyBy的算法。原理就是用Flink的算法反向生成一波新Key，新Key和老Key形成对应关系，在KeyBy的时候传新Key，也算是实现了我的指定算子执行的意愿。代码修改如下，主要是重写了KEY生成逻辑

Key重生成代码逻辑如下，参考了网上一篇FlinK KeyBy分布不均匀问题的总结思考代码的思路，但是他的解释不是很详细，我在基础之上结合我的业务场景丰富了描述，再次重写了逻辑。

import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.runtime.state.KeyGroupRangeAssignment;

import java.util.HashMap;
import java.util.Map;

public class ReBalanceKeySelector<T extends DataDTO> implements KeySelector<T, Integer> {
    /**
     * 再平衡结果集。arr[i],下标i代表真实的subTaskIndex，value等于Flink算法得到的newKey
     */
    private final Integer[] reBalanceKeys;

    public ReBalanceKeySelector(int parallelism) {
        reBalanceKeys = createReBalanceKeys(parallelism);
    }

    @Override
    public Integer getKey(T value) {
        return reBalanceKeys[value.getKeyGroup()];
    }

    public static Integer[] createReBalanceKeys(int parallelism) {
        //Flink算出来的最大并行度，注意这个并不等于任务并行度，取这个值是因为Flink的KEYBY运算用的着，我们不用关心其内在意义
        int maxParallelism = KeyGroupRangeAssignment.computeDefaultMaxParallelism(parallelism);
        //记录每个oldKey对应Flink计算出来的SubTaskIndex
        Map<Integer, Integer> keySubIndexMap = new HashMap<>();
        //anyOne初始为0无特殊意义，只是代表要循环足够多的次数，保证parallelism个并行度下所有SubTaskIndex都有对应的newKey
        int anyOne = 0;
        while (keySubIndexMap.size() < parallelism) {
            int subtaskIndex = KeyGroupRangeAssignment.assignKeyToParallelOperator(anyOne, maxParallelism, parallelism);
            if (!keySubIndexMap.containsKey(subtaskIndex)) {
                keySubIndexMap.put(subtaskIndex, anyOne);
            }
            anyOne++;
        }
        //map中key值升序排列无需二次排序，直接输出值数组即可
        return keySubIndexMap.values().toArray(new Integer[0]);
    }

    /**
     * 测试代码
     */
    public static void main(String[] args) {
        int parallelism = 30;
        int maxParallelism = KeyGroupRangeAssignment.computeDefaultMaxParallelism(parallelism);
        Integer[] reBalanceKeys = createReBalanceKeys(parallelism);
        //这里的i模拟的是我在初始化表中设置的KEY_GROUP
        for (int i = 0; i < 100; i++) {
            int newKey = reBalanceKeys[i % parallelism];
            int subtaskIndex = KeyGroupRangeAssignment.assignKeyToParallelOperator(newKey, maxParallelism, parallelism);
            System.out.printf("初始化信息表中的Key:%d,Flink运算得到的Key:%d,Flink最终子任务槽位:%d \n", i, newKey, subtaskIndex);
        }
    }
}