Examples with WakeupException org.apache.kafka.common.errors.WakeupException used on opensource projects

Example 51 with WakeupException

use of org.apache.kafka.common.errors.WakeupException in project flink by apache.

the class KafkaConsumerThread method reassignPartitions.

// ------------------------------------------------------------------------
/**
 * Reestablishes the assigned partitions for the consumer. The reassigned partitions consists of
 * the provided new partitions and whatever partitions was already previously assigned to the
 * consumer.
 *
 * <p>The reassignment process is protected against wakeup calls, so that after this method
 * returns, the consumer is either untouched or completely reassigned with the correct offset
 * positions.
 *
 * <p>If the consumer was already woken-up prior to a reassignment resulting in an interruption
 * any time during the reassignment, the consumer is guaranteed to roll back as if it was
 * untouched. On the other hand, if there was an attempt to wakeup the consumer during the
 * reassignment, the wakeup call is "buffered" until the reassignment completes.
 *
 * <p>This method is exposed for testing purposes.
 */
@VisibleForTesting
void reassignPartitions(List<KafkaTopicPartitionState<T, TopicPartition>> newPartitions) throws Exception {
    if (newPartitions.size() == 0) {
        return;
    }
    hasAssignedPartitions = true;
    boolean reassignmentStarted = false;
    // since the reassignment may introduce several Kafka blocking calls that cannot be
    // interrupted,
    // the consumer needs to be isolated from external wakeup calls in setOffsetsToCommit() and
    // shutdown()
    // until the reassignment is complete.
    final KafkaConsumer<byte[], byte[]> consumerTmp;
    synchronized (consumerReassignmentLock) {
        consumerTmp = this.consumer;
        this.consumer = null;
    }
    final Map<TopicPartition, Long> oldPartitionAssignmentsToPosition = new HashMap<>();
    try {
        for (TopicPartition oldPartition : consumerTmp.assignment()) {
            oldPartitionAssignmentsToPosition.put(oldPartition, consumerTmp.position(oldPartition));
        }
        final List<TopicPartition> newPartitionAssignments = new ArrayList<>(newPartitions.size() + oldPartitionAssignmentsToPosition.size());
        newPartitionAssignments.addAll(oldPartitionAssignmentsToPosition.keySet());
        newPartitionAssignments.addAll(convertKafkaPartitions(newPartitions));
        // reassign with the new partitions
        consumerTmp.assign(newPartitionAssignments);
        reassignmentStarted = true;
        // old partitions should be seeked to their previous position
        for (Map.Entry<TopicPartition, Long> oldPartitionToPosition : oldPartitionAssignmentsToPosition.entrySet()) {
            consumerTmp.seek(oldPartitionToPosition.getKey(), oldPartitionToPosition.getValue());
        }
        // replace those with actual offsets, according to what the sentinel value represent.
        for (KafkaTopicPartitionState<T, TopicPartition> newPartitionState : newPartitions) {
            if (newPartitionState.getOffset() == KafkaTopicPartitionStateSentinel.EARLIEST_OFFSET) {
                consumerTmp.seekToBeginning(Collections.singletonList(newPartitionState.getKafkaPartitionHandle()));
                newPartitionState.setOffset(consumerTmp.position(newPartitionState.getKafkaPartitionHandle()) - 1);
            } else if (newPartitionState.getOffset() == KafkaTopicPartitionStateSentinel.LATEST_OFFSET) {
                consumerTmp.seekToEnd(Collections.singletonList(newPartitionState.getKafkaPartitionHandle()));
                newPartitionState.setOffset(consumerTmp.position(newPartitionState.getKafkaPartitionHandle()) - 1);
            } else if (newPartitionState.getOffset() == KafkaTopicPartitionStateSentinel.GROUP_OFFSET) {
                // the KafkaConsumer by default will automatically seek the consumer position
                // to the committed group offset, so we do not need to do it.
                newPartitionState.setOffset(consumerTmp.position(newPartitionState.getKafkaPartitionHandle()) - 1);
            } else {
                consumerTmp.seek(newPartitionState.getKafkaPartitionHandle(), newPartitionState.getOffset() + 1);
            }
        }
    } catch (WakeupException e) {
        synchronized (consumerReassignmentLock) {
            this.consumer = consumerTmp;
            // we do a full roll back so that it is as if it was left untouched
            if (reassignmentStarted) {
                this.consumer.assign(new ArrayList<>(oldPartitionAssignmentsToPosition.keySet()));
                for (Map.Entry<TopicPartition, Long> oldPartitionToPosition : oldPartitionAssignmentsToPosition.entrySet()) {
                    this.consumer.seek(oldPartitionToPosition.getKey(), oldPartitionToPosition.getValue());
                }
            }
            // no need to restore the wakeup state in this case,
            // since only the last wakeup call is effective anyways
            hasBufferedWakeup = false;
            // again
            for (KafkaTopicPartitionState<T, TopicPartition> newPartition : newPartitions) {
                unassignedPartitionsQueue.add(newPartition);
            }
            // this signals the main fetch loop to continue through the loop
            throw new AbortedReassignmentException();
        }
    }
    // reassignment complete; expose the reassigned consumer
    synchronized (consumerReassignmentLock) {
        this.consumer = consumerTmp;
        // restore wakeup state for the consumer if necessary
        if (hasBufferedWakeup) {
            this.consumer.wakeup();
            hasBufferedWakeup = false;
        }
    }
}