Flink MySQL两阶段提交示例代码

最近在工作中使用到Flink，正好也涉及到往MySQL进行最终的数据写入。且程序的模式是精准一次，故要配合checkpoint进行两阶段提交。

在网上找了大部分的案例代码，但是发现部署后或多或少都会有一些问题。

在浏览了大部分的案例代码后，目前本人在工作中的两阶段提交示例代码如下，这里仅做案例展示使用，具体的业务逻辑需要各位替换成自己的。

Flink使用的版本是1.16.1

业务主类

public class CanbusStreamApp {
    public static void main(String[] args) throws Exception {   
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();        
        // 一分钟一次checkpoint，模式为精准一次
        env.enableCheckpointing(60000, CheckpointingMode.EXACTLY_ONCE);
        // 状态后端相关设置，HashMapStateBackend表示程序运行时状态存储在内存中
        env.setStateBackend(new HashMapStateBackend());
        CheckpointConfig checkpointConfig = env.getCheckpointConfig();
        // 程序做checkpoint时，状态落地到文件系统
        checkpointConfig.setCheckpointStorage(new FileSystemCheckpointStorage(checkpointPath));
        // 程序取消时，checkpoint是否保留，默认是保留
        checkpointConfig.setExternalizedCheckpointCleanup(CheckpointConfig.ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION);
        // 程序失败重启策略
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(3, Time.seconds(10)));

        // 中间业务过程省略。。。
        // 将结果数据写入mysql
        resultData.addSink(new YRCanbusTwoPhaseCommitSinkFunction()).name("MySQL Sink");
        env.execute("YRCanbus Stream");
    }
}

两阶段提交类

package cn.gwm.dp.function.canbus;

import cn.gwm.dp.entity.canbus.YRcanbusStore;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.typeutils.base.VoidSerializer;
import org.apache.flink.api.java.typeutils.runtime.kryo.KryoSerializer;
import org.apache.flink.streaming.api.functions.sink.TwoPhaseCommitSinkFunction;

import java.sql.SQLException;

/**
 * @Author: Spring
 * @Description:
 * @Date: Created on 13:50 2024/5/4
 */
public class YRCanbusTwoPhaseCommitSinkFunction extends TwoPhaseCommitSinkFunction<YRcanbusStore, YRCanbusTransaction, Void> {

    public YRCanbusTwoPhaseCommitSinkFunction() {
        super(new KryoSerializer<>(YRCanbusTransaction.class, new ExecutionConfig()), VoidSerializer.INSTANCE);
    }

    @Override
    protected void invoke(YRCanbusTransaction yrCanbusTransaction, YRcanbusStore yRcanbusStore, Context context) throws Exception {
        yrCanbusTransaction.store(yRcanbusStore);
    }

    @Override
    protected YRCanbusTransaction beginTransaction() throws Exception {
        return new YRCanbusTransaction();
    }

    @Override
    protected void preCommit(YRCanbusTransaction yrCanbusTransaction) throws Exception {

    }

    @Override
    protected void commit(YRCanbusTransaction yrCanbusTransaction) {
        try {
            yrCanbusTransaction.commit();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    protected void abort(YRCanbusTransaction yrCanbusTransaction) {
        try {
            yrCanbusTransaction.rollback();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }
}

其中，YRCanbusTransaction是实现的具体的事务类，YRcanbusStore是上游输出到当前的数据，它是一个自定义的实体类。

具体事务类

package cn.gwm.dp.function.canbus;

import cn.gwm.dp.constant.Constant;
import cn.gwm.dp.entity.canbus.YRcanbusStore;
import cn.gwm.dp.util.DruidUtil;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;

/**
 * 两阶段提交事务
 *
 * @author Spring
 * @date 2024/5/25 8:52
 */
public class YRCanbusTransaction {
    private transient Connection connection;
    private List<YRcanbusStore> list = new ArrayList<>();

    public void store(YRcanbusStore yrcanbusStore) {
        list.add(yrcanbusStore);
    }

    public void commit() throws SQLException {
        connection = DruidUtil.getConnection();
        connection.setAutoCommit(false);
        for (YRcanbusStore yrcanbusStore : list) {
            String platform = yrcanbusStore.getPlatform();
            String type = yrcanbusStore.getType();
            String day = yrcanbusStore.getDay();
            String hour = yrcanbusStore.getHour();
            String sql = Constant.UPDATE_LATEST_DATA;
            PreparedStatement preparedStatement = connection.prepareStatement(sql);
            preparedStatement.setString(1, platform);
            preparedStatement.setString(2, type);
            preparedStatement.setString(3, day);
            preparedStatement.setString(4, hour);
            preparedStatement.execute();
        }
        connection.commit();
        connection.close();
    }

    public void rollback() throws SQLException {
        connection.rollback();
        connection.close();
    }
}

commit方法中，请将具体的写入MySQL的逻辑替换为你自己的业务逻辑。我的代码中，列表中存储的只有一条数据，所以没有做批量提交，如果你同时写入的数据有很多，请自行修改为批量提交的语句。

对于MySQL连接的获取，我使用的是连接池，获取连接的方法封装到了DruidUtil中。关于连接池，网上的资料有很多，这里我就不赘述了。

以上就是两阶段提交的主体代码，目前测试通过，运行暂未发现异常。各位可根据自己的业务需要自行修改使用。如果有任何使用问题，也欢迎留言评论交流。