99. Flink面试高频题：Flink两阶段提交

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction.SourceContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


import java.util.ArrayList;
import java.util.List;
import java.util.Random;


public class TwoPhaseCommitExample {


    private static final Logger LOG = LoggerFactory.getLogger(TwoPhaseCommitExample.class);


    public static void main(String[] args) throws Exception {


        final StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();


        // 开启检查点，每5000ms进行一次
        env.enableCheckpointing(5000);
        env.getCheckpointConfig().setCheckpointingMode(CheckpointingMode.EXACTLY_ONCE);


        // 数据源
        env.addSource(new SourceFunction<String>() {
            private volatile boolean running = true;
            private String[] data = new String[] {"1", "2", "3", "4", "5", "6", "7", "8", "9", "10"};


            @Override
            public void run(SourceContext<String> ctx) throws Exception {
                final Random rand = new Random(System.currentTimeMillis());
                while(running){
                    int idx = rand.nextInt(data.length);
                    ctx.collect(data[idx]);
                    Thread.sleep(1000);
                }
            }


            @Override
            public void cancel() {
                running = false;
            }
        })
        // 进行转换操作
        .keyBy(new KeySelector<String, String>() {
            @Override
            public String getKey(String value) throws Exception {
                return value;
            }
        })
        .map(new TransactionMapper())
        // 写入输出
        .addSink(new TransactionSink());


        env.execute("TwoPhaseCommitExample");
    }


    /**
     * 事务对象
     */
    public static class Transaction {
        private String id;
        private String data;
        private TransactionState state;


        public Transaction() {}


        public Transaction(String id, String data) {
            this.id = id;
            this.data = data;
            this.state = TransactionState.New;
        }


        public String getId() {
            return id;
        }


        public void setId(String id) {
            this.id = id;
        }


        public String getData() {
            return data;
        }


        public void setData(String data) {
            this.data = data;
        }


        public TransactionState getState(){
            return state;
        }


        public void setState(TransactionState state){
            this.state = state;
        }


        @Override
        public String toString() {
            return "Transaction{" +
                    "id='" + id + '\'' +
                    ", data='" + data + '\'' +
                    ", state=" + state +
                    '}';
        }
    }


    /**
     * 事务状态枚举类
     */
    public static enum TransactionState {
        New,
        Prepared,
        Committed,
        Aborted
    }


    /**
     * 事务转换操作
     */
    public static class TransactionMapper extends RichMapFunction<String, Transaction> {
        private transient ValueState<Transaction> state;


        @Override
        public void open(Configuration config) throws Exception {
            super.open(config);
            ValueStateDescriptor<Transaction> descriptor =
                    new ValueStateDescriptor<>("transaction", Transaction.class);
            state = getRuntimeContext().getState(descriptor);
        }


        @Override
        public Transaction map(String value) throws Exception {
            Transaction transaction = state.value();
            if(transaction == null){
                transaction = new Transaction(value, "");
                state.update(transaction);
            }
            transaction.setData(value);
            transaction.setState(TransactionState.Prepared);
            return transaction;
        }
    }


    /**
     * 事务输出操作
     */
    public static class TransactionSink extends RichSinkFunction<Transaction> {
        private static final long serialVersionUID = 1L;
        private static final Logger LOG = LoggerFactory.getLogger(TransactionSink.class);


        private List<Transaction> transactions = new ArrayList<>();
        private transient ValueState<Byte> state;


        @Override
        public void open(Configuration config) throws Exception {
            super.open(config);
            ValueStateDescriptor<Byte> descriptor =
                    new ValueStateDescriptor<>("state", Byte.class);
            state = getRuntimeContext().getState(descriptor);
        }


        @Override
        public void invoke(Transaction transaction) throws Exception {
            switch (transaction.getState()) {
                case Prepared:
                    transactions.add(transaction);
                    break;
                case Aborted:
                    transaction.setState(TransactionState.Aborted);
                    break;
                case Committed:
                    for (Transaction t : transactions) {
                        t.setState(TransactionState.Committed);
                    }
                    state.update((byte)0);
                    transactions.clear();
                    break;
                default:
                    break;
            }
        }
    }
}