IncrementalCooperativeAssignor.java

/*
 * 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.connect.runtime.distributed;

import java.util.Arrays;
import java.util.Map.Entry;

import org.apache.kafka.common.utils.ExponentialBackoff;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.ConnectorsAndTasks;
import org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.WorkerLoad;
import org.apache.kafka.connect.util.ConnectUtils;
import org.apache.kafka.connect.storage.ClusterConfigState;
import org.apache.kafka.connect.util.ConnectorTaskId;
import org.slf4j.Logger;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.TreeSet;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import static org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember;
import static org.apache.kafka.connect.runtime.distributed.ConnectProtocol.Assignment;
import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG;
import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2;
import static org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.LeaderState;
import static org.apache.kafka.connect.util.ConnectUtils.combineCollections;
import static org.apache.kafka.connect.util.ConnectUtils.transformValues;

/**
 * An assignor that computes a distribution of connectors and tasks according to the incremental
 * cooperative strategy for rebalancing. Note that this class is NOT thread-safe.
 * @see <a href="https://cwiki.apache.org/confluence/display/KAFKA/KIP-415%3A+Incremental+Cooperative+Rebalancing+in+Kafka+Connect">
 * KIP-415 for a description of the assignment policy. </a>
 *
 */
public class IncrementalCooperativeAssignor implements ConnectAssignor {
    private final Logger log;
    private final Time time;
    private final int maxDelay;
    private ConnectorsAndTasks previousAssignment;
    private final ConnectorsAndTasks previousRevocation;
    private boolean revokedInPrevious;
    protected final Set<String> candidateWorkersForReassignment;
    protected long scheduledRebalance;
    protected int delay;
    protected int previousGenerationId;
    protected Set<String> previousMembers;

    private final ExponentialBackoff consecutiveRevokingRebalancesBackoff;

    private int numSuccessiveRevokingRebalances;

    public IncrementalCooperativeAssignor(LogContext logContext, Time time, int maxDelay) {
        this.log = logContext.logger(IncrementalCooperativeAssignor.class);
        this.time = time;
        this.maxDelay = maxDelay;
        this.previousAssignment = ConnectorsAndTasks.EMPTY;
        this.previousRevocation = new ConnectorsAndTasks.Builder().build();
        this.scheduledRebalance = 0;
        this.revokedInPrevious = false;
        this.candidateWorkersForReassignment = new LinkedHashSet<>();
        this.delay = 0;
        this.previousGenerationId = -1;
        this.previousMembers = Collections.emptySet();
        this.numSuccessiveRevokingRebalances = 0;
        // By default, initial interval is 1. The only corner case is when the user has set maxDelay to 0
        // in which case, the exponential backoff delay should be 0 which would return the backoff delay to be 0 always
        this.consecutiveRevokingRebalancesBackoff = new ExponentialBackoff(maxDelay == 0 ? 0 : 1, 40, maxDelay, 0);
    }

    @Override
    public Map<String, ByteBuffer> performAssignment(String leaderId, String protocol,
                                                     List<JoinGroupResponseMember> allMemberMetadata,
                                                     WorkerCoordinator coordinator) {
        log.debug("Performing task assignment");

        Map<String, ExtendedWorkerState> memberConfigs = new HashMap<>();
        for (JoinGroupResponseMember member : allMemberMetadata) {
            memberConfigs.put(
                    member.memberId(),
                    IncrementalCooperativeConnectProtocol.deserializeMetadata(ByteBuffer.wrap(member.metadata())));
        }
        log.debug("Member configs: {}", memberConfigs);

        // The new config offset is the maximum seen by any member. We always perform assignment using this offset,
        // even if some members have fallen behind. The config offset used to generate the assignment is included in
        // the response so members that have fallen behind will not use the assignment until they have caught up.
        long maxOffset = memberConfigs.values().stream().map(ExtendedWorkerState::offset).max(Long::compare).get();
        log.debug("Max config offset root: {}, local snapshot config offsets root: {}",
                  maxOffset, coordinator.configSnapshot().offset());

        short protocolVersion = ConnectProtocolCompatibility.fromProtocol(protocol).protocolVersion();

        Long leaderOffset = ensureLeaderConfig(maxOffset, coordinator);
        if (leaderOffset == null) {
            Map<String, ExtendedAssignment> assignments = fillAssignments(
                    memberConfigs.keySet(), Assignment.CONFIG_MISMATCH,
                    leaderId, memberConfigs.get(leaderId).url(), maxOffset,
                    ClusterAssignment.EMPTY, 0, protocolVersion);
            return serializeAssignments(assignments, protocolVersion);
        }
        return performTaskAssignment(leaderId, leaderOffset, memberConfigs, coordinator, protocolVersion);
    }

    private Long ensureLeaderConfig(long maxOffset, WorkerCoordinator coordinator) {
        // If this leader is behind some other members, we can't do assignment
        if (coordinator.configSnapshot().offset() < maxOffset) {
            // We might be able to take a new snapshot to catch up immediately and avoid another round of syncing here.
            // Alternatively, if this node has already passed the maximum reported by any other member of the group, it
            // is also safe to use this newer state.
            ClusterConfigState updatedSnapshot = coordinator.configFreshSnapshot();
            if (updatedSnapshot.offset() < maxOffset) {
                log.info("Was selected to perform assignments, but do not have latest config found in sync request. "
                         + "Returning an empty configuration to trigger re-sync.");
                return null;
            } else {
                coordinator.configSnapshot(updatedSnapshot);
                return updatedSnapshot.offset();
            }
        }
        return maxOffset;
    }

    /**
     * Performs task assignment based on the incremental cooperative connect protocol.
     * Read more on the design and implementation in:
     * <a href="https://cwiki.apache.org/confluence/display/KAFKA/KIP-415%3A+Incremental+Cooperative+Rebalancing+in+Kafka+Connect">
     * KIP-415</a>
     *
     * @param leaderId the ID of the group leader
     * @param maxOffset the latest known offset of the configuration topic
     * @param memberConfigs the metadata of all the members of the group as gather in the current
     * round of rebalancing
     * @param coordinator the worker coordinator instance that provide the configuration snapshot
     * and get assigned the leader state during this assignment
     * @param protocolVersion the Connect subprotocol version
     * @return the serialized assignment of tasks to the whole group, including assigned or
     * revoked tasks
     */
    protected Map<String, ByteBuffer> performTaskAssignment(String leaderId, long maxOffset,
                                                            Map<String, ExtendedWorkerState> memberConfigs,
                                                            WorkerCoordinator coordinator, short protocolVersion) {
        log.debug("Performing task assignment during generation: {} with memberId: {}",
                coordinator.generationId(), coordinator.memberId());
        Map<String, ConnectorsAndTasks> memberAssignments = transformValues(
                memberConfigs,
                memberConfig -> new ConnectorsAndTasks.Builder()
                        .with(memberConfig.assignment().connectors(), memberConfig.assignment().tasks())
                        .build()
        );
        ClusterAssignment clusterAssignment = performTaskAssignment(
                coordinator.configSnapshot(),
                coordinator.lastCompletedGenerationId(),
                coordinator.generationId(),
                memberAssignments
        );

        coordinator.leaderState(new LeaderState(memberConfigs, clusterAssignment.allAssignedConnectors(), clusterAssignment.allAssignedTasks()));

        Map<String, ExtendedAssignment> assignments =
                fillAssignments(memberConfigs.keySet(), Assignment.NO_ERROR, leaderId,
                        memberConfigs.get(leaderId).url(), maxOffset,
                        clusterAssignment,
                        delay, protocolVersion);

        log.debug("Actual assignments: {}", assignments);
        return serializeAssignments(assignments, protocolVersion);
    }

    // Visible for testing
    ClusterAssignment performTaskAssignment(
            ClusterConfigState configSnapshot,
            int lastCompletedGenerationId,
            int currentGenerationId,
            Map<String, ConnectorsAndTasks> memberAssignments
    ) {
        // Base set: The previous assignment of connectors-and-tasks is a standalone snapshot that
        // can be used to calculate derived sets
        log.debug("Previous assignments: {}", previousAssignment);
        if (previousGenerationId != lastCompletedGenerationId) {
            log.debug("Clearing the view of previous assignments due to generation mismatch between "
                    + "previous generation ID {} and last completed generation ID {}. This can "
                    + "happen if the leader fails to sync the assignment within a rebalancing round. "
                    + "The following view of previous assignments might be outdated and will be "
                    + "ignored by the leader in the current computation of new assignments. "
                    + "Possibly outdated previous assignments: {}",
                    previousGenerationId, lastCompletedGenerationId, previousAssignment);
            this.previousAssignment = ConnectorsAndTasks.EMPTY;
        }

        Set<String> configuredConnectors = new TreeSet<>(configSnapshot.connectors());
        Set<ConnectorTaskId> configuredTasks = combineCollections(configuredConnectors, configSnapshot::tasks, Collectors.toSet());

        // Base set: The set of configured connectors-and-tasks is a standalone snapshot that can
        // be used to calculate derived sets
        ConnectorsAndTasks configured = new ConnectorsAndTasks.Builder()
                .with(configuredConnectors, configuredTasks).build();
        log.debug("Configured assignments: {}", configured);

        // Base set: The set of active connectors-and-tasks is a standalone snapshot that can be
        // used to calculate derived sets
        ConnectorsAndTasks activeAssignments = assignment(memberAssignments);
        log.debug("Active assignments: {}", activeAssignments);

        // This means that a previous revocation did not take effect. In this case, reset
        // appropriately and be ready to re-apply revocation of tasks
        if (!previousRevocation.isEmpty()) {
            if (previousRevocation.connectors().stream().anyMatch(c -> activeAssignments.connectors().contains(c))
                    || previousRevocation.tasks().stream().anyMatch(t -> activeAssignments.tasks().contains(t))) {
                previousAssignment = activeAssignments;
            }
            previousRevocation.connectors().clear();
            previousRevocation.tasks().clear();
        }

        // Derived set: The set of deleted connectors-and-tasks is a derived set from the set
        // difference of previous - configured
        ConnectorsAndTasks deleted = diff(previousAssignment, configured);
        log.debug("Deleted assignments: {}", deleted);

        // The connectors and tasks that are currently running on more than one worker each
        ConnectorsAndTasks duplicated = duplicatedAssignments(memberAssignments);
        log.trace("Duplicated assignments: {}", duplicated);

        // Derived set: The set of lost or unaccounted connectors-and-tasks is a derived set from
        // the set difference of previous - active - deleted
        ConnectorsAndTasks lostAssignments = diff(previousAssignment, activeAssignments, deleted);
        log.debug("Lost assignments: {}", lostAssignments);

        // Derived set: The set of new connectors-and-tasks is a derived set from the set
        // difference of configured - previous - active
        ConnectorsAndTasks created = diff(configured, previousAssignment, activeAssignments);
        log.debug("Created: {}", created);

        // A collection of the current assignment excluding the connectors-and-tasks to be deleted
        List<WorkerLoad> currentWorkerAssignment = workerAssignment(memberAssignments, deleted);

        Map<String, ConnectorsAndTasks.Builder> toRevoke = new HashMap<>();

        Map<String, ConnectorsAndTasks> deletedToRevoke = intersection(deleted, memberAssignments);
        log.debug("Deleted connectors and tasks to revoke from each worker: {}", deletedToRevoke);
        addAll(toRevoke, deletedToRevoke);

        // Revoking redundant connectors/tasks if the workers have duplicate assignments
        Map<String, ConnectorsAndTasks> duplicatedToRevoke = intersection(duplicated, memberAssignments);
        log.debug("Duplicated connectors and tasks to revoke from each worker: {}", duplicatedToRevoke);
        addAll(toRevoke, duplicatedToRevoke);

        // Compute the assignment that will be applied across the cluster after this round of rebalance
        // Later on, new submissions and lost-and-reassigned connectors and tasks will be added to these assignments,
        // and load-balancing revocations will be removed from them.
        List<WorkerLoad> nextWorkerAssignment = workerLoads(memberAssignments);
        removeAll(nextWorkerAssignment, deletedToRevoke);
        removeAll(nextWorkerAssignment, duplicatedToRevoke);

        // Collect the lost assignments that are ready to be reassigned because the workers that were
        // originally responsible for them appear to have left the cluster instead of rejoining within
        // the scheduled rebalance delay. These assignments will be re-allocated to the existing workers
        // in the cluster later on
        ConnectorsAndTasks.Builder lostAssignmentsToReassignBuilder = new ConnectorsAndTasks.Builder();
        handleLostAssignments(lostAssignments, lostAssignmentsToReassignBuilder, nextWorkerAssignment);
        ConnectorsAndTasks lostAssignmentsToReassign = lostAssignmentsToReassignBuilder.build();

        // Do not revoke resources for re-assignment while a delayed rebalance is active
        if (delay == 0) {
            Map<String, ConnectorsAndTasks> loadBalancingRevocations =
                    performLoadBalancingRevocations(configured, nextWorkerAssignment);

            // If this round and the previous round involved revocation, we will calculate a delay for
            // the next round when revoking rebalance would be allowed. Note that delay could be 0, in which
            // case we would always revoke.
            if (revokedInPrevious && !loadBalancingRevocations.isEmpty()) {
                numSuccessiveRevokingRebalances++; // Should we consider overflow for this?
                log.debug("Consecutive revoking rebalances observed. Computing delay and next scheduled rebalance.");
                delay = (int) consecutiveRevokingRebalancesBackoff.backoff(numSuccessiveRevokingRebalances);
                if (delay != 0) {
                    scheduledRebalance = time.milliseconds() + delay;
                    log.debug("Skipping revocations in the current round with a delay of {}ms. Next scheduled rebalance:{}",
                            delay, scheduledRebalance);
                } else {
                    log.debug("Revoking assignments immediately since scheduled.rebalance.max.delay.ms is set to 0");
                    addAll(toRevoke, loadBalancingRevocations);
                    // Remove all newly-revoked connectors and tasks from the next assignment, both to
                    // ensure that they are not included in the assignments during this round, and to produce
                    // an accurate allocation of all newly-created and lost-and-reassigned connectors and tasks
                    // that will have to be distributed across the cluster during this round
                    removeAll(nextWorkerAssignment, loadBalancingRevocations);
                }
            } else if (!loadBalancingRevocations.isEmpty()) {
                // We had a revocation in this round but not in the previous round. Let's store that state.
                log.debug("Performing allocation-balancing revocation immediately as no revocations took place during the previous rebalance");
                addAll(toRevoke, loadBalancingRevocations);
                removeAll(nextWorkerAssignment, loadBalancingRevocations);
                revokedInPrevious = true;
            } else if (revokedInPrevious) {
                // No revocations in this round but the previous round had one. Probably the workers
                // have converged to a balanced load. We can reset the rebalance clock
                log.debug("Previous round had revocations but this round didn't. Probably, the cluster has reached a " +
                        "balanced load. Resetting the exponential backoff clock");
                revokedInPrevious = false;
                numSuccessiveRevokingRebalances = 0;
            } else {
                // no-op
                log.debug("No revocations in previous and current round.");
            }
        } else {
            log.debug("Delayed rebalance is active. Delaying {}ms before revoking connectors and tasks: {}", delay, toRevoke);
            revokedInPrevious = false;
        }

        // The complete set of connectors and tasks that should be newly-assigned during this round
        ConnectorsAndTasks toAssign = new ConnectorsAndTasks.Builder()
                .addAll(created)
                .addAll(lostAssignmentsToReassign)
                .build();

        assignConnectors(nextWorkerAssignment, toAssign.connectors());
        assignTasks(nextWorkerAssignment, toAssign.tasks());

        Map<String, Collection<String>> nextConnectorAssignments = nextWorkerAssignment.stream()
                .collect(Collectors.toMap(
                        WorkerLoad::worker,
                        WorkerLoad::connectors
                ));
        Map<String, Collection<ConnectorTaskId>> nextTaskAssignments = nextWorkerAssignment.stream()
                .collect(Collectors.toMap(
                        WorkerLoad::worker,
                        WorkerLoad::tasks
                ));

        Map<String, Collection<String>> currentConnectorAssignments =
                currentWorkerAssignment.stream().collect(Collectors.toMap(WorkerLoad::worker, WorkerLoad::connectors));
        Map<String, Collection<ConnectorTaskId>> currentTaskAssignments =
                currentWorkerAssignment.stream().collect(Collectors.toMap(WorkerLoad::worker, WorkerLoad::tasks));
        Map<String, Collection<String>> incrementalConnectorAssignments =
                diff(nextConnectorAssignments, currentConnectorAssignments);
        Map<String, Collection<ConnectorTaskId>> incrementalTaskAssignments =
                diff(nextTaskAssignments, currentTaskAssignments);

        Map<String, ConnectorsAndTasks> revoked = buildAll(toRevoke);

        previousAssignment = computePreviousAssignment(revoked, nextConnectorAssignments, nextTaskAssignments, lostAssignments);
        previousGenerationId = currentGenerationId;
        previousMembers = memberAssignments.keySet();

        log.debug("Incremental connector assignments: {}", incrementalConnectorAssignments);
        log.debug("Incremental task assignments: {}", incrementalTaskAssignments);

        Map<String, Collection<String>> revokedConnectors = transformValues(revoked, ConnectorsAndTasks::connectors);
        Map<String, Collection<ConnectorTaskId>> revokedTasks = transformValues(revoked, ConnectorsAndTasks::tasks);

        return new ClusterAssignment(
                incrementalConnectorAssignments,
                incrementalTaskAssignments,
                revokedConnectors,
                revokedTasks,
                diff(nextConnectorAssignments, revokedConnectors),
                diff(nextTaskAssignments, revokedTasks)
        );
    }

    private ConnectorsAndTasks computePreviousAssignment(Map<String, ConnectorsAndTasks> toRevoke,
                                                         Map<String, Collection<String>> connectorAssignments,
                                                         Map<String, Collection<ConnectorTaskId>> taskAssignments,
                                                         ConnectorsAndTasks lostAssignments) {
        ConnectorsAndTasks previousAssignment = new ConnectorsAndTasks.Builder().with(
                ConnectUtils.combineCollections(connectorAssignments.values()),
                ConnectUtils.combineCollections(taskAssignments.values())
        ).build();

        for (ConnectorsAndTasks revoked : toRevoke.values()) {
            previousAssignment.connectors().removeAll(revoked.connectors());
            previousAssignment.tasks().removeAll(revoked.tasks());
            previousRevocation.connectors().addAll(revoked.connectors());
            previousRevocation.tasks().addAll(revoked.tasks());
        }

        // Depends on the previous assignment's collections being sets at the moment.
        // TODO: make it independent
        previousAssignment.connectors().addAll(lostAssignments.connectors());
        previousAssignment.tasks().addAll(lostAssignments.tasks());

        return previousAssignment;
    }

    private ConnectorsAndTasks duplicatedAssignments(Map<String, ConnectorsAndTasks> memberAssignments) {
        Map<String, Long> connectorInstanceCounts = combineCollections(
                memberAssignments.values(),
                ConnectorsAndTasks::connectors,
                Collectors.groupingBy(Function.identity(), Collectors.counting())
        );
        Set<String> duplicatedConnectors = connectorInstanceCounts
                .entrySet().stream()
                .filter(entry -> entry.getValue() > 1L)
                .map(Entry::getKey)
                .collect(Collectors.toSet());

        Map<ConnectorTaskId, Long> taskInstanceCounts = combineCollections(
                memberAssignments.values(),
                ConnectorsAndTasks::tasks,
                Collectors.groupingBy(Function.identity(), Collectors.counting())
        );
        Set<ConnectorTaskId> duplicatedTasks = taskInstanceCounts
                .entrySet().stream()
                .filter(entry -> entry.getValue() > 1L)
                .map(Entry::getKey)
                .collect(Collectors.toSet());

        return new ConnectorsAndTasks.Builder().with(duplicatedConnectors, duplicatedTasks).build();
    }

    // visible for testing
    protected void handleLostAssignments(ConnectorsAndTasks lostAssignments,
                                         ConnectorsAndTasks.Builder lostAssignmentsToReassign,
                                         List<WorkerLoad> completeWorkerAssignment) {
        // There are no lost assignments and there have been no successive revoking rebalances
        if (lostAssignments.isEmpty() && !revokedInPrevious) {
            resetDelay();
            return;
        }

        final long now = time.milliseconds();
        log.debug("Found the following connectors and tasks missing from previous assignments: "
                + lostAssignments);

        Set<String> activeMembers = completeWorkerAssignment.stream()
                .map(WorkerLoad::worker)
                .collect(Collectors.toSet());
        if (scheduledRebalance <= 0 && activeMembers.containsAll(previousMembers)) {
            log.debug("No worker seems to have departed the group during the rebalance. The "
                    + "missing assignments that the leader is detecting are probably due to some "
                    + "workers failing to receive the new assignments in the previous rebalance. "
                    + "Will reassign missing tasks as new tasks");
            lostAssignmentsToReassign.addAll(lostAssignments);
            return;
        } else if (maxDelay == 0) {
            log.debug("Scheduled rebalance delays are disabled ({} = 0); "
                    + "reassigning all lost connectors and tasks immediately",
                    SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG
            );
            lostAssignmentsToReassign.addAll(lostAssignments);
            return;
        }

        if (scheduledRebalance > 0 && now >= scheduledRebalance) {
            // delayed rebalance expired and it's time to assign resources
            log.debug("Delayed rebalance expired. Reassigning lost tasks");
            List<WorkerLoad> candidateWorkerLoad = Collections.emptyList();
            if (!candidateWorkersForReassignment.isEmpty()) {
                candidateWorkerLoad = pickCandidateWorkerForReassignment(completeWorkerAssignment);
            }

            if (!candidateWorkerLoad.isEmpty()) {
                log.debug("Assigning lost tasks to {} candidate workers: {}", 
                        candidateWorkerLoad.size(),
                        candidateWorkerLoad.stream().map(WorkerLoad::worker).collect(Collectors.joining(",")));
                Iterator<WorkerLoad> candidateWorkerIterator = candidateWorkerLoad.iterator();
                for (String connector : lostAssignments.connectors()) {
                    // Loop over the candidate workers as many times as it takes
                    if (!candidateWorkerIterator.hasNext()) {
                        candidateWorkerIterator = candidateWorkerLoad.iterator();
                    }
                    WorkerLoad worker = candidateWorkerIterator.next();
                    log.debug("Assigning connector id {} to member {}", connector, worker.worker());
                    worker.assign(connector);
                }
                candidateWorkerIterator = candidateWorkerLoad.iterator();
                for (ConnectorTaskId task : lostAssignments.tasks()) {
                    if (!candidateWorkerIterator.hasNext()) {
                        candidateWorkerIterator = candidateWorkerLoad.iterator();
                    }
                    WorkerLoad worker = candidateWorkerIterator.next();
                    log.debug("Assigning task id {} to member {}", task, worker.worker());
                    worker.assign(task);
                }
            } else {
                log.debug("No single candidate worker was found to assign lost tasks. Treating lost tasks as new tasks");
                lostAssignmentsToReassign.addAll(lostAssignments);
            }
            resetDelay();
            // Resetting the flag as now we can permit successive revoking rebalances.
            // since we have gone through the full rebalance delay
            revokedInPrevious = false;
        } else {
            candidateWorkersForReassignment
                    .addAll(candidateWorkersForReassignment(completeWorkerAssignment));
            if (now < scheduledRebalance) {
                // a delayed rebalance is in progress, but it's not yet time to reassign
                // unaccounted resources
                delay = calculateDelay(now);
                log.debug("Delayed rebalance in progress. Task reassignment is postponed. New computed rebalance delay: {}", delay);
            } else {
                // This means scheduledRebalance == 0
                // We could also extract the current minimum delay from the group, to make
                // independent of consecutive leader failures, but this optimization is skipped
                // at the moment
                delay = maxDelay;
                log.debug("Resetting rebalance delay to the max: {}. scheduledRebalance: {} now: {} diff scheduledRebalance - now: {}",
                        delay, scheduledRebalance, now, scheduledRebalance - now);
            }
            scheduledRebalance = now + delay;
        }
    }

    private void resetDelay() {
        candidateWorkersForReassignment.clear();
        scheduledRebalance = 0;
        if (delay != 0) {
            log.debug("Resetting delay from previous value: {} to 0", delay);
        }
        delay = 0;
    }

    private Set<String> candidateWorkersForReassignment(List<WorkerLoad> completeWorkerAssignment) {
        return completeWorkerAssignment.stream()
                .filter(WorkerLoad::isEmpty)
                .map(WorkerLoad::worker)
                .collect(Collectors.toSet());
    }

    private List<WorkerLoad> pickCandidateWorkerForReassignment(List<WorkerLoad> completeWorkerAssignment) {
        Map<String, WorkerLoad> activeWorkers = completeWorkerAssignment.stream()
                .collect(Collectors.toMap(WorkerLoad::worker, Function.identity()));
        return candidateWorkersForReassignment.stream()
                .map(activeWorkers::get)
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
    }

    private Map<String, ExtendedAssignment> fillAssignments(Collection<String> members, short error,
                                                            String leaderId, String leaderUrl, long maxOffset,
                                                            ClusterAssignment clusterAssignment,
                                                            int delay, short protocolVersion) {
        Map<String, ExtendedAssignment> groupAssignment = new HashMap<>();
        for (String member : members) {
            Collection<String> connectorsToStart = clusterAssignment.newlyAssignedConnectors(member);
            Collection<ConnectorTaskId> tasksToStart = clusterAssignment.newlyAssignedTasks(member);
            Collection<String> connectorsToStop = clusterAssignment.newlyRevokedConnectors(member);
            Collection<ConnectorTaskId> tasksToStop = clusterAssignment.newlyRevokedTasks(member);
            ExtendedAssignment assignment =
                    new ExtendedAssignment(protocolVersion, error, leaderId, leaderUrl, maxOffset,
                            connectorsToStart, tasksToStart, connectorsToStop, tasksToStop, delay);
            log.debug("Filling assignment: {} -> {}", member, assignment);
            groupAssignment.put(member, assignment);
        }
        log.debug("Finished assignment");
        return groupAssignment;
    }

    /**
     * From a map of workers to assignment object generate the equivalent map of workers to byte
     * buffers of serialized assignments.
     *
     * @param assignments the map of worker assignments
     * @return the serialized map of assignments to workers
     */
    protected Map<String, ByteBuffer> serializeAssignments(Map<String, ExtendedAssignment> assignments, short protocolVersion) {
        boolean sessioned = protocolVersion >= CONNECT_PROTOCOL_V2;
        return assignments.entrySet()
                .stream()
                .collect(Collectors.toMap(
                    Map.Entry::getKey,
                    e -> IncrementalCooperativeConnectProtocol.serializeAssignment(e.getValue(), sessioned)));
    }

