org.apache.kafka.clients.consumer.internals.SubscriptionState Maven / Gradle / Ivy
/**
* Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements. See the NOTICE
* file distributed with this work for additional information regarding copyright ownership. The ASF licenses this file
* to You under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the
* License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
* an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
* specific language governing permissions and limitations under the License.
*/
package org.apache.kafka.clients.consumer.internals;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.clients.consumer.OffsetResetStrategy;
import org.apache.kafka.common.TopicPartition;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.regex.Pattern;
/**
* A class for tracking the topics, partitions, and offsets for the consumer. A partition
* is "assigned" either directly with {@link #assignFromUser(Collection)} (manual assignment)
* or with {@link #assignFromSubscribed(Collection)} (automatic assignment from subscription).
*
* Once assigned, the partition is not considered "fetchable" until its initial position has
* been set with {@link #seek(TopicPartition, long)}. Fetchable partitions track a fetch
* position which is used to set the offset of the next fetch, and a consumed position
* which is the last offset that has been returned to the user. You can suspend fetching
* from a partition through {@link #pause(TopicPartition)} without affecting the fetched/consumed
* offsets. The partition will remain unfetchable until the {@link #resume(TopicPartition)} is
* used. You can also query the pause state independently with {@link #isPaused(TopicPartition)}.
*
* Note that pause state as well as fetch/consumed positions are not preserved when partition
* assignment is changed whether directly by the user or through a group rebalance.
*
* This class also maintains a cache of the latest commit position for each of the assigned
* partitions. This is updated through {@link #committed(TopicPartition, OffsetAndMetadata)} and can be used
* to set the initial fetch position (e.g. {@link Fetcher#resetOffset(TopicPartition)}.
*/
public class SubscriptionState {
/* the pattern user has requested */
private Pattern subscribedPattern;
/* the list of topics the user has requested */
private final Set subscription;
/* the list of topics the group has subscribed to (set only for the leader on join group completion) */
private final Set groupSubscription;
/* the list of partitions the user has requested */
private final Set userAssignment;
/* the list of partitions currently assigned */
private final Map assignment;
/* do we need to request a partition assignment from the coordinator? */
private boolean needsPartitionAssignment;
/* do we need to request the latest committed offsets from the coordinator? */
private boolean needsFetchCommittedOffsets;
/* Default offset reset strategy */
private final OffsetResetStrategy defaultResetStrategy;
/* Listener to be invoked when assignment changes */
private ConsumerRebalanceListener listener;
private static final String SUBSCRIPTION_EXCEPTION_MESSAGE =
"Subscription to topics, partitions and pattern are mutually exclusive";
public SubscriptionState(OffsetResetStrategy defaultResetStrategy) {
this.defaultResetStrategy = defaultResetStrategy;
this.subscription = new HashSet<>();
this.userAssignment = new HashSet<>();
this.assignment = new HashMap<>();
this.groupSubscription = new HashSet<>();
this.needsPartitionAssignment = false;
this.needsFetchCommittedOffsets = true; // initialize to true for the consumers to fetch offset upon starting up
this.subscribedPattern = null;
}
public void subscribe(Collection topics, ConsumerRebalanceListener listener) {
if (listener == null)
throw new IllegalArgumentException("RebalanceListener cannot be null");
if (!this.userAssignment.isEmpty() || this.subscribedPattern != null)
throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
this.listener = listener;
changeSubscription(topics);
}
public void changeSubscription(Collection topicsToSubscribe) {
if (!this.subscription.equals(new HashSet<>(topicsToSubscribe))) {
this.subscription.clear();
this.subscription.addAll(topicsToSubscribe);
this.groupSubscription.addAll(topicsToSubscribe);
this.needsPartitionAssignment = true;
// Remove any assigned partitions which are no longer subscribed to
for (Iterator it = assignment.keySet().iterator(); it.hasNext(); ) {
TopicPartition tp = it.next();
if (!subscription.contains(tp.topic()))
it.remove();
}
}
}
/**
* Add topics to the current group subscription. This is used by the group leader to ensure
* that it receives metadata updates for all topics that the group is interested in.
* @param topics The topics to add to the group subscription
*/
public void groupSubscribe(Collection topics) {
if (!this.userAssignment.isEmpty())
throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
this.groupSubscription.addAll(topics);
}
public void needReassignment() {
this.groupSubscription.retainAll(subscription);
this.needsPartitionAssignment = true;
}
/**
* Change the assignment to the specified partitions provided by the user,
* note this is different from {@link #assignFromSubscribed(Collection)}
* whose input partitions are provided from the subscribed topics.
*/
public void assignFromUser(Collection partitions) {
if (!this.subscription.isEmpty() || this.subscribedPattern != null)
throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
this.userAssignment.clear();
this.userAssignment.addAll(partitions);
for (TopicPartition partition : partitions)
if (!assignment.containsKey(partition))
addAssignedPartition(partition);
this.assignment.keySet().retainAll(this.userAssignment);
this.needsPartitionAssignment = false;
}
/**
* Change the assignment to the specified partitions returned from the coordinator,
* note this is different from {@link #assignFromUser(Collection)} which directly set the assignment from user inputs
*/
public void assignFromSubscribed(Collection assignments) {
for (TopicPartition tp : assignments)
if (!this.subscription.contains(tp.topic()))
throw new IllegalArgumentException("Assigned partition " + tp + " for non-subscribed topic.");
this.assignment.clear();
for (TopicPartition tp: assignments)
addAssignedPartition(tp);
this.needsPartitionAssignment = false;
}
public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
if (listener == null)
throw new IllegalArgumentException("RebalanceListener cannot be null");
if (!this.subscription.isEmpty() || !this.userAssignment.isEmpty())
throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
this.listener = listener;
this.subscribedPattern = pattern;
}
public boolean hasPatternSubscription() {
return subscribedPattern != null;
}
public void unsubscribe() {
this.subscription.clear();
this.userAssignment.clear();
this.assignment.clear();
this.needsPartitionAssignment = true;
this.subscribedPattern = null;
}
public Pattern getSubscribedPattern() {
return this.subscribedPattern;
}
public Set subscription() {
return this.subscription;
}
public Set pausedPartitions() {
HashSet paused = new HashSet<>();
for (Map.Entry entry : assignment.entrySet()) {
final TopicPartition tp = entry.getKey();
final TopicPartitionState state = entry.getValue();
if (state.paused) {
paused.add(tp);
}
}
return paused;
}
/**
* Get the subscription for the group. For the leader, this will include the union of the
* subscriptions of all group members. For followers, it is just that member's subscription.
* This is used when querying topic metadata to detect the metadata changes which would
* require rebalancing. The leader fetches metadata for all topics in the group so that it
* can do the partition assignment (which requires at least partition counts for all topics
* to be assigned).
* @return The union of all subscribed topics in the group if this member is the leader
* of the current generation; otherwise it returns the same set as {@link #subscription()}
*/
public Set groupSubscription() {
return this.groupSubscription;
}
private TopicPartitionState assignedState(TopicPartition tp) {
TopicPartitionState state = this.assignment.get(tp);
if (state == null)
throw new IllegalStateException("No current assignment for partition " + tp);
return state;
}
public void committed(TopicPartition tp, OffsetAndMetadata offset) {
assignedState(tp).committed(offset);
}
public OffsetAndMetadata committed(TopicPartition tp) {
return assignedState(tp).committed;
}
public void needRefreshCommits() {
this.needsFetchCommittedOffsets = true;
}
public boolean refreshCommitsNeeded() {
return this.needsFetchCommittedOffsets;
}
public void commitsRefreshed() {
this.needsFetchCommittedOffsets = false;
}
public void seek(TopicPartition tp, long offset) {
assignedState(tp).seek(offset);
}
public Set assignedPartitions() {
return this.assignment.keySet();
}
public Set fetchablePartitions() {
Set fetchable = new HashSet<>();
for (Map.Entry entry : assignment.entrySet()) {
if (entry.getValue().isFetchable())
fetchable.add(entry.getKey());
}
return fetchable;
}
public boolean partitionsAutoAssigned() {
return !this.subscription.isEmpty();
}
public void position(TopicPartition tp, long offset) {
assignedState(tp).position(offset);
}
public Long position(TopicPartition tp) {
return assignedState(tp).position;
}
public Map allConsumed() {
Map allConsumed = new HashMap<>();
for (Map.Entry entry : assignment.entrySet()) {
TopicPartitionState state = entry.getValue();
if (state.hasValidPosition())
allConsumed.put(entry.getKey(), new OffsetAndMetadata(state.position));
}
return allConsumed;
}
public void needOffsetReset(TopicPartition partition, OffsetResetStrategy offsetResetStrategy) {
assignedState(partition).awaitReset(offsetResetStrategy);
}
public void needOffsetReset(TopicPartition partition) {
needOffsetReset(partition, defaultResetStrategy);
}
public boolean hasDefaultOffsetResetPolicy() {
return defaultResetStrategy != OffsetResetStrategy.NONE;
}
public boolean isOffsetResetNeeded(TopicPartition partition) {
return assignedState(partition).awaitingReset();
}
public OffsetResetStrategy resetStrategy(TopicPartition partition) {
return assignedState(partition).resetStrategy;
}
public boolean hasAllFetchPositions() {
for (TopicPartitionState state : assignment.values())
if (!state.hasValidPosition())
return false;
return true;
}
public Set missingFetchPositions() {
Set missing = new HashSet<>();
for (Map.Entry entry : assignment.entrySet())
if (!entry.getValue().hasValidPosition())
missing.add(entry.getKey());
return missing;
}
public boolean partitionAssignmentNeeded() {
return this.needsPartitionAssignment;
}
public boolean isAssigned(TopicPartition tp) {
return assignment.containsKey(tp);
}
public boolean isPaused(TopicPartition tp) {
return isAssigned(tp) && assignedState(tp).paused;
}
public boolean isFetchable(TopicPartition tp) {
return isAssigned(tp) && assignedState(tp).isFetchable();
}
public void pause(TopicPartition tp) {
assignedState(tp).pause();
}
public void resume(TopicPartition tp) {
assignedState(tp).resume();
}
private void addAssignedPartition(TopicPartition tp) {
this.assignment.put(tp, new TopicPartitionState());
}
public ConsumerRebalanceListener listener() {
return listener;
}
private static class TopicPartitionState {
private Long position; // last consumed position
private OffsetAndMetadata committed; // last committed position
private boolean paused; // whether this partition has been paused by the user
private OffsetResetStrategy resetStrategy; // the strategy to use if the offset needs resetting
public TopicPartitionState() {
this.paused = false;
this.position = null;
this.committed = null;
this.resetStrategy = null;
}
private void awaitReset(OffsetResetStrategy strategy) {
this.resetStrategy = strategy;
this.position = null;
}
public boolean awaitingReset() {
return resetStrategy != null;
}
public boolean hasValidPosition() {
return position != null;
}
private void seek(long offset) {
this.position = offset;
this.resetStrategy = null;
}
private void position(long offset) {
if (!hasValidPosition())
throw new IllegalStateException("Cannot set a new position without a valid current position");
this.position = offset;
}
private void committed(OffsetAndMetadata offset) {
this.committed = offset;
}
private void pause() {
this.paused = true;
}
private void resume() {
this.paused = false;
}
private boolean isFetchable() {
return !paused && hasValidPosition();
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy