From 91400976890ee84bde9c28cc864a2de98c18a31a Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Mon, 1 Jun 2020 17:58:13 -0700 Subject: [PATCH 01/23] WIP --- .../internals/StreamsPartitionAssignor.java | 72 +++++++++++++++++++ .../internals/assignment/ClientState.java | 5 ++ 2 files changed, 77 insertions(+) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index a6fbdfbb6ddc0..d1836e88275b8 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -16,6 +16,9 @@ */ package org.apache.kafka.streams.processor.internals; +import java.util.Queue; +import java.util.SortedSet; +import java.util.TreeSet; import org.apache.kafka.clients.admin.Admin; import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo; import org.apache.kafka.clients.consumer.ConsumerGroupMetadata; @@ -1249,6 +1252,75 @@ private Set previousConsumersOfTaskPartitions(final Set return previousConsumers; } + /** + * @param taskIds the set of tasks to be distributed + * @param consumers the set of consumers to receive tasks + * @return a map of task assignments keyed by the consumer id + */ + private static Map> assignTasksToClients(final Collection taskIds, + final Set statefulTasks, + final Set consumers, + final ClientState state) { + final Map> assignment = new HashMap<>(); + for (final String consumer : consumers) { + assignment.put(consumer, new ArrayList<>()); + } + + final int minStatefulTasksPerThread = (int) Math.floor(((double) statefulTasks.size()) / consumers.size()); + final SortedSet unassignedStatefulTasks = new TreeSet<>(statefulTasks); + + final Queue unfilledConsumers = new LinkedList<>(); + + // First assign stateful tasks to previous owner, up to the min expected tasks/thread + for (final String consumer : consumers) { + final List threadAssignment = assignment.get(consumer); + + int i = 0; + for (final TaskId task : state.previousStatefulTasksForConsumer(consumer)) { + if (i < minStatefulTasksPerThread) { + threadAssignment.add(task); + unassignedStatefulTasks.remove(task); + } else { + break; + } + ++i; + } + + if (threadAssignment.size() < minStatefulTasksPerThread) { + unfilledConsumers.offer(consumer); + } + } + + // Next interleave remaining unassigned tasks amongst unfilled consumers + while (!unassignedStatefulTasks.isEmpty()) { + final Iterator unassignedTaskIter = unassignedStatefulTasks.iterator(); + while (unassignedTaskIter.hasNext()) { + final TaskId task = unassignedTaskIter.next(); + unassignedTaskIter.remove(); + final String consumer = unfilledConsumers.poll(); + if (consumer != null) { + final List threadAssignment = assignment.get(consumer); + + threadAssignment.add(task); + if (threadAssignment.size() < minStatefulTasksPerThread) { + unfilledConsumers.offer(consumer); + } + } else { + break; + } + } + } + + // At this point all consumers are at the minimum capacity, so just give one unassigned task to each + while (!unassignedStatefulTasks.isEmpty()) { + + } + + + // distribute stateless tasks + return assignment; + } + /** * Generate an assignment that attempts to maximize load balance without regard for stickiness, by spreading * tasks of the same groupId (subtopology) over different consumers. diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java index 4c96ade222e48..27b1dd9323152 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java @@ -301,6 +301,11 @@ long lagFor(final TaskId task) { } } + public Set previousStatefulTasksForConsumer(final String memberId) { + //TODO + return null; + } + boolean hasUnfulfilledQuota(final int tasksPerThread) { return activeTasks.size() < capacity * tasksPerThread; } From 508388975488180e8240cc9edf2955ca5fbeabc4 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Mon, 1 Jun 2020 19:45:03 -0700 Subject: [PATCH 02/23] implemented stickiness --- .../internals/StreamsPartitionAssignor.java | 377 ++++-------------- .../internals/assignment/ClientState.java | 29 +- 2 files changed, 105 insertions(+), 301 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index d1836e88275b8..cc93ac0ce8133 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -16,9 +16,11 @@ */ package org.apache.kafka.streams.processor.internals; +import java.util.PriorityQueue; import java.util.Queue; import java.util.SortedSet; import java.util.TreeSet; +import java.util.function.Function; import org.apache.kafka.clients.admin.Admin; import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo; import org.apache.kafka.clients.consumer.ConsumerGroupMetadata; @@ -383,19 +385,7 @@ public GroupAssignment assign(final Cluster metadata, final GroupSubscription gr // compute the assignment of tasks to threads within each client and build the final group assignment - final Map assignment; - if (versionProbing) { - assignment = versionProbingAssignment( - clientMetadataMap, - partitionsForTask, - partitionsByHost, - standbyPartitionsByHost, - allOwnedPartitions, - minReceivedMetadataVersion, - minSupportedMetadataVersion - ); - } else { - assignment = computeNewAssignment( + final Map assignment = computeNewAssignment( clientMetadataMap, partitionsForTask, partitionsByHost, @@ -403,9 +393,9 @@ public GroupAssignment assign(final Cluster metadata, final GroupSubscription gr allOwnedPartitions, minReceivedMetadataVersion, minSupportedMetadataVersion, + versionProbing, probingRebalanceNeeded - ); - } + ); return new GroupAssignment(assignment); } @@ -881,8 +871,9 @@ private Map computeNewAssignment(final Map allOwnedPartitions, final int minUserMetadataVersion, final int minSupportedMetadataVersion, + final boolean versionProbing, final boolean shouldTriggerProbingRebalance) { - boolean rebalanceRequired = shouldTriggerProbingRebalance; + boolean rebalanceRequired = shouldTriggerProbingRebalance || versionProbing; final Map assignment = new HashMap<>(); // within the client, distribute tasks to its owned consumers @@ -891,40 +882,39 @@ private Map computeNewAssignment(final Map consumers = clientMetadata.consumers; - Map> activeTaskAssignments; - - // Try to avoid triggering another rebalance by giving active tasks back to their previous owners within a - // client, without violating load balance. If we already know another rebalance will be required, or the - // client had no owned partitions, try to balance the workload as evenly as possible by interleaving tasks - if (rebalanceRequired || state.ownedPartitions().isEmpty()) { - activeTaskAssignments = interleaveConsumerTasksByGroupId(state.activeTasks(), consumers); - } else if ((activeTaskAssignments = tryStickyAndBalancedTaskAssignmentWithinClient(state, consumers, partitionsForTask, allOwnedPartitions)) - .equals(Collections.emptyMap())) { - rebalanceRequired = true; - activeTaskAssignments = interleaveConsumerTasksByGroupId(state.activeTasks(), consumers); - } - final Map> interleavedStandby = - interleaveConsumerTasksByGroupId(state.standbyTasks(), consumers); + final Map> activeTaskAssignment = assignTasksToThreads( + state.statefulActiveTasks(), + state.statelessActiveTasks(), + consumers, + state::previousStatefulActiveTasksForConsumer + ); + + final Map> standbyTaskAssignment = assignTasksToThreads( + state.standbyTasks(), + Collections.emptySet(), + consumers, + state::previousStandbyTasksForConsumer + ); // Arbitrarily choose the leader's client to be responsible for triggering the probing rebalance - final boolean encodeNextProbingRebalanceTime = clientId.equals(taskManager.processId()) && shouldTriggerProbingRebalance; + final boolean encodeNextProbingRebalanceTime = shouldTriggerProbingRebalance && clientId.equals(taskManager.processId()); - final boolean followupRebalanceScheduled = addClientAssignments( + final boolean tasksRevoked = addClientAssignments( assignment, clientMetadata, partitionsForTask, partitionsByHostState, standbyPartitionsByHost, allOwnedPartitions, - activeTaskAssignments, - interleavedStandby, + activeTaskAssignment, + standbyTaskAssignment, minUserMetadataVersion, minSupportedMetadataVersion, - false, - encodeNextProbingRebalanceTime); + encodeNextProbingRebalanceTime + ); - if (followupRebalanceScheduled) { + if (tasksRevoked || encodeNextProbingRebalanceTime) { rebalanceRequired = true; log.debug("Requested client {} to schedule a followup rebalance", clientId); } @@ -941,57 +931,9 @@ private Map computeNewAssignment(final Map versionProbingAssignment(final Map clientsMetadata, - final Map> partitionsForTask, - final Map> partitionsByHost, - final Map> standbyPartitionsByHost, - final Set allOwnedPartitions, - final int minUserMetadataVersion, - final int minSupportedMetadataVersion) { - final Map assignment = new HashMap<>(); - - // Since we know another rebalance will be triggered anyway, just try and generate a balanced assignment - // (without violating cooperative protocol) now so that on the second rebalance we can just give tasks - // back to their previous owners - // within the client, distribute tasks to its owned consumers - for (final ClientMetadata clientMetadata : clientsMetadata.values()) { - final ClientState state = clientMetadata.state; - - final Map> interleavedActive = - interleaveConsumerTasksByGroupId(state.activeTasks(), clientMetadata.consumers); - final Map> interleavedStandby = - interleaveConsumerTasksByGroupId(state.standbyTasks(), clientMetadata.consumers); - - addClientAssignments( - assignment, - clientMetadata, - partitionsForTask, - partitionsByHost, - standbyPartitionsByHost, - allOwnedPartitions, - interleavedActive, - interleavedStandby, - minUserMetadataVersion, - minSupportedMetadataVersion, - true, - false); - } - - log.info("Finished unstable assignment of tasks, a followup rebalance will be scheduled due to version probing."); - - return assignment; - } - /** * Adds the encoded assignment for each StreamThread consumer in the client to the overall assignment map - * @return true if this client has been told to schedule a followup rebalance + * @return true if a followup rebalance will be required due to revoekd tasks */ private boolean addClientAssignments(final Map assignment, final ClientMetadata clientMetadata, @@ -1003,9 +945,10 @@ private boolean addClientAssignments(final Map assignment, final Map> standbyTaskAssignments, final int minUserMetadataVersion, final int minSupportedMetadataVersion, - final boolean versionProbing, final boolean probingRebalanceNeeded) { - boolean rebalanceRequested = probingRebalanceNeeded || versionProbing; + boolean followupRebalanceRequiredForRevokedTasks = false; + + // We only want to encode a scheduled probing rebalance for a single member in this client boolean shouldEncodeProbingRebalance = probingRebalanceNeeded; // Loop through the consumers and build their assignment @@ -1023,7 +966,8 @@ private boolean addClientAssignments(final Map assignment, clientMetadata.state, activeTasksForConsumer, partitionsForTask, - allOwnedPartitions); + allOwnedPartitions + ); final Map> standbyTaskMap = buildStandbyTaskMap(standbyTaskAssignments.get(consumer), partitionsForTask); @@ -1039,11 +983,15 @@ private boolean addClientAssignments(final Map assignment, ); if (tasksRevoked) { - // TODO: once KAFKA-9821 is resolved we can leave it to the client to trigger this rebalance + // TODO: once KAFKA-10078 is resolved we can leave it to the client to trigger this rebalance log.debug("Requesting followup rebalance be scheduled immediately due to tasks changing ownership."); info.setNextRebalanceTime(0L); - rebalanceRequested = true; - } else if (shouldEncodeProbingRebalance) { + followupRebalanceRequiredForRevokedTasks = true; + // Don't bother to schedule a probing rebalance if an immediate one is already scheduled + shouldEncodeProbingRebalance = false; + } + + if (shouldEncodeProbingRebalance) { final long nextRebalanceTimeMs = time.milliseconds() + probingRebalanceIntervalMs(); log.debug("Requesting followup rebalance be scheduled for {} ms to probe for caught-up replica tasks.", nextRebalanceTimeMs); info.setNextRebalanceTime(nextRebalanceTimeMs); @@ -1058,7 +1006,7 @@ private boolean addClientAssignments(final Map assignment, ) ); } - return rebalanceRequested; + return followupRebalanceRequiredForRevokedTasks; } /** @@ -1081,7 +1029,7 @@ private boolean populateActiveTaskAndPartitionsLists(final List for (final TaskId taskId : activeTasksForConsumer) { final List assignedPartitionsForTask = new ArrayList<>(); for (final TopicPartition partition : partitionsForTask.get(taskId)) { - final String oldOwner = clientState.ownedPartitions().get(partition); + 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 @@ -1124,150 +1072,21 @@ private static Map> buildStandbyTaskMap(final Collec } /** - * Generates an assignment that tries to satisfy two conditions: no active task previously owned by a consumer - * be assigned to another (ie nothing gets revoked), and the number of tasks is evenly distributed throughout - * the client. - *

- * If it is impossible to satisfy both constraints we abort early and return an empty map so we can use a - * different assignment strategy that tries to distribute tasks of a single subtopology across different threads. - * - * @param state state for this client - * @param consumers the consumers in this client - * @param partitionsForTask mapping from task to its associated partitions - * @param allOwnedPartitions set of all partitions claimed as owned by the group - * @return task assignment for the consumers of this client - * empty map if it is not possible to generate a balanced assignment without moving a task to a new consumer - */ - Map> tryStickyAndBalancedTaskAssignmentWithinClient(final ClientState state, - final Set consumers, - final Map> partitionsForTask, - final Set allOwnedPartitions) { - final Map> assignments = new HashMap<>(); - final LinkedList newTasks = new LinkedList<>(); - final Set unfilledConsumers = new HashSet<>(consumers); - - final int maxTasksPerClient = (int) Math.ceil(((double) state.activeTaskCount()) / consumers.size()); - - // initialize task list for consumers - for (final String consumer : consumers) { - assignments.put(consumer, new ArrayList<>()); - } - - for (final TaskId task : state.activeTasks()) { - final Set previousConsumers = previousConsumersOfTaskPartitions(partitionsForTask.get(task), state.ownedPartitions(), allOwnedPartitions); - - // If this task's partitions were owned by different consumers, we can't avoid revoking partitions - if (previousConsumers.size() > 1) { - log.warn("The partitions of task {} were claimed as owned by different StreamThreads. " + - "This indicates the mapping from partitions to tasks has changed!", task); - return Collections.emptyMap(); - } - - // If this is a new task, or its old consumer no longer exists, it can be freely (re)assigned - if (previousConsumers.isEmpty()) { - log.debug("Task {} was not previously owned by any consumers still in the group. It's owner may " + - "have died or it may be a new task", task); - newTasks.add(task); - } else { - final String consumer = previousConsumers.iterator().next(); - - // If the previous consumer was from another client, these partitions will have to be revoked - if (!consumers.contains(consumer)) { - log.debug("This client was assigned a task {} whose partition(s) were previously owned by another " + - "client, falling back to an interleaved assignment since a rebalance is inevitable.", task); - return Collections.emptyMap(); - } - - // If this consumer previously owned more tasks than it has capacity for, some must be revoked - if (assignments.get(consumer).size() >= maxTasksPerClient) { - log.debug("Cannot create a sticky and balanced assignment as this client's consumers owned more " + - "previous tasks than it has capacity for during this assignment, falling back to interleaved " + - "assignment since a realance is inevitable."); - return Collections.emptyMap(); - } - - assignments.get(consumer).add(task); - - // If we have now reached capacity, remove it from set of consumers who still need more tasks - if (assignments.get(consumer).size() == maxTasksPerClient) { - unfilledConsumers.remove(consumer); - } - } - } - - // Interleave any remaining tasks by groupId among the consumers with remaining capacity. For further - // explanation, see the javadocs for #interleaveConsumerTasksByGroupId - Collections.sort(newTasks); - while (!newTasks.isEmpty()) { - if (unfilledConsumers.isEmpty()) { - throw new IllegalStateException("Some tasks could not be distributed"); - } - - final Iterator consumerIt = unfilledConsumers.iterator(); - - // Loop through the unfilled consumers and distribute tasks until newTasks is empty - while (consumerIt.hasNext()) { - final String consumer = consumerIt.next(); - final List consumerAssignment = assignments.get(consumer); - final TaskId task = newTasks.poll(); - if (task == null) { - break; - } - - consumerAssignment.add(task); - if (consumerAssignment.size() == maxTasksPerClient) { - consumerIt.remove(); - } - } - } - - return assignments; - } - - /** - * Get the previous consumer for the partitions of a task - * - * @param taskPartitions the TopicPartitions for a single given task - * @param clientOwnedPartitions the partitions owned by all consumers in a client - * @param allOwnedPartitions all partitions claimed as owned by any consumer in any client - * @return set of consumer(s) that previously owned the partitions in this task - * empty set signals that it is a new task, or its previous owner is no longer in the group - */ - private Set previousConsumersOfTaskPartitions(final Set taskPartitions, - final Map clientOwnedPartitions, - final Set allOwnedPartitions) { - // this "foreignConsumer" indicates a partition was owned by someone from another client -- we don't really care who - final String foreignConsumer = ""; - final Set previousConsumers = new HashSet<>(); - - for (final TopicPartition tp : taskPartitions) { - final String currentPartitionConsumer = clientOwnedPartitions.get(tp); - if (currentPartitionConsumer != null) { - previousConsumers.add(currentPartitionConsumer); - } else if (allOwnedPartitions.contains(tp)) { - previousConsumers.add(foreignConsumer); - } - } - - return previousConsumers; - } - - /** - * @param taskIds the set of tasks to be distributed - * @param consumers the set of consumers to receive tasks - * @return a map of task assignments keyed by the consumer id + * Generate an assignment that tries to preserve thread-level stickiness of stateful tasks without violating + * balance. The stateful and total task load are both balanced across threads. Tasks without previous owners + * will be interleaved by group id to spread subtopologies across threads and further balance the workload. */ - private static Map> assignTasksToClients(final Collection taskIds, - final Set statefulTasks, + private static Map> assignTasksToThreads(final Set statefulTasksToAssign, + final Set statelessTasksToAssign, final Set consumers, - final ClientState state) { + final Function> previousTasksForConsumer) { final Map> assignment = new HashMap<>(); for (final String consumer : consumers) { assignment.put(consumer, new ArrayList<>()); } - final int minStatefulTasksPerThread = (int) Math.floor(((double) statefulTasks.size()) / consumers.size()); - final SortedSet unassignedStatefulTasks = new TreeSet<>(statefulTasks); + final int minStatefulTasksPerThread = (int) Math.floor(((double) statefulTasksToAssign.size()) / consumers.size()); + final PriorityQueue unassignedStatefulTasks = new PriorityQueue<>(statefulTasksToAssign); final Queue unfilledConsumers = new LinkedList<>(); @@ -1276,14 +1095,14 @@ private static Map> assignTasksToClients(final Collection threadAssignment = assignment.get(consumer); int i = 0; - for (final TaskId task : state.previousStatefulTasksForConsumer(consumer)) { - if (i < minStatefulTasksPerThread) { - threadAssignment.add(task); - unassignedStatefulTasks.remove(task); - } else { - break; - } - ++i; + for (final TaskId task : previousTasksForConsumer.apply(consumer)) { + if (i < minStatefulTasksPerThread) { + threadAssignment.add(task); + unassignedStatefulTasks.remove(task); + } else { + break; + } + ++i; } if (threadAssignment.size() < minStatefulTasksPerThread) { @@ -1292,73 +1111,39 @@ private static Map> assignTasksToClients(final Collection unassignedTaskIter = unassignedStatefulTasks.iterator(); - while (unassignedTaskIter.hasNext()) { - final TaskId task = unassignedTaskIter.next(); - unassignedTaskIter.remove(); + while (!unfilledConsumers.isEmpty()) { + final TaskId task = unassignedStatefulTasks.poll(); + if (task != null) { final String consumer = unfilledConsumers.poll(); - if (consumer != null) { - final List threadAssignment = assignment.get(consumer); - - threadAssignment.add(task); - if (threadAssignment.size() < minStatefulTasksPerThread) { - unfilledConsumers.offer(consumer); - } - } else { - break; + final List threadAssignment = assignment.get(consumer); + threadAssignment.add(task); + if (threadAssignment.size() < minStatefulTasksPerThread) { + unfilledConsumers.offer(consumer); } + } else { + throw new IllegalStateException("Ran out of unassigned stateful tasks but some members were not at capacity"); } } - // At this point all consumers are at the minimum capacity, so just give one unassigned task to each - while (!unassignedStatefulTasks.isEmpty()) { - + // At this point all consumers are at the minimum capacity, so just add the remaining stateful tasks and + // distribute them with the stateless tasks + final LinkedList unassignedTasks = new LinkedList<>(statelessTasksToAssign); + Collections.sort(unassignedTasks); + for (final TaskId task : unassignedStatefulTasks) { + unassignedTasks.addFirst(task); } + // Reset the unfilled consumers and just cycle through them until we've assigned all tasks + unfilledConsumers.addAll(consumers); - // distribute stateless tasks - return assignment; - } - - /** - * Generate an assignment that attempts to maximize load balance without regard for stickiness, by spreading - * tasks of the same groupId (subtopology) over different consumers. - * - * @param taskIds the set of tasks to be distributed - * @param consumers the set of consumers to receive tasks - * @return a map of task assignments keyed by the consumer id - */ - static Map> interleaveConsumerTasksByGroupId(final Collection taskIds, - final Set consumers) { - // First we make a sorted list of the tasks, grouping them by groupId - final LinkedList sortedTasks = new LinkedList<>(taskIds); - Collections.sort(sortedTasks); - - // Initialize the assignment map and task list for each consumer. We use a TreeMap here for a consistent - // ordering of the consumers in the hope they will end up with the same set of tasks in subsequent assignments - final Map> taskIdsForConsumerAssignment = new TreeMap<>(); - for (final String consumer : consumers) { - taskIdsForConsumerAssignment.put(consumer, new ArrayList<>()); + while (!unassignedTasks.isEmpty()) { + final TaskId task = unassignedTasks.poll(); + final String consumer = unfilledConsumers.poll(); + assignment.get(consumer).add(task); + unfilledConsumers.offer(consumer); } - // We loop until the tasks have all been assigned, removing them from the list when they are given to a - // consumer. To interleave the tasks, we loop through the consumers and give each one task from the head - // of the list. When we finish going through the list of consumers we start over at the beginning of the - // consumers list, continuing until we run out of tasks. - while (!sortedTasks.isEmpty()) { - for (final Map.Entry> consumerTaskIds : taskIdsForConsumerAssignment.entrySet()) { - final List taskIdList = consumerTaskIds.getValue(); - final TaskId taskId = sortedTasks.poll(); - - // Check for null here as we may run out of tasks before giving every consumer exactly the same number - if (taskId == null) { - break; - } - taskIdList.add(taskId); - } - } - return taskIdsForConsumerAssignment; + return assignment; } private void validateMetadataVersions(final int receivedAssignmentMetadataVersion, diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java index 27b1dd9323152..3305c7ecd328a 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.streams.processor.internals.assignment; +import java.util.stream.Collectors; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.streams.processor.TaskId; import org.apache.kafka.streams.processor.internals.Task; @@ -143,7 +144,7 @@ public void assignActiveTasks(final Collection tasks) { activeTasks.addAll(tasks); } - void assignActive(final TaskId task) { + public void assignActive(final TaskId task) { assertNotAssigned(task); activeTasks.add(task); } @@ -232,8 +233,9 @@ Set previousAssignedTasks() { return union(() -> new HashSet<>(prevActiveTasks.size() + prevStandbyTasks.size()), prevActiveTasks, prevStandbyTasks); } - public Map ownedPartitions() { - return unmodifiableMap(ownedPartitions); + // May return null + public String previousOwnerForPartition(final TopicPartition partition) { + return ownedPartitions.get(partition); } public void addOwnedPartitions(final Collection ownedPartitions, final String consumer) { @@ -301,7 +303,20 @@ long lagFor(final TaskId task) { } } - public Set previousStatefulTasksForConsumer(final String memberId) { + public Set statefulActiveTasks() { + return activeTasks.stream().filter(this::isStateful).collect(Collectors.toSet()); + } + + public Set statelessActiveTasks() { + return activeTasks.stream().filter(task -> !isStateful(task)).collect(Collectors.toSet()); + } + + public Set previousStatefulActiveTasksForConsumer(final String memberId) { + //TODO + return null; + } + + public Set previousStandbyTasksForConsumer(final String memberId) { //TODO return null; } @@ -345,12 +360,16 @@ public String toString() { "]"; } + private boolean isStateful(final TaskId task) { + return taskLagTotals.containsKey(task); + } + private void initializePrevActiveTasksFromOwnedPartitions(final Map taskForPartitionMap) { // there are three cases where we need to construct some or all of the prevTasks from the ownedPartitions: // 1) COOPERATIVE clients on version 2.4-2.5 do not encode active tasks at all and rely on ownedPartitions // 2) future client during version probing, when we can't decode the future subscription info's prev tasks // 3) stateless tasks are not encoded in the task lags, and must be figured out from the ownedPartitions - for (final Map.Entry partitionEntry : ownedPartitions().entrySet()) { + for (final Map.Entry partitionEntry : ownedPartitions.entrySet()) { final TopicPartition tp = partitionEntry.getKey(); final TaskId task = taskForPartitionMap.get(tp); if (task != null) { From be118634a1a51f6144adf9dc9055c4626bd7d88b Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Mon, 1 Jun 2020 20:27:17 -0700 Subject: [PATCH 03/23] compiling tests --- .../internals/StreamsPartitionAssignor.java | 18 ++++------ .../internals/assignment/ClientState.java | 28 ++++++++++++--- .../StreamsPartitionAssignorTest.java | 34 ++++++++++++------- .../internals/assignment/ClientStateTest.java | 16 ++++----- .../TaskAssignorConvergenceTest.java | 2 +- 5 files changed, 60 insertions(+), 38 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index cc93ac0ce8133..7f5a57ef3c53d 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -18,8 +18,6 @@ import java.util.PriorityQueue; import java.util.Queue; -import java.util.SortedSet; -import java.util.TreeSet; import java.util.function.Function; import org.apache.kafka.clients.admin.Admin; import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo; @@ -59,13 +57,11 @@ import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; -import java.util.Iterator; import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; -import java.util.TreeMap; import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; @@ -138,8 +134,8 @@ void addConsumer(final String consumerMemberId, final List owned state.addOwnedPartitions(ownedPartitions, consumerMemberId); } - void addPreviousTasksAndOffsetSums(final Map taskOffsetSums) { - state.addPreviousTasksAndOffsetSums(taskOffsetSums); + void addPreviousTasksAndOffsetSums(final String consumerId, final Map taskOffsetSums) { + state.addPreviousTasksAndOffsetSums(consumerId, taskOffsetSums); } @Override @@ -323,7 +319,7 @@ public GroupAssignment assign(final Cluster metadata, final GroupSubscription gr // add the consumer and any info in its subscription to the client clientMetadata.addConsumer(consumerId, subscription.ownedPartitions()); allOwnedPartitions.addAll(subscription.ownedPartitions()); - clientMetadata.addPreviousTasksAndOffsetSums(info.taskOffsetSums()); + clientMetadata.addPreviousTasksAndOffsetSums(consumerId, info.taskOffsetSums()); } final boolean versionProbing = @@ -1076,10 +1072,10 @@ private static Map> buildStandbyTaskMap(final Collec * balance. The stateful and total task load are both balanced across threads. Tasks without previous owners * will be interleaved by group id to spread subtopologies across threads and further balance the workload. */ - private static Map> assignTasksToThreads(final Set statefulTasksToAssign, - final Set statelessTasksToAssign, - final Set consumers, - final Function> previousTasksForConsumer) { + static Map> assignTasksToThreads(final Set statefulTasksToAssign, + final Set statelessTasksToAssign, + final Set consumers, + final Function> previousTasksForConsumer) { final Map> assignment = new HashMap<>(); for (final String consumer : consumers) { assignment.put(consumer, new ArrayList<>()); diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java index 3305c7ecd328a..604022c003005 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java @@ -49,6 +49,7 @@ public class ClientState { private final Set prevActiveTasks; private final Set prevStandbyTasks; + private final Map> consumerToPreviousTaskIds; private final Map ownedPartitions; private final Map taskOffsetSums; // contains only stateful tasks we previously owned private final Map taskLagTotals; // contains lag for all stateful tasks in the app topology @@ -64,6 +65,7 @@ public ClientState() { standbyTasks = new TreeSet<>(); prevActiveTasks = new TreeSet<>(); prevStandbyTasks = new TreeSet<>(); + consumerToPreviousTaskIds = new TreeMap<>(); ownedPartitions = new TreeMap<>(TOPIC_PARTITION_COMPARATOR); taskOffsetSums = new TreeMap<>(); taskLagTotals = new TreeMap<>(); @@ -74,6 +76,7 @@ private ClientState(final Set activeTasks, final Set standbyTasks, final Set prevActiveTasks, final Set prevStandbyTasks, + final Map> consumerToPreviousTaskIds, final SortedMap ownedPartitions, final Map taskOffsetSums, final Map taskLagTotals, @@ -82,6 +85,7 @@ private ClientState(final Set activeTasks, this.standbyTasks = standbyTasks; this.prevActiveTasks = prevActiveTasks; this.prevStandbyTasks = prevStandbyTasks; + this.consumerToPreviousTaskIds = consumerToPreviousTaskIds; this.ownedPartitions = ownedPartitions; this.taskOffsetSums = taskOffsetSums; this.taskLagTotals = taskLagTotals; @@ -96,6 +100,7 @@ public ClientState(final Set previousActiveTasks, standbyTasks = new TreeSet<>(); prevActiveTasks = unmodifiableSet(new TreeSet<>(previousActiveTasks)); prevStandbyTasks = unmodifiableSet(new TreeSet<>(previousStandbyTasks)); + consumerToPreviousTaskIds = new TreeMap<>(); ownedPartitions = new TreeMap<>(TOPIC_PARTITION_COMPARATOR); taskOffsetSums = emptyMap(); this.taskLagTotals = unmodifiableMap(taskLagTotals); @@ -110,6 +115,7 @@ public ClientState copy() { new TreeSet<>(standbyTasks), new TreeSet<>(prevActiveTasks), new TreeSet<>(prevStandbyTasks), + new TreeMap<>(consumerToPreviousTaskIds), newOwnedPartitions, new TreeMap<>(taskOffsetSums), new TreeMap<>(taskLagTotals), @@ -244,8 +250,9 @@ public void addOwnedPartitions(final Collection ownedPartitions, } } - public void addPreviousTasksAndOffsetSums(final Map taskOffsetSums) { + public void addPreviousTasksAndOffsetSums(final String consumerId, final Map taskOffsetSums) { this.taskOffsetSums.putAll(taskOffsetSums); + consumerToPreviousTaskIds.put(consumerId, taskOffsetSums.keySet()); } public void initializePrevTasks(final Map taskForPartitionMap) { @@ -312,13 +319,24 @@ public Set statelessActiveTasks() { } public Set previousStatefulActiveTasksForConsumer(final String memberId) { - //TODO - return null; + final Set prevTasks = new HashSet<>(); + for (final TaskId task : consumerToPreviousTaskIds.get(memberId)) { + if (isStateful(task) && prevActiveTasks.contains(task)) { + prevTasks.add(task); + } + } + return prevTasks; } + // Equivalent to previousStatefulStandbyTasksForConsumer since all standbys are by definition stateful public Set previousStandbyTasksForConsumer(final String memberId) { - //TODO - return null; + final Set prevTasks = new HashSet<>(); + for (final TaskId task : consumerToPreviousTaskIds.get(memberId)) { + if (prevStandbyTasks.contains(task)) { + prevTasks.add(task); + } + } + return prevTasks; } boolean hasUnfulfilledQuota(final int tasksPerThread) { diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index 350a59801e017..e28c3af05e76a 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -88,6 +88,7 @@ import static org.apache.kafka.common.utils.Utils.mkEntry; import static org.apache.kafka.common.utils.Utils.mkMap; import static org.apache.kafka.common.utils.Utils.mkSet; +import static org.apache.kafka.streams.processor.internals.StreamsPartitionAssignor.assignTasksToThreads; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.EMPTY_CHANGELOG_END_OFFSETS; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.EMPTY_TASKS; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.EMPTY_TASK_OFFSET_SUMS; @@ -325,10 +326,9 @@ public void shouldUseCooperativeRebalancingProtocol() { @Test public void shouldProduceStickyAndBalancedAssignmentWhenNothingChanges() { - configureDefault(); final ClientState state = new ClientState(); - final List allTasks = asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, - TASK_1_3); + final List allTasks = + asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, TASK_1_3); final Map> previousAssignment = mkMap( mkEntry(CONSUMER_1, asList(TASK_0_0, TASK_1_1, TASK_1_3)), @@ -336,26 +336,28 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNothingChanges() { mkEntry(CONSUMER_3, asList(TASK_0_1, TASK_0_2, TASK_1_2)) ); - for (final Map.Entry> entry : previousAssignment.entrySet()) { - for (final TaskId task : entry.getValue()) { - state.addOwnedPartitions(partitionsForTask.get(task), entry.getKey()); - } - } + state.addPreviousTasksAndOffsetSums(CONSUMER_1, getTaskOffsetSums(asList(TASK_0_0, TASK_1_1, TASK_1_3), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_2, getTaskOffsetSums(asList(TASK_0_3, TASK_1_0), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_3, getTaskOffsetSums(asList(TASK_0_1, TASK_0_2, TASK_1_2), EMPTY_TASKS)); + state.initializePrevTasks(emptyMap()); + state.computeTaskLags(UUID_1, getTaskEndOffsetSums(allTasks)); final Set consumers = mkSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); state.assignActiveTasks(allTasks); assertEquivalentAssignment( previousAssignment, - partitionAssignor.tryStickyAndBalancedTaskAssignmentWithinClient( - state, + assignTasksToThreads( + state.activeTasks(), + emptySet(), consumers, - partitionsForTask, - emptySet() + state::previousStatefulActiveTasksForConsumer ) ); } + /* + @Test public void shouldProduceStickyAndBalancedAssignmentWhenNewTasksAreAdded() { configureDefault(); @@ -485,6 +487,7 @@ public void shouldInterleaveTasksByGroupId() { assertThat(interleavedTaskIds, equalTo(assignment)); } +*/ @Test public void testEagerSubscription() { @@ -2019,10 +2022,15 @@ private static SubscriptionInfo getInfoForOlderVersion(final int version, } // Stub offset sums for when we only care about the prev/standby task sets, not the actual offsets - private static Map getTaskOffsetSums(final Set activeTasks, final Set standbyTasks) { + private static Map getTaskOffsetSums(final Collection activeTasks, final Collection standbyTasks) { final Map taskOffsetSums = activeTasks.stream().collect(Collectors.toMap(t -> t, t -> Task.LATEST_OFFSET)); taskOffsetSums.putAll(standbyTasks.stream().collect(Collectors.toMap(t -> t, t -> 0L))); return taskOffsetSums; } + // Stub end offsets sums for situations where we don't really care about computing exact lags + private static Map getTaskEndOffsetSums(final Collection allStatefulTasks) { + return allStatefulTasks.stream().collect(Collectors.toMap(t -> t, t-> Long.MAX_VALUE)); + } + } diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java index 8e5fa36ecd463..d0e95cf4d2c0d 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java @@ -300,7 +300,7 @@ public void shouldNotHaveUnfulfilledQuotaWhenActiveTaskSizeGreaterEqualThanCapac @Test public void shouldAddTasksWithLatestOffsetToPrevActiveTasks() { final Map taskOffsetSums = Collections.singletonMap(TASK_0_1, Task.LATEST_OFFSET); - client.addPreviousTasksAndOffsetSums(taskOffsetSums); + client.addPreviousTasksAndOffsetSums("c1", taskOffsetSums); client.initializePrevTasks(Collections.emptyMap()); assertThat(client.prevActiveTasks(), equalTo(Collections.singleton(TASK_0_1))); assertThat(client.previousAssignedTasks(), equalTo(Collections.singleton(TASK_0_1))); @@ -313,7 +313,7 @@ public void shouldAddTasksInOffsetSumsMapToPrevStandbyTasks() { mkEntry(TASK_0_1, 0L), mkEntry(TASK_0_2, 100L) ); - client.addPreviousTasksAndOffsetSums(taskOffsetSums); + client.addPreviousTasksAndOffsetSums("c1", taskOffsetSums); client.initializePrevTasks(Collections.emptyMap()); assertThat(client.prevStandbyTasks(), equalTo(mkSet(TASK_0_1, TASK_0_2))); assertThat(client.previousAssignedTasks(), equalTo(mkSet(TASK_0_1, TASK_0_2))); @@ -330,7 +330,7 @@ public void shouldComputeTaskLags() { mkEntry(TASK_0_1, 500L), mkEntry(TASK_0_2, 100L) ); - client.addPreviousTasksAndOffsetSums(taskOffsetSums); + client.addPreviousTasksAndOffsetSums("c1", taskOffsetSums); client.computeTaskLags(null, allTaskEndOffsetSums); assertThat(client.lagFor(TASK_0_1), equalTo(500L)); @@ -341,7 +341,7 @@ public void shouldComputeTaskLags() { public void shouldReturnEndOffsetSumForLagOfTaskWeDidNotPreviouslyOwn() { final Map taskOffsetSums = Collections.emptyMap(); final Map allTaskEndOffsetSums = Collections.singletonMap(TASK_0_1, 500L); - client.addPreviousTasksAndOffsetSums(taskOffsetSums); + client.addPreviousTasksAndOffsetSums("c1", taskOffsetSums); client.computeTaskLags(null, allTaskEndOffsetSums); assertThat(client.lagFor(TASK_0_1), equalTo(500L)); } @@ -350,7 +350,7 @@ public void shouldReturnEndOffsetSumForLagOfTaskWeDidNotPreviouslyOwn() { public void shouldReturnLatestOffsetForLagOfPreviousActiveRunningTask() { final Map taskOffsetSums = Collections.singletonMap(TASK_0_1, Task.LATEST_OFFSET); final Map allTaskEndOffsetSums = Collections.singletonMap(TASK_0_1, 500L); - client.addPreviousTasksAndOffsetSums(taskOffsetSums); + client.addPreviousTasksAndOffsetSums("c1", taskOffsetSums); client.computeTaskLags(null, allTaskEndOffsetSums); assertThat(client.lagFor(TASK_0_1), equalTo(Task.LATEST_OFFSET)); } @@ -359,7 +359,7 @@ public void shouldReturnLatestOffsetForLagOfPreviousActiveRunningTask() { public void shouldReturnUnknownOffsetSumForLagOfTaskWithUnknownOffset() { final Map taskOffsetSums = Collections.singletonMap(TASK_0_1, UNKNOWN_OFFSET_SUM); final Map allTaskEndOffsetSums = Collections.singletonMap(TASK_0_1, 500L); - client.addPreviousTasksAndOffsetSums(taskOffsetSums); + client.addPreviousTasksAndOffsetSums("c1", taskOffsetSums); client.computeTaskLags(null, allTaskEndOffsetSums); assertThat(client.lagFor(TASK_0_1), equalTo(UNKNOWN_OFFSET_SUM)); } @@ -368,7 +368,7 @@ public void shouldReturnUnknownOffsetSumForLagOfTaskWithUnknownOffset() { public void shouldReturnEndOffsetSumIfOffsetSumIsGreaterThanEndOffsetSum() { final Map taskOffsetSums = Collections.singletonMap(TASK_0_1, 5L); final Map allTaskEndOffsetSums = Collections.singletonMap(TASK_0_1, 1L); - client.addPreviousTasksAndOffsetSums(taskOffsetSums); + client.addPreviousTasksAndOffsetSums("c1", taskOffsetSums); client.computeTaskLags(null, allTaskEndOffsetSums); assertThat(client.lagFor(TASK_0_1), equalTo(1L)); } @@ -385,7 +385,7 @@ public void shouldThrowIllegalStateExceptionIfTaskLagsMapIsNotEmpty() { public void shouldThrowIllegalStateExceptionOnLagForUnknownTask() { final Map taskOffsetSums = Collections.singletonMap(TASK_0_1, 0L); final Map allTaskEndOffsetSums = Collections.singletonMap(TASK_0_1, 500L); - client.addPreviousTasksAndOffsetSums(taskOffsetSums); + client.addPreviousTasksAndOffsetSums("c1", taskOffsetSums); client.computeTaskLags(null, allTaskEndOffsetSums); assertThrows(IllegalStateException.class, () -> client.lagFor(TASK_0_2)); } diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java index 8253bcd3ade0c..68c9dfec107de 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/TaskAssignorConvergenceTest.java @@ -185,7 +185,7 @@ private void prepareForNextRebalance() { } newClientState.addPreviousActiveTasks(clientState.activeTasks()); newClientState.addPreviousStandbyTasks(clientState.standbyTasks()); - newClientState.addPreviousTasksAndOffsetSums(taskOffsetSums); + newClientState.addPreviousTasksAndOffsetSums("consumer", taskOffsetSums); newClientState.computeTaskLags(uuid, statefulTaskEndOffsetSums); newClientStates.put(uuid, newClientState); } From dc296fc0380105691bc768b0f4b128fee987b89a Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Mon, 1 Jun 2020 21:18:12 -0700 Subject: [PATCH 04/23] improve stickiness to the max --- .../internals/StreamsPartitionAssignor.java | 66 ++++++++++++++----- 1 file changed, 49 insertions(+), 17 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 7f5a57ef3c53d..0d206edb4d693 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.streams.processor.internals; +import java.util.Iterator; import java.util.PriorityQueue; import java.util.Queue; import java.util.function.Function; @@ -1084,7 +1085,9 @@ static Map> assignTasksToThreads(final Set stateful final int minStatefulTasksPerThread = (int) Math.floor(((double) statefulTasksToAssign.size()) / consumers.size()); final PriorityQueue unassignedStatefulTasks = new PriorityQueue<>(statefulTasksToAssign); - final Queue unfilledConsumers = new LinkedList<>(); + final PriorityQueue consumersToFill = new PriorityQueue<>(); + // keep track of tasks that we have to skip during the first pass in case we can reassign them later + final Map unassignedTaskToPreviousOwner = new HashMap<>(); // First assign stateful tasks to previous owner, up to the min expected tasks/thread for (final String consumer : consumers) { @@ -1096,47 +1099,76 @@ static Map> assignTasksToThreads(final Set stateful threadAssignment.add(task); unassignedStatefulTasks.remove(task); } else { - break; + unassignedTaskToPreviousOwner.put(task, consumer); } ++i; } if (threadAssignment.size() < minStatefulTasksPerThread) { - unfilledConsumers.offer(consumer); + consumersToFill.offer(consumer); } } // Next interleave remaining unassigned tasks amongst unfilled consumers - while (!unfilledConsumers.isEmpty()) { + while (!consumersToFill.isEmpty()) { final TaskId task = unassignedStatefulTasks.poll(); if (task != null) { - final String consumer = unfilledConsumers.poll(); + final String consumer = consumersToFill.poll(); final List threadAssignment = assignment.get(consumer); threadAssignment.add(task); if (threadAssignment.size() < minStatefulTasksPerThread) { - unfilledConsumers.offer(consumer); + consumersToFill.offer(consumer); } } else { throw new IllegalStateException("Ran out of unassigned stateful tasks but some members were not at capacity"); } } - // At this point all consumers are at the minimum capacity, so just add the remaining stateful tasks and - // distribute them with the stateless tasks - final LinkedList unassignedTasks = new LinkedList<>(statelessTasksToAssign); - Collections.sort(unassignedTasks); + // At this point all consumers are at the min capacity, so there may be up to N - 1 unassigned + // stateful tasks remaining that should now be distributed over the consumers + consumersToFill.addAll(consumers); + + // Go over the tasks we skipped earlier and assign them to their previous owner when possible + for (final Map.Entry taskEntry : unassignedTaskToPreviousOwner.entrySet()) { + final TaskId task = taskEntry.getKey(); + final String consumer = taskEntry.getValue(); + if (consumersToFill.contains(consumer)) { + assignment.get(consumer).add(task); + unassignedStatefulTasks.remove(task); + consumersToFill.remove(consumer); + } + } + + // Now just distribute the remaining unassigned tasks over the consumers still at min capacity for (final TaskId task : unassignedStatefulTasks) { - unassignedTasks.addFirst(task); + final String consumer = consumersToFill.poll(); + assignment.get(consumer).add(task); + } + + // There must be at least one consumer still at min capacity while all the others are at min + // capacity + 1, so start distributing stateless tasks to get all consumers back to the same count + final List unassignedStatelessTasks = new ArrayList<>(statelessTasksToAssign); + Collections.sort(unassignedStatelessTasks); + + final Iterator unassignedStatelessTasksIter = unassignedStatelessTasks.iterator(); + while (unassignedStatelessTasksIter.hasNext()) { + final TaskId task = unassignedStatelessTasksIter.next(); + final String consumer = consumersToFill.poll(); + if (consumer != null) { + assignment.get(consumer).add(task); + } else { + break; + } } - // Reset the unfilled consumers and just cycle through them until we've assigned all tasks - unfilledConsumers.addAll(consumers); + // Now just distribute tasks while circling through all the consumers + consumersToFill.addAll(consumers); - while (!unassignedTasks.isEmpty()) { - final TaskId task = unassignedTasks.poll(); - final String consumer = unfilledConsumers.poll(); + while (unassignedStatelessTasksIter.hasNext()) { + final TaskId task = unassignedStatelessTasksIter.next(); + final String consumer = consumersToFill.poll(); assignment.get(consumer).add(task); - unfilledConsumers.offer(consumer); + consumersToFill.offer(consumer); } return assignment; From 4b8e96cb899d3b6ee6c91766ae2722e7c30ee5e4 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Mon, 1 Jun 2020 21:30:00 -0700 Subject: [PATCH 05/23] debugging last test --- .../internals/StreamsPartitionAssignor.java | 4 +- .../StreamsPartitionAssignorTest.java | 139 +++++------------- 2 files changed, 37 insertions(+), 106 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 0d206edb4d693..0714b0dcda490 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -1073,8 +1073,8 @@ private static Map> buildStandbyTaskMap(final Collec * balance. The stateful and total task load are both balanced across threads. Tasks without previous owners * will be interleaved by group id to spread subtopologies across threads and further balance the workload. */ - static Map> assignTasksToThreads(final Set statefulTasksToAssign, - final Set statelessTasksToAssign, + static Map> assignTasksToThreads(final Collection statefulTasksToAssign, + final Collection statelessTasksToAssign, final Set consumers, final Function> previousTasksForConsumer) { final Map> assignment = new HashMap<>(); diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index e28c3af05e76a..6cdc1002b903d 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -326,7 +326,6 @@ public void shouldUseCooperativeRebalancingProtocol() { @Test public void shouldProduceStickyAndBalancedAssignmentWhenNothingChanges() { - final ClientState state = new ClientState(); final List allTasks = asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, TASK_1_3); @@ -336,19 +335,18 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNothingChanges() { mkEntry(CONSUMER_3, asList(TASK_0_1, TASK_0_2, TASK_1_2)) ); + final ClientState state = new ClientState(); + final Set consumers = mkSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); state.addPreviousTasksAndOffsetSums(CONSUMER_1, getTaskOffsetSums(asList(TASK_0_0, TASK_1_1, TASK_1_3), EMPTY_TASKS)); state.addPreviousTasksAndOffsetSums(CONSUMER_2, getTaskOffsetSums(asList(TASK_0_3, TASK_1_0), EMPTY_TASKS)); state.addPreviousTasksAndOffsetSums(CONSUMER_3, getTaskOffsetSums(asList(TASK_0_1, TASK_0_2, TASK_1_2), EMPTY_TASKS)); state.initializePrevTasks(emptyMap()); state.computeTaskLags(UUID_1, getTaskEndOffsetSums(allTasks)); - final Set consumers = mkSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); - state.assignActiveTasks(allTasks); - assertEquivalentAssignment( previousAssignment, assignTasksToThreads( - state.activeTasks(), + allTasks, emptySet(), consumers, state::previousStatefulActiveTasksForConsumer @@ -356,15 +354,10 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNothingChanges() { ); } - /* - @Test public void shouldProduceStickyAndBalancedAssignmentWhenNewTasksAreAdded() { - configureDefault(); - final ClientState state = new ClientState(); - - final Set allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, - TASK_1_3); + final List allTasks = + asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, TASK_1_3); final Map> previousAssignment = mkMap( mkEntry(CONSUMER_1, new ArrayList<>(asList(TASK_0_0, TASK_1_1, TASK_1_3))), @@ -372,122 +365,60 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNewTasksAreAdded() { mkEntry(CONSUMER_3, new ArrayList<>(asList(TASK_0_1, TASK_0_2, TASK_1_2))) ); - for (final Map.Entry> entry : previousAssignment.entrySet()) { - for (final TaskId task : entry.getValue()) { - state.addOwnedPartitions(partitionsForTask.get(task), entry.getKey()); - } - } - + final ClientState state = new ClientState(); final Set consumers = mkSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); + state.addPreviousTasksAndOffsetSums(CONSUMER_1, getTaskOffsetSums(asList(TASK_0_0, TASK_1_1, TASK_1_3), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_2, getTaskOffsetSums(asList(TASK_0_3, TASK_1_0), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_3, getTaskOffsetSums(asList(TASK_0_1, TASK_0_2, TASK_1_2), EMPTY_TASKS)); + state.initializePrevTasks(emptyMap()); + state.computeTaskLags(UUID_1, getTaskEndOffsetSums(allTasks)); - // We should be able to add a new task without sacrificing stickyness + // We should be able to add a new task without sacrificing stickiness final TaskId newTask = TASK_2_0; allTasks.add(newTask); state.assignActiveTasks(allTasks); final Map> newAssignment = - partitionAssignor.tryStickyAndBalancedTaskAssignmentWithinClient(state, consumers, partitionsForTask, emptySet()); + assignTasksToThreads( + allTasks, + emptySet(), + consumers, + state::previousStatefulActiveTasksForConsumer + ); previousAssignment.get(CONSUMER_2).add(newTask); assertEquivalentAssignment(previousAssignment, newAssignment); } @Test - public void shouldReturnEmptyMapWhenStickyAndBalancedAssignmentIsNotPossibleBecauseNewConsumerJoined() { - configureDefault(); - final ClientState state = new ClientState(); - - final List allTasks = asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, - TASK_1_3); + public void shouldInterleaveTasksByGroupIdDuringNewAssignment() { + final List allTasks = + asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, TASK_2_0, TASK_2_1); - final Map> previousAssignment = mkMap( - mkEntry(CONSUMER_1, asList(TASK_0_0, TASK_1_1, TASK_1_3)), - mkEntry(CONSUMER_2, asList(TASK_0_3, TASK_1_0)), - mkEntry(CONSUMER_3, asList(TASK_0_1, TASK_0_2, TASK_1_2)) + final Map> assignment = mkMap( + mkEntry(CONSUMER_1, new ArrayList<>(asList(TASK_0_0, TASK_0_3, TASK_1_2))), + mkEntry(CONSUMER_2, new ArrayList<>(asList(TASK_0_1, TASK_1_0, TASK_2_0))), + mkEntry(CONSUMER_3, new ArrayList<>(asList(TASK_0_2, TASK_1_1, TASK_2_1))) ); - for (final Map.Entry> entry : previousAssignment.entrySet()) { - for (final TaskId task : entry.getValue()) { - state.addOwnedPartitions(partitionsForTask.get(task), entry.getKey()); - } - } - - // If we add a new consumer here, we cannot produce an assignment that is both sticky and balanced - final Set consumers = mkSet(CONSUMER_1, CONSUMER_2, CONSUMER_3, CONSUMER_4); - state.assignActiveTasks(allTasks); - - assertThat(partitionAssignor.tryStickyAndBalancedTaskAssignmentWithinClient(state, consumers, partitionsForTask, emptySet()), - equalTo(emptyMap())); - } - - @Test - public void shouldReturnEmptyMapWhenStickyAndBalancedAssignmentIsNotPossibleBecauseOtherClientOwnedPartition() { - configureDefault(); final ClientState state = new ClientState(); - - final List allTasks = asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, - TASK_1_3); - - final Map> previousAssignment = mkMap( - mkEntry(CONSUMER_1, new ArrayList<>(asList(TASK_1_1, TASK_1_3))), - mkEntry(CONSUMER_2, new ArrayList<>(asList(TASK_0_3, TASK_1_0))), - mkEntry(CONSUMER_3, new ArrayList<>(asList(TASK_0_1, TASK_0_2, TASK_1_2))) - ); - - for (final Map.Entry> entry : previousAssignment.entrySet()) { - for (final TaskId task : entry.getValue()) { - state.addOwnedPartitions(partitionsForTask.get(task), entry.getKey()); - } - } - - // Add the partitions of TASK_0_0 to allOwnedPartitions but not c1's ownedPartitions/previousAssignment - final Set allOwnedPartitions = new HashSet<>(partitionsForTask.get(TASK_0_0)); - final Set consumers = mkSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); - state.assignActiveTasks(allTasks); - - assertThat(partitionAssignor.tryStickyAndBalancedTaskAssignmentWithinClient(state, consumers, partitionsForTask, allOwnedPartitions), - equalTo(emptyMap())); - } - - @Test - public void shouldInterleaveTasksByGroupId() { - final TaskId taskIdA0 = new TaskId(0, 0); - final TaskId taskIdA1 = new TaskId(0, 1); - final TaskId taskIdA2 = new TaskId(0, 2); - final TaskId taskIdA3 = new TaskId(0, 3); - - final TaskId taskIdB0 = new TaskId(1, 0); - final TaskId taskIdB1 = new TaskId(1, 1); - final TaskId taskIdB2 = new TaskId(1, 2); - - final TaskId taskIdC0 = new TaskId(2, 0); - final TaskId taskIdC1 = new TaskId(2, 1); + state.addPreviousTasksAndOffsetSums(CONSUMER_1, emptyMap()); + state.addPreviousTasksAndOffsetSums(CONSUMER_2, emptyMap()); + state.addPreviousTasksAndOffsetSums(CONSUMER_3, emptyMap()); - final String c1 = "c1"; - final String c2 = "c2"; - final String c3 = "c3"; - - final Set consumers = mkSet(c1, c2, c3); - - final List expectedSubList1 = asList(taskIdA0, taskIdA3, taskIdB2); - final List expectedSubList2 = asList(taskIdA1, taskIdB0, taskIdC0); - final List expectedSubList3 = asList(taskIdA2, taskIdB1, taskIdC1); - - final Map> assignment = new HashMap<>(); - assignment.put(c1, expectedSubList1); - assignment.put(c2, expectedSubList2); - assignment.put(c3, expectedSubList3); - - final List tasks = asList(taskIdC0, taskIdC1, taskIdB0, taskIdB1, taskIdB2, taskIdA0, taskIdA1, taskIdA2, taskIdA3); - Collections.shuffle(tasks); + Collections.shuffle(allTasks); final Map> interleavedTaskIds = - StreamsPartitionAssignor.interleaveConsumerTasksByGroupId(tasks, consumers); + assignTasksToThreads( + allTasks, + emptySet(), + consumers, + state::previousStatefulActiveTasksForConsumer + ); assertThat(interleavedTaskIds, equalTo(assignment)); } -*/ @Test public void testEagerSubscription() { From 8300d491ae64f4f8ab4fba5b2e9baa1ffadd9a7c Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Mon, 1 Jun 2020 22:21:53 -0700 Subject: [PATCH 06/23] use sorted set for test --- .../internals/StreamsPartitionAssignor.java | 12 +++++++----- .../internals/StreamsPartitionAssignorTest.java | 14 ++++++++------ 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 0714b0dcda490..d038e0ecce182 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -19,6 +19,8 @@ import java.util.Iterator; import java.util.PriorityQueue; import java.util.Queue; +import java.util.SortedSet; +import java.util.TreeSet; import java.util.function.Function; import org.apache.kafka.clients.admin.Admin; import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo; @@ -114,7 +116,7 @@ public int hashCode() { private static class ClientMetadata { private final HostInfo hostInfo; - private final Set consumers; + private final SortedSet consumers; private final ClientState state; ClientMetadata(final String endPoint) { @@ -123,7 +125,7 @@ private static class ClientMetadata { hostInfo = HostInfo.buildFromEndpoint(endPoint); // initialize the consumer memberIds - consumers = new HashSet<>(); + consumers = new TreeSet<>(); // initialize the client state state = new ClientState(); @@ -878,7 +880,7 @@ private Map computeNewAssignment(final Map consumers = clientMetadata.consumers; + final SortedSet consumers = clientMetadata.consumers; final Map> activeTaskAssignment = assignTasksToThreads( state.statefulActiveTasks(), @@ -1075,7 +1077,7 @@ private static Map> buildStandbyTaskMap(final Collec */ static Map> assignTasksToThreads(final Collection statefulTasksToAssign, final Collection statelessTasksToAssign, - final Set consumers, + final SortedSet consumers, final Function> previousTasksForConsumer) { final Map> assignment = new HashMap<>(); for (final String consumer : consumers) { @@ -1085,7 +1087,7 @@ static Map> assignTasksToThreads(final Collection s final int minStatefulTasksPerThread = (int) Math.floor(((double) statefulTasksToAssign.size()) / consumers.size()); final PriorityQueue unassignedStatefulTasks = new PriorityQueue<>(statefulTasksToAssign); - final PriorityQueue consumersToFill = new PriorityQueue<>(); + final Queue consumersToFill = new LinkedList<>(); // keep track of tasks that we have to skip during the first pass in case we can reassign them later final Map unassignedTaskToPreviousOwner = new HashMap<>(); diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index 6cdc1002b903d..17955a3faf3f0 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.streams.processor.internals; +import java.util.SortedSet; import org.apache.kafka.clients.admin.Admin; import org.apache.kafka.clients.admin.AdminClient; import org.apache.kafka.clients.admin.ListOffsetsResult; @@ -88,6 +89,7 @@ import static org.apache.kafka.common.utils.Utils.mkEntry; import static org.apache.kafka.common.utils.Utils.mkMap; import static org.apache.kafka.common.utils.Utils.mkSet; +import static org.apache.kafka.common.utils.Utils.mkSortedSet; import static org.apache.kafka.streams.processor.internals.StreamsPartitionAssignor.assignTasksToThreads; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.EMPTY_CHANGELOG_END_OFFSETS; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.EMPTY_TASKS; @@ -336,7 +338,7 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNothingChanges() { ); final ClientState state = new ClientState(); - final Set consumers = mkSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); + final SortedSet consumers = mkSortedSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); state.addPreviousTasksAndOffsetSums(CONSUMER_1, getTaskOffsetSums(asList(TASK_0_0, TASK_1_1, TASK_1_3), EMPTY_TASKS)); state.addPreviousTasksAndOffsetSums(CONSUMER_2, getTaskOffsetSums(asList(TASK_0_3, TASK_1_0), EMPTY_TASKS)); state.addPreviousTasksAndOffsetSums(CONSUMER_3, getTaskOffsetSums(asList(TASK_0_1, TASK_0_2, TASK_1_2), EMPTY_TASKS)); @@ -357,7 +359,7 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNothingChanges() { @Test public void shouldProduceStickyAndBalancedAssignmentWhenNewTasksAreAdded() { final List allTasks = - asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, TASK_1_3); + new ArrayList<>(asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, TASK_1_3)); final Map> previousAssignment = mkMap( mkEntry(CONSUMER_1, new ArrayList<>(asList(TASK_0_0, TASK_1_1, TASK_1_3))), @@ -366,10 +368,10 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNewTasksAreAdded() { ); final ClientState state = new ClientState(); - final Set consumers = mkSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); - state.addPreviousTasksAndOffsetSums(CONSUMER_1, getTaskOffsetSums(asList(TASK_0_0, TASK_1_1, TASK_1_3), EMPTY_TASKS)); + final SortedSet consumers = mkSortedSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); + state.addPreviousTasksAndOffsetSums(CONSUMER_1, getTaskOffsetSums(asList(TASK_0_1, TASK_0_2, TASK_1_2), EMPTY_TASKS)); state.addPreviousTasksAndOffsetSums(CONSUMER_2, getTaskOffsetSums(asList(TASK_0_3, TASK_1_0), EMPTY_TASKS)); - state.addPreviousTasksAndOffsetSums(CONSUMER_3, getTaskOffsetSums(asList(TASK_0_1, TASK_0_2, TASK_1_2), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_3, getTaskOffsetSums(asList(TASK_0_0, TASK_1_1, TASK_1_3), EMPTY_TASKS)); state.initializePrevTasks(emptyMap()); state.computeTaskLags(UUID_1, getTaskEndOffsetSums(allTasks)); @@ -402,7 +404,7 @@ public void shouldInterleaveTasksByGroupIdDuringNewAssignment() { ); final ClientState state = new ClientState(); - final Set consumers = mkSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); + final SortedSet consumers = mkSortedSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); state.addPreviousTasksAndOffsetSums(CONSUMER_1, emptyMap()); state.addPreviousTasksAndOffsetSums(CONSUMER_2, emptyMap()); state.addPreviousTasksAndOffsetSums(CONSUMER_3, emptyMap()); From d40dd5ae74c3a0b3fc22b60fd0380c44a44d3f7b Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Tue, 2 Jun 2020 10:44:40 -0700 Subject: [PATCH 07/23] checkstyle --- .../streams/processor/internals/StreamsPartitionAssignor.java | 2 +- .../processor/internals/StreamsPartitionAssignorTest.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index d038e0ecce182..adb31fcef3bea 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -932,7 +932,7 @@ private Map computeNewAssignment(final Map assignment, final ClientMetadata clientMetadata, diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index 17955a3faf3f0..5ce36008aed72 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -1963,7 +1963,7 @@ private static Map getTaskOffsetSums(final Collection acti // Stub end offsets sums for situations where we don't really care about computing exact lags private static Map getTaskEndOffsetSums(final Collection allStatefulTasks) { - return allStatefulTasks.stream().collect(Collectors.toMap(t -> t, t-> Long.MAX_VALUE)); + return allStatefulTasks.stream().collect(Collectors.toMap(t -> t, t -> Long.MAX_VALUE)); } } From 5cb6f40bd138fbfa57ac60fccafeb18a4e12c45e Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Tue, 2 Jun 2020 12:33:33 -0700 Subject: [PATCH 08/23] add unit tests --- .../internals/assignment/ClientStateTest.java | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java index d0e95cf4d2c0d..8008ceabb36ed 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java @@ -32,6 +32,7 @@ import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_1; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_2; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_3; +import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_1; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.hasActiveTasks; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.hasStandbyTasks; import static org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo.UNKNOWN_OFFSET_SUM; @@ -307,6 +308,46 @@ public void shouldAddTasksWithLatestOffsetToPrevActiveTasks() { assertTrue(client.prevStandbyTasks().isEmpty()); } + @Test + public void shouldGetPreviousStatefulActiveTasksForConsumer() { + client.addPreviousTasksAndOffsetSums("c1", Collections.singletonMap(TASK_0_0, Task.LATEST_OFFSET)); + client.addPreviousTasksAndOffsetSums("c2", Collections.singletonMap(TASK_0_1, 0L)); + client.addPreviousTasksAndOffsetSums("c3", Collections.singletonMap(TASK_0_2, Task.LATEST_OFFSET)); + + client.initializePrevTasks(Collections.emptyMap()); + client.computeTaskLags( + UUID_1, + mkMap( + mkEntry(TASK_0_0, 100L), + mkEntry(TASK_0_1, 100L) + ) + ); + + assertThat(client.previousStatefulActiveTasksForConsumer("c1"), equalTo(Collections.singleton(TASK_0_0))); + assertTrue(client.previousStatefulActiveTasksForConsumer("c2").isEmpty()); + assertTrue(client.previousStatefulActiveTasksForConsumer("c3").isEmpty()); + } + + @Test + public void shouldGetPreviousStandbyTasksForConsumer() { + client.addPreviousTasksAndOffsetSums("c1", Collections.singletonMap(TASK_0_0, Task.LATEST_OFFSET)); + client.addPreviousTasksAndOffsetSums("c2", Collections.singletonMap(TASK_0_1, 0L)); + client.addPreviousTasksAndOffsetSums("c3", Collections.singletonMap(TASK_0_2, Task.LATEST_OFFSET)); + + client.initializePrevTasks(Collections.emptyMap()); + client.computeTaskLags( + UUID_1, + mkMap( + mkEntry(TASK_0_0, 100L), + mkEntry(TASK_0_1, 100L) + ) + ); + + assertTrue(client.previousStandbyTasksForConsumer("c1").isEmpty()); + assertThat(client.previousStandbyTasksForConsumer("c2"), equalTo(Collections.singleton(TASK_0_1))); + assertTrue(client.previousStandbyTasksForConsumer("c3").isEmpty()); + } + @Test public void shouldAddTasksInOffsetSumsMapToPrevStandbyTasks() { final Map taskOffsetSums = mkMap( From 72e64ef6089a6ebd7c197e4a0f383c252c9e09f8 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Tue, 2 Jun 2020 13:25:29 -0700 Subject: [PATCH 09/23] filter for statefulness in SPA --- .../internals/StreamsPartitionAssignor.java | 14 ++++++++------ .../internals/assignment/ClientState.java | 4 ++-- .../internals/StreamsPartitionAssignorTest.java | 6 +++--- .../internals/assignment/ClientStateTest.java | 6 +++--- 4 files changed, 16 insertions(+), 14 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index adb31fcef3bea..b925747a7e3b5 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -886,7 +886,7 @@ private Map computeNewAssignment(final Map> standbyTaskAssignment = assignTasksToThreads( @@ -1097,11 +1097,13 @@ static Map> assignTasksToThreads(final Collection s int i = 0; for (final TaskId task : previousTasksForConsumer.apply(consumer)) { - if (i < minStatefulTasksPerThread) { - threadAssignment.add(task); - unassignedStatefulTasks.remove(task); - } else { - unassignedTaskToPreviousOwner.put(task, consumer); + if (unassignedStatefulTasks.contains(task)) { + if (i < minStatefulTasksPerThread) { + threadAssignment.add(task); + unassignedStatefulTasks.remove(task); + } else { + unassignedTaskToPreviousOwner.put(task, consumer); + } } ++i; } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java index 604022c003005..1e21f6ace21fa 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java @@ -318,10 +318,10 @@ public Set statelessActiveTasks() { return activeTasks.stream().filter(task -> !isStateful(task)).collect(Collectors.toSet()); } - public Set previousStatefulActiveTasksForConsumer(final String memberId) { + public Set previousActiveTasksForConsumer(final String memberId) { final Set prevTasks = new HashSet<>(); for (final TaskId task : consumerToPreviousTaskIds.get(memberId)) { - if (isStateful(task) && prevActiveTasks.contains(task)) { + if (prevActiveTasks.contains(task)) { prevTasks.add(task); } } diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index 5ce36008aed72..168f60f71c9f6 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -351,7 +351,7 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNothingChanges() { allTasks, emptySet(), consumers, - state::previousStatefulActiveTasksForConsumer + state::previousActiveTasksForConsumer ) ); } @@ -385,7 +385,7 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNewTasksAreAdded() { allTasks, emptySet(), consumers, - state::previousStatefulActiveTasksForConsumer + state::previousActiveTasksForConsumer ); previousAssignment.get(CONSUMER_2).add(newTask); @@ -416,7 +416,7 @@ public void shouldInterleaveTasksByGroupIdDuringNewAssignment() { allTasks, emptySet(), consumers, - state::previousStatefulActiveTasksForConsumer + state::previousActiveTasksForConsumer ); assertThat(interleavedTaskIds, equalTo(assignment)); diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java index 8008ceabb36ed..7bf139b23c16d 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java @@ -323,9 +323,9 @@ public void shouldGetPreviousStatefulActiveTasksForConsumer() { ) ); - assertThat(client.previousStatefulActiveTasksForConsumer("c1"), equalTo(Collections.singleton(TASK_0_0))); - assertTrue(client.previousStatefulActiveTasksForConsumer("c2").isEmpty()); - assertTrue(client.previousStatefulActiveTasksForConsumer("c3").isEmpty()); + assertThat(client.previousActiveTasksForConsumer("c1"), equalTo(Collections.singleton(TASK_0_0))); + assertTrue(client.previousActiveTasksForConsumer("c2").isEmpty()); + assertTrue(client.previousActiveTasksForConsumer("c3").isEmpty()); } @Test From acb7a4c0993b7fd7c455b734aa895583341c965d Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Tue, 2 Jun 2020 14:52:52 -0700 Subject: [PATCH 10/23] fixing up tests --- .../internals/StreamsPartitionAssignor.java | 127 +++++++++--------- .../internals/assignment/ClientState.java | 43 +++--- .../StreamsPartitionAssignorTest.java | 27 ++-- .../internals/assignment/ClientStateTest.java | 50 ++++--- 4 files changed, 123 insertions(+), 124 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index b925747a7e3b5..f07d24a9d1025 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -21,7 +21,6 @@ import java.util.Queue; import java.util.SortedSet; import java.util.TreeSet; -import java.util.function.Function; import org.apache.kafka.clients.admin.Admin; import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo; import org.apache.kafka.clients.consumer.ConsumerGroupMetadata; @@ -886,14 +885,14 @@ private Map computeNewAssignment(final Map> standbyTaskAssignment = assignTasksToThreads( state.standbyTasks(), Collections.emptySet(), consumers, - state::previousStandbyTasksForConsumer + state ); // Arbitrarily choose the leader's client to be responsible for triggering the probing rebalance @@ -1078,12 +1077,17 @@ private static Map> buildStandbyTaskMap(final Collec static Map> assignTasksToThreads(final Collection statefulTasksToAssign, final Collection statelessTasksToAssign, final SortedSet consumers, - final Function> previousTasksForConsumer) { + final ClientState state) { final Map> assignment = new HashMap<>(); for (final String consumer : consumers) { assignment.put(consumer, new ArrayList<>()); } + final List unassignedStatelessTasks = new ArrayList<>(statelessTasksToAssign); + Collections.sort(unassignedStatelessTasks); + + final Iterator unassignedStatelessTasksIter = unassignedStatelessTasks.iterator(); + final int minStatefulTasksPerThread = (int) Math.floor(((double) statefulTasksToAssign.size()) / consumers.size()); final PriorityQueue unassignedStatefulTasks = new PriorityQueue<>(statefulTasksToAssign); @@ -1091,77 +1095,76 @@ static Map> assignTasksToThreads(final Collection s // keep track of tasks that we have to skip during the first pass in case we can reassign them later final Map unassignedTaskToPreviousOwner = new HashMap<>(); - // First assign stateful tasks to previous owner, up to the min expected tasks/thread - for (final String consumer : consumers) { - final List threadAssignment = assignment.get(consumer); - - int i = 0; - for (final TaskId task : previousTasksForConsumer.apply(consumer)) { - if (unassignedStatefulTasks.contains(task)) { - if (i < minStatefulTasksPerThread) { - threadAssignment.add(task); - unassignedStatefulTasks.remove(task); - } else { - unassignedTaskToPreviousOwner.put(task, consumer); + if (!unassignedStatefulTasks.isEmpty()) { + // First assign stateful tasks to previous owner, up to the min expected tasks/thread + for (final String consumer : consumers) { + final List threadAssignment = assignment.get(consumer); + + int i = 0; + for (final TaskId task : state.previousTasksForConsumer(consumer)) { + if (unassignedStatefulTasks.contains(task)) { + if (i < minStatefulTasksPerThread) { + threadAssignment.add(task); + unassignedStatefulTasks.remove(task); + } else { + unassignedTaskToPreviousOwner.put(task, consumer); + } } + ++i; } - ++i; - } - if (threadAssignment.size() < minStatefulTasksPerThread) { - consumersToFill.offer(consumer); - } - } - - // Next interleave remaining unassigned tasks amongst unfilled consumers - while (!consumersToFill.isEmpty()) { - final TaskId task = unassignedStatefulTasks.poll(); - if (task != null) { - final String consumer = consumersToFill.poll(); - final List threadAssignment = assignment.get(consumer); - threadAssignment.add(task); if (threadAssignment.size() < minStatefulTasksPerThread) { consumersToFill.offer(consumer); } - } else { - throw new IllegalStateException("Ran out of unassigned stateful tasks but some members were not at capacity"); } - } - // At this point all consumers are at the min capacity, so there may be up to N - 1 unassigned - // stateful tasks remaining that should now be distributed over the consumers - consumersToFill.addAll(consumers); + // Next interleave remaining unassigned tasks amongst unfilled consumers + while (!consumersToFill.isEmpty()) { + final TaskId task = unassignedStatefulTasks.poll(); + if (task != null) { + final String consumer = consumersToFill.poll(); + final List threadAssignment = assignment.get(consumer); + threadAssignment.add(task); + if (threadAssignment.size() < minStatefulTasksPerThread) { + consumersToFill.offer(consumer); + } + } else { + throw new IllegalStateException("Ran out of unassigned stateful tasks but some members were not at capacity"); + } + } - // Go over the tasks we skipped earlier and assign them to their previous owner when possible - for (final Map.Entry taskEntry : unassignedTaskToPreviousOwner.entrySet()) { - final TaskId task = taskEntry.getKey(); - final String consumer = taskEntry.getValue(); - if (consumersToFill.contains(consumer)) { - assignment.get(consumer).add(task); - unassignedStatefulTasks.remove(task); - consumersToFill.remove(consumer); + // At this point all consumers are at the min capacity, so there may be up to N - 1 unassigned + // stateful tasks remaining that should now be distributed over the consumers + consumersToFill.addAll(consumers); + + // Go over the tasks we skipped earlier and assign them to their previous owner when possible + for (final Map.Entry taskEntry : unassignedTaskToPreviousOwner.entrySet()) { + final TaskId task = taskEntry.getKey(); + final String consumer = taskEntry.getValue(); + if (consumersToFill.contains(consumer)) { + assignment.get(consumer).add(task); + unassignedStatefulTasks.remove(task); + consumersToFill.remove(consumer); + } } - } - // Now just distribute the remaining unassigned tasks over the consumers still at min capacity - for (final TaskId task : unassignedStatefulTasks) { - final String consumer = consumersToFill.poll(); - assignment.get(consumer).add(task); - } + // Now just distribute the remaining unassigned tasks over the consumers still at min capacity + for (final TaskId task : unassignedStatefulTasks) { + final String consumer = consumersToFill.poll(); + assignment.get(consumer).add(task); + } - // There must be at least one consumer still at min capacity while all the others are at min - // capacity + 1, so start distributing stateless tasks to get all consumers back to the same count - final List unassignedStatelessTasks = new ArrayList<>(statelessTasksToAssign); - Collections.sort(unassignedStatelessTasks); - final Iterator unassignedStatelessTasksIter = unassignedStatelessTasks.iterator(); - while (unassignedStatelessTasksIter.hasNext()) { - final TaskId task = unassignedStatelessTasksIter.next(); - final String consumer = consumersToFill.poll(); - if (consumer != null) { - assignment.get(consumer).add(task); - } else { - break; + // There must be at least one consumer still at min capacity while all the others are at min + // capacity + 1, so start distributing stateless tasks to get all consumers back to the same count + while (unassignedStatelessTasksIter.hasNext()) { + final String consumer = consumersToFill.poll(); + if (consumer != null) { + final TaskId task = unassignedStatelessTasksIter.next(); + assignment.get(consumer).add(task); + } else { + break; + } } } diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java index 1e21f6ace21fa..2ca74dd4f430f 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java @@ -16,6 +16,7 @@ */ package org.apache.kafka.streams.processor.internals.assignment; +import java.util.SortedSet; import java.util.stream.Collectors; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.streams.processor.TaskId; @@ -37,6 +38,7 @@ import static java.util.Collections.unmodifiableMap; import static java.util.Collections.unmodifiableSet; import static java.util.Comparator.comparing; +import static java.util.Comparator.comparingLong; import static org.apache.kafka.common.utils.Utils.union; import static org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo.UNKNOWN_OFFSET_SUM; @@ -49,7 +51,7 @@ public class ClientState { private final Set prevActiveTasks; private final Set prevStandbyTasks; - private final Map> consumerToPreviousTaskIds; + private final Map> consumerToPreviousStatefulTaskIds; private final Map ownedPartitions; private final Map taskOffsetSums; // contains only stateful tasks we previously owned private final Map taskLagTotals; // contains lag for all stateful tasks in the app topology @@ -65,7 +67,7 @@ public ClientState() { standbyTasks = new TreeSet<>(); prevActiveTasks = new TreeSet<>(); prevStandbyTasks = new TreeSet<>(); - consumerToPreviousTaskIds = new TreeMap<>(); + consumerToPreviousStatefulTaskIds = new TreeMap<>(); ownedPartitions = new TreeMap<>(TOPIC_PARTITION_COMPARATOR); taskOffsetSums = new TreeMap<>(); taskLagTotals = new TreeMap<>(); @@ -76,7 +78,7 @@ private ClientState(final Set activeTasks, final Set standbyTasks, final Set prevActiveTasks, final Set prevStandbyTasks, - final Map> consumerToPreviousTaskIds, + final Map> consumerToPreviousStatefulTaskIds, final SortedMap ownedPartitions, final Map taskOffsetSums, final Map taskLagTotals, @@ -85,7 +87,7 @@ private ClientState(final Set activeTasks, this.standbyTasks = standbyTasks; this.prevActiveTasks = prevActiveTasks; this.prevStandbyTasks = prevStandbyTasks; - this.consumerToPreviousTaskIds = consumerToPreviousTaskIds; + this.consumerToPreviousStatefulTaskIds = consumerToPreviousStatefulTaskIds; this.ownedPartitions = ownedPartitions; this.taskOffsetSums = taskOffsetSums; this.taskLagTotals = taskLagTotals; @@ -100,7 +102,7 @@ public ClientState(final Set previousActiveTasks, standbyTasks = new TreeSet<>(); prevActiveTasks = unmodifiableSet(new TreeSet<>(previousActiveTasks)); prevStandbyTasks = unmodifiableSet(new TreeSet<>(previousStandbyTasks)); - consumerToPreviousTaskIds = new TreeMap<>(); + consumerToPreviousStatefulTaskIds = new TreeMap<>(); ownedPartitions = new TreeMap<>(TOPIC_PARTITION_COMPARATOR); taskOffsetSums = emptyMap(); this.taskLagTotals = unmodifiableMap(taskLagTotals); @@ -115,7 +117,7 @@ public ClientState copy() { new TreeSet<>(standbyTasks), new TreeSet<>(prevActiveTasks), new TreeSet<>(prevStandbyTasks), - new TreeMap<>(consumerToPreviousTaskIds), + new TreeMap<>(consumerToPreviousStatefulTaskIds), newOwnedPartitions, new TreeMap<>(taskOffsetSums), new TreeMap<>(taskLagTotals), @@ -252,7 +254,7 @@ public void addOwnedPartitions(final Collection ownedPartitions, public void addPreviousTasksAndOffsetSums(final String consumerId, final Map taskOffsetSums) { this.taskOffsetSums.putAll(taskOffsetSums); - consumerToPreviousTaskIds.put(consumerId, taskOffsetSums.keySet()); + consumerToPreviousStatefulTaskIds.put(consumerId, taskOffsetSums.keySet()); } public void initializePrevTasks(final Map taskForPartitionMap) { @@ -318,25 +320,18 @@ public Set statelessActiveTasks() { return activeTasks.stream().filter(task -> !isStateful(task)).collect(Collectors.toSet()); } - public Set previousActiveTasksForConsumer(final String memberId) { - final Set prevTasks = new HashSet<>(); - for (final TaskId task : consumerToPreviousTaskIds.get(memberId)) { - if (prevActiveTasks.contains(task)) { - prevTasks.add(task); - } - } - return prevTasks; - } + // Return a list of that consumer's previous tasks in increasing lag order + public SortedSet previousTasksForConsumer(final String memberId) { + final Set prevTasks = consumerToPreviousStatefulTaskIds.get(memberId); - // Equivalent to previousStatefulStandbyTasksForConsumer since all standbys are by definition stateful - public Set previousStandbyTasksForConsumer(final String memberId) { - final Set prevTasks = new HashSet<>(); - for (final TaskId task : consumerToPreviousTaskIds.get(memberId)) { - if (prevStandbyTasks.contains(task)) { - prevTasks.add(task); - } + // If we were unable to fetch the end offsets and could not compute lags, just return in task order + if (taskLagTotals.isEmpty()) { + return new TreeSet<>(prevTasks); + } else { + final SortedSet prevTasksByLag = new TreeSet<>(comparingLong(this::lagFor).thenComparing(TaskId::compareTo)); + prevTasksByLag.addAll(prevTasks); + return prevTasksByLag; } - return prevTasks; } boolean hasUnfulfilledQuota(final int tasksPerThread) { diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index 168f60f71c9f6..91e955438f4bb 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -128,8 +128,7 @@ public class StreamsPartitionAssignorTest { private static final String CONSUMER_1 = "consumer1"; private static final String CONSUMER_2 = "consumer2"; private static final String CONSUMER_3 = "consumer3"; - private static final String CONSUMER_4 = "consumer4"; - + private final Set allTopics = mkSet("topic1", "topic2"); private final TopicPartition t1p0 = new TopicPartition("topic1", 0); @@ -351,7 +350,7 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNothingChanges() { allTasks, emptySet(), consumers, - state::previousActiveTasksForConsumer + state ) ); } @@ -369,9 +368,9 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNewTasksAreAdded() { final ClientState state = new ClientState(); final SortedSet consumers = mkSortedSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); - state.addPreviousTasksAndOffsetSums(CONSUMER_1, getTaskOffsetSums(asList(TASK_0_1, TASK_0_2, TASK_1_2), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_1, getTaskOffsetSums(asList(TASK_0_0, TASK_1_1, TASK_1_3), EMPTY_TASKS)); state.addPreviousTasksAndOffsetSums(CONSUMER_2, getTaskOffsetSums(asList(TASK_0_3, TASK_1_0), EMPTY_TASKS)); - state.addPreviousTasksAndOffsetSums(CONSUMER_3, getTaskOffsetSums(asList(TASK_0_0, TASK_1_1, TASK_1_3), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_3, getTaskOffsetSums(asList(TASK_0_1, TASK_0_2, TASK_1_2), EMPTY_TASKS)); state.initializePrevTasks(emptyMap()); state.computeTaskLags(UUID_1, getTaskEndOffsetSums(allTasks)); @@ -385,7 +384,7 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNewTasksAreAdded() { allTasks, emptySet(), consumers, - state::previousActiveTasksForConsumer + state ); previousAssignment.get(CONSUMER_2).add(newTask); @@ -416,7 +415,7 @@ public void shouldInterleaveTasksByGroupIdDuringNewAssignment() { allTasks, emptySet(), consumers, - state::previousActiveTasksForConsumer + state ); assertThat(interleavedTaskIds, equalTo(assignment)); @@ -573,16 +572,6 @@ public void shouldAssignEvenlyAcrossConsumersOneClientMultipleThreads() { final List topics = asList("topic1", "topic2"); - final TaskId taskIdA0 = new TaskId(0, 0); - final TaskId taskIdA1 = new TaskId(0, 1); - final TaskId taskIdA2 = new TaskId(0, 2); - final TaskId taskIdA3 = new TaskId(0, 3); - - final TaskId taskIdB0 = new TaskId(1, 0); - final TaskId taskIdB1 = new TaskId(1, 1); - final TaskId taskIdB2 = new TaskId(1, 2); - final TaskId taskIdB3 = new TaskId(1, 3); - configureDefault(); subscriptions.put("consumer10", @@ -605,12 +594,12 @@ public void shouldAssignEvenlyAcrossConsumersOneClientMultipleThreads() { // the first consumer final AssignmentInfo info10 = AssignmentInfo.decode(assignments.get("consumer10").userData()); - final List expectedInfo10TaskIds = asList(taskIdA0, taskIdA2, taskIdB0, taskIdB2); + final List expectedInfo10TaskIds = asList(TASK_0_0, TASK_0_2, TASK_1_0, TASK_1_2); assertEquals(expectedInfo10TaskIds, info10.activeTasks()); // the second consumer final AssignmentInfo info11 = AssignmentInfo.decode(assignments.get("consumer11").userData()); - final List expectedInfo11TaskIds = asList(taskIdA1, taskIdA3, taskIdB1, taskIdB3); + final List expectedInfo11TaskIds = asList(TASK_0_1, TASK_0_3, TASK_1_1, TASK_1_3); assertEquals(expectedInfo11TaskIds, info11.activeTasks()); } diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java index 7bf139b23c16d..9b68e477f4f3d 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java @@ -28,6 +28,7 @@ import static org.apache.kafka.common.utils.Utils.mkEntry; import static org.apache.kafka.common.utils.Utils.mkMap; import static org.apache.kafka.common.utils.Utils.mkSet; +import static org.apache.kafka.common.utils.Utils.mkSortedSet; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_0; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_1; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_2; @@ -309,43 +310,54 @@ public void shouldAddTasksWithLatestOffsetToPrevActiveTasks() { } @Test - public void shouldGetPreviousStatefulActiveTasksForConsumer() { - client.addPreviousTasksAndOffsetSums("c1", Collections.singletonMap(TASK_0_0, Task.LATEST_OFFSET)); - client.addPreviousTasksAndOffsetSums("c2", Collections.singletonMap(TASK_0_1, 0L)); - client.addPreviousTasksAndOffsetSums("c3", Collections.singletonMap(TASK_0_2, Task.LATEST_OFFSET)); + public void shouldReturnPreviousStatefulTasksForConsumer() { + client.addPreviousTasksAndOffsetSums("c1", Collections.singletonMap(TASK_0_1, Task.LATEST_OFFSET)); + client.addPreviousTasksAndOffsetSums("c2", Collections.singletonMap(TASK_0_2, 0L)); + client.addPreviousTasksAndOffsetSums("c3", Collections.emptyMap()); client.initializePrevTasks(Collections.emptyMap()); client.computeTaskLags( UUID_1, mkMap( - mkEntry(TASK_0_0, 100L), - mkEntry(TASK_0_1, 100L) + mkEntry(TASK_0_1, 1_000L), + mkEntry(TASK_0_2, 1_000L) ) ); - assertThat(client.previousActiveTasksForConsumer("c1"), equalTo(Collections.singleton(TASK_0_0))); - assertTrue(client.previousActiveTasksForConsumer("c2").isEmpty()); - assertTrue(client.previousActiveTasksForConsumer("c3").isEmpty()); + assertThat(client.previousTasksForConsumer("c1"), equalTo(mkSortedSet(TASK_0_1))); + assertThat(client.previousTasksForConsumer("c2"), equalTo(mkSortedSet(TASK_0_2))); + assertTrue(client.previousTasksForConsumer("c3").isEmpty()); } @Test - public void shouldGetPreviousStandbyTasksForConsumer() { - client.addPreviousTasksAndOffsetSums("c1", Collections.singletonMap(TASK_0_0, Task.LATEST_OFFSET)); - client.addPreviousTasksAndOffsetSums("c2", Collections.singletonMap(TASK_0_1, 0L)); - client.addPreviousTasksAndOffsetSums("c3", Collections.singletonMap(TASK_0_2, Task.LATEST_OFFSET)); - + public void shouldReturnPreviousStatefulTasksForConsumerWhenLagIsNotComputed() { + client.addPreviousTasksAndOffsetSums("c1", Collections.singletonMap(TASK_0_1, 1000L)); client.initializePrevTasks(Collections.emptyMap()); + + assertThat(client.previousTasksForConsumer("c1"), equalTo(mkSortedSet(TASK_0_1))); + } + + @Test + public void shouldReturnPreviousStatefulTasksForConsumerInIncreasingLagOrder() { + client.addPreviousTasksAndOffsetSums("c1", mkMap( + mkEntry(TASK_0_1, 100L), + mkEntry(TASK_0_2, 0L), + mkEntry(TASK_0_3, Task.LATEST_OFFSET) + )); + client.computeTaskLags( UUID_1, mkMap( - mkEntry(TASK_0_0, 100L), - mkEntry(TASK_0_1, 100L) + mkEntry(TASK_0_1, 1_000L), + mkEntry(TASK_0_2, 1_000L), + mkEntry(TASK_0_3, 1_000L) + ) ); - assertTrue(client.previousStandbyTasksForConsumer("c1").isEmpty()); - assertThat(client.previousStandbyTasksForConsumer("c2"), equalTo(Collections.singleton(TASK_0_1))); - assertTrue(client.previousStandbyTasksForConsumer("c3").isEmpty()); + client.initializePrevTasks(Collections.emptyMap()); + + assertThat(client.previousTasksForConsumer("c1"), equalTo(mkSortedSet(TASK_0_3, TASK_0_2, TASK_0_1))); } @Test From f580e63fb889115ef66987656459baefcce3473c Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Wed, 3 Jun 2020 11:16:05 -0700 Subject: [PATCH 11/23] bump log to INFO --- .../processor/internals/StreamsPartitionAssignor.java | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index f07d24a9d1025..578cbbcd4d0ed 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -982,16 +982,14 @@ private boolean addClientAssignments(final Map assignment, if (tasksRevoked) { // TODO: once KAFKA-10078 is resolved we can leave it to the client to trigger this rebalance - log.debug("Requesting followup rebalance be scheduled immediately due to tasks changing ownership."); + log.info("Requesting followup rebalance be scheduled immediately due to tasks changing ownership."); info.setNextRebalanceTime(0L); followupRebalanceRequiredForRevokedTasks = true; // Don't bother to schedule a probing rebalance if an immediate one is already scheduled shouldEncodeProbingRebalance = false; - } - - if (shouldEncodeProbingRebalance) { + } else if (shouldEncodeProbingRebalance) { final long nextRebalanceTimeMs = time.milliseconds() + probingRebalanceIntervalMs(); - log.debug("Requesting followup rebalance be scheduled for {} ms to probe for caught-up replica tasks.", nextRebalanceTimeMs); + log.info("Requesting followup rebalance be scheduled for {} ms to probe for caught-up replica tasks.", nextRebalanceTimeMs); info.setNextRebalanceTime(nextRebalanceTimeMs); shouldEncodeProbingRebalance = false; } From 7c18e59541ed5a5a482766f7bf30d4ba253ac2f8 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Thu, 4 Jun 2020 09:54:04 -0700 Subject: [PATCH 12/23] pull sorting into SPA --- .../internals/StreamsPartitionAssignor.java | 9 ++++- .../internals/assignment/ClientState.java | 27 ++++++------- .../internals/assignment/ClientStateTest.java | 38 +++++++------------ 3 files changed, 33 insertions(+), 41 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 578cbbcd4d0ed..5ae6e944dd9e5 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -70,6 +70,7 @@ import java.util.function.Supplier; import java.util.stream.Collectors; +import static java.util.Comparator.comparingLong; import static java.util.UUID.randomUUID; import static org.apache.kafka.streams.processor.internals.ClientUtils.fetchEndOffsets; import static org.apache.kafka.streams.processor.internals.assignment.StreamsAssignmentProtocolVersions.EARLIEST_PROBEABLE_VERSION; @@ -1099,7 +1100,7 @@ static Map> assignTasksToThreads(final Collection s final List threadAssignment = assignment.get(consumer); int i = 0; - for (final TaskId task : state.previousTasksForConsumer(consumer)) { + for (final TaskId task : getPreviousTasksByLag(state, consumer)) { if (unassignedStatefulTasks.contains(task)) { if (i < minStatefulTasksPerThread) { threadAssignment.add(task); @@ -1179,6 +1180,12 @@ static Map> assignTasksToThreads(final Collection s return assignment; } + private static SortedSet getPreviousTasksByLag(final ClientState state, final String consumer) { + final SortedSet prevTasksByLag = new TreeSet<>(comparingLong(state::lagFor).thenComparing(TaskId::compareTo)); + prevTasksByLag.addAll(state.previousTasksForConsumer(consumer)); + return prevTasksByLag; + } + private void validateMetadataVersions(final int receivedAssignmentMetadataVersion, final int latestCommonlySupportedVersion) { diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java index 2ca74dd4f430f..cff70d4f9d3a7 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java @@ -302,14 +302,19 @@ public void computeTaskLags(final UUID uuid, final Map allTaskEndO * @return end offset sum - offset sum * Task.LATEST_OFFSET if this was previously an active running task on this client */ - long lagFor(final TaskId task) { - final Long totalLag = taskLagTotals.get(task); + public long lagFor(final TaskId task) { + final Long totalLag; + if (taskLagTotals.isEmpty()) { + // If we couldn't compute the task lags due to failure to fetch offsets, just return a flat constant + totalLag = 0L; + } else { + totalLag = taskLagTotals.get(task); + } if (totalLag == null) { throw new IllegalStateException("Tried to lookup lag for unknown task " + task); - } else { - return totalLag; } + return totalLag; } public Set statefulActiveTasks() { @@ -320,18 +325,8 @@ public Set statelessActiveTasks() { return activeTasks.stream().filter(task -> !isStateful(task)).collect(Collectors.toSet()); } - // Return a list of that consumer's previous tasks in increasing lag order - public SortedSet previousTasksForConsumer(final String memberId) { - final Set prevTasks = consumerToPreviousStatefulTaskIds.get(memberId); - - // If we were unable to fetch the end offsets and could not compute lags, just return in task order - if (taskLagTotals.isEmpty()) { - return new TreeSet<>(prevTasks); - } else { - final SortedSet prevTasksByLag = new TreeSet<>(comparingLong(this::lagFor).thenComparing(TaskId::compareTo)); - prevTasksByLag.addAll(prevTasks); - return prevTasksByLag; - } + public Set previousTasksForConsumer(final String memberId) { + return consumerToPreviousStatefulTaskIds.get(memberId); } boolean hasUnfulfilledQuota(final int tasksPerThread) { diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java index 9b68e477f4f3d..b64e6718130b3 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java @@ -28,7 +28,6 @@ import static org.apache.kafka.common.utils.Utils.mkEntry; import static org.apache.kafka.common.utils.Utils.mkMap; import static org.apache.kafka.common.utils.Utils.mkSet; -import static org.apache.kafka.common.utils.Utils.mkSortedSet; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_0; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_1; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_2; @@ -46,7 +45,6 @@ import static org.junit.Assert.assertTrue; public class ClientStateTest { - private final ClientState client = new ClientState(1); private final ClientState zeroCapacityClient = new ClientState(0); @@ -324,40 +322,22 @@ public void shouldReturnPreviousStatefulTasksForConsumer() { ) ); - assertThat(client.previousTasksForConsumer("c1"), equalTo(mkSortedSet(TASK_0_1))); - assertThat(client.previousTasksForConsumer("c2"), equalTo(mkSortedSet(TASK_0_2))); + assertThat(client.previousTasksForConsumer("c1"), equalTo(mkSet(TASK_0_1))); + assertThat(client.previousTasksForConsumer("c2"), equalTo(mkSet(TASK_0_2))); assertTrue(client.previousTasksForConsumer("c3").isEmpty()); } @Test - public void shouldReturnPreviousStatefulTasksForConsumerWhenLagIsNotComputed() { - client.addPreviousTasksAndOffsetSums("c1", Collections.singletonMap(TASK_0_1, 1000L)); - client.initializePrevTasks(Collections.emptyMap()); - - assertThat(client.previousTasksForConsumer("c1"), equalTo(mkSortedSet(TASK_0_1))); - } - - @Test - public void shouldReturnPreviousStatefulTasksForConsumerInIncreasingLagOrder() { + public void shouldReturnPreviousTasksForConsumer() { client.addPreviousTasksAndOffsetSums("c1", mkMap( mkEntry(TASK_0_1, 100L), mkEntry(TASK_0_2, 0L), mkEntry(TASK_0_3, Task.LATEST_OFFSET) )); - client.computeTaskLags( - UUID_1, - mkMap( - mkEntry(TASK_0_1, 1_000L), - mkEntry(TASK_0_2, 1_000L), - mkEntry(TASK_0_3, 1_000L) - - ) - ); - client.initializePrevTasks(Collections.emptyMap()); - assertThat(client.previousTasksForConsumer("c1"), equalTo(mkSortedSet(TASK_0_3, TASK_0_2, TASK_0_1))); + assertThat(client.previousTasksForConsumer("c1"), equalTo(mkSet(TASK_0_3, TASK_0_2, TASK_0_1))); } @Test @@ -434,6 +414,16 @@ public void shouldThrowIllegalStateExceptionIfTaskLagsMapIsNotEmpty() { assertThrows(IllegalStateException.class, () -> client.computeTaskLags(null, allTaskEndOffsetSums)); } + @Test + public void shouldReturnZeroForAllTasksIfLagNotComputed() { + client.addPreviousTasksAndOffsetSums("c1", Collections.singletonMap(TASK_0_1, 0L)); + client.addPreviousTasksAndOffsetSums("c2", Collections.singletonMap(TASK_0_2, Task.LATEST_OFFSET)); + client.addPreviousTasksAndOffsetSums("c3", Collections.singletonMap(TASK_0_3, 500L)); + assertThat(client.lagFor(TASK_0_1), equalTo(0L)); + assertThat(client.lagFor(TASK_0_2), equalTo(0L)); + assertThat(client.lagFor(TASK_0_3), equalTo(0L)); + } + @Test public void shouldThrowIllegalStateExceptionOnLagForUnknownTask() { final Map taskOffsetSums = Collections.singletonMap(TASK_0_1, 0L); From 96a6c23d23b4c8b4018f4bb1bf044ce924f73f90 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Thu, 4 Jun 2020 09:54:35 -0700 Subject: [PATCH 13/23] checkstyle --- .../streams/processor/internals/assignment/ClientState.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java index cff70d4f9d3a7..fa33d45d7ab95 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java @@ -16,7 +16,6 @@ */ package org.apache.kafka.streams.processor.internals.assignment; -import java.util.SortedSet; import java.util.stream.Collectors; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.streams.processor.TaskId; @@ -38,7 +37,6 @@ import static java.util.Collections.unmodifiableMap; import static java.util.Collections.unmodifiableSet; import static java.util.Comparator.comparing; -import static java.util.Comparator.comparingLong; import static org.apache.kafka.common.utils.Utils.union; import static org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo.UNKNOWN_OFFSET_SUM; From d26d46290fdd806ef85abd989507495598b2854e Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Thu, 4 Jun 2020 11:39:38 -0700 Subject: [PATCH 14/23] use UNKNOWN_OFFSET_SUM, always initialize lag --- .../processor/internals/StreamsPartitionAssignor.java | 7 +++---- .../processor/internals/assignment/ClientState.java | 9 +-------- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 5ae6e944dd9e5..7fa94fd67d3b8 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -76,6 +76,7 @@ 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 { @@ -767,7 +768,7 @@ private boolean populateClientStatesMap(final Map clientState allTaskEndOffsetSums = computeEndOffsetSumsByTask(endOffsets, changelogsByStatefulTask, allNewlyCreatedChangelogPartitions); fetchEndOffsetsSuccessful = true; } catch (final StreamsException e) { - allTaskEndOffsetSums = null; + allTaskEndOffsetSums = changelogsByStatefulTask.keySet().stream().collect(Collectors.toMap(t -> t, t -> UNKNOWN_OFFSET_SUM)); fetchEndOffsetsSuccessful = false; } @@ -776,9 +777,7 @@ private boolean populateClientStatesMap(final Map clientState final ClientState state = entry.getValue().state; state.initializePrevTasks(taskForPartition); - if (fetchEndOffsetsSuccessful) { - state.computeTaskLags(uuid, allTaskEndOffsetSums); - } + state.computeTaskLags(uuid, allTaskEndOffsetSums); clientStates.put(uuid, state); } return fetchEndOffsetsSuccessful; diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java index fa33d45d7ab95..616cd42735101 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/assignment/ClientState.java @@ -301,14 +301,7 @@ public void computeTaskLags(final UUID uuid, final Map allTaskEndO * Task.LATEST_OFFSET if this was previously an active running task on this client */ public long lagFor(final TaskId task) { - final Long totalLag; - if (taskLagTotals.isEmpty()) { - // If we couldn't compute the task lags due to failure to fetch offsets, just return a flat constant - totalLag = 0L; - } else { - totalLag = taskLagTotals.get(task); - } - + final Long totalLag = taskLagTotals.get(task); if (totalLag == null) { throw new IllegalStateException("Tried to lookup lag for unknown task " + task); } From 6df10ca2f7c1447a246a18a4cc6fdd2906af7cf5 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Thu, 4 Jun 2020 11:40:40 -0700 Subject: [PATCH 15/23] remove no longer relevant new test --- .../internals/assignment/ClientStateTest.java | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java index b64e6718130b3..d50c00eae2617 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/ClientStateTest.java @@ -414,16 +414,6 @@ public void shouldThrowIllegalStateExceptionIfTaskLagsMapIsNotEmpty() { assertThrows(IllegalStateException.class, () -> client.computeTaskLags(null, allTaskEndOffsetSums)); } - @Test - public void shouldReturnZeroForAllTasksIfLagNotComputed() { - client.addPreviousTasksAndOffsetSums("c1", Collections.singletonMap(TASK_0_1, 0L)); - client.addPreviousTasksAndOffsetSums("c2", Collections.singletonMap(TASK_0_2, Task.LATEST_OFFSET)); - client.addPreviousTasksAndOffsetSums("c3", Collections.singletonMap(TASK_0_3, 500L)); - assertThat(client.lagFor(TASK_0_1), equalTo(0L)); - assertThat(client.lagFor(TASK_0_2), equalTo(0L)); - assertThat(client.lagFor(TASK_0_3), equalTo(0L)); - } - @Test public void shouldThrowIllegalStateExceptionOnLagForUnknownTask() { final Map taskOffsetSums = Collections.singletonMap(TASK_0_1, 0L); From fc98357ad9d2e7ee1e37e2098548c52320a379ad Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Thu, 4 Jun 2020 11:46:22 -0700 Subject: [PATCH 16/23] fix/rename counter --- .../processor/internals/StreamsPartitionAssignor.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 7fa94fd67d3b8..90a35b160bbc3 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -1098,17 +1098,17 @@ static Map> assignTasksToThreads(final Collection s for (final String consumer : consumers) { final List threadAssignment = assignment.get(consumer); - int i = 0; + int assignedTasks = 0; for (final TaskId task : getPreviousTasksByLag(state, consumer)) { if (unassignedStatefulTasks.contains(task)) { - if (i < minStatefulTasksPerThread) { + if (assignedTasks < minStatefulTasksPerThread) { threadAssignment.add(task); unassignedStatefulTasks.remove(task); + ++assignedTasks; } else { unassignedTaskToPreviousOwner.put(task, consumer); } } - ++i; } if (threadAssignment.size() < minStatefulTasksPerThread) { From b5cfd406a33deb19d18f85f76860055d78b38351 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Thu, 4 Jun 2020 11:59:58 -0700 Subject: [PATCH 17/23] some cleanup --- .../internals/StreamsPartitionAssignor.java | 57 ++++++++++--------- 1 file changed, 31 insertions(+), 26 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 90a35b160bbc3..8f77802dd203a 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -1132,36 +1132,41 @@ static Map> assignTasksToThreads(final Collection s } // At this point all consumers are at the min capacity, so there may be up to N - 1 unassigned - // stateful tasks remaining that should now be distributed over the consumers - consumersToFill.addAll(consumers); - - // Go over the tasks we skipped earlier and assign them to their previous owner when possible - for (final Map.Entry taskEntry : unassignedTaskToPreviousOwner.entrySet()) { - final TaskId task = taskEntry.getKey(); - final String consumer = taskEntry.getValue(); - if (consumersToFill.contains(consumer)) { - assignment.get(consumer).add(task); - unassignedStatefulTasks.remove(task); - consumersToFill.remove(consumer); + // stateful tasks still remaining that should now be distributed over the consumers + if (!unassignedStatefulTasks.isEmpty()) { + consumersToFill.addAll(consumers); + + // Go over the tasks we skipped earlier and assign them to their previous owner when possible + for (final Map.Entry taskEntry : unassignedTaskToPreviousOwner.entrySet()) { + final TaskId task = taskEntry.getKey(); + final String consumer = taskEntry.getValue(); + if (consumersToFill.contains(consumer)) { + assignment.get(consumer).add(task); + unassignedStatefulTasks.remove(task); + // Remove this consumer since we know it is now at minCapacity + 1 + consumersToFill.remove(consumer); + } } - } - // Now just distribute the remaining unassigned tasks over the consumers still at min capacity - for (final TaskId task : unassignedStatefulTasks) { - final String consumer = consumersToFill.poll(); - assignment.get(consumer).add(task); - } + // Now just distribute the remaining unassigned stateful tasks over the consumers still at min capacity + for (final TaskId task : unassignedStatefulTasks) { + final String consumer = consumersToFill.poll(); + final List threadAssignment = assignment.get(consumer); + threadAssignment.add(task); + } - // There must be at least one consumer still at min capacity while all the others are at min - // capacity + 1, so start distributing stateless tasks to get all consumers back to the same count - while (unassignedStatelessTasksIter.hasNext()) { - final String consumer = consumersToFill.poll(); - if (consumer != null) { - final TaskId task = unassignedStatelessTasksIter.next(); - assignment.get(consumer).add(task); - } else { - break; + // There must be at least one consumer still at min capacity while all the others are at min + // capacity + 1, so start distributing stateless tasks to get all consumers back to the same count + while (unassignedStatelessTasksIter.hasNext()) { + final String consumer = consumersToFill.poll(); + if (consumer != null) { + final TaskId task = unassignedStatelessTasksIter.next(); + unassignedStatelessTasksIter.remove(); + assignment.get(consumer).add(task); + } else { + break; + } } } } From f83baf95aa467a0b1236e7c10ef0ff9a8fced345 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Thu, 4 Jun 2020 12:01:28 -0700 Subject: [PATCH 18/23] fixindentation --- .../internals/StreamsPartitionAssignor.java | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 8f77802dd203a..cd03fb505715a 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -882,17 +882,17 @@ private Map computeNewAssignment(final Map consumers = clientMetadata.consumers; final Map> activeTaskAssignment = assignTasksToThreads( - state.statefulActiveTasks(), - state.statelessActiveTasks(), - consumers, - state + state.statefulActiveTasks(), + state.statelessActiveTasks(), + consumers, + state ); final Map> standbyTaskAssignment = assignTasksToThreads( - state.standbyTasks(), - Collections.emptySet(), - consumers, - state + state.standbyTasks(), + Collections.emptySet(), + consumers, + state ); // Arbitrarily choose the leader's client to be responsible for triggering the probing rebalance From f0c5a1954114381f09fb65caff9b5c9a6c60791c Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Thu, 4 Jun 2020 13:30:55 -0700 Subject: [PATCH 19/23] simplify --- .../streams/processor/internals/StreamsPartitionAssignor.java | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index cd03fb505715a..11d95d6cf2968 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -1098,13 +1098,11 @@ static Map> assignTasksToThreads(final Collection s for (final String consumer : consumers) { final List threadAssignment = assignment.get(consumer); - int assignedTasks = 0; for (final TaskId task : getPreviousTasksByLag(state, consumer)) { if (unassignedStatefulTasks.contains(task)) { - if (assignedTasks < minStatefulTasksPerThread) { + if (threadAssignment.size() < minStatefulTasksPerThread) { threadAssignment.add(task); unassignedStatefulTasks.remove(task); - ++assignedTasks; } else { unassignedTaskToPreviousOwner.put(task, consumer); } From ee384a9a9e5697e479010989a90f76a0140f6d66 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Thu, 4 Jun 2020 14:01:14 -0700 Subject: [PATCH 20/23] fix bug caught in test --- .../streams/processor/internals/StreamsPartitionAssignor.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java index 11d95d6cf2968..3f2cc87c51db5 100644 --- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java +++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java @@ -1138,7 +1138,7 @@ static Map> assignTasksToThreads(final Collection s for (final Map.Entry taskEntry : unassignedTaskToPreviousOwner.entrySet()) { final TaskId task = taskEntry.getKey(); final String consumer = taskEntry.getValue(); - if (consumersToFill.contains(consumer)) { + if (consumersToFill.contains(consumer) && unassignedStatefulTasks.contains(task)) { assignment.get(consumer).add(task); unassignedStatefulTasks.remove(task); // Remove this consumer since we know it is now at minCapacity + 1 From 168051c69252a4a5e6e64e5162fa8c4ae59ec3ce Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Mon, 8 Jun 2020 10:02:18 -0700 Subject: [PATCH 21/23] Adding unit tests --- .../StreamsPartitionAssignorTest.java | 101 +++++++++++++++--- 1 file changed, 84 insertions(+), 17 deletions(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index 91e955438f4bb..be6eda8c914b4 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -114,6 +114,7 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.is; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -128,6 +129,7 @@ public class StreamsPartitionAssignorTest { private static final String CONSUMER_1 = "consumer1"; private static final String CONSUMER_2 = "consumer2"; private static final String CONSUMER_3 = "consumer3"; + private static final String CONSUMER_4 = "consumer4"; private final Set allTopics = mkSet("topic1", "topic2"); @@ -148,23 +150,6 @@ public class StreamsPartitionAssignorTest { private final TopicPartition t4p2 = new TopicPartition("topic4", 2); private final TopicPartition t4p3 = new TopicPartition("topic4", 3); - private final Map> partitionsForTask = mkMap( - mkEntry(TASK_0_0, mkSet(t1p0, t2p0)), - mkEntry(TASK_0_1, mkSet(t1p1, t2p1)), - mkEntry(TASK_0_2, mkSet(t1p2, t2p2)), - mkEntry(TASK_0_3, mkSet(t1p3, t2p3)), - - mkEntry(TASK_1_0, mkSet(t3p0)), - mkEntry(TASK_1_1, mkSet(t3p1)), - mkEntry(TASK_1_2, mkSet(t3p2)), - mkEntry(TASK_1_3, mkSet(t3p3)), - - mkEntry(TASK_2_0, mkSet(t4p0)), - mkEntry(TASK_2_1, mkSet(t4p1)), - mkEntry(TASK_2_2, mkSet(t4p2)), - mkEntry(TASK_2_3, mkSet(t4p3)) - ); - private final List infos = asList( new PartitionInfo("topic1", 0, Node.noNode(), new Node[0], new Node[0]), new PartitionInfo("topic1", 1, Node.noNode(), new Node[0], new Node[0]), @@ -391,6 +376,88 @@ public void shouldProduceStickyAndBalancedAssignmentWhenNewTasksAreAdded() { assertEquivalentAssignment(previousAssignment, newAssignment); } + @Test + public void shouldProduceMaximallyStickyAssignmentWhenMemberLeaves() { + final List allTasks = + asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, TASK_1_3); + + final Map> previousAssignment = mkMap( + mkEntry(CONSUMER_1, asList(TASK_0_0, TASK_1_1, TASK_1_3)), + mkEntry(CONSUMER_2, asList(TASK_0_3, TASK_1_0)), + mkEntry(CONSUMER_3, asList(TASK_0_1, TASK_0_2, TASK_1_2)) + ); + + final ClientState state = new ClientState(); + final SortedSet consumers = mkSortedSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); + state.addPreviousTasksAndOffsetSums(CONSUMER_1, getTaskOffsetSums(asList(TASK_0_0, TASK_1_1, TASK_1_3), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_2, getTaskOffsetSums(asList(TASK_0_3, TASK_1_0), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_3, getTaskOffsetSums(asList(TASK_0_1, TASK_0_2, TASK_1_2), EMPTY_TASKS)); + state.initializePrevTasks(emptyMap()); + state.computeTaskLags(UUID_1, getTaskEndOffsetSums(allTasks)); + + // Consumer 3 leaves the group + consumers.remove(CONSUMER_3); + + final Map> assignment = assignTasksToThreads( + allTasks, + emptySet(), + consumers, + state + ); + + // Each member should have all of its previous tasks reassigned plus some of consumer 3's tasks + // We should give one of its tasks to consumer 1, and two of its tasks to consumer 2 + assertTrue(assignment.get(CONSUMER_1).containsAll(previousAssignment.get(CONSUMER_1))); + assertTrue(assignment.get(CONSUMER_2).containsAll(previousAssignment.get(CONSUMER_2))); + + assertThat(assignment.get(CONSUMER_1).size(), equalTo(4)); + assertThat(assignment.get(CONSUMER_2).size(), equalTo(4)); + } + + @Test + public void shouldProduceStickyEnoughAssignmentWhenNewMemberJoins() { + final List allTasks = + asList(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_1_0, TASK_1_1, TASK_1_2, TASK_1_3); + + final Map> previousAssignment = mkMap( + mkEntry(CONSUMER_1, asList(TASK_0_0, TASK_1_1, TASK_1_3)), + mkEntry(CONSUMER_2, asList(TASK_0_3, TASK_1_0)), + mkEntry(CONSUMER_3, asList(TASK_0_1, TASK_0_2, TASK_1_2)) + ); + + final ClientState state = new ClientState(); + final SortedSet consumers = mkSortedSet(CONSUMER_1, CONSUMER_2, CONSUMER_3); + state.addPreviousTasksAndOffsetSums(CONSUMER_1, getTaskOffsetSums(asList(TASK_0_0, TASK_1_1, TASK_1_3), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_2, getTaskOffsetSums(asList(TASK_0_3, TASK_1_0), EMPTY_TASKS)); + state.addPreviousTasksAndOffsetSums(CONSUMER_3, getTaskOffsetSums(asList(TASK_0_1, TASK_0_2, TASK_1_2), EMPTY_TASKS)); + + // Consumer 4 joins the group + consumers.add(CONSUMER_4); + state.addPreviousTasksAndOffsetSums(CONSUMER_4, getTaskOffsetSums(EMPTY_TASKS, EMPTY_TASKS)); + + state.initializePrevTasks(emptyMap()); + state.computeTaskLags(UUID_1, getTaskEndOffsetSums(allTasks)); + + final Map> assignment = assignTasksToThreads( + allTasks, + emptySet(), + consumers, + state + ); + + // we should move one task each from consumer 1 and consumer 3 to the new member, and none from consumer 2 + assertTrue(previousAssignment.get(CONSUMER_1).containsAll(assignment.get(CONSUMER_1))); + assertTrue(previousAssignment.get(CONSUMER_3).containsAll(assignment.get(CONSUMER_3))); + + assertTrue(assignment.get(CONSUMER_2).containsAll(previousAssignment.get(CONSUMER_2))); + + + assertThat(assignment.get(CONSUMER_1).size(), equalTo(2)); + assertThat(assignment.get(CONSUMER_2).size(), equalTo(2)); + assertThat(assignment.get(CONSUMER_3).size(), equalTo(2)); + assertThat(assignment.get(CONSUMER_4).size(), equalTo(2)); + } + @Test public void shouldInterleaveTasksByGroupIdDuringNewAssignment() { final List allTasks = From 6b0f6fbe8ea8fd23bcbecbd6f56036bb1e3fb4b6 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Mon, 8 Jun 2020 10:07:31 -0700 Subject: [PATCH 22/23] checkstyle --- .../processor/internals/StreamsPartitionAssignorTest.java | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index be6eda8c914b4..ae813483e7017 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -104,8 +104,6 @@ import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_1_3; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_2_0; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_2_1; -import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_2_2; -import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_2_3; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_1; import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.UUID_2; import static org.apache.kafka.streams.processor.internals.assignment.StreamsAssignmentProtocolVersions.LATEST_SUPPORTED_VERSION; @@ -114,7 +112,6 @@ import static org.hamcrest.CoreMatchers.equalTo; import static org.hamcrest.CoreMatchers.not; import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.contains; import static org.hamcrest.Matchers.is; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -434,7 +431,7 @@ public void shouldProduceStickyEnoughAssignmentWhenNewMemberJoins() { // Consumer 4 joins the group consumers.add(CONSUMER_4); state.addPreviousTasksAndOffsetSums(CONSUMER_4, getTaskOffsetSums(EMPTY_TASKS, EMPTY_TASKS)); - + state.initializePrevTasks(emptyMap()); state.computeTaskLags(UUID_1, getTaskEndOffsetSums(allTasks)); From e921175af03a61f20c24398540c1bdcf5a94fbb1 Mon Sep 17 00:00:00 2001 From: ableegoldman Date: Mon, 8 Jun 2020 10:29:41 -0700 Subject: [PATCH 23/23] remove unused topic partitions --- .../processor/internals/StreamsPartitionAssignorTest.java | 4 ---- 1 file changed, 4 deletions(-) diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java index ae813483e7017..03ab1a7103e79 100644 --- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java +++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignorTest.java @@ -142,10 +142,6 @@ public class StreamsPartitionAssignorTest { private final TopicPartition t3p1 = new TopicPartition("topic3", 1); private final TopicPartition t3p2 = new TopicPartition("topic3", 2); private final TopicPartition t3p3 = new TopicPartition("topic3", 3); - private final TopicPartition t4p0 = new TopicPartition("topic4", 0); - private final TopicPartition t4p1 = new TopicPartition("topic4", 1); - private final TopicPartition t4p2 = new TopicPartition("topic4", 2); - private final TopicPartition t4p3 = new TopicPartition("topic4", 3); private final List infos = asList( new PartitionInfo("topic1", 0, Node.noNode(), new Node[0], new Node[0]),