    private static ConnectorsAndTasks diff(ConnectorsAndTasks base,
                                           ConnectorsAndTasks... toSubtract) {
        Collection<String> connectors = new TreeSet<>(base.connectors());
        Collection<ConnectorTaskId> tasks = new TreeSet<>(base.tasks());
        for (ConnectorsAndTasks sub : toSubtract) {
            connectors.removeAll(sub.connectors());
            tasks.removeAll(sub.tasks());
        }
        return new ConnectorsAndTasks.Builder().with(connectors, tasks).build();
    }

    private static <T> Map<String, Collection<T>> diff(Map<String, Collection<T>> base,
                                                       Map<String, Collection<T>> toSubtract) {
        Map<String, Collection<T>> incremental = new HashMap<>();
        for (Map.Entry<String, Collection<T>> entry : base.entrySet()) {
            List<T> values = new ArrayList<>(entry.getValue());
            values.removeAll(toSubtract.getOrDefault(entry.getKey(), Collections.emptySet()));
            incremental.put(entry.getKey(), values);
        }
        return incremental;
    }

    private ConnectorsAndTasks assignment(Map<String, ConnectorsAndTasks> memberAssignments) {
        log.debug("Received assignments: {}", memberAssignments);
        return new ConnectorsAndTasks.Builder().with(
                ConnectUtils.combineCollections(memberAssignments.values(), ConnectorsAndTasks::connectors),
                ConnectUtils.combineCollections(memberAssignments.values(), ConnectorsAndTasks::tasks)
        ).build();
    }

    /**
     * Revoke connectors and tasks from each worker in the cluster until no worker is running more than it
     * would be with a perfectly-balanced assignment.
     * @param configured the set of configured connectors and tasks across the entire cluster
     * @param workers the workers in the cluster, whose assignments should not include any deleted or duplicated connectors or tasks
     *                that are already due to be revoked from the worker in this rebalance
     * @return which connectors and tasks should be revoked from which workers; never null, but may be empty
     * if no load-balancing revocations are necessary or possible
     */
    private Map<String, ConnectorsAndTasks> performLoadBalancingRevocations(
            ConnectorsAndTasks configured,
            Collection<WorkerLoad> workers
    ) {
        if (log.isTraceEnabled()) {
            workers.forEach(wl -> log.trace(
                    "Per worker current load size; worker: {} connectors: {} tasks: {}",
                    wl.worker(), wl.connectorsSize(), wl.tasksSize()));
        }

        if (workers.stream().allMatch(WorkerLoad::isEmpty)) {
            log.trace("No load-balancing revocations required; all workers are either new "
                    + "or will have all currently-assigned connectors and tasks revoked during this round"
            );
            return Collections.emptyMap();
        }
        if (configured.isEmpty()) {
            log.trace("No load-balancing revocations required; no connectors are currently configured on this cluster");
            return Collections.emptyMap();
        }

        Map<String, ConnectorsAndTasks.Builder> result = new HashMap<>();

        Map<String, Set<String>> connectorRevocations = loadBalancingRevocations(
                "connector",
                configured.connectors().size(),
                workers,
                WorkerLoad::connectors
        );
        Map<String, Set<ConnectorTaskId>> taskRevocations = loadBalancingRevocations(
                "task",
                configured.tasks().size(),
                workers,
                WorkerLoad::tasks
        );

        connectorRevocations.forEach((worker, revoked) ->
                result.computeIfAbsent(worker, w -> new ConnectorsAndTasks.Builder()).addConnectors(revoked)
        );
        taskRevocations.forEach((worker, revoked) ->
                result.computeIfAbsent(worker, w -> new ConnectorsAndTasks.Builder()).addTasks(revoked)
        );

        return buildAll(result);
    }

    private <E> Map<String, Set<E>> loadBalancingRevocations(
            String allocatedResourceName,
            int totalToAllocate,
            Collection<WorkerLoad> workers,
            Function<WorkerLoad, Collection<E>> workerAllocation
    ) {
        int totalWorkers = workers.size();
        // The minimum instances of this resource that should be assigned to each worker
        int minAllocatedPerWorker = totalToAllocate / totalWorkers;
        // How many workers are going to have to be allocated exactly one extra instance
        // (since the total number to allocate may not be a perfect multiple of the number of workers)
        int workersToAllocateExtra = totalToAllocate % totalWorkers;
        // Useful function to determine exactly how many instances of the resource a given worker is currently allocated
        Function<WorkerLoad, Integer> workerAllocationSize = workerAllocation.andThen(Collection::size);

        long workersAllocatedMinimum = workers.stream()
                .map(workerAllocationSize)
                .filter(n -> n == minAllocatedPerWorker)
                .count();
        long workersAllocatedSingleExtra = workers.stream()
                .map(workerAllocationSize)
                .filter(n -> n == minAllocatedPerWorker + 1)
                .count();
        if (workersAllocatedSingleExtra == workersToAllocateExtra
                && workersAllocatedMinimum + workersAllocatedSingleExtra == totalWorkers) {
            log.trace(
                    "No load-balancing {} revocations required; the current allocations, when combined with any newly-created {}s, should be balanced",
                    allocatedResourceName,
                    allocatedResourceName
            );
            return Collections.emptyMap();
        }

        Map<String, Set<E>> result = new HashMap<>();
        // How many workers we've allocated a single extra resource instance to
        int allocatedExtras = 0;
        // Calculate how many (and which) connectors/tasks to revoke from each worker here
        for (WorkerLoad worker : workers) {
            int currentAllocationSizeForWorker = workerAllocationSize.apply(worker);
            if (currentAllocationSizeForWorker <= minAllocatedPerWorker) {
                // This worker isn't allocated more than the minimum; no need to revoke anything
                continue;
            }
            int maxAllocationForWorker;
            if (allocatedExtras < workersToAllocateExtra) {
                // We'll allocate one of the extra resource instances to this worker
                allocatedExtras++;
                if (currentAllocationSizeForWorker == minAllocatedPerWorker + 1) {
                    // If the worker's running exactly one more than the minimum, and we're allowed to
                    // allocate an extra to it, there's no need to revoke anything
                    continue;
                }
                maxAllocationForWorker = minAllocatedPerWorker + 1;
            } else {
                maxAllocationForWorker = minAllocatedPerWorker;
            }

            Set<E> revokedFromWorker = new LinkedHashSet<>();
            result.put(worker.worker(), revokedFromWorker);

            Iterator<E> currentWorkerAllocation = workerAllocation.apply(worker).iterator();
            // Revoke resources from the worker until it isn't allocated any more than it should be
            for (int numRevoked = 0; currentAllocationSizeForWorker - numRevoked > maxAllocationForWorker; numRevoked++) {
                if (!currentWorkerAllocation.hasNext()) {
                    // Should never happen, but better to log a warning and move on than die and fail the whole rebalance if it does
                    log.warn(
                            "Unexpectedly ran out of {}s to revoke from worker {} while performing load-balancing revocations; " +
                                    "worker appears to still be allocated {} instances, which is more than the intended allocation of {}",
                            allocatedResourceName,
                            worker.worker(),
                            workerAllocationSize.apply(worker),
                            maxAllocationForWorker
                    );
                    break;
                }
                E revocation = currentWorkerAllocation.next();
                revokedFromWorker.add(revocation);
            }
        }
        return result;
    }

    private int calculateDelay(long now) {
        long diff = scheduledRebalance - now;
        return diff > 0 ? (int) Math.min(diff, maxDelay) : 0;
    }

    /**
     * Perform a round-robin assignment of connectors to workers with existing worker load. This
     * assignment tries to balance the load between workers, by assigning connectors to workers
     * that have equal load, starting with the least loaded workers.
     *
     * @param workerAssignment the current worker assignment; assigned connectors are added to this list
     * @param connectors the connectors to be assigned
     */
    protected void assignConnectors(List<WorkerLoad> workerAssignment, Collection<String> connectors) {
        workerAssignment.sort(WorkerLoad.connectorComparator());
        WorkerLoad first = workerAssignment.get(0);

        Iterator<String> load = connectors.iterator();
        while (load.hasNext()) {
            int firstLoad = first.connectorsSize();
            int upTo = IntStream.range(0, workerAssignment.size())
                    .filter(i -> workerAssignment.get(i).connectorsSize() > firstLoad)
                    .findFirst()
                    .orElse(workerAssignment.size());
            for (WorkerLoad worker : workerAssignment.subList(0, upTo)) {
                String connector = load.next();
                log.debug("Assigning connector {} to {}", connector, worker.worker());
                worker.assign(connector);
                if (!load.hasNext()) {
                    break;
                }
            }
        }
    }

    /**
     * Perform a round-robin assignment of tasks to workers with existing worker load. This
     * assignment tries to balance the load between workers, by assigning tasks to workers that
     * have equal load, starting with the least loaded workers.
     *
     * @param workerAssignment the current worker assignment; assigned tasks are added to this list
     * @param tasks the tasks to be assigned
     */
    protected void assignTasks(List<WorkerLoad> workerAssignment, Collection<ConnectorTaskId> tasks) {
        workerAssignment.sort(WorkerLoad.taskComparator());
        WorkerLoad first = workerAssignment.get(0);

        Iterator<ConnectorTaskId> load = tasks.iterator();
        while (load.hasNext()) {
            int firstLoad = first.tasksSize();
            int upTo = IntStream.range(0, workerAssignment.size())
                    .filter(i -> workerAssignment.get(i).tasksSize() > firstLoad)
                    .findFirst()
                    .orElse(workerAssignment.size());
            for (WorkerLoad worker : workerAssignment.subList(0, upTo)) {
                ConnectorTaskId task = load.next();
                log.debug("Assigning task {} to {}", task, worker.worker());
                worker.assign(task);
                if (!load.hasNext()) {
                    break;
                }
            }
        }
    }

    private static List<WorkerLoad> workerAssignment(Map<String, ConnectorsAndTasks> memberAssignments,
                                                     ConnectorsAndTasks toExclude) {
        ConnectorsAndTasks ignore = new ConnectorsAndTasks.Builder()
                .with(toExclude.connectors(), toExclude.tasks())
                .build();

        return memberAssignments.entrySet().stream()
                .map(e -> new WorkerLoad.Builder(e.getKey()).with(
                        e.getValue().connectors().stream()
                                .filter(v -> !ignore.connectors().contains(v))
                                .collect(Collectors.toList()),
                        e.getValue().tasks().stream()
                                .filter(v -> !ignore.tasks().contains(v))
                                .collect(Collectors.toList())
                        ).build()
                ).collect(Collectors.toList());
    }

    private static void addAll(Map<String, ConnectorsAndTasks.Builder> base, Map<String, ConnectorsAndTasks> toAdd) {
        toAdd.forEach((worker, assignment) -> base
                .computeIfAbsent(worker, w -> new ConnectorsAndTasks.Builder())
                .addAll(assignment)
        );
    }

    private static <K> Map<K, ConnectorsAndTasks> buildAll(Map<K, ConnectorsAndTasks.Builder> builders) {
        return transformValues(builders, ConnectorsAndTasks.Builder::build);
    }

