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.NoOffsetForPartitionException;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.clients.consumer.OffsetResetStrategy;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.internals.PartitionStates;
import org.apache.kafka.common.requests.IsolationLevel;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
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(Set)} (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.
*/
public class SubscriptionState {
private static final String SUBSCRIPTION_EXCEPTION_MESSAGE =
"Subscription to topics, partitions and pattern are mutually exclusive";
private enum SubscriptionType {
NONE, AUTO_TOPICS, AUTO_PATTERN, USER_ASSIGNED
}
/* the type of subscription */
private SubscriptionType subscriptionType;
/* the pattern user has requested */
private Pattern subscribedPattern;
/* the list of topics the user has requested */
private 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 partitions that are currently assigned, note that the order of partition matters (see FetchBuilder for more details) */
private final PartitionStates assignment;
/* Default offset reset strategy */
private final OffsetResetStrategy defaultResetStrategy;
/* Listeners provide a hook for internal state cleanup (e.g. metrics) on assignment changes */
private final List listeners = new ArrayList<>();
/* User-provided listener to be invoked when assignment changes */
private ConsumerRebalanceListener rebalanceListener;
public SubscriptionState(OffsetResetStrategy defaultResetStrategy) {
this.defaultResetStrategy = defaultResetStrategy;
this.subscription = Collections.emptySet();
this.assignment = new PartitionStates<>();
this.groupSubscription = new HashSet<>();
this.subscribedPattern = null;
this.subscriptionType = SubscriptionType.NONE;
}
/**
* This method sets the subscription type if it is not already set (i.e. when it is NONE),
* or verifies that the subscription type is equal to the give type when it is set (i.e.
* when it is not NONE)
* @param type The given subscription type
*/
private void setSubscriptionType(SubscriptionType type) {
if (this.subscriptionType == SubscriptionType.NONE)
this.subscriptionType = type;
else if (this.subscriptionType != type)
throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
}
public void subscribe(Set topics, ConsumerRebalanceListener listener) {
if (listener == null)
throw new IllegalArgumentException("RebalanceListener cannot be null");
setSubscriptionType(SubscriptionType.AUTO_TOPICS);
this.rebalanceListener = listener;
changeSubscription(topics);
}
public void subscribeFromPattern(Set topics) {
if (subscriptionType != SubscriptionType.AUTO_PATTERN)
throw new IllegalArgumentException("Attempt to subscribe from pattern while subscription type set to " +
subscriptionType);
changeSubscription(topics);
}
private void changeSubscription(Set topicsToSubscribe) {
if (!this.subscription.equals(topicsToSubscribe)) {
this.subscription = topicsToSubscribe;
this.groupSubscription.addAll(topicsToSubscribe);
}
}
/**
* 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.subscriptionType == SubscriptionType.USER_ASSIGNED)
throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
this.groupSubscription.addAll(topics);
}
/**
* Reset the group's subscription to only contain topics subscribed by this consumer.
*/
public void resetGroupSubscription() {
this.groupSubscription.retainAll(subscription);
}
/**
* 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(Set partitions) {
setSubscriptionType(SubscriptionType.USER_ASSIGNED);
if (!this.assignment.partitionSet().equals(partitions)) {
fireOnAssignment(partitions);
Map partitionToState = new HashMap<>();
for (TopicPartition partition : partitions) {
TopicPartitionState state = assignment.stateValue(partition);
if (state == null)
state = new TopicPartitionState();
partitionToState.put(partition, state);
}
this.assignment.set(partitionToState);
}
}
/**
* Change the assignment to the specified partitions returned from the coordinator,
* note this is different from {@link #assignFromUser(Set)} which directly set the assignment from user inputs
*/
public void assignFromSubscribed(Collection assignments) {
if (!this.partitionsAutoAssigned())
throw new IllegalArgumentException("Attempt to dynamically assign partitions while manual assignment in use");
Map assignedPartitionStates = partitionToStateMap(assignments);
fireOnAssignment(assignedPartitionStates.keySet());
if (this.subscribedPattern != null) {
for (TopicPartition tp : assignments) {
if (!this.subscribedPattern.matcher(tp.topic()).matches())
throw new IllegalArgumentException("Assigned partition " + tp + " for non-subscribed topic regex pattern; subscription pattern is " + this.subscribedPattern);
}
} else {
for (TopicPartition tp : assignments)
if (!this.subscription.contains(tp.topic()))
throw new IllegalArgumentException("Assigned partition " + tp + " for non-subscribed topic; subscription is " + this.subscription);
}
this.assignment.set(assignedPartitionStates);
}
public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
if (listener == null)
throw new IllegalArgumentException("RebalanceListener cannot be null");
setSubscriptionType(SubscriptionType.AUTO_PATTERN);
this.rebalanceListener = listener;
this.subscribedPattern = pattern;
}
public boolean hasPatternSubscription() {
return this.subscriptionType == SubscriptionType.AUTO_PATTERN;
}
public boolean hasNoSubscriptionOrUserAssignment() {
return this.subscriptionType == SubscriptionType.NONE;
}
public void unsubscribe() {
this.subscription = Collections.emptySet();
this.assignment.clear();
this.subscribedPattern = null;
this.subscriptionType = SubscriptionType.NONE;
fireOnAssignment(Collections.emptySet());
}
public Pattern subscribedPattern() {
return this.subscribedPattern;
}
public Set subscription() {
return this.subscription;
}
public Set pausedPartitions() {
HashSet paused = new HashSet<>();
for (PartitionStates.PartitionState state : assignment.partitionStates()) {
if (state.value().paused) {
paused.add(state.topicPartition());
}
}
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.stateValue(tp);
if (state == null)
throw new IllegalStateException("No current assignment for partition " + tp);
return state;
}
public void seek(TopicPartition tp, long offset) {
assignedState(tp).seek(offset);
}
/**
* @return an unmodifiable view of the currently assigned partitions
*/
public Set assignedPartitions() {
return this.assignment.partitionSet();
}
/**
* Provides the number of assigned partitions in a thread safe manner.
* @return the number of assigned partitions.
*/
public int numAssignedPartitions() {
return this.assignment.size();
}
public List fetchablePartitions() {
List fetchable = new ArrayList<>(assignment.size());
for (PartitionStates.PartitionState state : assignment.partitionStates()) {
if (state.value().isFetchable())
fetchable.add(state.topicPartition());
}
return fetchable;
}
public boolean partitionsAutoAssigned() {
return this.subscriptionType == SubscriptionType.AUTO_TOPICS || this.subscriptionType == SubscriptionType.AUTO_PATTERN;
}
public void position(TopicPartition tp, long offset) {
assignedState(tp).position(offset);
}
public Long position(TopicPartition tp) {
return assignedState(tp).position;
}
public Long partitionLag(TopicPartition tp, IsolationLevel isolationLevel) {
TopicPartitionState topicPartitionState = assignedState(tp);
if (isolationLevel == IsolationLevel.READ_COMMITTED)
return topicPartitionState.lastStableOffset == null ? null : topicPartitionState.lastStableOffset - topicPartitionState.position;
else
return topicPartitionState.highWatermark == null ? null : topicPartitionState.highWatermark - topicPartitionState.position;
}
public Long partitionLead(TopicPartition tp) {
TopicPartitionState topicPartitionState = assignedState(tp);
return topicPartitionState.logStartOffset == null ? null : topicPartitionState.position - topicPartitionState.logStartOffset;
}
public void updateHighWatermark(TopicPartition tp, long highWatermark) {
assignedState(tp).highWatermark = highWatermark;
}
public void updateLogStartOffset(TopicPartition tp, long logStartOffset) {
assignedState(tp).logStartOffset = logStartOffset;
}
public void updateLastStableOffset(TopicPartition tp, long lastStableOffset) {
assignedState(tp).lastStableOffset = lastStableOffset;
}
public Map allConsumed() {
Map allConsumed = new HashMap<>();
for (PartitionStates.PartitionState state : assignment.partitionStates()) {
if (state.value().hasValidPosition())
allConsumed.put(state.topicPartition(), new OffsetAndMetadata(state.value().position));
}
return allConsumed;
}
public void requestOffsetReset(TopicPartition partition, OffsetResetStrategy offsetResetStrategy) {
assignedState(partition).reset(offsetResetStrategy);
}
public void requestOffsetReset(TopicPartition partition) {
requestOffsetReset(partition, defaultResetStrategy);
}
public void setResetPending(Set partitions, long nextAllowResetTimeMs) {
for (TopicPartition partition : partitions) {
assignedState(partition).setResetPending(nextAllowResetTimeMs);
}
}
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 (PartitionStates.PartitionState state : assignment.partitionStates()) {
if (!state.value().hasValidPosition())
return false;
}
return true;
}
public Set missingFetchPositions() {
Set missing = new HashSet<>();
for (PartitionStates.PartitionState state : assignment.partitionStates()) {
if (state.value().isMissingPosition())
missing.add(state.topicPartition());
}
return missing;
}
public void resetMissingPositions() {
final Set partitionsWithNoOffsets = new HashSet<>();
for (PartitionStates.PartitionState state : assignment.partitionStates()) {
TopicPartition tp = state.topicPartition();
TopicPartitionState partitionState = state.value();
if (partitionState.isMissingPosition()) {
if (defaultResetStrategy == OffsetResetStrategy.NONE)
partitionsWithNoOffsets.add(tp);
else
partitionState.reset(defaultResetStrategy);
}
}
if (!partitionsWithNoOffsets.isEmpty())
throw new NoOffsetForPartitionException(partitionsWithNoOffsets);
}
public Set partitionsNeedingReset(long nowMs) {
Set partitions = new HashSet<>();
for (PartitionStates.PartitionState state : assignment.partitionStates()) {
TopicPartitionState partitionState = state.value();
if (partitionState.awaitingReset() && partitionState.isResetAllowed(nowMs))
partitions.add(state.topicPartition());
}
return partitions;
}
public boolean isAssigned(TopicPartition tp) {
return assignment.contains(tp);
}
public boolean isPaused(TopicPartition tp) {
return isAssigned(tp) && assignedState(tp).paused;
}
public boolean isFetchable(TopicPartition tp) {
return isAssigned(tp) && assignedState(tp).isFetchable();
}
public boolean hasValidPosition(TopicPartition tp) {
return isAssigned(tp) && assignedState(tp).hasValidPosition();
}
public void pause(TopicPartition tp) {
assignedState(tp).pause();
}
public void resume(TopicPartition tp) {
assignedState(tp).resume();
}
public void resetFailed(Set partitions, long nextRetryTimeMs) {
for (TopicPartition partition : partitions)
assignedState(partition).resetFailed(nextRetryTimeMs);
}
public void movePartitionToEnd(TopicPartition tp) {
assignment.moveToEnd(tp);
}
public ConsumerRebalanceListener rebalanceListener() {
return rebalanceListener;
}
public void addListener(Listener listener) {
listeners.add(listener);
}
public void fireOnAssignment(Set assignment) {
for (Listener listener : listeners)
listener.onAssignment(assignment);
}
private static Map partitionToStateMap(Collection assignments) {
Map map = new HashMap<>(assignments.size());
for (TopicPartition tp : assignments)
map.put(tp, new TopicPartitionState());
return map;
}
private static class TopicPartitionState {
private Long position; // last consumed position
private Long highWatermark; // the high watermark from last fetch
private Long logStartOffset; // the log start offset
private Long lastStableOffset;
private boolean paused; // whether this partition has been paused by the user
private OffsetResetStrategy resetStrategy; // the strategy to use if the offset needs resetting
private Long nextAllowedRetryTimeMs;
TopicPartitionState() {
this.paused = false;
this.position = null;
this.highWatermark = null;
this.logStartOffset = null;
this.lastStableOffset = null;
this.resetStrategy = null;
this.nextAllowedRetryTimeMs = null;
}
private void reset(OffsetResetStrategy strategy) {
this.resetStrategy = strategy;
this.position = null;
this.nextAllowedRetryTimeMs = null;
}
private boolean isResetAllowed(long nowMs) {
return nextAllowedRetryTimeMs == null || nowMs >= nextAllowedRetryTimeMs;
}
private boolean awaitingReset() {
return resetStrategy != null;
}
private void setResetPending(long nextAllowedRetryTimeMs) {
this.nextAllowedRetryTimeMs = nextAllowedRetryTimeMs;
}
private void resetFailed(long nextAllowedRetryTimeMs) {
this.nextAllowedRetryTimeMs = nextAllowedRetryTimeMs;
}
private boolean hasValidPosition() {
return position != null;
}
private boolean isMissingPosition() {
return !hasValidPosition() && !awaitingReset();
}
private void seek(long offset) {
this.position = offset;
this.resetStrategy = null;
this.nextAllowedRetryTimeMs = 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 pause() {
this.paused = true;
}
private void resume() {
this.paused = false;
}
private boolean isFetchable() {
return !paused && hasValidPosition();
}
}
public interface Listener {
/**
* Fired after a new assignment is received (after a group rebalance or when the user manually changes the
* assignment).
*
* @param assignment The topic partitions assigned to the consumer
*/
void onAssignment(Set assignment);
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy