Flink源码走读（二）：Flink+Kafka实现端到端Exactly Once语义

@Override
public void snapshotState(FunctionSnapshotContext context) throws Exception {
	// this is like the pre-commit of a 2-phase-commit transaction
	// we are ready to commit and remember the transaction

	checkState(currentTransactionHolder != null, "bug: no transaction object when performing state snapshot");

	long checkpointId = context.getCheckpointId();
	LOG.debug("{} - checkpoint {} triggered, flushing transaction '{}'", name(), context.getCheckpointId(), currentTransactionHolder);

	preCommit(currentTransactionHolder.handle);
	pendingCommitTransactions.put(checkpointId, currentTransactionHolder);
	LOG.debug("{} - stored pending transactions {}", name(), pendingCommitTransactions);

	currentTransactionHolder = beginTransactionInternal();
	LOG.debug("{} - started new transaction '{}'", name(), currentTransactionHolder);

	state.clear();
	state.add(new State<>(
		this.currentTransactionHolder,
		new ArrayList<>(pendingCommitTransactions.values()),
		userContext));
}

@Override
public final void notifyCheckpointComplete(long checkpointId) throws Exception {
	// the following scenarios are possible here
	//
	//  (1) there is exactly one transaction from the latest checkpoint that
	//      was triggered and completed. That should be the common case.
	//      Simply commit that transaction in that case.
	//
	//  (2) there are multiple pending transactions because one previous
	//      checkpoint was skipped. That is a rare case, but can happen
	//      for example when:
	//
	//        - the master cannot persist the metadata of the last
	//          checkpoint (temporary outage in the storage system) but
	//          could persist a successive checkpoint (the one notified here)
	//
	//        - other tasks could not persist their status during
	//          the previous checkpoint, but did not trigger a failure because they
	//          could hold onto their state and could successfully persist it in
	//          a successive checkpoint (the one notified here)
	//
	//      In both cases, the prior checkpoint never reach a committed state, but
	//      this checkpoint is always expected to subsume the prior one and cover all
	//      changes since the last successful one. As a consequence, we need to commit
	//      all pending transactions.
	//
	//  (3) Multiple transactions are pending, but the checkpoint complete notification
	//      relates not to the latest. That is possible, because notification messages
	//      can be delayed (in an extreme case till arrive after a succeeding checkpoint
	//      was triggered) and because there can be concurrent overlapping checkpoints
	//      (a new one is started before the previous fully finished).
	//
	// ==> There should never be a case where we have no pending transaction here
	//
	Iterator<Map.Entry<Long, TransactionHolder<TXN>>> pendingTransactionIterator = pendingCommitTransactions.entrySet().iterator();
	checkState(pendingTransactionIterator.hasNext(), "checkpoint completed, but no transaction pending");
	Throwable firstError = null;

	while (pendingTransactionIterator.hasNext()) {
		Map.Entry<Long, TransactionHolder<TXN>> entry = pendingTransactionIterator.next();
		Long pendingTransactionCheckpointId = entry.getKey();
		TransactionHolder<TXN> pendingTransaction = entry.getValue();
		if (pendingTransactionCheckpointId > checkpointId) {
			continue;
		}

		LOG.info("{} - checkpoint {} complete, committing transaction {} from checkpoint {}",
			name(), checkpointId, pendingTransaction, pendingTransactionCheckpointId);

		logWarningIfTimeoutAlmostReached(pendingTransaction);
		try {
			commit(pendingTransaction.handle);
		} catch (Throwable t) {
			if (firstError == null) {
					firstError = t;
			}
		}

		LOG.debug("{} - committed checkpoint transaction {}", name(), pendingTransaction);

		pendingTransactionIterator.remove();
	}

	if (firstError != null) {
		throw new FlinkRuntimeException("Committing one of transactions failed, logging first encountered failure",
			firstError);
	}
}

// ------ methods that should be implemented in child class to support two phase commit algorithm ------

/**
 * Write value within a transaction.
 */
protected abstract void invoke(TXN transaction, IN value, Context context) throws Exception;

/**
 * Method that starts a new transaction.
 *
 * @return newly created transaction.
 */
protected abstract TXN beginTransaction() throws Exception;

/**
 * Pre commit previously created transaction. Pre commit must make all of the necessary steps to prepare the
 * transaction for a commit that might happen in the future. After this point the transaction might still be
 * aborted, but underlying implementation must ensure that commit calls on already pre committed transactions
 * will always succeed.
 *
 * <p>Usually implementation involves flushing the data.
 */
