/*
 * 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.streams.processor.internals;

import org.apache.kafka.clients.admin.Admin;
import org.apache.kafka.clients.admin.ListOffsetsResult;
import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.Configurable;
import org.apache.kafka.common.KafkaException;
import org.apache.kafka.common.Node;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.TimeoutException;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.streams.errors.MissingSourceTopicException;
import org.apache.kafka.streams.errors.StreamsException;
import org.apache.kafka.streams.errors.TaskAssignmentException;
import org.apache.kafka.streams.processor.TaskId;
import org.apache.kafka.streams.processor.assignment.ApplicationState;
import org.apache.kafka.streams.processor.assignment.AssignmentConfigs;
import org.apache.kafka.streams.processor.assignment.ProcessId;
import org.apache.kafka.streams.processor.assignment.TaskAssignmentUtils;
import org.apache.kafka.streams.processor.assignment.TaskAssignor.AssignmentError;
import org.apache.kafka.streams.processor.assignment.TaskAssignor.TaskAssignment;
import org.apache.kafka.streams.processor.assignment.TaskInfo;
import org.apache.kafka.streams.processor.assignment.TaskTopicPartition;
import org.apache.kafka.streams.processor.internals.InternalTopologyBuilder.TopicsInfo;
import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
import org.apache.kafka.streams.processor.internals.assignment.AssignmentInfo;
import org.apache.kafka.streams.processor.internals.assignment.AssignorConfiguration;
import org.apache.kafka.streams.processor.internals.assignment.AssignorConfiguration.AssignmentListener;
import org.apache.kafka.streams.processor.internals.assignment.AssignorError;
import org.apache.kafka.streams.processor.internals.assignment.ClientState;
import org.apache.kafka.streams.processor.internals.assignment.CopartitionedTopicsEnforcer;
import org.apache.kafka.streams.processor.internals.assignment.DefaultApplicationState;
import org.apache.kafka.streams.processor.internals.assignment.DefaultTaskInfo;
import org.apache.kafka.streams.processor.internals.assignment.DefaultTaskTopicPartition;
import org.apache.kafka.streams.processor.internals.assignment.FallbackPriorTaskAssignor;
import org.apache.kafka.streams.processor.internals.assignment.LegacyStickyTaskAssignor;
import org.apache.kafka.streams.processor.internals.assignment.LegacyTaskAssignor;
import org.apache.kafka.streams.processor.internals.assignment.RackAwareTaskAssignor;
import org.apache.kafka.streams.processor.internals.assignment.RackUtils;
import org.apache.kafka.streams.processor.internals.assignment.ReferenceContainer;
import org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo;
import org.apache.kafka.streams.state.HostInfo;

import org.slf4j.Logger;

import java.nio.ByteBuffer;
import java.time.Instant;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.PriorityQueue;
import java.util.Queue;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

import static java.util.Collections.singletonList;
import static java.util.Collections.unmodifiableSet;
import static java.util.Map.Entry.comparingByKey;
import static org.apache.kafka.common.utils.Utils.filterMap;
import static org.apache.kafka.streams.processor.internals.ClientUtils.fetchCommittedOffsets;
import static org.apache.kafka.streams.processor.internals.ClientUtils.fetchEndOffsetsResult;
import static org.apache.kafka.streams.processor.internals.assignment.StreamsAssignmentProtocolVersions.EARLIEST_PROBEABLE_VERSION;
import static org.apache.kafka.streams.processor.internals.assignment.StreamsAssignmentProtocolVersions.LATEST_SUPPORTED_VERSION;
import static org.apache.kafka.streams.processor.internals.assignment.StreamsAssignmentProtocolVersions.UNKNOWN;
import static org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo.UNKNOWN_OFFSET_SUM;

public class StreamsPartitionAssignor implements ConsumerPartitionAssignor, Configurable {

    private Logger log;
    private String logPrefix;

    private static class AssignedPartition implements Comparable<AssignedPartition> {

        private final TaskId taskId;
        private final TopicPartition partition;

        AssignedPartition(final TaskId taskId, final TopicPartition partition) {
            this.taskId = taskId;
            this.partition = partition;
        }

        @Override
        public int compareTo(final AssignedPartition that) {
            return PARTITION_COMPARATOR.compare(partition, that.partition);
        }

        @Override
        public boolean equals(final Object o) {
            if (!(o instanceof AssignedPartition)) {
                return false;
            }
            final AssignedPartition other = (AssignedPartition) o;
            return compareTo(other) == 0;
        }

        @Override
        public int hashCode() {
            // Only partition is important for compareTo, equals and hashCode.
            return partition.hashCode();
        }
    }

    public static class ClientMetadata {

        private final HostInfo hostInfo;
        private final ClientState state;
        private final SortedSet<String> consumers;
        private final Optional<String> rackId;

        ClientMetadata(final ProcessId processId, final String endPoint, final Map<String, String> clientTags, final Optional<String> rackId) {

            // get the host info, or null if no endpoint is configured (ie endPoint == null)
            hostInfo = HostInfo.buildFromEndpoint(endPoint);

            // initialize the consumer memberIds
            consumers = new TreeSet<>();

            // initialize the client state with client tags
            state = new ClientState(processId, clientTags);

            this.rackId = rackId;
        }

        void addConsumer(final String consumerMemberId, final List<TopicPartition> ownedPartitions) {
            consumers.add(consumerMemberId);
            state.incrementCapacity();
            state.addOwnedPartitions(ownedPartitions, consumerMemberId);
        }

        void addPreviousTasksAndOffsetSums(final String consumerId, final Map<TaskId, Long> taskOffsetSums) {
            state.addPreviousTasksAndOffsetSums(consumerId, taskOffsetSums);
        }

        public ClientState state() {
            return state;
        }

        public HostInfo hostInfo() {
            return hostInfo;
        }

        public Optional<String> rackId() {
            return rackId;
        }

        @Override
        public String toString() {
            return "ClientMetadata{" +
                "hostInfo=" + hostInfo +
                ", consumers=" + consumers +
                ", state=" + state +
                '}';
        }
    }

    @FunctionalInterface
    public interface UserTaskAssignmentListener {
        void onAssignmentComputed(GroupAssignment assignment, GroupSubscription subscription);
    }

    // keep track of any future consumers in a "dummy" Client since we can't decipher their subscription
    private static final ProcessId FUTURE_ID = ProcessId.randomProcessId();

    protected static final Comparator<TopicPartition> PARTITION_COMPARATOR =
        Comparator.comparing(TopicPartition::topic).thenComparingInt(TopicPartition::partition);

    private String userEndPoint;
    private AssignmentConfigs assignmentConfigs;

    // for the main consumer, we need to use a supplier to break a cyclic setup dependency
    private Supplier<Consumer<byte[], byte[]>> mainConsumerSupplier;
    private Admin adminClient;
    private TaskManager taskManager;
    private StreamsMetadataState streamsMetadataState;
    private PartitionGrouper partitionGrouper;
    private AtomicInteger assignmentErrorCode;
    private AtomicLong nextScheduledRebalanceMs;
    private Queue<StreamsException> nonFatalExceptionsToHandle;
    private Time time;

    // since live upgrades from 2.3 (or earlier) to 4.0 or above are no longer supported, we can always
    // start with the latest supported metadata version since version probing will take
    // care of downgrading it if/when necessary
    protected int usedSubscriptionMetadataVersion = LATEST_SUPPORTED_VERSION;

    private InternalTopicManager internalTopicManager;
    private CopartitionedTopicsEnforcer copartitionedTopicsEnforcer;
    private AssignmentListener assignmentListener;

    private Supplier<Optional<org.apache.kafka.streams.processor.assignment.TaskAssignor>>
        customTaskAssignorSupplier;
    private Supplier<LegacyTaskAssignor> legacyTaskAssignorSupplier;
    private byte uniqueField;
    private Map<String, String> clientTags;

    /**
     * We need to have the PartitionAssignor and its StreamThread to be mutually accessible since the former needs
     * latter's cached metadata while sending subscriptions, and the latter needs former's returned assignment when
     * adding tasks.
     *
     * @throws KafkaException if the stream thread is not specified
     */
    @Override
    public void configure(final Map<String, ?> configs) {
        final AssignorConfiguration assignorConfiguration = new AssignorConfiguration(configs);

        logPrefix = assignorConfiguration.logPrefix();
        log = new LogContext(logPrefix).logger(getClass());

        final ReferenceContainer referenceContainer = assignorConfiguration.referenceContainer();
        mainConsumerSupplier = () -> Objects.requireNonNull(referenceContainer.mainConsumer, "Main consumer was not specified");
        adminClient = Objects.requireNonNull(referenceContainer.adminClient, "Admin client was not specified");
        taskManager = Objects.requireNonNull(referenceContainer.taskManager, "TaskManager was not specified");
        streamsMetadataState = Objects.requireNonNull(referenceContainer.streamsMetadataState, "StreamsMetadataState was not specified");
        assignmentErrorCode = referenceContainer.assignmentErrorCode;
        nextScheduledRebalanceMs = referenceContainer.nextScheduledRebalanceMs;
        nonFatalExceptionsToHandle = referenceContainer.nonFatalExceptionsToHandle;
        time = Objects.requireNonNull(referenceContainer.time, "Time was not specified");
        assignmentConfigs = assignorConfiguration.assignmentConfigs();
        partitionGrouper = new PartitionGrouper();
        userEndPoint = assignorConfiguration.userEndPoint();
        internalTopicManager = assignorConfiguration.internalTopicManager();
        copartitionedTopicsEnforcer = assignorConfiguration.copartitionedTopicsEnforcer();
        customTaskAssignorSupplier = assignorConfiguration::customTaskAssignor;
        legacyTaskAssignorSupplier = assignorConfiguration::taskAssignor;
        assignmentListener = assignorConfiguration.assignmentListener();
        uniqueField = 0;
        clientTags = referenceContainer.clientTags;
    }

    @Override
    public String name() {
        return "stream";
    }

    @Override
    public List<RebalanceProtocol> supportedProtocols() {
        return singletonList(RebalanceProtocol.COOPERATIVE);
    }

    @Override
    public ByteBuffer subscriptionUserData(final Set<String> topics) {
        // Adds the following information to subscription
        // 1. Client ProcessId (a UUID assigned to an instance of KafkaStreams)
        // 2. Map from task id to its overall lag
        // 3. Unique Field to ensure a rebalance when a thread rejoins by forcing the user data to be different

        handleRebalanceStart(topics);
        uniqueField++;

        final Set<String> currentNamedTopologies = taskManager.topologyMetadata().namedTopologiesView();

        // If using NamedTopologies, filter out any that are no longer recognized/have been removed
        final Map<TaskId, Long> taskOffsetSums = taskManager.topologyMetadata().hasNamedTopologies() ?
            filterMap(taskManager.taskOffsetSums(), t -> currentNamedTopologies.contains(t.getKey().topologyName())) :
            taskManager.taskOffsetSums();

        return new SubscriptionInfo(
            usedSubscriptionMetadataVersion,
            LATEST_SUPPORTED_VERSION,
            taskManager.processId(),
            userEndPoint,
            taskOffsetSums,
            uniqueField,
            assignmentErrorCode.get(),
            clientTags
        ).encode();
    }

    private Map<String, Assignment> errorAssignment(final Map<ProcessId, ClientMetadata> clientsMetadata,
                                                    final int errorCode) {
        final Map<String, Assignment> assignment = new HashMap<>();
        for (final ClientMetadata clientMetadata : clientsMetadata.values()) {
            for (final String consumerId : clientMetadata.consumers) {
                assignment.put(consumerId, new Assignment(
                    Collections.emptyList(),
                    new AssignmentInfo(LATEST_SUPPORTED_VERSION,
                        Collections.emptyList(),
                        Collections.emptyMap(),
                        Collections.emptyMap(),
                        Collections.emptyMap(),
                        errorCode).encode()
                ));
            }
        }
        return assignment;
    }

    /*
     * This assigns tasks to consumer clients in the following steps.
     *
     * 0. decode the subscriptions to assemble the metadata for each client and check for version probing
     *
     * 1. check all repartition source topics and use internal topic manager to make sure
     *    they have been created with the right number of partitions. Also verify and/or create
     *    any changelog topics with the correct number of partitions.
     *
     * 2. use the partition grouper to generate tasks along with their assigned partitions, then use
     *    the configured TaskAssignor to construct the mapping of tasks to clients.
     *
     * 3. construct the global mapping of host to partitions to enable query routing.
     *
     * 4. within each client, assign tasks to consumer clients.
     */
    @Override
    public GroupAssignment assign(final Cluster metadata, final GroupSubscription groupSubscription) {
        final Map<String, Subscription> subscriptions = groupSubscription.groupSubscription();

        // ---------------- Step Zero ---------------- //

        // construct the client metadata from the decoded subscription info

        final Map<ProcessId, ClientMetadata> clientMetadataMap = new HashMap<>();
        final Set<TopicPartition> allOwnedPartitions = new HashSet<>();
        final Map<ProcessId, Map<String, Optional<String>>> racksForProcessConsumer = new HashMap<>();

        int minReceivedMetadataVersion = LATEST_SUPPORTED_VERSION;
        int minSupportedMetadataVersion = LATEST_SUPPORTED_VERSION;

        boolean shutdownRequested = false;
        boolean assignmentErrorFound = false;
        int futureMetadataVersion = UNKNOWN;
        for (final Map.Entry<String, Subscription> entry : subscriptions.entrySet()) {
            final String consumerId = entry.getKey();
            final Subscription subscription = entry.getValue();
            final SubscriptionInfo info = SubscriptionInfo.decode(subscription.userData());
            final int usedVersion = info.version();
            if (info.errorCode() == AssignorError.SHUTDOWN_REQUESTED.code()) {
                shutdownRequested = true;
            }

            minReceivedMetadataVersion = updateMinReceivedVersion(usedVersion, minReceivedMetadataVersion);
            minSupportedMetadataVersion = updateMinSupportedVersion(info.latestSupportedVersion(), minSupportedMetadataVersion);

            final ProcessId processId;
            if (usedVersion > LATEST_SUPPORTED_VERSION) {
                futureMetadataVersion = usedVersion;
                processId = FUTURE_ID;
                if (!clientMetadataMap.containsKey(FUTURE_ID)) {
                    clientMetadataMap.put(FUTURE_ID, new ClientMetadata(FUTURE_ID, null, Collections.emptyMap(), subscription.rackId()));
                }
            } else {
                processId = info.processId();
            }

            racksForProcessConsumer.computeIfAbsent(processId, kv -> new HashMap<>()).put(consumerId, subscription.rackId());

            ClientMetadata clientMetadata = clientMetadataMap.get(processId);

            // create the new client metadata if necessary
            if (clientMetadata == null) {
                clientMetadata = new ClientMetadata(info.processId(), info.userEndPoint(), info.clientTags(), subscription.rackId());
                clientMetadataMap.put(info.processId(), clientMetadata);
            }

            // add the consumer and any info in its subscription to the client
            clientMetadata.addConsumer(consumerId, subscription.ownedPartitions());
            final int prevSize = allOwnedPartitions.size();
            allOwnedPartitions.addAll(subscription.ownedPartitions());
            if (allOwnedPartitions.size() < prevSize + subscription.ownedPartitions().size()) {
                assignmentErrorFound = true;
            }
            clientMetadata.addPreviousTasksAndOffsetSums(consumerId, info.taskOffsetSums());
        }

        if (assignmentErrorFound) {
            log.warn("The previous assignment contains a partition more than once. " +
                "\t Mapping: {}", subscriptions);
        }

        try {
            log.debug("Constructed client metadata {} from the member subscriptions.", clientMetadataMap);

            // ---------------- Step One ---------------- //

            if (shutdownRequested) {
                return new GroupAssignment(errorAssignment(clientMetadataMap, AssignorError.SHUTDOWN_REQUESTED.code()));
            }

            // parse the topology to determine the repartition source topics,
            // making sure they are created with the number of partitions as
            // the maximum of the depending sub-topologies source topics' number of partitions
            final RepartitionTopics repartitionTopics = prepareRepartitionTopics(metadata);
            final Map<TopicPartition, PartitionInfo> allRepartitionTopicPartitions = repartitionTopics.topicPartitionsInfo();

            final Cluster fullMetadata = metadata.withPartitions(allRepartitionTopicPartitions);
            log.debug("Created repartition topics {} from the parsed topology.", allRepartitionTopicPartitions.values());

            // ---------------- Step Two ---------------- //

            // construct the assignment of tasks to clients

            final Map<Subtopology, TopicsInfo> topicGroups =
                taskManager.topologyMetadata().subtopologyTopicsInfoMapExcluding(repartitionTopics.topologiesWithMissingInputTopics());

            final Set<String> allSourceTopics = new HashSet<>();
            final Map<Subtopology, Set<String>> sourceTopicsByGroup = new HashMap<>();
            for (final Map.Entry<Subtopology, TopicsInfo> entry : topicGroups.entrySet()) {
                allSourceTopics.addAll(entry.getValue().sourceTopics);
                sourceTopicsByGroup.put(entry.getKey(), entry.getValue().sourceTopics);
            }

            // get the tasks as partition groups from the partition grouper
            final Map<TaskId, Set<TopicPartition>> partitionsForTask =
                partitionGrouper.partitionGroups(sourceTopicsByGroup, fullMetadata);

            final Set<TaskId> statefulTasks = new HashSet<>();

            final boolean versionProbing =
                checkMetadataVersions(minReceivedMetadataVersion, minSupportedMetadataVersion, futureMetadataVersion);
            final UserTaskAssignmentListener userTaskAssignmentListener = assignTasksToClients(fullMetadata, groupSubscription,
                allSourceTopics, topicGroups, clientMetadataMap, partitionsForTask, racksForProcessConsumer, statefulTasks);

            // ---------------- Step Three ---------------- //

            // construct the global partition assignment per host map

            final Map<HostInfo, Set<TopicPartition>> partitionsByHost = new HashMap<>();
            final Map<HostInfo, Set<TopicPartition>> standbyPartitionsByHost = new HashMap<>();
            if (minReceivedMetadataVersion >= 2) {
                populatePartitionsByHostMaps(partitionsByHost, standbyPartitionsByHost, partitionsForTask, clientMetadataMap);
            }

            // ---------------- Step Four ---------------- //

            // compute the assignment of tasks to threads within each client and build the final group assignment
            final Map<String, Assignment> assignment = computeNewAssignment(
                statefulTasks,
                clientMetadataMap,
                partitionsForTask,
                partitionsByHost,
                standbyPartitionsByHost,
                allOwnedPartitions,
                minReceivedMetadataVersion,
                minSupportedMetadataVersion,
                versionProbing
            );

            final GroupAssignment groupAssignment = new GroupAssignment(assignment);
            userTaskAssignmentListener.onAssignmentComputed(groupAssignment, groupSubscription);
            return groupAssignment;
        } catch (final MissingSourceTopicException e) {
            log.error("Caught an error in the task assignment. Returning an error assignment.", e);
            return new GroupAssignment(
                errorAssignment(clientMetadataMap, AssignorError.INCOMPLETE_SOURCE_TOPIC_METADATA.code())
            );
        } catch (final TaskAssignmentException e) {
            log.error("Caught an error in the task assignment. Returning an error assignment.", e);
            return new GroupAssignment(
                errorAssignment(clientMetadataMap, AssignorError.ASSIGNMENT_ERROR.code())
            );
        }
    }

    /**
     *
     * @param clientMetadataMap the map of process id to client metadata used to build an immutable
     *                          {@code ApplicationState}
     * @return The {@code ApplicationState} needed by the TaskAssigner to compute new task
     *         assignments.
     */
    private ApplicationState buildApplicationState(final TopologyMetadata topologyMetadata,
                                                   final Map<ProcessId, ClientMetadata> clientMetadataMap,
                                                   final Map<Subtopology, TopicsInfo> topicGroups,
                                                   final Cluster cluster) {
        final Map<Subtopology, Set<String>> sourceTopicsByGroup = new HashMap<>();
        final Map<Subtopology, Set<String>> changelogTopicsByGroup = new HashMap<>();
        for (final Map.Entry<Subtopology, TopicsInfo> entry : topicGroups.entrySet()) {
            final Set<String> sourceTopics = entry.getValue().sourceTopics;
            final Set<String> changelogTopics = entry.getValue().changelogTopics();
            sourceTopicsByGroup.put(entry.getKey(), sourceTopics);
            changelogTopicsByGroup.put(entry.getKey(), changelogTopics);
        }

        final Map<TaskId, Set<TopicPartition>> changelogPartitionsForTask = new HashMap<>();
        final Map<TaskId, Set<TopicPartition>> sourcePartitionsForTask =
            partitionGrouper.partitionGroups(sourceTopicsByGroup, changelogTopicsByGroup,
                changelogPartitionsForTask, cluster);

        if (!sourcePartitionsForTask.keySet().equals(changelogPartitionsForTask.keySet())) {
            log.error("Partition grouper returned {} tasks for source topics but {} tasks for changelog topics",
                sourcePartitionsForTask.size(), changelogPartitionsForTask.size());
            throw new TaskAssignmentException("Partition grouper returned conflicting information about the "
                                              + "tasks for source topics vs changelog topics.");
        }

        final Set<DefaultTaskTopicPartition> topicsRequiringRackInfo = new HashSet<>();
        final AtomicBoolean rackInformationFetched = new AtomicBoolean(false);
        final Runnable fetchRackInformation = () -> {
            if (!rackInformationFetched.get()) {
                RackUtils.annotateTopicPartitionsWithRackInfo(cluster, internalTopicManager, topicsRequiringRackInfo);
                rackInformationFetched.set(true);
            }
        };

        final Map<TaskId, Set<TaskTopicPartition>> topicPartitionsForTask = new HashMap<>();
        final Set<TaskId> logicalTaskIds = unmodifiableSet(sourcePartitionsForTask.keySet());
        logicalTaskIds.forEach(taskId -> {
            final Set<TaskTopicPartition> topicPartitions = new HashSet<>();

            for (final TopicPartition topicPartition : sourcePartitionsForTask.get(taskId)) {
                final boolean isSource = true;
                final boolean isChangelog = changelogPartitionsForTask.get(taskId).contains(topicPartition);
                final DefaultTaskTopicPartition racklessTopicPartition = new DefaultTaskTopicPartition(
                    topicPartition, isSource, isChangelog, fetchRackInformation);
                topicsRequiringRackInfo.add(racklessTopicPartition);
                topicPartitions.add(racklessTopicPartition);
            }

            for (final TopicPartition topicPartition : changelogPartitionsForTask.get(taskId)) {
                final boolean isSource = sourcePartitionsForTask.get(taskId).contains(topicPartition);
                final boolean isChangelog = true;
                final DefaultTaskTopicPartition racklessTopicPartition = new DefaultTaskTopicPartition(
                    topicPartition, isSource, isChangelog, fetchRackInformation);
                topicsRequiringRackInfo.add(racklessTopicPartition);
                topicPartitions.add(racklessTopicPartition);
            }

            topicPartitionsForTask.put(taskId, topicPartitions);
        });

        final Map<TaskId, TaskInfo> logicalTasks = logicalTaskIds.stream().collect(Collectors.toMap(
            Function.identity(),
            taskId -> {
                final Set<String> stateStoreNames = topologyMetadata
                    .stateStoreNamesForSubtopology(taskId.topologyName(), taskId.subtopology());
                final Set<TaskTopicPartition> topicPartitions = topicPartitionsForTask.get(taskId);
                return new DefaultTaskInfo(
                    taskId,
                    !stateStoreNames.isEmpty(),
                    stateStoreNames,
                    topicPartitions
                );
            }
        ));

        return new DefaultApplicationState(
            assignmentConfigs,
            logicalTasks,
            clientMetadataMap
        );
    }

    private void processStreamsPartitionAssignment(final org.apache.kafka.streams.processor.assignment.TaskAssignor assignor,
                                                   final TaskAssignment taskAssignment,
                                                   final AssignmentError assignmentError,
                                                   final Map<ProcessId, ClientMetadata> clientMetadataMap,
                                                   final GroupSubscription groupSubscription) {
        if (assignmentError == AssignmentError.UNKNOWN_PROCESS_ID || assignmentError == AssignmentError.UNKNOWN_TASK_ID) {
            assignor.onAssignmentComputed(new GroupAssignment(Collections.emptyMap()), groupSubscription, assignmentError);
            log.error("Rebalance failed due to task assignor returning assignment with error {}, " +
                      "assignor callback will receive empty GroupAssignment due to this error", assignmentError);
            throw new StreamsException("Task assignment with " + assignor.getClass().getName() +
                                       " returned a fatal error: " + assignmentError);
        }

        taskAssignment.assignment().forEach(kafkaStreamsAssignment -> {
            final ProcessId processId = kafkaStreamsAssignment.processId();
            final ClientMetadata clientMetadata = clientMetadataMap.get(processId);
            clientMetadata.state.setAssignedTasks(kafkaStreamsAssignment);
            if (kafkaStreamsAssignment.followupRebalanceDeadline().isPresent()) {
                clientMetadata.state.setFollowupRebalanceDeadline(kafkaStreamsAssignment.followupRebalanceDeadline().get());
            }
        });
    }

    /**
     * Verify the subscription versions are within the expected bounds and check for version probing.
     *
     * @return whether this was a version probing rebalance
     */
    private boolean checkMetadataVersions(final int minReceivedMetadataVersion,
                                          final int minSupportedMetadataVersion,
                                          final int futureMetadataVersion) {
        final boolean versionProbing;

        if (futureMetadataVersion == UNKNOWN) {
            versionProbing = false;
        } else if (minReceivedMetadataVersion >= EARLIEST_PROBEABLE_VERSION) {
            versionProbing = true;
            log.info("Received a future (version probing) subscription (version: {})."
                         + " Sending assignment back (with supported version {}).",
                futureMetadataVersion,
                minSupportedMetadataVersion);

        } else {
            throw new TaskAssignmentException(
                "Received a future (version probing) subscription (version: " + futureMetadataVersion
                    + ") and an incompatible pre Kafka 2.0 subscription (version: " + minReceivedMetadataVersion
                    + ") at the same time."
            );
        }

        if (minReceivedMetadataVersion < LATEST_SUPPORTED_VERSION) {
            log.info("Downgrade metadata to version {}. Latest supported version is {}.",
                minReceivedMetadataVersion,
                LATEST_SUPPORTED_VERSION);
        }
        if (minSupportedMetadataVersion < LATEST_SUPPORTED_VERSION) {
            log.info("Downgrade latest supported metadata to version {}. Latest supported version is {}.",
                minSupportedMetadataVersion,
                LATEST_SUPPORTED_VERSION);
        }
        return versionProbing;
    }

    /**
     * Computes and assembles all repartition topic metadata then creates the topics if necessary. Also verifies
     * that all user input topics of each topology have been created ahead of time. If any such source topics are
     * missing from a NamedTopology, the assignor will skip distributing its tasks until they have been created
     * and invoke the exception handler (without killing the thread) once for each topology to alert the user of
     * the missing topics.
     * <p>
     * For regular applications without named topologies, the assignor will instead send a shutdown signal to
     * all clients so the user can identify and resolve the problem.
     *
     * @return application metadata such as partition info of repartition topics, missing external topics, etc
     */
    private RepartitionTopics prepareRepartitionTopics(final Cluster metadata) {
        final RepartitionTopics repartitionTopics = new RepartitionTopics(
            taskManager.topologyMetadata(),
            internalTopicManager,
            copartitionedTopicsEnforcer,
            metadata,
            logPrefix
        );
        repartitionTopics.setup();
        final boolean isMissingInputTopics = !repartitionTopics.missingSourceTopicExceptions().isEmpty();
        if (isMissingInputTopics) {
            if (!taskManager.topologyMetadata().hasNamedTopologies()) {
                final String errorMsg = String.format("Missing source topics. %s", repartitionTopics.missingSourceTopics());
                log.error(errorMsg);
                throw new MissingSourceTopicException(errorMsg);
            } else {
                nonFatalExceptionsToHandle.addAll(repartitionTopics.missingSourceTopicExceptions());
            }
        }
        return repartitionTopics;
    }


    /**
     * Populates the taskForPartition and tasksForTopicGroup maps, and checks that partitions are assigned to exactly
     * one task.
     *
     * @param taskForPartition a map from partition to the corresponding task. Populated here.
     * @param tasksForTopicGroup a map from the topicGroupId to the set of corresponding tasks. Populated here.
     * @param allSourceTopics a set of all source topics in the topology
     * @param partitionsForTask a map from task to the set of input partitions
     * @param fullMetadata the cluster metadata
     */
    private void populateTasksForMaps(final Map<TopicPartition, TaskId> taskForPartition,
                                      final Map<Subtopology, Set<TaskId>> tasksForTopicGroup,
                                      final Set<String> allSourceTopics,
                                      final Map<TaskId, Set<TopicPartition>> partitionsForTask,
                                      final Cluster fullMetadata) {
        // check if all partitions are assigned, and there are no duplicates of partitions in multiple tasks
        final Set<TopicPartition> allAssignedPartitions = new HashSet<>();
        for (final Map.Entry<TaskId, Set<TopicPartition>> entry : partitionsForTask.entrySet()) {
            final TaskId id = entry.getKey();
            final Set<TopicPartition> partitions = entry.getValue();

            for (final TopicPartition partition : partitions) {
                taskForPartition.put(partition, id);
                if (allAssignedPartitions.contains(partition)) {
                    log.warn("Partition {} is assigned to more than one tasks: {}", partition, partitionsForTask);
                }
            }
            allAssignedPartitions.addAll(partitions);

            tasksForTopicGroup.computeIfAbsent(new Subtopology(id.subtopology(), id.topologyName()), k -> new HashSet<>()).add(id);
        }

        checkAllPartitions(allSourceTopics, partitionsForTask, allAssignedPartitions, fullMetadata);
    }

    // Logs a warning if any partitions are not assigned to a task, or a task has no assigned partitions
    private void checkAllPartitions(final Set<String> allSourceTopics,
                                    final Map<TaskId, Set<TopicPartition>> partitionsForTask,
                                    final Set<TopicPartition> allAssignedPartitions,
                                    final Cluster fullMetadata) {
        for (final String topic : allSourceTopics) {
            final List<PartitionInfo> partitionInfoList = fullMetadata.partitionsForTopic(topic);
            if (partitionInfoList.isEmpty()) {
                log.warn("No partitions found for topic {}", topic);
            } else {
                for (final PartitionInfo partitionInfo : partitionInfoList) {
                    final TopicPartition partition = new TopicPartition(partitionInfo.topic(),
                        partitionInfo.partition());
                    if (!allAssignedPartitions.contains(partition)) {
                        log.warn("Partition {} is not assigned to any tasks: {}"
                                     + " Possible causes of a partition not getting assigned"
                                     + " is that another topic defined in the topology has not been"
                                     + " created when starting your streams application,"
                                     + " resulting in no tasks created for this topology at all.", partition,
                            partitionsForTask);
                    }
                }
            }
        }
    }

    /**
     * Assigns a set of tasks to each client (Streams instance) using the configured task assignor, and also
     * populate the stateful tasks that have been assigned to the clients
     */
    private UserTaskAssignmentListener assignTasksToClients(final Cluster fullMetadata,
                                                            final GroupSubscription groupSubscription,
                                                            final Set<String> allSourceTopics,
                                                            final Map<Subtopology, TopicsInfo> topicGroups,
                                                            final Map<ProcessId, ClientMetadata> clientMetadataMap,
                                                            final Map<TaskId, Set<TopicPartition>> partitionsForTask,
                                                            final Map<ProcessId, Map<String, Optional<String>>> racksForProcessConsumer,
                                                            final Set<TaskId> statefulTasks) {
        if (!statefulTasks.isEmpty()) {
            throw new TaskAssignmentException("The stateful tasks should not be populated before assigning tasks to clients");
        }

        final Map<TopicPartition, TaskId> taskForPartition = new HashMap<>();
        final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = new HashMap<>();
        populateTasksForMaps(taskForPartition, tasksForTopicGroup, allSourceTopics, partitionsForTask, fullMetadata);

        final ChangelogTopics changelogTopics = new ChangelogTopics(
            internalTopicManager,
            topicGroups,
            tasksForTopicGroup,
            logPrefix
        );
        changelogTopics.setup();

        final Map<ProcessId, ClientState> clientStates = new HashMap<>();
        final boolean lagComputationSuccessful =
            populateClientStatesMap(clientStates, clientMetadataMap, taskForPartition, changelogTopics);


        log.info("{} client nodes and {} consumers participating in this rebalance: \n{}.",
                 clientStates.size(),
                 clientStates.values().stream().map(ClientState::capacity).reduce(Integer::sum).orElse(0),
                 clientStates.entrySet().stream()
                     .sorted(comparingByKey())
                     .map(entry -> entry.getKey() + ": " + entry.getValue().consumers())
                     .collect(Collectors.joining(Utils.NL)));

        final Set<TaskId> allTasks = partitionsForTask.keySet();
        statefulTasks.addAll(changelogTopics.statefulTaskIds());

        log.info("Assigning stateful tasks: {}\n"
                     + "and stateless tasks: {}",
                 statefulTasks,
                 allTasks.stream().filter(t -> !statefulTasks.contains(t)).collect(Collectors.toSet()));
        log.debug("Assigning tasks and {} standby replicas to client nodes {}",
                  numStandbyReplicas(), clientStates);

        final Optional<org.apache.kafka.streams.processor.assignment.TaskAssignor> userTaskAssignor =
            customTaskAssignorSupplier.get();
        final UserTaskAssignmentListener customTaskAssignmentListener;
        if (userTaskAssignor.isPresent()) {
            final ApplicationState applicationState = buildApplicationState(
                taskManager.topologyMetadata(),
                clientMetadataMap,
                topicGroups,
                fullMetadata
            );
            final org.apache.kafka.streams.processor.assignment.TaskAssignor assignor = userTaskAssignor.get();
            final TaskAssignment taskAssignment = assignor.assign(applicationState);
            final AssignmentError assignmentError = TaskAssignmentUtils.validateTaskAssignment(applicationState, taskAssignment);
            processStreamsPartitionAssignment(assignor, taskAssignment, assignmentError, clientMetadataMap, groupSubscription);
            customTaskAssignmentListener = (assignment, subscription) -> {
                assignor.onAssignmentComputed(assignment, subscription, assignmentError);
                if (assignmentError != AssignmentError.NONE) {
                    log.error("Rebalance failed due to task assignor returning assignment with error {}", assignmentError);
                    throw new StreamsException("Task assignment with " + assignor.getClass().getName() +
                                               " returned an error: " + assignmentError);
                }
            };
        } else {
            customTaskAssignmentListener = (assignment, subscription) -> { };
            final LegacyTaskAssignor taskAssignor = createTaskAssignor(lagComputationSuccessful);
            final RackAwareTaskAssignor rackAwareTaskAssignor = new RackAwareTaskAssignor(
                fullMetadata,
                partitionsForTask,
                changelogTopics.changelogPartionsForTask(),
                tasksForTopicGroup,
                racksForProcessConsumer,
                internalTopicManager,
                assignmentConfigs,
                time
            );
            final boolean probingRebalanceNeeded = taskAssignor.assign(clientStates,
                allTasks,
                statefulTasks,
                rackAwareTaskAssignor,
                assignmentConfigs);
            if (probingRebalanceNeeded) {
                // Arbitrarily choose the leader's client to be responsible for triggering the probing rebalance,
                // note once we pick the first consumer within the process to trigger probing rebalance, other consumer
                // would not set to trigger any more.
                final ClientMetadata rebalanceClientMetadata = clientMetadataMap.get(taskManager.processId());
                if (rebalanceClientMetadata != null) {
                    final Instant rebalanceDeadline = Instant.ofEpochMilli(time.milliseconds() + probingRebalanceIntervalMs());
                    rebalanceClientMetadata.state.setFollowupRebalanceDeadline(rebalanceDeadline);
                }
            }
        }

        // Break this up into multiple logs to make sure the summary info gets through, which helps avoid
        // info loss for example due to long line truncation with large apps
        log.info("Assigned {} total tasks including {} stateful tasks to {} client nodes.",
                 allTasks.size(),
                 statefulTasks.size(),
                 clientStates.size());
        log.info("Assignment of tasks to nodes: {}",
                 clientStates.entrySet().stream()
                     .sorted(comparingByKey())
                     .map(entry -> entry.getKey() + "=" + entry.getValue().currentAssignment())
                     .collect(Collectors.joining(Utils.NL)));
        return customTaskAssignmentListener;
    }

    private LegacyTaskAssignor createTaskAssignor(final boolean lagComputationSuccessful) {
        final LegacyTaskAssignor taskAssignor = legacyTaskAssignorSupplier.get();
        if (taskAssignor instanceof LegacyStickyTaskAssignor) {
            // special case: to preserve pre-existing behavior, we invoke the LegacyStickyTaskAssignor
            // whether or not lag computation failed.
            return taskAssignor;
        } else if (lagComputationSuccessful) {
            return taskAssignor;
        } else {
            log.info("Failed to fetch end offsets for changelogs, will return previous assignment to clients and "
                         + "trigger another rebalance to retry.");
            return new FallbackPriorTaskAssignor();
        }
    }

    /**
     * Builds a map from client to state, and readies each ClientState for assignment by adding any missing prev tasks
     * and computing the per-task overall lag based on the fetched end offsets for each changelog.
     *
     * @param clientStates a map from each client to its state, including offset lags. Populated by this method.
     * @param clientMetadataMap a map from each client to its full metadata
     * @param taskForPartition map from topic partition to its corresponding task
     * @param changelogTopics object that manages changelog topics
     *
     * @return whether we were able to successfully fetch the changelog end offsets and compute each client's lag
     */
    private boolean populateClientStatesMap(final Map<ProcessId, ClientState> clientStates,
                                            final Map<ProcessId, ClientMetadata> clientMetadataMap,
                                            final Map<TopicPartition, TaskId> taskForPartition,
                                            final ChangelogTopics changelogTopics) {
        boolean fetchEndOffsetsSuccessful;
        Map<TaskId, Long> allTaskEndOffsetSums;
        try {
            // Make the listOffsets request first so it can  fetch the offsets for non-source changelogs
            // asynchronously while we use the blocking Consumer#committed call to fetch source-changelog offsets;
            // note that we would need to wrap all exceptions as Streams exception with partition-level fine-grained
            // error messages
            final ListOffsetsResult endOffsetsResult =
                fetchEndOffsetsResult(changelogTopics.preExistingNonSourceTopicBasedPartitions(), adminClient);

            final Map<TopicPartition, Long> sourceChangelogEndOffsets =
                fetchCommittedOffsets(changelogTopics.preExistingSourceTopicBasedPartitions(), mainConsumerSupplier.get());

            final Map<TopicPartition, ListOffsetsResultInfo> endOffsets = ClientUtils.getEndOffsets(
                endOffsetsResult, changelogTopics.preExistingNonSourceTopicBasedPartitions());

            allTaskEndOffsetSums = computeEndOffsetSumsByTask(
                endOffsets,
                sourceChangelogEndOffsets,
                changelogTopics
            );
            fetchEndOffsetsSuccessful = true;
        } catch (final StreamsException | TimeoutException e) {
            log.info("Failed to retrieve all end offsets for changelogs, and hence could not calculate the per-task lag; " +
                "this is not a fatal error but would cause the assignor to fallback to a naive algorithm", e);
            allTaskEndOffsetSums = changelogTopics.statefulTaskIds().stream().collect(Collectors.toMap(t -> t, t -> UNKNOWN_OFFSET_SUM));
            fetchEndOffsetsSuccessful = false;
        }

        for (final Map.Entry<ProcessId, ClientMetadata> entry : clientMetadataMap.entrySet()) {
            final ProcessId processId = entry.getKey();
            final ClientState state = entry.getValue().state;
            state.initializePrevTasks(taskForPartition, taskManager.topologyMetadata().hasNamedTopologies());

            state.computeTaskLags(processId, allTaskEndOffsetSums);
            clientStates.put(processId, state);
        }

        return fetchEndOffsetsSuccessful;
    }

    /**
     * @param endOffsets the listOffsets result from the adminClient
     * @param sourceChangelogEndOffsets the end (committed) offsets of optimized source changelogs
     * @param changelogTopics object that manages changelog topics
     *
     * @return Map from stateful task to its total end offset summed across all changelog partitions
     */
    private Map<TaskId, Long> computeEndOffsetSumsByTask(final Map<TopicPartition, ListOffsetsResultInfo> endOffsets,
                                                         final Map<TopicPartition, Long> sourceChangelogEndOffsets,
                                                         final ChangelogTopics changelogTopics) {

        final Map<TaskId, Long> taskEndOffsetSums = new HashMap<>();
        for (final TaskId taskId : changelogTopics.statefulTaskIds()) {
            taskEndOffsetSums.put(taskId, 0L);
            for (final TopicPartition changelogPartition : changelogTopics.preExistingPartitionsFor(taskId)) {
                final long changelogPartitionEndOffset;
                if (sourceChangelogEndOffsets.containsKey(changelogPartition)) {
                    changelogPartitionEndOffset = sourceChangelogEndOffsets.get(changelogPartition);
                } else if (endOffsets.containsKey(changelogPartition)) {
                    changelogPartitionEndOffset = endOffsets.get(changelogPartition).offset();
                } else {
                    log.debug("Fetched offsets did not contain the changelog {} of task {}", changelogPartition, taskId);
                    throw new IllegalStateException("Could not get end offset for " + changelogPartition);
                }
                final long newEndOffsetSum = taskEndOffsetSums.get(taskId) + changelogPartitionEndOffset;
                if (newEndOffsetSum < 0) {
                    taskEndOffsetSums.put(taskId, Long.MAX_VALUE);
                    break;
                } else {
                    taskEndOffsetSums.put(taskId, newEndOffsetSum);
                }
            }
        }
        return taskEndOffsetSums;
    }

    /**
     * Populates the global partitionsByHost and standbyPartitionsByHost maps that are sent to each member
     *
     * @param partitionsByHost a map from host to the set of partitions hosted there. Populated here.
     * @param standbyPartitionsByHost a map from host to the set of standby partitions hosted there. Populated here.
     * @param partitionsForTask a map from task to its set of assigned partitions
     * @param clientMetadataMap a map from client to its metadata and state
     */
    private void populatePartitionsByHostMaps(final Map<HostInfo, Set<TopicPartition>> partitionsByHost,
                                              final Map<HostInfo, Set<TopicPartition>> standbyPartitionsByHost,
                                              final Map<TaskId, Set<TopicPartition>> partitionsForTask,
                                              final Map<ProcessId, ClientMetadata> clientMetadataMap) {
        for (final Map.Entry<ProcessId, ClientMetadata> entry : clientMetadataMap.entrySet()) {
            final HostInfo hostInfo = entry.getValue().hostInfo;

            // if application server is configured, also include host state map
            if (hostInfo != null) {
                final Set<TopicPartition> topicPartitions = new HashSet<>();
                final Set<TopicPartition> standbyPartitions = new HashSet<>();
                final ClientState state = entry.getValue().state;

                for (final TaskId id : state.activeTasks()) {
                    topicPartitions.addAll(partitionsForTask.get(id));
                }

                for (final TaskId id : state.standbyTasks()) {
                    standbyPartitions.addAll(partitionsForTask.get(id));
                }

                partitionsByHost.put(hostInfo, topicPartitions);
                standbyPartitionsByHost.put(hostInfo, standbyPartitions);
            }
        }
    }

    /**
     * Computes the assignment of tasks to threads within each client and assembles the final assignment to send out.
     *
     * @return the final assignment for each StreamThread consumer
     */
    private Map<String, Assignment> computeNewAssignment(final Set<TaskId> statefulTasks,
                                                         final Map<ProcessId, ClientMetadata> clientsMetadata,
                                                         final Map<TaskId, Set<TopicPartition>> partitionsForTask,
                                                         final Map<HostInfo, Set<TopicPartition>> partitionsByHostState,
                                                         final Map<HostInfo, Set<TopicPartition>> standbyPartitionsByHost,
                                                         final Set<TopicPartition> allOwnedPartitions,
                                                         final int minUserMetadataVersion,
                                                         final int minSupportedMetadataVersion,
                                                         final boolean versionProbing) {
        boolean rebalanceRequired = versionProbing;
        final Map<String, Assignment> assignment = new HashMap<>();

        // within the client, distribute tasks to its owned consumers
        for (final Map.Entry<ProcessId, ClientMetadata> clientEntry : clientsMetadata.entrySet()) {
            final ProcessId clientId = clientEntry.getKey();
            final ClientMetadata clientMetadata = clientEntry.getValue();
            final ClientState state = clientMetadata.state;
            final SortedSet<String> consumers = clientMetadata.consumers;
            final Map<String, Integer> threadTaskCounts = new HashMap<>();

            final Map<String, List<TaskId>> activeTaskStatefulAssignment = assignTasksToThreads(
                state.statefulActiveTasks(),
                true,
                consumers,
                state,
                threadTaskCounts
            );

            final Map<String, List<TaskId>> standbyTaskAssignment = assignTasksToThreads(
                state.standbyTasks(),
                true,
                consumers,
                state,
                threadTaskCounts
            );

            final Map<String, List<TaskId>> activeTaskStatelessAssignment = assignTasksToThreads(
                state.statelessActiveTasks(),
                false,
                consumers,
                state,
                threadTaskCounts
            );

            // Combine activeTaskStatefulAssignment and activeTaskStatelessAssignment together into
            // activeTaskStatelessAssignment
            final Map<String, List<TaskId>> activeTaskAssignment = activeTaskStatefulAssignment;
            for (final Map.Entry<String, List<TaskId>> threadEntry : activeTaskStatelessAssignment.entrySet()) {
                activeTaskAssignment.get(threadEntry.getKey()).addAll(threadEntry.getValue());
            }

            final boolean isNextProbingRebalanceEncoded = clientMetadata.state.followupRebalanceDeadline().isPresent();

            final boolean tasksRevoked = addClientAssignments(
                statefulTasks,
                assignment,
                clientMetadata,
                partitionsForTask,
                partitionsByHostState,
                standbyPartitionsByHost,
                allOwnedPartitions,
                activeTaskAssignment,
                standbyTaskAssignment,
                minUserMetadataVersion,
                minSupportedMetadataVersion
            );

            if (tasksRevoked || isNextProbingRebalanceEncoded) {
                rebalanceRequired = true;
                log.debug("Requested client {} to schedule a followup rebalance", clientId);
            }

            log.info("Client {} per-consumer assignment:\n" +
                "\tprev owned active {}\n" +
                "\tprev owned standby {}\n" +
                "\tassigned active {}\n" +
                "\trevoking active {}\n" +
                "\tassigned standby {}\n",
                clientId,
                clientMetadata.state.prevOwnedActiveTasksByConsumer(),
                clientMetadata.state.prevOwnedStandbyByConsumer(),
                clientMetadata.state.assignedActiveTasksByConsumer(),
                clientMetadata.state.revokingActiveTasksByConsumer(),
                clientMetadata.state.assignedStandbyTasksByConsumer());
        }

        if (rebalanceRequired) {
            assignmentListener.onAssignmentComplete(false);
            log.info("Finished unstable assignment of tasks, a followup rebalance will be scheduled.");
        } else {
            assignmentListener.onAssignmentComplete(true);
            log.info("Finished stable assignment of tasks, no followup rebalances required.");
        }

        return assignment;
    }

    /**
     * Adds the encoded assignment for each StreamThread consumer in the client to the overall assignment map
     * @return true if a followup rebalance will be required due to revoked tasks
     */
    private boolean addClientAssignments(final Set<TaskId> statefulTasks,
                                         final Map<String, Assignment> assignment,
                                         final ClientMetadata clientMetadata,
                                         final Map<TaskId, Set<TopicPartition>> partitionsForTask,
                                         final Map<HostInfo, Set<TopicPartition>> partitionsByHostState,
                                         final Map<HostInfo, Set<TopicPartition>> standbyPartitionsByHost,
                                         final Set<TopicPartition> allOwnedPartitions,
                                         final Map<String, List<TaskId>> activeTaskAssignments,
                                         final Map<String, List<TaskId>> standbyTaskAssignments,
                                         final int minUserMetadataVersion,
                                         final int minSupportedMetadataVersion) {
        boolean followupRebalanceRequiredForRevokedTasks = false;

        // We only want to encode a scheduled probing rebalance for a single member in this client
        final Optional<Instant> followupRebalanceDeadline = clientMetadata.state.followupRebalanceDeadline();
        boolean shouldEncodeProbingRebalance = followupRebalanceDeadline.isPresent();

        // Loop through the consumers and build their assignment
        for (final String consumer : clientMetadata.consumers) {
            final List<TaskId> activeTasksForConsumer = activeTaskAssignments.get(consumer);

            // These will be filled in by populateActiveTaskAndPartitionsLists below
            final List<TopicPartition> activePartitionsList = new ArrayList<>();
            final List<TaskId> assignedActiveList = new ArrayList<>();

            final Set<TaskId> activeTasksRemovedPendingRevokation = populateActiveTaskAndPartitionsLists(
                activePartitionsList,
                assignedActiveList,
                consumer,
                clientMetadata.state,
                activeTasksForConsumer,
                partitionsForTask,
                allOwnedPartitions
            );

            final Map<TaskId, Set<TopicPartition>> standbyTaskMap = buildStandbyTaskMap(
                    consumer,
                    standbyTaskAssignments.get(consumer),
                    activeTasksRemovedPendingRevokation,
                    statefulTasks,
                    partitionsForTask,
                    clientMetadata.state
                );

            final AssignmentInfo info = new AssignmentInfo(
                minUserMetadataVersion,
                minSupportedMetadataVersion,
                assignedActiveList,
                standbyTaskMap,
                partitionsByHostState,
                standbyPartitionsByHost,
                AssignorError.NONE.code()
            );

            if (!activeTasksRemovedPendingRevokation.isEmpty()) {
                // TODO: once KAFKA-10078 is resolved we can leave it to the client to trigger this rebalance
                log.info("Requesting followup rebalance be scheduled immediately by {} due to tasks changing ownership.", consumer);
                info.setNextRebalanceTime(0L);
                followupRebalanceRequiredForRevokedTasks = true;
                // Don't bother to schedule a probing rebalance if an immediate one is already scheduled
                shouldEncodeProbingRebalance = false;
            } else if (shouldEncodeProbingRebalance) {
                final long nextRebalanceTimeMs = followupRebalanceDeadline.get().toEpochMilli();
                log.info("Requesting followup rebalance be scheduled by {} for {} to probe for caught-up replica tasks.",
                        consumer, Utils.toLogDateTimeFormat(nextRebalanceTimeMs));
                info.setNextRebalanceTime(nextRebalanceTimeMs);
                shouldEncodeProbingRebalance = false;
            }

            assignment.put(
                consumer,
                new Assignment(
                    activePartitionsList,
                    info.encode()
                )
            );
        }
        return followupRebalanceRequiredForRevokedTasks;
    }

    /**
     * Populates the lists of active tasks and active task partitions for the consumer with a 1:1 mapping between them
     * such that the nth task corresponds to the nth partition in the list. This means tasks with multiple partitions
     * will be repeated in the list.
     */
    private Set<TaskId> populateActiveTaskAndPartitionsLists(final List<TopicPartition> activePartitionsList,
                                                             final List<TaskId> assignedActiveList,
                                                             final String consumer,
                                                             final ClientState clientState,
                                                             final List<TaskId> activeTasksForConsumer,
                                                             final Map<TaskId, Set<TopicPartition>> partitionsForTask,
                                                             final Set<TopicPartition> allOwnedPartitions) {
        final List<AssignedPartition> assignedPartitions = new ArrayList<>();
        final Set<TaskId> removedActiveTasks = new TreeSet<>();

        for (final TaskId taskId : activeTasksForConsumer) {
            // Populate the consumer for assigned tasks without considering revocation,
            // this is for debugging purposes only
            clientState.assignActiveToConsumer(taskId, consumer);

            final List<AssignedPartition> assignedPartitionsForTask = new ArrayList<>();
            for (final TopicPartition partition : partitionsForTask.get(taskId)) {
                final String oldOwner = clientState.previousOwnerForPartition(partition);
                final boolean newPartitionForConsumer = oldOwner == null || !oldOwner.equals(consumer);

                // If the partition is new to this consumer but is still owned by another, remove from the assignment
                // until it has been revoked and can safely be reassigned according to the COOPERATIVE protocol
                if (newPartitionForConsumer && allOwnedPartitions.contains(partition)) {
                    log.info(
                        "Removing task {} from {} active assignment until it is safely revoked in followup rebalance",
                        taskId,
                        consumer
                    );
                    removedActiveTasks.add(taskId);

                    clientState.revokeActiveFromConsumer(taskId, consumer);

                    // Clear the assigned partitions list for this task if any partition can not safely be assigned,
                    // so as not to encode a partial task
                    assignedPartitionsForTask.clear();

                    // This has no effect on the assignment, as we'll never consult the ClientState again, but
                    // it does perform a useful assertion that the task was actually assigned.
                    clientState.unassignActive(taskId);
                    break;
                } else {
                    assignedPartitionsForTask.add(new AssignedPartition(taskId, partition));
                }
            }
            // assignedPartitionsForTask will either contain all partitions for the task or be empty, so just add all
            assignedPartitions.addAll(assignedPartitionsForTask);
        }

        // Add one copy of a task for each corresponding partition, so the receiver can determine the task <-> tp mapping
        Collections.sort(assignedPartitions);
        for (final AssignedPartition partition : assignedPartitions) {
            assignedActiveList.add(partition.taskId);
            activePartitionsList.add(partition.partition);
        }
        return removedActiveTasks;
    }

    /**
     * @return map from task id to its assigned partitions for all standby tasks
     */
    private Map<TaskId, Set<TopicPartition>> buildStandbyTaskMap(final String consumer,
                                                                 final Iterable<TaskId> standbyTasks,
                                                                 final Iterable<TaskId> revokedTasks,
                                                                 final Set<TaskId> allStatefulTasks,
                                                                 final Map<TaskId, Set<TopicPartition>> partitionsForTask,
                                                                 final ClientState clientState) {
        final Map<TaskId, Set<TopicPartition>> standbyTaskMap = new HashMap<>();

        for (final TaskId task : standbyTasks) {
            clientState.assignStandbyToConsumer(task, consumer);
            standbyTaskMap.put(task, partitionsForTask.get(task));
        }

        for (final TaskId task : revokedTasks) {
            // If this task is stateful and already owned by the consumer, but can't (yet) be assigned as an active
            // task during this rebalance as it must be revoked from another consumer first, place a temporary
            // standby task here until it can receive the active task to avoid closing the state store (and losing
            // all of the accumulated state in the case of in-memory stores)
            if (clientState.previouslyOwnedStandby(task) && allStatefulTasks.contains(task)) {
                log.info("Adding removed stateful active task {} as a standby for {} until it is revoked and can "
                             + "be transitioned to active in a followup rebalance", task, consumer);

                // This has no effect on the assignment, as we'll never consult the ClientState again, but
                // it does perform a useful assertion that the it's legal to assign this task as a standby to this instance
                clientState.assignStandbyToConsumer(task, consumer);
                clientState.assignStandby(task);

                standbyTaskMap.put(task, partitionsForTask.get(task));
            }
        }
        return standbyTaskMap;
    }

    /**
     * Generate an assignment that tries to preserve thread-level stickiness for stateful tasks without violating
     * balance. The tasks are balanced across threads. Stateful tasks without previous owners will be interleaved by
     * group id to spread subtopologies across threads and further balance the workload.
     * Stateless tasks are simply spread across threads without taking into account previous ownership.
     * threadLoad is a map that keeps track of task load per thread across multiple calls so active and standby
     * tasks are evenly distributed
     */
    static Map<String, List<TaskId>> assignTasksToThreads(final Collection<TaskId> tasksToAssign,
                                                          final boolean isStateful,
                                                          final SortedSet<String> consumers,
                                                          final ClientState state,
                                                          final Map<String, Integer> threadLoad) {
        final Map<String, List<TaskId>> assignment = new HashMap<>();
        for (final String consumer : consumers) {
            assignment.put(consumer, new ArrayList<>());
        }

        final int totalTasks = threadLoad.values().stream().reduce(tasksToAssign.size(), Integer::sum);

        final int minTasksPerThread = (int) Math.floor(((double) totalTasks) / consumers.size());
        final PriorityQueue<TaskId> unassignedTasks = new PriorityQueue<>(tasksToAssign);

        final Queue<String> consumersToFill = new LinkedList<>();
        // keep track of tasks that we have to skip during the first pass in case we can reassign them later
        // using tree-map to make sure the iteration ordering over keys are preserved
        final Map<TaskId, String> unassignedTaskToPreviousOwner = new TreeMap<>();

        if (!unassignedTasks.isEmpty()) {
            // First assign tasks to previous owner, up to the min expected tasks/thread if these are stateful
            for (final String consumer : consumers) {
                final List<TaskId> threadAssignment = assignment.get(consumer);
                // The number of tasks we have to assign here to hit minTasksPerThread
                final int tasksTargetCount = minTasksPerThread - threadLoad.getOrDefault(consumer, 0);

                if (isStateful) {
                    for (final TaskId task : state.prevTasksByLag(consumer)) {
                        if (unassignedTasks.contains(task)) {
                            if (threadAssignment.size() < tasksTargetCount) {
                                threadAssignment.add(task);
                                unassignedTasks.remove(task);
                            } else {
                                unassignedTaskToPreviousOwner.put(task, consumer);
                            }
                        }
                    }
                }

                if (threadAssignment.size() < tasksTargetCount) {
                    consumersToFill.offer(consumer);
                }
            }

            // Next interleave remaining unassigned tasks amongst unfilled consumers
            while (!consumersToFill.isEmpty()) {
                final TaskId task = unassignedTasks.poll();
                if (task != null) {
                    final String consumer = consumersToFill.poll();
                    final List<TaskId> threadAssignment = assignment.get(consumer);
                    threadAssignment.add(task);
                    final int threadTaskCount = threadAssignment.size() + threadLoad.getOrDefault(consumer, 0);
                    if (threadTaskCount < minTasksPerThread) {
                        consumersToFill.offer(consumer);
                    }
                } else {
                    throw new TaskAssignmentException("Ran out of unassigned stateful tasks but some members were not at capacity");
                }
            }

            // At this point all consumers are at the min or min + 1 capacity.
            // The min + 1 case can occur for standbys where there's fewer standbys than consumers and after assigning
            // the active tasks some consumers already have min + 1 one tasks assigned.
            // The tasks still remaining should now be distributed over the consumers that are still at min capacity
            if (!unassignedTasks.isEmpty()) {
                for (final String consumer : consumers) {
                    final int taskCount = assignment.get(consumer).size() + threadLoad.getOrDefault(consumer, 0);
                    if (taskCount == minTasksPerThread) {
                        consumersToFill.add(consumer);
                    }
                }

                // Go over the tasks we skipped earlier and assign them to their previous owner when possible
                for (final Map.Entry<TaskId, String> taskEntry : unassignedTaskToPreviousOwner.entrySet()) {
                    final TaskId task = taskEntry.getKey();
                    final String consumer = taskEntry.getValue();
                    if (consumersToFill.contains(consumer) && unassignedTasks.contains(task)) {
                        assignment.get(consumer).add(task);
                        unassignedTasks.remove(task);
                        // Remove this consumer since we know it is now at minCapacity + 1
                        consumersToFill.remove(consumer);
                    }
                }

                // Now just distribute the remaining unassigned stateful tasks over the consumers still at min capacity
                for (final TaskId task : unassignedTasks) {
                    final String consumer = consumersToFill.poll();
                    final List<TaskId> threadAssignment = assignment.get(consumer);
                    threadAssignment.add(task);
                }
            }
        }
        // Update threadLoad
        for (final Map.Entry<String, List<TaskId>> taskEntry : assignment.entrySet()) {
            final String consumer = taskEntry.getKey();
            final int totalCount = threadLoad.getOrDefault(consumer, 0) + taskEntry.getValue().size();
            threadLoad.put(consumer, totalCount);
        }

        return assignment;
    }

    private void validateMetadataVersions(final int receivedAssignmentMetadataVersion,
                                          final int latestCommonlySupportedVersion) {

        if (receivedAssignmentMetadataVersion > usedSubscriptionMetadataVersion) {
            log.error("Leader sent back an assignment with version {} which was greater than our used version {}",
                receivedAssignmentMetadataVersion, usedSubscriptionMetadataVersion);
            throw new TaskAssignmentException(
                "Sent a version " + usedSubscriptionMetadataVersion
                    + " subscription but got an assignment with higher version "
                    + receivedAssignmentMetadataVersion + "."
            );
        }

        if (latestCommonlySupportedVersion > LATEST_SUPPORTED_VERSION) {
            log.error("Leader sent back assignment with commonly supported version {} that is greater than our "
                + "actual latest supported version {}", latestCommonlySupportedVersion, LATEST_SUPPORTED_VERSION);
            throw new TaskAssignmentException("Can't upgrade to metadata version greater than we support");
        }
    }

    // Returns true if subscription version was changed, indicating version probing and need to rebalance again
    protected boolean maybeUpdateSubscriptionVersion(final int receivedAssignmentMetadataVersion,
                                                     final int latestCommonlySupportedVersion) {
        if (receivedAssignmentMetadataVersion >= EARLIEST_PROBEABLE_VERSION) {
            // If the latest commonly supported version is now greater than our used version, this indicates we have just
            // completed the rolling upgrade and can now update our subscription version for the final rebalance
            if (latestCommonlySupportedVersion > usedSubscriptionMetadataVersion) {
                log.info(
                    "Sent a version {} subscription and group's latest commonly supported version is {} (successful "
                        +
                        "version probing and end of rolling upgrade). Upgrading subscription metadata version to " +
                        "{} for next rebalance.",
                    usedSubscriptionMetadataVersion,
                    latestCommonlySupportedVersion,
                    latestCommonlySupportedVersion
                );
                usedSubscriptionMetadataVersion = latestCommonlySupportedVersion;
                return true;
            }

            // If we received a lower version than we sent, someone else in the group still hasn't upgraded. We
            // should downgrade our subscription until everyone is on the latest version
            if (receivedAssignmentMetadataVersion < usedSubscriptionMetadataVersion) {
                log.info(
                    "Sent a version {} subscription and got version {} assignment back (successful version probing). "
                        +
                        "Downgrade subscription metadata to commonly supported version {} and trigger new rebalance.",
                    usedSubscriptionMetadataVersion,
                    receivedAssignmentMetadataVersion,
                    latestCommonlySupportedVersion
                );
                usedSubscriptionMetadataVersion = latestCommonlySupportedVersion;
                return true;
            }
        } else {
            log.debug("Received an assignment version {} that is less than the earliest version that allows version " +
                "probing {}. If this is not during a rolling upgrade from version 2.0 or below, this is an error.",
                receivedAssignmentMetadataVersion, EARLIEST_PROBEABLE_VERSION);
        }

        return false;
    }

    @Override
    public void onAssignment(final Assignment assignment, final ConsumerGroupMetadata metadata) {
        final List<TopicPartition> partitions = new ArrayList<>(assignment.partitions());
        partitions.sort(PARTITION_COMPARATOR);

        final AssignmentInfo info = AssignmentInfo.decode(assignment.userData());
        if (info.errCode() != AssignorError.NONE.code()) {
            // set flag to shutdown streams app
            assignmentErrorCode.set(info.errCode());
            return;
        }
        /*
         * latestCommonlySupportedVersion belongs to [usedSubscriptionMetadataVersion, LATEST_SUPPORTED_VERSION]
         * receivedAssignmentMetadataVersion belongs to [EARLIEST_PROBEABLE_VERSION, usedSubscriptionMetadataVersion]
         *
         * usedSubscriptionMetadataVersion will be downgraded to receivedAssignmentMetadataVersion during a rolling
         * bounce upgrade with version probing.
         *
         * usedSubscriptionMetadataVersion will be upgraded to latestCommonlySupportedVersion when all members have
         * been bounced and it is safe to use the latest version.
         */
        final int receivedAssignmentMetadataVersion = info.version();
        final int latestCommonlySupportedVersion = info.commonlySupportedVersion();

        validateMetadataVersions(receivedAssignmentMetadataVersion, latestCommonlySupportedVersion);

        // version 1 field
        final Map<TaskId, Set<TopicPartition>> activeTasks;
        // version 2 fields
        final Map<TopicPartition, PartitionInfo> topicToPartitionInfo;
        final Map<HostInfo, Set<TopicPartition>> partitionsByHost;
        final Map<HostInfo, Set<TopicPartition>> standbyPartitionsByHost;
        final long encodedNextScheduledRebalanceMs;

        switch (receivedAssignmentMetadataVersion) {
            case 1:
                validateActiveTaskEncoding(partitions, info, logPrefix);

                activeTasks = getActiveTasks(partitions, info);
                partitionsByHost = Collections.emptyMap();
                standbyPartitionsByHost = Collections.emptyMap();
                topicToPartitionInfo = Collections.emptyMap();
                encodedNextScheduledRebalanceMs = Long.MAX_VALUE;
                break;
            case 2:
            case 3:
            case 4:
            case 5:
                validateActiveTaskEncoding(partitions, info, logPrefix);

                activeTasks = getActiveTasks(partitions, info);
                partitionsByHost = info.partitionsByHost();
                standbyPartitionsByHost = Collections.emptyMap();
                topicToPartitionInfo = getTopicPartitionInfo(partitionsByHost);
                encodedNextScheduledRebalanceMs = Long.MAX_VALUE;
                break;
            case 6:
                validateActiveTaskEncoding(partitions, info, logPrefix);

                activeTasks = getActiveTasks(partitions, info);
                partitionsByHost = info.partitionsByHost();
                standbyPartitionsByHost = info.standbyPartitionByHost();
                topicToPartitionInfo = getTopicPartitionInfo(partitionsByHost);
                encodedNextScheduledRebalanceMs = Long.MAX_VALUE;
                break;
            case 7:
            case 8:
            case 9:
            case 10:
            case 11:
                validateActiveTaskEncoding(partitions, info, logPrefix);

                activeTasks = getActiveTasks(partitions, info);
                partitionsByHost = info.partitionsByHost();
                standbyPartitionsByHost = info.standbyPartitionByHost();
                topicToPartitionInfo = getTopicPartitionInfo(partitionsByHost);
                encodedNextScheduledRebalanceMs = info.nextRebalanceMs();
                break;
            default:
                throw new IllegalStateException(
                    "This code should never be reached."
                        + " Please file a bug report at https://issues.apache.org/jira/projects/KAFKA/"
                );
        }

        maybeScheduleFollowupRebalance(
            encodedNextScheduledRebalanceMs,
            receivedAssignmentMetadataVersion,
            latestCommonlySupportedVersion,
            partitionsByHost.keySet()
        );

        streamsMetadataState.onChange(partitionsByHost, standbyPartitionsByHost, topicToPartitionInfo);

        // we do not capture any exceptions but just let the exception thrown from consumer.poll directly
        // since when stream thread captures it, either we close all tasks as dirty or we close thread
        taskManager.handleAssignment(activeTasks, info.standbyTasks());
    }

    private void maybeScheduleFollowupRebalance(final long encodedNextScheduledRebalanceMs,
                                                final int receivedAssignmentMetadataVersion,
                                                final int latestCommonlySupportedVersion,
                                                final Set<HostInfo> groupHostInfo) {
        if (maybeUpdateSubscriptionVersion(receivedAssignmentMetadataVersion, latestCommonlySupportedVersion)) {
            log.info("Requested to schedule immediate rebalance due to version probing.");
            nextScheduledRebalanceMs.set(0L);
        } else if (!verifyHostInfo(groupHostInfo)) {
            log.info("Requested to schedule immediate rebalance to update group with new host endpoint = {}.", userEndPoint);
            nextScheduledRebalanceMs.set(0L);
        } else if (encodedNextScheduledRebalanceMs == 0L) {
            log.info("Requested to schedule immediate rebalance for new tasks to be safely revoked from current owner.");
            nextScheduledRebalanceMs.set(0L);
        } else if (encodedNextScheduledRebalanceMs < Long.MAX_VALUE) {
            log.info(
                "Requested to schedule next probing rebalance at {} to try for a more balanced assignment.",
                Utils.toLogDateTimeFormat(encodedNextScheduledRebalanceMs)
            );
            nextScheduledRebalanceMs.set(encodedNextScheduledRebalanceMs);
        } else {
            log.info("No followup rebalance was requested, resetting the rebalance schedule.");
            nextScheduledRebalanceMs.set(Long.MAX_VALUE);
        }
    }

    /**
     * Verify that this client's host info was included in the map returned in the assignment, and trigger a
     * rebalance if not. This may be necessary when using static membership, as a rejoining client will be handed
     * back its original assignment to avoid an unnecessary rebalance. If the client's endpoint has changed, we need
     * to force a rebalance for the other members in the group to get the updated host info for this client.
     *
     * @param groupHostInfo the HostInfo of all clients in the group
     * @return false if the current host info does not match that in the group assignment
     */
    private boolean verifyHostInfo(final Set<HostInfo> groupHostInfo) {
        if (userEndPoint != null && !groupHostInfo.isEmpty()) {
            final HostInfo myHostInfo = HostInfo.buildFromEndpoint(userEndPoint);

            return groupHostInfo.contains(myHostInfo);
        } else {
            return true;
        }
    }

    // protected for upgrade test
    protected static Map<TaskId, Set<TopicPartition>> getActiveTasks(final List<TopicPartition> partitions, final AssignmentInfo info) {
        final Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>();
        for (int i = 0; i < partitions.size(); i++) {
            final TopicPartition partition = partitions.get(i);
            final TaskId id = info.activeTasks().get(i);
            activeTasks.computeIfAbsent(id, k1 -> new HashSet<>()).add(partition);
        }
        return activeTasks;
    }

    static Map<TopicPartition, PartitionInfo> getTopicPartitionInfo(final Map<HostInfo, Set<TopicPartition>> partitionsByHost) {
        final Map<TopicPartition, PartitionInfo> topicToPartitionInfo = new HashMap<>();
        for (final Set<TopicPartition> value : partitionsByHost.values()) {
            for (final TopicPartition topicPartition : value) {
                topicToPartitionInfo.put(
                    topicPartition,
                    new PartitionInfo(
                        topicPartition.topic(),
                        topicPartition.partition(),
                        null,
                        new Node[0],
                        new Node[0]
                    )
                );
            }
        }
        return topicToPartitionInfo;
    }

    private static void validateActiveTaskEncoding(final List<TopicPartition> partitions, final AssignmentInfo info, final String logPrefix) {
        // the number of assigned partitions should be the same as number of active tasks, which
        // could be duplicated if one task has more than one assigned partitions
        if (partitions.size() != info.activeTasks().size()) {
            throw new TaskAssignmentException(
                String.format(
                    "%sNumber of assigned partitions %d is not equal to "
                        + "the number of active taskIds %d, assignmentInfo=%s",
                    logPrefix, partitions.size(),
                    info.activeTasks().size(), info
                )
            );
        }
    }

    private int updateMinReceivedVersion(final int usedVersion, final int minReceivedMetadataVersion) {
        return Math.min(usedVersion, minReceivedMetadataVersion);
    }

    private int updateMinSupportedVersion(final int supportedVersion, final int minSupportedMetadataVersion) {
        if (supportedVersion < minSupportedMetadataVersion) {
            log.debug("Downgrade the current minimum supported version {} to the smaller seen supported version {}",
                minSupportedMetadataVersion, supportedVersion);
            return supportedVersion;
        } else {
            log.debug("Current minimum supported version remains at {}, last seen supported version was {}",
                minSupportedMetadataVersion, supportedVersion);
            return minSupportedMetadataVersion;
        }
    }

    // following functions are for test only
    void setInternalTopicManager(final InternalTopicManager internalTopicManager) {
        this.internalTopicManager = internalTopicManager;
    }

    protected String userEndPoint() {
        return userEndPoint;
    }

    protected TaskManager taskManager() {
        return taskManager;
    }

    protected byte uniqueField() {
        return uniqueField;
    }

    protected Map<String, String> clientTags() {
        return clientTags;
    }

    protected void handleRebalanceStart(final Set<String> topics) {
        taskManager.handleRebalanceStart(topics);
    }

    long acceptableRecoveryLag() {
        return assignmentConfigs.acceptableRecoveryLag();
    }

    int maxWarmupReplicas() {
        return assignmentConfigs.maxWarmupReplicas();
    }

    int numStandbyReplicas() {
        return assignmentConfigs.numStandbyReplicas();
    }

    long probingRebalanceIntervalMs() {
        return assignmentConfigs.probingRebalanceIntervalMs();
    }

}
