com.huaweicloud.dis.adapter.common.consumer.SubscriptionState Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of huaweicloud-dis-kafka-adapter-common Show documentation
Huawei Cloud Data Ingestion Service Kafka Adapter
The newest version!
/*
 * Copyright 2002-2010 the original author or authors.
 *
 * Licensed 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 com.huaweicloud.dis.adapter.common.consumer;


/**
 * base on apache common client 0.10.0.1
 */

import com.huaweicloud.dis.adapter.common.model.DisOffsetAndMetadata;
import com.huaweicloud.dis.adapter.common.model.DisOffsetResetStrategy;
import com.huaweicloud.dis.adapter.common.model.StreamPartition;

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;


public class SubscriptionState {

    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 final Set subscription;

    private final Set subscriptionIds;

    /* 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 DisOffsetResetStrategy defaultResetStrategy;

    /* Listener to be invoked when assignment changes */
    private DisConsumerRebalanceListener listener;

    private static final String SUBSCRIPTION_EXCEPTION_MESSAGE =
            "Subscription to topics, partitions and pattern are mutually exclusive";

    /**
     * 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 SubscriptionState(DisOffsetResetStrategy defaultResetStrategy) {
        this.defaultResetStrategy = defaultResetStrategy;
        this.subscription = new HashSet<>();
        this.subscriptionIds = 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;
        this.subscriptionType = SubscriptionType.NONE;
    }

    public void subscribe(Collection topics, DisConsumerRebalanceListener listener) {
        if (listener == null)
            throw new IllegalArgumentException("RebalanceListener cannot be null");

        setSubscriptionType(SubscriptionType.AUTO_TOPICS);

        this.listener = listener;

        changeSubscription(topics);
    }

    public void subscribe(Collection topicNames, Collection topicIds, DisConsumerRebalanceListener listener) {
        if (listener == null)
            throw new IllegalArgumentException("RebalanceListener cannot be null");

        setSubscriptionType(SubscriptionType.AUTO_TOPICS);
        this.listener = listener;
        Collection topics = new HashSet<>();
        if (topicIds != null) {
            subscriptionIds.addAll(topicIds);
            topics.addAll(topicIds);
        }
        if (topicNames != null) {
            topics.addAll(topicNames);
            subscription.addAll(topicNames);
        }
        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(); ) {
                StreamPartition tp = it.next();
                if (!subscription.contains(tp.stream()))
                    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.subscriptionType == SubscriptionType.USER_ASSIGNED)
            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) {
        setSubscriptionType(SubscriptionType.USER_ASSIGNED);

        this.userAssignment.clear();
        this.userAssignment.addAll(partitions);

        for (StreamPartition partition : partitions)
            if (!assignment.containsKey(partition))
                addAssignedPartition(partition);

        this.assignment.keySet().retainAll(this.userAssignment);

        this.needsPartitionAssignment = false;
        this.needsFetchCommittedOffsets = true;
    }

    /**
     * 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 (StreamPartition tp : assignments)
            if (!this.subscription.contains(tp.stream()) && !this.subscription.contains(tp.getStreamId()))
                throw new IllegalArgumentException("Assigned partition " + tp + " for non-subscribed topic.");
        this.assignment.clear();
        for (StreamPartition tp : assignments)
            addAssignedPartition(tp);
        this.needsPartitionAssignment = false;
    }

    public void subscribe(Pattern pattern, DisConsumerRebalanceListener listener) {
        if (listener == null)
            throw new IllegalArgumentException("RebalanceListener cannot be null");

        setSubscriptionType(SubscriptionType.AUTO_PATTERN);

        this.listener = listener;
        this.subscribedPattern = pattern;
    }

    public boolean hasPatternSubscription() {
        return this.subscriptionType == SubscriptionType.AUTO_PATTERN;
    }

    public void unsubscribe() {
        this.subscription.clear();
        this.userAssignment.clear();
        this.assignment.clear();
        this.needsPartitionAssignment = true;
        this.subscribedPattern = null;
        this.subscriptionType = SubscriptionType.NONE;
    }


    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 StreamPartition tp = entry.getKey();
            final StreamPartitionState 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 StreamPartitionState assignedState(StreamPartition tp) {
        StreamPartitionState state = this.assignment.get(tp);
        if (state == null)
            throw new IllegalStateException("No current assignment for partition " + tp);
        return state;
    }

    public void committed(StreamPartition tp, DisOffsetAndMetadata offset) {
        assignedState(tp).committed(offset);
    }

    public DisOffsetAndMetadata committed(StreamPartition 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(StreamPartition 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.subscriptionType == SubscriptionType.AUTO_TOPICS || this.subscriptionType == SubscriptionType.AUTO_PATTERN;
    }

    public void position(StreamPartition tp, long offset) {
        assignedState(tp).position(offset);
    }

    public Long position(StreamPartition tp) {
        return assignedState(tp).position;
    }

    public Map allConsumed() {
        Map allConsumed = new HashMap<>();
        for (Map.Entry entry : assignment.entrySet()) {
            StreamPartitionState state = entry.getValue();
            if (state.hasValidPosition())
                allConsumed.put(entry.getKey(), new DisOffsetAndMetadata(state.position));
        }
        return allConsumed;
    }

    public void needOffsetReset(StreamPartition partition, DisOffsetResetStrategy disOffsetResetStrategy) {
        assignedState(partition).awaitReset(disOffsetResetStrategy);
    }

    public void needOffsetReset(StreamPartition partition) {
        needOffsetReset(partition, defaultResetStrategy);
    }

    public boolean hasDefaultOffsetResetPolicy() {
        return defaultResetStrategy != DisOffsetResetStrategy.NONE;
    }

    public boolean isOffsetResetNeeded(StreamPartition partition) {
        return assignedState(partition).awaitingReset();
    }

    public DisOffsetResetStrategy resetStrategy(StreamPartition partition) {
        return assignedState(partition).resetStrategy;
    }

    public boolean hasAllFetchPositions() {
        for (StreamPartitionState 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 Set getSubscriptionIds() {
        return subscriptionIds;
    }

    public Set getSubscriptionStreams() {
        return subscription;
    }

    public boolean partitionAssignmentNeeded() {
        return this.needsPartitionAssignment;
    }

    public boolean isAssigned(StreamPartition tp) {
        return assignment.containsKey(tp);
    }

    public boolean isPaused(StreamPartition tp) {
        return isAssigned(tp) && assignedState(tp).paused;
    }

    public boolean isFetchable(StreamPartition tp) {
        return isAssigned(tp) && assignedState(tp).isFetchable();
    }

    public void pause(StreamPartition tp) {
        assignedState(tp).pause();
    }

    public void resume(StreamPartition tp) {
        assignedState(tp).resume();
    }

    private void addAssignedPartition(StreamPartition tp) {
        this.assignment.put(tp, new StreamPartitionState());
    }

    public DisConsumerRebalanceListener listener() {
        return listener;
    }

    private static class StreamPartitionState {
        private Long position; // last consumed position
        private DisOffsetAndMetadata committed;  // last committed position
        private boolean paused;  // whether this partition has been paused by the user
        private DisOffsetResetStrategy resetStrategy;  // the strategy to use if the offset needs resetting

        public StreamPartitionState() {
            this.paused = false;
            this.position = null;
            this.committed = null;
            this.resetStrategy = null;
        }

        private void awaitReset(DisOffsetResetStrategy 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(DisOffsetAndMetadata offset) {
            this.committed = offset;
        }

        private void pause() {
            this.paused = true;
        }

        private void resume() {
            this.paused = false;
        }

        private boolean isFetchable() {
            return !paused && hasValidPosition();
        }

    }

}