protected abstract void preCommit(TXN transaction) throws Exception;

/**
 * Commit a pre-committed transaction. If this method fail, Flink application will be
 * restarted and {@link TwoPhaseCommitSinkFunction#recoverAndCommit(Object)} will be called again for the
 * same transaction.
 */
protected abstract void commit(TXN transaction);

/**
 * Abort a transaction.
 */
protected abstract void abort(TXN transaction);

@Override
public void invoke(KafkaTransactionState transaction, IN next, Context context) throws FlinkKafka011Exception {
	checkErroneous();

	byte[] serializedKey = schema.serializeKey(next);
	byte[] serializedValue = schema.serializeValue(next);
	String targetTopic = schema.getTargetTopic(next);
	if (targetTopic == null) {
		targetTopic = defaultTopicId;
	}

	Long timestamp = null;
	if (this.writeTimestampToKafka) {
		timestamp = context.timestamp();
	}

	ProducerRecord<byte[], byte[]> record;
	int[] partitions = topicPartitionsMap.get(targetTopic);
	if (null == partitions) {
		partitions = getPartitionsByTopic(targetTopic, transaction.producer);
		topicPartitionsMap.put(targetTopic, partitions);
	}
	if (flinkKafkaPartitioner != null) {
		record = new ProducerRecord<>(
			targetTopic,
			flinkKafkaPartitioner.partition(next, serializedKey, serializedValue, targetTopic, partitions),
			timestamp,
			serializedKey,
			serializedValue);
	} else {
		record = new ProducerRecord<>(targetTopic, null, timestamp, serializedKey, serializedValue);
	}
	pendingRecords.incrementAndGet();
	transaction.producer.send(record, callback);
}

@Override
protected KafkaTransactionState beginTransaction() throws FlinkKafka011Exception {
	switch (semantic) {
		case EXACTLY_ONCE:
			FlinkKafkaProducer<byte[], byte[]> producer = createTransactionalProducer();
			producer.beginTransaction();
			return new KafkaTransactionState(producer.getTransactionalId(), producer);
		case AT_LEAST_ONCE:
		case NONE:
			// Do not create new producer on each beginTransaction() if it is not necessary
			final KafkaTransactionState currentTransaction = currentTransaction();
			if (currentTransaction != null && currentTransaction.producer != null) {
				return new KafkaTransactionState(currentTransaction.producer);
			}
			return new KafkaTransactionState(initNonTransactionalProducer(true));
		default:
			throw new UnsupportedOperationException("Not implemented semantic");
	}
}

@Override
protected void preCommit(KafkaTransactionState transaction) throws FlinkKafka011Exception {
	switch (semantic) {
		case EXACTLY_ONCE:
		case AT_LEAST_ONCE:
			flush(transaction);
			break;
		case NONE:
			break;
		default:
			throw new UnsupportedOperationException("Not implemented semantic");
	}
	checkErroneous();
}

@Override
protected void commit(KafkaTransactionState transaction) {
	if (transaction.isTransactional()) {
		try {
			transaction.producer.commitTransaction();
		} finally {
			recycleTransactionalProducer(transaction.producer);
		}
	}
}

@Override
protected void abort(KafkaTransactionState transaction) {
	if (transaction.isTransactional()) {
		transaction.producer.abortTransaction();
		recycleTransactionalProducer(transaction.producer);
	}
}

/**
 * Pool of available transactional ids.
 */
private final BlockingDeque<String> availableTransactionalIds = new LinkedBlockingDeque<>();
	
/**
 * For each checkpoint we create new {@link FlinkKafkaProducer} so that new transactions will not clash
 * with transactions created during previous checkpoints ({@code producer.initTransactions()} assures that we
 * obtain new producerId and epoch counters).
 */
private FlinkKafkaProducer<byte[], byte[]> createTransactionalProducer() throws FlinkKafka011Exception {
	String transactionalId = availableTransactionalIds.poll();
	if (transactionalId == null) {
		throw new FlinkKafka011Exception(
			FlinkKafka011ErrorCode.PRODUCERS_POOL_EMPTY,
			"Too many ongoing snapshots. Increase kafka producers pool size or decrease number of concurrent checkpoints.");
	}
	FlinkKafkaProducer<byte[], byte[]> producer = initTransactionalProducer(transactionalId, true);
	producer.initTransactions();
	return producer;
}

private void recycleTransactionalProducer(FlinkKafkaProducer<byte[], byte[]> producer) {
	availableTransactionalIds.add(producer.getTransactionalId());
	producer.close();
}