Flink+Kafka 实现原理

// TwoPhaseCommitSinkFunction
public void snapshotState(FunctionSnapshotContext context) throws Exception {
    ...
    preCommit(currentTransactionHolder.handle);
    ...
    currentTransactionHolder = beginTransactionInternal();
    ...
}
// beginTransaction
private TransactionHolder<TXN> beginTransactionInternal() throws Exception {
    return new TransactionHolder<>(beginTransaction(), clock.millis());
}
// notifyCheckpointComplete
public final void notifyCheckpointComplete(long checkpointId){
    commit(pendingTransaction.handle);
}


// FlinkKafkaProducer
public class FlinkKafkaProducer<IN>
	extends TwoPhaseCommitSinkFunction{
        ...
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            super.snapshotState(context);
            ...
            nextTransactionalIdHintState.add(new FlinkKafkaProducer.NextTransactionalIdHint(
                getRuntimeContext().getNumberOfParallelSubtasks(),
                nextFreeTransactionalId));
		}
         // preCommit
        protected void preCommit(FlinkKafkaProducer.KafkaTransactionState transaction) {
			...
            // transaction.producer.flush();
            flush(transaction); 
        }
        // beginTransaction
        protected FlinkKafkaProducer.KafkaTransactionState beginTransaction() {
            				FlinkKafkaInternalProducer<byte[], byte[]> producer = createTransactionalProducer();
				producer.beginTransaction();
				return new FlinkKafkaProducer.KafkaTransactionState(producer.getTransactionalId(), producer);
        }
         // invoke 
        public void invoke(FlinkKafkaProducer.KafkaTransactionState transaction, IN next, Context context){
            record = new ProducerRecord<>(
                targetTopic,
                flinkKafkaPartitioner.partition(next, serializedKey, serializedValue, targetTopic, partitions),
                timestamp,
                serializedKey,
                serializedValue);
            ...
            transaction.producer.send(record, callback);
        }
        // commit
        protected void commit(FlinkKafkaProducer.KafkaTransactionState transaction) {
            // kafkaProducer.commitTransaction();
            transaction.producer.commitTransaction();
        }
        // abort
        protected void abort(FlinkKafkaProducer.KafkaTransactionState transaction) {
            if (transaction.isTransactional()) {
                transaction.producer.abortTransaction();
                recycleTransactionalProducer(transaction.producer);
            }
	}
    }

// FlinkKafkaProducer
public void close() throws FlinkKafkaException {
    currentTransaction.producer.close();
}

// FlinkKafkaConsumerBase
// 运行时针对单个  Operator 实例(并行度) 保存的offset
private transient volatile TreeMap<KafkaTopicPartition, Long> restoredState;

/** Accessor for state in the operator state backend. */
// offset state checkpoint 时状态，ListState(OperateState) 
// 一个Flink 并行度可能消费多个 Kafka 分区 (Spark 中任务并行度和Kafka 分区一致)
private transient ListState<Tuple2<KafkaTopicPartition, Long>> unionOffsetStates;

// FlinkKafkaConsumerBase
// ------------------------------------------------------------------------
//  Checkpoint and restore
// ------------------------------------------------------------------------

@Override
public final void initializeState(FunctionInitializationContext context){
    // 获取 state
    		OperatorStateStore stateStore = context.getOperatorStateStore();
    // subTask 获取全量 state，包含所有分区的 offset
		this.unionOffsetStates = stateStore.getUnionListState(new ListStateDescriptor<>(
			OFFSETS_STATE_NAME,
			TypeInformation.of(new TypeHint<Tuple2<KafkaTopicPartition, Long>>() {})));
    	// populate actual holder for restored state
		for (Tuple2<KafkaTopicPartition, Long> kafkaOffset : unionOffsetStates.get()) {
				restoredState.put(kafkaOffset.f0, kafkaOffset.f1);
			}
}

// FlinkKafkaConsumerBase
public void open(Configuration configuration) {
    ...
final List<KafkaTopicPartition> allPartitions = partitionDiscoverer.discoverPartitions();
    ...
    for (Map.Entry<KafkaTopicPartition, Long> restoredStateEntry : restoredState.entrySet()) {
        		// restoredFromOldState 指从flink 1.2 恢复，当前版本1.9，永远为 false
				if (!restoredFromOldState) {
				// seed the partition discoverer with the union state while filtering out
				// restored partitions that should not be subscribed by this subtask
                    // 分配分区
					if (KafkaTopicPartitionAssigner.assign(
						restoredStateEntry.getKey(), getRuntimeContext().getNumberOfParallelSubtasks())
							== getRuntimeContext().getIndexOfThisSubtask()){
                    // 恢复当前subTask 负责分区的 offset 
					subscribedPartitionsToStartOffsets.put(restoredStateEntry.getKey(), 					restoredStateEntry.getValue());
					}
				} else {// when restoring from older 1.1 / 1.2 state, the restored state would not be the union state;
				// in this case, just use the restored state as the subscribed partitions
				subscribedPartitionsToStartOffsets.put(restoredStateEntry.getKey(), restoredStateEntry.getValue());
				}
			}
    ...
}

// KafkaTopicPartitionAssigner
/* 
* Utility for assigning Kafka partitions to consumer subtasks.
* @param partition the Kafka partition
* @param numParallelSubtasks total number of parallel subtasks
*
* @return index of the target subtask that the Kafka partition should be assigned to.
*/
public static int assign(KafkaTopicPartition partition, int numParallelSubtasks) {
    int startIndex = ((partition.getTopic().hashCode() * 31) & 0x7FFFFFFF) % numParallelSubtasks;
    // here, the assumption is that the id of Kafka partitions are always ascending
    // starting from 0, and therefore can be used directly as the offset clockwise from the start index
    return (startIndex + partition.getPartition()) % numParallelSubtasks;
}

public final void initializeState(FunctionInitializationContext context){
    // subTask 获取全量 state，包含所有分区的 offset
		this.unionOffsetStates = stateStore.getUnionListState(new ListStateDescriptor<>(
			OFFSETS_STATE_NAME,
			TypeInformation.of(new TypeHint<Tuple2<KafkaTopicPartition, Long>>() {})));
    ...
}
public void open(Configuration configuration) {
    ...
    // 第一次获取所有分区数：可以感知 Kafka 的分区变化
    // 后续分区分配如常，因为subTask 得到全分区的 offset
final List<KafkaTopicPartition> allPartitions = partitionDiscoverer.discoverPartitions();
    ...
}    

// FlinkKafkaConsumerBase
public void run(SourceContext<T> sourceContext) throws Exception {
    ...
    // KEY_PARTITION_DISCOVERY_INTERVAL_MILLIS 参数判断
    if (discoveryIntervalMillis == PARTITION_DISCOVERY_DISABLED) {
        kafkaFetcher.runFetchLoop();
    } else {
        runWithPartitionDiscovery();
    }
}
private void runWithPartitionDiscovery() throws Exception {
    final AtomicReference<Exception> discoveryLoopErrorRef = new AtomicReference<>();
    // 创建 发现Kafka 新分区 线程
    createAndStartDiscoveryLoop(discoveryLoopErrorRef);
	// 消费 Kafka 数据
    kafkaFetcher.runFetchLoop();
    ...
}
private void createAndStartDiscoveryLoop {
    // 间隔 discoveryIntervalMillis 去检测发现新分区添加到 kafkaFetcher
		discoveryLoopThread = new Thread(() -> {
			try {
	   // --------------------- partition discovery loop ---------------------
		// throughout the loop, we always eagerly check if we are still running before
	// performing the next operation, so that we can escape the loop as soon as possible
				while (running) {

					final List<KafkaTopicPartition> discoveredPartitions;
					try {
                        // 发现新分区
                        // discoverPartitions 第一次调用会返回当前 subTask 消费的所有分区
                        // 非第一次调用，则返回当前 subTask 消费的 新分区
						discoveredPartitions = partitionDiscoverer.discoverPartitions();
					} 
                    ...
		// no need to add the discovered partitions if we were closed during the meantime
					if (running && !discoveredPartitions.isEmpty()) {
                        // 添加新分区去消费
						kafkaFetcher.addDiscoveredPartitions(discoveredPartitions);
					}
					// do not waste any time sleeping if we're not running anymore
					if (running && discoveryIntervalMillis != 0) {
						try {
							Thread.sleep(discoveryIntervalMillis);
						} catch (InterruptedException iex) {
		// may be interrupted if the consumer was canceled midway; simply escape the loop
							break;
						}
					}
				}
			}
            ...
		discoveryLoopThread.start();
	}