    private static List<WorkerLoad> workerLoads(Map<String, ConnectorsAndTasks> memberAssignments) {
        return memberAssignments.entrySet().stream()
                .map(e -> new WorkerLoad.Builder(e.getKey()).with(e.getValue().connectors(), e.getValue().tasks()).build())
                .collect(Collectors.toList());
    }

    private static void removeAll(List<WorkerLoad> workerLoads, Map<String, ConnectorsAndTasks> toRemove) {
        workerLoads.forEach(workerLoad -> {
            String worker = workerLoad.worker();
            ConnectorsAndTasks toRemoveFromWorker = toRemove.getOrDefault(worker, ConnectorsAndTasks.EMPTY);
            workerLoad.connectors().removeAll(toRemoveFromWorker.connectors());
            workerLoad.tasks().removeAll(toRemoveFromWorker.tasks());
        });
    }

    private static Map<String, ConnectorsAndTasks> intersection(ConnectorsAndTasks connectorsAndTasks, Map<String, ConnectorsAndTasks> assignments) {
        return transformValues(assignments, assignment -> {
            Collection<String> connectors = new HashSet<>(assignment.connectors());
            connectors.retainAll(connectorsAndTasks.connectors());
            Collection<ConnectorTaskId> tasks = new HashSet<>(assignment.tasks());
            tasks.retainAll(connectorsAndTasks.tasks());
            return new ConnectorsAndTasks.Builder().with(connectors, tasks).build();
        });
    }

    static class ClusterAssignment {

        private final Map<String, Collection<String>> newlyAssignedConnectors;
        private final Map<String, Collection<ConnectorTaskId>> newlyAssignedTasks;
        private final Map<String, Collection<String>> newlyRevokedConnectors;
        private final Map<String, Collection<ConnectorTaskId>> newlyRevokedTasks;
        private final Map<String, Collection<String>> allAssignedConnectors;
        private final Map<String, Collection<ConnectorTaskId>> allAssignedTasks;
        private final Set<String> allWorkers;

        public static final ClusterAssignment EMPTY = new ClusterAssignment(
                Collections.emptyMap(),
                Collections.emptyMap(),
                Collections.emptyMap(),
                Collections.emptyMap(),
                Collections.emptyMap(),
                Collections.emptyMap()
        );

        public ClusterAssignment(
                Map<String, Collection<String>> newlyAssignedConnectors,
                Map<String, Collection<ConnectorTaskId>> newlyAssignedTasks,
                Map<String, Collection<String>> newlyRevokedConnectors,
                Map<String, Collection<ConnectorTaskId>> newlyRevokedTasks,
                Map<String, Collection<String>> allAssignedConnectors,
                Map<String, Collection<ConnectorTaskId>> allAssignedTasks
        ) {
            this.newlyAssignedConnectors = newlyAssignedConnectors;
            this.newlyAssignedTasks = newlyAssignedTasks;
            this.newlyRevokedConnectors = newlyRevokedConnectors;
            this.newlyRevokedTasks = newlyRevokedTasks;
            this.allAssignedConnectors = allAssignedConnectors;
            this.allAssignedTasks = allAssignedTasks;
            this.allWorkers = combineCollections(
                    Arrays.asList(newlyAssignedConnectors, newlyAssignedTasks, newlyRevokedConnectors, newlyRevokedTasks, allAssignedConnectors, allAssignedTasks),
                    Map::keySet,
                    Collectors.toSet()
            );
        }

        public Map<String, Collection<String>> newlyAssignedConnectors() {
            return newlyAssignedConnectors;
        }

        public Collection<String> newlyAssignedConnectors(String worker) {
            return newlyAssignedConnectors.getOrDefault(worker, Collections.emptySet());
        }

        public Map<String, Collection<ConnectorTaskId>> newlyAssignedTasks() {
            return newlyAssignedTasks;
        }

        public Collection<ConnectorTaskId> newlyAssignedTasks(String worker) {
            return newlyAssignedTasks.getOrDefault(worker, Collections.emptySet());
        }

        public Map<String, Collection<String>> newlyRevokedConnectors() {
            return newlyRevokedConnectors;
        }

        public Collection<String> newlyRevokedConnectors(String worker) {
            return newlyRevokedConnectors.getOrDefault(worker, Collections.emptySet());
        }

        public Map<String, Collection<ConnectorTaskId>> newlyRevokedTasks() {
            return newlyRevokedTasks;
        }

        public Collection<ConnectorTaskId> newlyRevokedTasks(String worker) {
            return newlyRevokedTasks.getOrDefault(worker, Collections.emptySet());
        }

        public Map<String, Collection<String>> allAssignedConnectors() {
            return allAssignedConnectors;
        }

        public Map<String, Collection<ConnectorTaskId>> allAssignedTasks() {
            return allAssignedTasks;
        }

        public Set<String> allWorkers() {
            return allWorkers;
        }

        @Override
        public String toString() {
            return "ClusterAssignment{"
                    + "newlyAssignedConnectors=" + newlyAssignedConnectors
                    + ", newlyAssignedTasks=" + newlyAssignedTasks
                    + ", newlyRevokedConnectors=" + newlyRevokedConnectors
                    + ", newlyRevokedTasks=" + newlyRevokedTasks
                    + ", allAssignedConnectors=" + allAssignedConnectors
                    + ", allAssignedTasks=" + allAssignedTasks
                    + ", allWorkers=" + allWorkers
                    + '}';
        }
    }

}