You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by lu...@apache.org on 2026/09/30 10:47:17 UTC
(kafka) branch trunk updated: KAFKA-20998: Add balanced task assignor for streams groups (#23419)
This is an automated email from the ASF dual-hosted git repository.
lucasbru pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new dbe3c9127fb KAFKA-20998: Add balanced task assignor for streams groups (#23419)
dbe3c9127fb is described below
commit dbe3c9127fbc8e32425c5ea04d0af8b7656e77a7
Author: Gavin Wang <z2...@uwaterloo.ca>
AuthorDate: Wed Sep 30 06:47:09 2026 -0400
KAFKA-20998: Add balanced task assignor for streams groups (#23419)
The streams rebalance protocol (KIP-1071) ships one built-in task
assignor, "sticky". This adds a second, "balanced", a translation of the
placement in the classic `HighAvailabilityTaskAssignor`:
- Stateful active tasks are dealt round-robin over the processes, then
moved between processes while a move reduces the per-member load skew.
- Standbys go to the least loaded process that does not hold the task,
and stateless tasks fill in by active task load.
- Within a process, tasks are spread over its members the way the
classic client spreads them over its stream threads
(`assignTasksToThreads`): stateful actives, then standbys, then
stateless actives, each round levelling the members' task counts. A
stateful task stays on the member that holds it in any role (active,
standby or warm-up) where that leaves no other member short.
Reported offsets are not read: warm-up tasks are the assignment
refiner's job (KAFKA-20665), and rack-aware standby placement is
KAFKA-20999. A shortfall of standby replicas is logged once per
assignment at INFO.
The assignor is not registered as a built-in yet, so
`streams.assignor.name=balanced` is still rejected and the docs are
unchanged. Registering it (appending it to
`GroupCoordinatorConfig.STREAMS_GROUP_BUILTIN_ASSIGNORS`) waits for an
optimized implementation, since the all-pairs skew loop takes 6-18 ms
for 1000 single-member processes (up to 25x the sticky assignor), and
for KIP-1071 to be amended from `highly_available` to `balanced`.
Testing: `BalancedTaskAssignorTest` covers distribution across processes
and members, member stickiness within a process (including standbys and
role changes), the standby shortfall, stateless placement, determinism,
and a seeded randomized check over 200 groups asserting validity, the
cross-process fixed point and the per-round levelling.
`StreamsStickyAssignorBenchmark` becomes `StreamsAssignorBenchmark` with
an `assignorType` parameter (run `BALANCED` with `-p
reportedOffsets=NONE`); the balanced assignor takes 1-5x the sticky
assignor's time except at 1000 single-member processes.
---
.../streams/assignor/BalancedTaskAssignor.java | 537 +++++++++++++++
.../streams/assignor/BalancedTaskAssignorTest.java | 753 +++++++++++++++++++++
...enchmark.java => StreamsAssignorBenchmark.java} | 33 +-
3 files changed, 1319 insertions(+), 4 deletions(-)
diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java
new file mode 100644
index 00000000000..93216357b91
--- /dev/null
+++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java
@@ -0,0 +1,537 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group.streams.assignor;
+
+import org.apache.kafka.coordinator.group.api.streams.assignor.GroupAssignment;
+import org.apache.kafka.coordinator.group.api.streams.assignor.GroupSpec;
+import org.apache.kafka.coordinator.group.api.streams.assignor.MemberAssignment;
+import org.apache.kafka.coordinator.group.api.streams.assignor.MemberAssignmentState;
+import org.apache.kafka.coordinator.group.api.streams.assignor.TaskAssignor;
+import org.apache.kafka.coordinator.group.api.streams.assignor.TaskAssignorException;
+import org.apache.kafka.coordinator.group.api.streams.assignor.TopologyDescriber;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.Deque;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.SortedSet;
+import java.util.TreeMap;
+import java.util.TreeSet;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * A task assignor that computes a <em>balanced</em> assignment: tasks of the same subtopology are spread over as
+ * many processes as possible, and the per-member task load is evened out across processes.
+ * <p>
+ * The assignment is the placement half of the "classic" protocol's {@code HighAvailabilityTaskAssignor}, translated
+ * to the streams rebalance protocol:
+ * <ol>
+ * <li>Stateful active tasks are dealt round-robin over the processes, in sorted order of task and process ID,
+ * and then moved between processes as long as a move reduces the skew of the per-member task load.</li>
+ * <li>Standby tasks are placed on the least loaded process that does not hold the task yet, and evened out the
+ * same way.</li>
+ * <li>Stateless active tasks fill in the gaps, going to the process with the lowest active task load.</li>
+ * <li>Within a process, the tasks are spread over its members in three rounds, as the classic client spreads a
+ * process's tasks over its stream threads: stateful active tasks, then standby tasks, then stateless active
+ * tasks. Each round levels the members' total task counts. A stateful task stays on the member that currently
+ * holds it, in whichever role, where that does not leave another member short; stateless tasks are dealt out
+ * without regard to their current member.</li>
+ * </ol>
+ * In contrast to the {@link StickyTaskAssignor}, the placement across processes does not depend on the previous
+ * assignment, so the assignment stays orderly across many membership changes at the price of moving more tasks.
+ * <p>
+ * The parts of the classic assignor that this assignor deliberately does <em>not</em> implement are not assignor
+ * concerns in the streams rebalance protocol: warm-up tasks are inserted by the group coordinator when it applies the
+ * target assignment (so the member's {@link MemberAssignmentState#taskOffsets()} and
+ * {@link MemberAssignmentState#taskEndOffsets()} are not read here, and {@link MemberAssignmentState#warmupTasks()}
+ * only tells the fan-out which member of a process holds a task's state), and rack-aware placement is tracked
+ * separately.
+ */
+public class BalancedTaskAssignor implements TaskAssignor {
+
+ private static final String BALANCED_ASSIGNOR_NAME = "balanced";
+ private static final Logger log = LoggerFactory.getLogger(BalancedTaskAssignor.class);
+
+ private static final Comparator<ProcessTasks> BY_ASSIGNED_LOAD =
+ Comparator.comparingDouble(ProcessTasks::assignedTaskLoad).thenComparing(ProcessTasks::processId);
+ private static final Comparator<ProcessTasks> BY_ACTIVE_LOAD =
+ Comparator.comparingDouble(ProcessTasks::activeTaskLoad).thenComparing(ProcessTasks::processId);
+
+ @Override
+ public String name() {
+ return BALANCED_ASSIGNOR_NAME;
+ }
+
+ @Override
+ public String toString() {
+ return name();
+ }
+
+ @Override
+ public GroupAssignment assign(final GroupSpec groupSpec, final TopologyDescriber topologyDescriber) throws TaskAssignorException {
+ final TreeMap<String, ProcessTasks> processes = initializeProcesses(groupSpec);
+
+ final SortedSet<TaskId> statefulTasks = new TreeSet<>();
+ final SortedSet<TaskId> statelessTasks = new TreeSet<>();
+ for (final String subtopology : topologyDescriber.subtopologies()) {
+ final SortedSet<TaskId> tasks = topologyDescriber.isStateful(subtopology) ? statefulTasks : statelessTasks;
+ final int numberOfPartitions = topologyDescriber.maxNumInputPartitions(subtopology);
+ for (int partition = 0; partition < numberOfPartitions; partition++) {
+ tasks.add(new TaskId(subtopology, partition));
+ }
+ }
+
+ if (processes.isEmpty()) {
+ if (!statefulTasks.isEmpty() || !statelessTasks.isEmpty()) {
+ throw new TaskAssignorException("No process available to assign active tasks.");
+ }
+ return new GroupAssignment(Map.of());
+ }
+
+ assignActiveStatefulTasks(processes.values(), statefulTasks);
+
+ final int numStandbyReplicas = groupSpec.configs().numStandbyReplicas();
+ if (numStandbyReplicas > 0) {
+ assignStandbyReplicaTasks(processes.values(), statefulTasks, numStandbyReplicas);
+ }
+
+ assignStatelessActiveTasks(processes.values(), statelessTasks);
+
+ return buildGroupAssignment(processes.values());
+ }
+
+ /**
+ * Groups the members by process. The processes are kept in sorted order of their ID, which makes the
+ * assignment deterministic for a given group.
+ */
+ private static TreeMap<String, ProcessTasks> initializeProcesses(final GroupSpec groupSpec) {
+ final TreeMap<String, ProcessTasks> processes = new TreeMap<>();
+ for (final String memberId : groupSpec.memberIds()) {
+ final String processId = groupSpec.memberMetadata(memberId).processId();
+ processes.computeIfAbsent(processId, ProcessTasks::new)
+ .addMember(memberId, groupSpec.memberAssignmentState(memberId));
+ }
+ return processes;
+ }
+
+ /**
+ * Deals the stateful tasks round-robin over the processes and then evens out the load. Iterating the tasks in
+ * subtopology order spreads the tasks of each subtopology over the processes, which is what makes the assignment
+ * balanced rather than merely even.
+ */
+ private static void assignActiveStatefulTasks(final Collection<ProcessTasks> processes,
+ final SortedSet<TaskId> statefulTasks) {
+ Iterator<ProcessTasks> processIterator = null;
+ for (final TaskId task : statefulTasks) {
+ if (processIterator == null || !processIterator.hasNext()) {
+ processIterator = processes.iterator();
+ }
+ processIterator.next().statefulActiveTasks.add(task);
+ }
+
+ balanceTasksOverProcesses(processes, process -> process.statefulActiveTasks);
+ }
+
+ private static void assignStandbyReplicaTasks(final Collection<ProcessTasks> processes,
+ final SortedSet<TaskId> statefulTasks,
+ final int numStandbyReplicas) {
+ // The queue orders the processes by their load as of the time they were offered. A polled process is offered
+ // again after its load changed, so the order stays valid.
+ final PriorityQueue<ProcessTasks> processesByLoad = new PriorityQueue<>(BY_ASSIGNED_LOAD);
+ processesByLoad.addAll(processes);
+
+ int tasksShortOfReplicas = 0;
+ int missingReplicas = 0;
+ for (final TaskId task : statefulTasks) {
+ int remainingReplicas = numStandbyReplicas;
+ while (remainingReplicas > 0) {
+ final ProcessTasks process = pollLeastLoadedProcessWithoutTask(processesByLoad, task);
+ if (process == null) {
+ break;
+ }
+ process.standbyTasks.add(task);
+ processesByLoad.add(process);
+ remainingReplicas--;
+ }
+
+ if (remainingReplicas > 0) {
+ tasksShortOfReplicas++;
+ missingReplicas += remainingReplicas;
+ }
+ }
+
+ if (tasksShortOfReplicas > 0) {
+ // Expected whenever the group runs on fewer processes than copies are configured, and repeated on every
+ // assignment of that group, so one line at INFO rather than a warning per task.
+ log.info("{} of {} stateful tasks got fewer than the configured {} standby replicas ({} replicas missing in "
+ + "total): the copies of a task must be on different processes, and the group runs on {} process(es). "
+ + "Add application instances to get the configured number of standby replicas.",
+ tasksShortOfReplicas, statefulTasks.size(), numStandbyReplicas, missingReplicas, processes.size());
+ }
+
+ balanceTasksOverProcesses(processes, process -> process.standbyTasks);
+ }
+
+ /**
+ * Polls the least loaded process that does not hold the task yet, or {@code null} if every process holds it.
+ * Processes that were skipped are offered back to the queue; the returned process is not, since the caller
+ * changes its load.
+ */
+ private static ProcessTasks pollLeastLoadedProcessWithoutTask(final PriorityQueue<ProcessTasks> processesByLoad,
+ final TaskId task) {
+ final List<ProcessTasks> skipped = new ArrayList<>();
+ ProcessTasks found = null;
+ while (!processesByLoad.isEmpty()) {
+ final ProcessTasks candidate = processesByLoad.poll();
+ if (candidate.hasTask(task)) {
+ skipped.add(candidate);
+ } else {
+ found = candidate;
+ break;
+ }
+ }
+ processesByLoad.addAll(skipped);
+ return found;
+ }
+
+ /**
+ * Stateless tasks carry no state to restore, so they are simply placed on the process with the lowest active task
+ * load, which fills in any imbalance the stateful placement left behind.
+ */
+ private static void assignStatelessActiveTasks(final Collection<ProcessTasks> processes,
+ final SortedSet<TaskId> statelessTasks) {
+ final PriorityQueue<ProcessTasks> processesByActiveLoad = new PriorityQueue<>(BY_ACTIVE_LOAD);
+ processesByActiveLoad.addAll(processes);
+
+ for (final TaskId task : statelessTasks) {
+ final ProcessTasks process = processesByActiveLoad.poll();
+ process.statelessActiveTasks.add(task);
+ processesByActiveLoad.add(process);
+ }
+ }
+
+ /**
+ * Moves tasks from more loaded to less loaded processes until no single move reduces the skew of the per-member
+ * task load any further. A task is never moved to a process that already holds it as active or standby task.
+ */
+ private static void balanceTasksOverProcesses(final Collection<ProcessTasks> processes,
+ final Function<ProcessTasks, SortedSet<TaskId>> tasksToBalance) {
+ boolean keepBalancing = true;
+ while (keepBalancing) {
+ keepBalancing = false;
+ for (final ProcessTasks source : processes) {
+ for (final ProcessTasks destination : processes) {
+ if (source == destination || !shouldMoveATask(source, destination)) {
+ continue;
+ }
+
+ final SortedSet<TaskId> destinationTasks = tasksToBalance.apply(destination);
+ // The moves change only the source's tasks, which the iterator removes in place.
+ final Iterator<TaskId> sourceIterator = tasksToBalance.apply(source).iterator();
+ while (shouldMoveATask(source, destination) && sourceIterator.hasNext()) {
+ final TaskId taskToMove = sourceIterator.next();
+ if (!destination.hasTask(taskToMove)) {
+ sourceIterator.remove();
+ destinationTasks.add(taskToMove);
+ keepBalancing = true;
+ }
+ }
+ }
+ }
+ }
+ }
+
+ /**
+ * A task should move from the source to the destination only if the destination is less loaded per member, and
+ * the move would reduce that skew without tipping the imbalance over to the other side.
+ */
+ private static boolean shouldMoveATask(final ProcessTasks source, final ProcessTasks destination) {
+ final double skew = source.assignedTaskLoad() - destination.assignedTaskLoad();
+
+ if (skew <= 0) {
+ return false;
+ }
+
+ final double proposedAssignedTasksPerMemberAtDestination =
+ (destination.assignedTaskCount() + 1.0) / destination.capacity();
+ final double proposedAssignedTasksPerMemberAtSource =
+ (source.assignedTaskCount() - 1.0) / source.capacity();
+ final double proposedSkew = proposedAssignedTasksPerMemberAtSource - proposedAssignedTasksPerMemberAtDestination;
+
+ if (proposedSkew < 0) {
+ // then the move would only create an imbalance in the other direction.
+ return false;
+ }
+ // we should only move a task if doing so would actually improve the skew.
+ return proposedSkew < skew;
+ }
+
+ /**
+ * Every member belongs to exactly one process and each process reports an assignment for each of its members,
+ * so the per-process assignments cover the whole group and never overlap.
+ */
+ private static GroupAssignment buildGroupAssignment(final Collection<ProcessTasks> processes) {
+ return new GroupAssignment(processes.stream()
+ .flatMap(process -> process.distributeTasksOverMembers().entrySet().stream())
+ .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
+ }
+
+ private static Map<String, Set<Integer>> toCompactedTaskIds(final Set<TaskId> taskIds) {
+ final Map<String, Set<Integer>> ret = new HashMap<>();
+ for (final TaskId taskId : taskIds) {
+ ret.computeIfAbsent(taskId.subtopologyId(), subtopologyId -> new HashSet<>())
+ .add(taskId.partition());
+ }
+ return ret;
+ }
+
+ /**
+ * The tasks placed on one process, and the members the process contributes. The assignment across processes
+ * works on this level; the tasks are spread over the members only once the placement is final.
+ */
+ private static final class ProcessTasks {
+ private final String processId;
+ private final TreeSet<String> memberIds = new TreeSet<>();
+ private final SortedSet<TaskId> statefulActiveTasks = new TreeSet<>();
+ private final SortedSet<TaskId> statelessActiveTasks = new TreeSet<>();
+ private final SortedSet<TaskId> standbyTasks = new TreeSet<>();
+ // The member of this process that currently holds a task in any role, so that a stateful task can stay on
+ // the member that has its state, as the classic client does through the offset sums each thread reports.
+ private final Map<TaskId, String> currentOwner = new HashMap<>();
+
+ private ProcessTasks(final String processId) {
+ this.processId = processId;
+ }
+
+ private String processId() {
+ return processId;
+ }
+
+ private void addMember(final String memberId, final MemberAssignmentState currentAssignment) {
+ memberIds.add(memberId);
+ recordCurrentOwner(memberId, currentAssignment.activeTasks());
+ recordCurrentOwner(memberId, currentAssignment.standbyTasks());
+ recordCurrentOwner(memberId, currentAssignment.warmupTasks());
+ }
+
+ private void recordCurrentOwner(final String memberId, final Map<String, Set<Integer>> tasks) {
+ for (final Map.Entry<String, Set<Integer>> entry : tasks.entrySet()) {
+ for (final int partition : entry.getValue()) {
+ // Two members of one process cannot hold the same task; the smaller member ID wins if they do.
+ currentOwner.merge(new TaskId(entry.getKey(), partition), memberId,
+ (existing, candidate) -> existing.compareTo(candidate) <= 0 ? existing : candidate);
+ }
+ }
+ }
+
+ private int capacity() {
+ return memberIds.size();
+ }
+
+ private int activeTaskCount() {
+ return statefulActiveTasks.size() + statelessActiveTasks.size();
+ }
+
+ private int assignedTaskCount() {
+ return activeTaskCount() + standbyTasks.size();
+ }
+
+ private double assignedTaskLoad() {
+ return ((double) assignedTaskCount()) / capacity();
+ }
+
+ private double activeTaskLoad() {
+ return ((double) activeTaskCount()) / capacity();
+ }
+
+ private boolean hasTask(final TaskId task) {
+ return statefulActiveTasks.contains(task) || statelessActiveTasks.contains(task) || standbyTasks.contains(task);
+ }
+
+ /**
+ * Spreads the process's tasks over its members in the three rounds of the classic client's
+ * {@code assignTasksToThreads}: stateful active tasks first, then standby tasks, then stateless active
+ * tasks. Each round levels the members' total task counts on top of what the previous rounds placed, so the
+ * stateful active tasks are even on their own and the standbys fill up the members with fewer of them. The
+ * stateless tasks come last and are dealt out without regard to their current member, since they have no
+ * state to keep close. As in the classic client, a member may therefore end up with more active tasks than a
+ * sibling that holds standbys instead.
+ *
+ * @return The assignment of every member of this process, including members that received no task.
+ */
+ private Map<String, MemberAssignment> distributeTasksOverMembers() {
+ final Map<String, Integer> taskCountByMember = new HashMap<>(capacity());
+ final Map<String, Set<TaskId>> activeTasksByMember = new HashMap<>(capacity());
+ final Map<String, Set<TaskId>> standbyTasksByMember = new HashMap<>(capacity());
+ for (final String memberId : memberIds) {
+ taskCountByMember.put(memberId, 0);
+ activeTasksByMember.put(memberId, new HashSet<>());
+ standbyTasksByMember.put(memberId, new HashSet<>());
+ }
+
+ distributeTasksOverMembers(statefulActiveTasks, currentOwner, taskCountByMember, activeTasksByMember);
+ distributeTasksOverMembers(standbyTasks, currentOwner, taskCountByMember, standbyTasksByMember);
+ distributeTasksOverMembers(statelessActiveTasks, Map.of(), taskCountByMember, activeTasksByMember);
+
+ return memberIds.stream().collect(Collectors.toMap(
+ Function.identity(),
+ memberId -> new MemberAssignment(
+ toCompactedTaskIds(activeTasksByMember.get(memberId)),
+ toCompactedTaskIds(standbyTasksByMember.get(memberId))
+ )
+ ));
+ }
+
+ /**
+ * Distributes the tasks of one round over the members so that every member ends the round at the same
+ * level, {@code (tasks placed so far + tasks of this round) / members}, or one above it:
+ * <ol>
+ * <li>A task stays on the member that currently owns it as long as that member is below the level.</li>
+ * <li>The members still below the level take the remaining tasks in turn, in member ID order, until
+ * they reach it.</li>
+ * <li>Fewer tasks than members are left. Each goes to a member at the level: its current owner if that
+ * member is one of them, otherwise the next such member in ID order.</li>
+ * </ol>
+ * Keeping the sticky pass below the level, rather than one above it, is what keeps a member that joins from
+ * being left with nothing while the others keep everything.
+ *
+ * @param tasksToDistribute The tasks of this round placed on the process.
+ * @param currentOwner The member of this process that currently holds a task, if any; empty for
+ * a round without stickiness.
+ * @param taskCountByMember The number of tasks each member holds so far, over all rounds; updated as
+ * tasks are distributed.
+ * @param distributedTasksByMember The tasks of this round's role each member receives; the output of this
+ * method.
+ */
+ private void distributeTasksOverMembers(final SortedSet<TaskId> tasksToDistribute,
+ final Map<TaskId, String> currentOwner,
+ final Map<String, Integer> taskCountByMember,
+ final Map<String, Set<TaskId>> distributedTasksByMember) {
+ if (tasksToDistribute.isEmpty()) {
+ return;
+ }
+ int tasksPlacedSoFar = 0;
+ for (final int count : taskCountByMember.values()) {
+ tasksPlacedSoFar += count;
+ }
+ final int level = (tasksPlacedSoFar + tasksToDistribute.size()) / capacity();
+
+ // Step 1. Tasks whose owner is at the level already are remembered, so the owner can still keep them in
+ // step 3.
+ final List<TaskId> unassignedTasks = new ArrayList<>();
+ final Map<TaskId, String> ownersOfSkippedTasks = new LinkedHashMap<>();
+ for (final TaskId task : tasksToDistribute) {
+ final String owner = currentOwner.get(task);
+ if (owner != null && taskCountByMember.get(owner) < level) {
+ place(task, owner, taskCountByMember, distributedTasksByMember);
+ } else {
+ unassignedTasks.add(task);
+ if (owner != null) {
+ ownersOfSkippedTasks.put(task, owner);
+ }
+ }
+ }
+
+ final int firstLeftoverTask = fillMembersBelowLevel(unassignedTasks, level, taskCountByMember, distributedTasksByMember);
+ placeLeftoverTasks(unassignedTasks.subList(firstLeftoverTask, unassignedTasks.size()), ownersOfSkippedTasks, level,
+ taskCountByMember, distributedTasksByMember);
+ }
+
+ /**
+ * Step 2 of a round: the members below the level take the unassigned tasks in turn, in member ID order,
+ * until every member has reached the level.
+ *
+ * @return The index of the first unassigned task that was not placed.
+ */
+ private int fillMembersBelowLevel(final List<TaskId> unassignedTasks,
+ final int level,
+ final Map<String, Integer> taskCountByMember,
+ final Map<String, Set<TaskId>> distributedTasksByMember) {
+ final Deque<String> membersToFill = new ArrayDeque<>();
+ for (final String memberId : memberIds) {
+ if (taskCountByMember.get(memberId) < level) {
+ membersToFill.add(memberId);
+ }
+ }
+ int nextTask = 0;
+ while (!membersToFill.isEmpty() && nextTask < unassignedTasks.size()) {
+ final String member = membersToFill.poll();
+ place(unassignedTasks.get(nextTask++), member, taskCountByMember, distributedTasksByMember);
+ if (taskCountByMember.get(member) < level) {
+ membersToFill.add(member);
+ }
+ }
+ return nextTask;
+ }
+
+ /**
+ * Step 3 of a round: every member is at the level or one above it, and fewer tasks are left than there are
+ * members at the level, since every round before this one levelled its tasks the same way. Each leftover
+ * task goes to its current owner if that member is at the level, otherwise to the next member at the level
+ * in ID order.
+ */
+ private void placeLeftoverTasks(final List<TaskId> leftoverTasks,
+ final Map<TaskId, String> ownersOfSkippedTasks,
+ final int level,
+ final Map<String, Integer> taskCountByMember,
+ final Map<String, Set<TaskId>> distributedTasksByMember) {
+ final Set<String> membersAtLevel = new LinkedHashSet<>();
+ for (final String memberId : memberIds) {
+ if (taskCountByMember.get(memberId) == level) {
+ membersAtLevel.add(memberId);
+ }
+ }
+ final List<TaskId> tasksWithoutOwnerAtLevel = new ArrayList<>();
+ for (final TaskId task : leftoverTasks) {
+ final String owner = ownersOfSkippedTasks.get(task);
+ if (owner != null && membersAtLevel.remove(owner)) {
+ place(task, owner, taskCountByMember, distributedTasksByMember);
+ } else {
+ tasksWithoutOwnerAtLevel.add(task);
+ }
+ }
+ for (final TaskId task : tasksWithoutOwnerAtLevel) {
+ final Iterator<String> nextMemberAtLevel = membersAtLevel.iterator();
+ final String member = nextMemberAtLevel.next();
+ nextMemberAtLevel.remove();
+ place(task, member, taskCountByMember, distributedTasksByMember);
+ }
+ }
+
+ private static void place(final TaskId task,
+ final String member,
+ final Map<String, Integer> taskCountByMember,
+ final Map<String, Set<TaskId>> distributedTasksByMember) {
+ distributedTasksByMember.get(member).add(task);
+ taskCountByMember.merge(member, 1, Integer::sum);
+ }
+ }
+}
diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignorTest.java
new file mode 100644
index 00000000000..cb61382e4bd
--- /dev/null
+++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignorTest.java
@@ -0,0 +1,753 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.coordinator.group.streams.assignor;
+
+import org.apache.kafka.common.utils.LogCaptureAppender;
+import org.apache.kafka.coordinator.group.api.streams.assignor.GroupAssignment;
+import org.apache.kafka.coordinator.group.api.streams.assignor.MemberAssignment;
+import org.apache.kafka.coordinator.group.api.streams.assignor.TaskAssignorException;
+import org.apache.kafka.coordinator.group.api.streams.assignor.TopologyDescriber;
+
+import org.apache.logging.log4j.Level;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.Optional;
+import java.util.Random;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.stream.Stream;
+
+import static org.apache.kafka.common.utils.Utils.mkEntry;
+import static org.apache.kafka.common.utils.Utils.mkMap;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class BalancedTaskAssignorTest {
+
+ private static final String SUBTOPOLOGY_1 = "test-subtopology1";
+ private static final String SUBTOPOLOGY_2 = "test-subtopology2";
+
+ private final BalancedTaskAssignor assignor = new BalancedTaskAssignor();
+
+ @Test
+ public void testToStringReturnsName() {
+ assertEquals("balanced", assignor.name());
+ assertEquals(assignor.name(), assignor.toString());
+ }
+
+ @Test
+ public void shouldThrowWhenThereAreTasksButNoMembers() {
+ assertThrows(TaskAssignorException.class, () -> assignor.assign(
+ new GroupSpecImpl(Map.of(), AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(3, SUBTOPOLOGY_1)
+ ));
+ }
+
+ @Test
+ public void shouldAssignNothingWhenThereAreNoTasks() {
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members("member1", "process1", "member2", "process2"), AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(0, SUBTOPOLOGY_1)
+ );
+
+ assertEquals(2, result.members().size());
+ assertEquals(new MemberAssignment(Map.of(), Map.of()), result.members().get("member1"));
+ assertEquals(new MemberAssignment(Map.of(), Map.of()), result.members().get("member2"));
+ }
+
+ @Test
+ public void shouldAssignActiveStatefulTasksEvenlyOverProcessesWhereNumberOfProcessesIntegralDivisorOfNumberOfTasks() {
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members("member1", "process1", "member2", "process2", "member3", "process3"), AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(6, SUBTOPOLOGY_1)
+ );
+
+ assertEquals(2, activeTaskCount(result, "member1"));
+ assertEquals(2, activeTaskCount(result, "member2"));
+ assertEquals(2, activeTaskCount(result, "member3"));
+ assertAllTasksAssignedOnce(result, statefulTopology(6, SUBTOPOLOGY_1));
+ assertNoStandbyTasks(result);
+ }
+
+ @Test
+ public void shouldAssignActiveStatefulTasksEvenlyOverMembersWhereNumberOfMembersIntegralDivisorOfNumberOfTasks() {
+ // process1 has two members, process2 one; so process1 gets twice as many tasks.
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members("member1_1", "process1", "member1_2", "process1", "member2_1", "process2"), AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(6, SUBTOPOLOGY_1)
+ );
+
+ assertEquals(2, activeTaskCount(result, "member1_1"));
+ assertEquals(2, activeTaskCount(result, "member1_2"));
+ assertEquals(2, activeTaskCount(result, "member2_1"));
+ assertAllTasksAssignedOnce(result, statefulTopology(6, SUBTOPOLOGY_1));
+ }
+
+ @Test
+ public void shouldAssignActiveStatefulTasksEvenlyOverProcessesWhereNumberOfProcessesNotIntegralDivisorOfNumberOfTasks() {
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members("member1", "process1", "member2", "process2"), AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(5, SUBTOPOLOGY_1)
+ );
+
+ final int count1 = activeTaskCount(result, "member1");
+ final int count2 = activeTaskCount(result, "member2");
+ assertEquals(5, count1 + count2);
+ assertTrue(Math.abs(count1 - count2) <= 1, "Expected an even split, got " + count1 + " and " + count2);
+ assertAllTasksAssignedOnce(result, statefulTopology(5, SUBTOPOLOGY_1));
+ }
+
+ @Test
+ public void shouldAssignActiveStatefulTasksEvenlyOverUnevenlyDistributedMembers() {
+ // process1 has three members, process2 has one. Round-robin would give each process four tasks, so the
+ // balancing step has to move two tasks to process1 to even out the per-member load.
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members(
+ "member1_1", "process1", "member1_2", "process1", "member1_3", "process1",
+ "member2_1", "process2"
+ ), AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(8, SUBTOPOLOGY_1)
+ );
+
+ assertEquals(2, activeTaskCount(result, "member1_1"));
+ assertEquals(2, activeTaskCount(result, "member1_2"));
+ assertEquals(2, activeTaskCount(result, "member1_3"));
+ assertEquals(2, activeTaskCount(result, "member2_1"));
+ assertAllTasksAssignedOnce(result, statefulTopology(8, SUBTOPOLOGY_1));
+ }
+
+ @Test
+ public void shouldAssignActiveStatefulTasksEvenlyOverProcessesWithMoreProcessesThanTasks() {
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members("member1", "process1", "member2", "process2", "member3", "process3", "member4", "process4"), AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(2, SUBTOPOLOGY_1)
+ );
+
+ // Which two of the four interchangeable processes receive a task is a tie-break, not a property.
+ assertEquals(List.of(0, 0, 1, 1), sortedActiveTaskCounts(result, "member1", "member2", "member3", "member4"));
+ assertAllTasksAssignedOnce(result, statefulTopology(2, SUBTOPOLOGY_1));
+ }
+
+ @Test
+ public void shouldSpreadTasksOfEachSubtopologyOverProcesses() {
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members("member1", "process1", "member2", "process2"), AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(2, SUBTOPOLOGY_1, SUBTOPOLOGY_2)
+ );
+
+ // Each process gets one partition of each subtopology, never both partitions of the same subtopology.
+ assertOnePartitionOfEachSubtopology(result, "member1", SUBTOPOLOGY_1, SUBTOPOLOGY_2);
+ assertOnePartitionOfEachSubtopology(result, "member2", SUBTOPOLOGY_1, SUBTOPOLOGY_2);
+ assertAllTasksAssignedOnce(result, statefulTopology(2, SUBTOPOLOGY_1, SUBTOPOLOGY_2));
+ }
+
+ @Test
+ public void shouldNotDependOnPreviousAssignmentAcrossProcessesUnlikeStickyAssignor() {
+ // The previous assignment groups the tasks by subtopology. The sticky assignor keeps that layout, because
+ // both members are below quota; the balanced assignor recomputes the interleaved layout regardless.
+ final Map<String, MemberMetadataAndStateImpl> members = mkMap(
+ mkEntry("member1", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0, 1))), Map.of())),
+ mkEntry("member2", memberWithTasks("process2", mkMap(mkEntry(SUBTOPOLOGY_2, Set.of(0, 1))), Map.of()))
+ );
+ final TopologyDescriber topology = statefulTopology(2, SUBTOPOLOGY_1, SUBTOPOLOGY_2);
+
+ final GroupAssignment sticky = new StickyTaskAssignor().assign(new GroupSpecImpl(members, AssignmentConfigsImpl.DEFAULT), topology);
+ assertEquals(mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0, 1))), sticky.members().get("member1").activeTasks());
+ assertEquals(mkMap(mkEntry(SUBTOPOLOGY_2, Set.of(0, 1))), sticky.members().get("member2").activeTasks());
+
+ final GroupAssignment balanced = assignor.assign(new GroupSpecImpl(members, AssignmentConfigsImpl.DEFAULT), topology);
+ assertOnePartitionOfEachSubtopology(balanced, "member1", SUBTOPOLOGY_1, SUBTOPOLOGY_2);
+ assertOnePartitionOfEachSubtopology(balanced, "member2", SUBTOPOLOGY_1, SUBTOPOLOGY_2);
+ assertAllTasksAssignedOnce(balanced, topology);
+ }
+
+ @Test
+ public void shouldKeepTasksOnTheirCurrentMemberWithinAProcess() {
+ // Both members of the process are at quota with their current tasks, so nothing needs to move within the
+ // process even though the tasks were dealt to the process in a different order.
+ final Map<String, MemberMetadataAndStateImpl> members = mkMap(
+ mkEntry("member1", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(1, 3))), Map.of())),
+ mkEntry("member2", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0, 2))), Map.of()))
+ );
+
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members, AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(4, SUBTOPOLOGY_1)
+ );
+
+ assertEquals(mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(1, 3))), result.members().get("member1").activeTasks());
+ assertEquals(mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0, 2))), result.members().get("member2").activeTasks());
+ }
+
+ @Test
+ public void shouldMoveTasksOffAnOverloadedMemberWithinAProcess() {
+ // member1 currently owns every task; it keeps two (its quota) and the other two move to member2.
+ final Map<String, MemberMetadataAndStateImpl> members = mkMap(
+ mkEntry("member1", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0, 1, 2, 3))), Map.of())),
+ mkEntry("member2", memberWithTasks("process1", Map.of(), Map.of()))
+ );
+
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members, AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(4, SUBTOPOLOGY_1)
+ );
+
+ // member1 keeps two of its current tasks -- which two is a tie-break -- and member2 receives the other two.
+ final Set<Integer> keptByMember1 = result.members().get("member1").activeTasks().get(SUBTOPOLOGY_1);
+ final Set<Integer> movedToMember2 = result.members().get("member2").activeTasks().get(SUBTOPOLOGY_1);
+ assertEquals(2, keptByMember1.size());
+ assertEquals(2, movedToMember2.size());
+ assertAllTasksAssignedOnce(result, statefulTopology(4, SUBTOPOLOGY_1));
+ }
+
+ @Test
+ public void shouldGiveAJoiningMemberTasksAlthoughTheOthersAreWithinTheirCeiling() {
+ // Four tasks over three members: two members may hold two tasks, but not both of the current owners, or the
+ // joining member3 would be left with nothing. Exactly one task moves; the other three stay where they are.
+ final Map<String, MemberMetadataAndStateImpl> members = mkMap(
+ mkEntry("member1", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0, 1))), Map.of())),
+ mkEntry("member2", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(2, 3))), Map.of())),
+ mkEntry("member3", memberWithTasks("process1", Map.of(), Map.of()))
+ );
+
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members, AssignmentConfigsImpl.DEFAULT),
+ statefulTopology(4, SUBTOPOLOGY_1)
+ );
+
+ assertEquals(List.of(1, 1, 2), sortedActiveTaskCounts(result, "member1", "member2", "member3"));
+ assertEquals(1, activeTaskCount(result, "member3"));
+ final Set<Integer> stayedOnMember1 = result.members().get("member1").activeTasks().getOrDefault(SUBTOPOLOGY_1, Set.of());
+ final Set<Integer> stayedOnMember2 = result.members().get("member2").activeTasks().getOrDefault(SUBTOPOLOGY_1, Set.of());
+ assertTrue(Set.of(0, 1).containsAll(stayedOnMember1), "member1 received a task it did not own: " + stayedOnMember1);
+ assertTrue(Set.of(2, 3).containsAll(stayedOnMember2), "member2 received a task it did not own: " + stayedOnMember2);
+ assertEquals(3, stayedOnMember1.size() + stayedOnMember2.size(), "only one task should have moved");
+ assertAllTasksAssignedOnce(result, statefulTopology(4, SUBTOPOLOGY_1));
+ }
+
+ @Test
+ public void shouldBalanceStatefulTasksOverMembersSeparatelyFromStatelessTasks() {
+ // member1 currently owns both stateful tasks. Keeping both would balance the active task counts with the
+ // two stateless tasks going to member2, but leave member1 with all the state; instead each member gets one
+ // stateful and one stateless task.
+ final Map<String, MemberMetadataAndStateImpl> members = mkMap(
+ mkEntry("member1", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0, 1))), Map.of())),
+ mkEntry("member2", memberWithTasks("process1", Map.of(), Map.of()))
+ );
+ final TopologyDescriber topology = new MixedTopologyDescriber(2, 2);
+
+ final GroupAssignment result = assignor.assign(new GroupSpecImpl(members, AssignmentConfigsImpl.DEFAULT), topology);
+
+ assertOnePartitionOfEachSubtopology(result, "member1", SUBTOPOLOGY_1, SUBTOPOLOGY_2);
+ assertOnePartitionOfEachSubtopology(result, "member2", SUBTOPOLOGY_1, SUBTOPOLOGY_2);
+ assertAllTasksAssignedOnce(result, topology);
+ }
+
+ @Test
+ public void shouldKeepStandbyTasksOnTheirCurrentMemberWithinAProcess() {
+ // The members already hold the placement the assignor computes, standbys included, so nothing moves. Without
+ // standby stickiness the standbys would be re-dealt in task order, giving member1 task 1 and member2 task 3.
+ final Map<String, MemberMetadataAndStateImpl> members = mkMap(
+ mkEntry("member1", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0))), mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(3))))),
+ mkEntry("member2", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(2))), mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(1))))),
+ mkEntry("member3", memberWithTasks("process2", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(1))), mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(2))))),
+ mkEntry("member4", memberWithTasks("process2", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(3))), mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0)))))
+ );
+
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members, AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(1)),
+ statefulTopology(4, SUBTOPOLOGY_1)
+ );
+
+ members.forEach((memberId, member) -> assertEquals(
+ new MemberAssignment(member.activeTasks(), member.standbyTasks()),
+ result.members().get(memberId),
+ memberId
+ ));
+ }
+
+ @Test
+ public void shouldKeepATaskOnItsMemberWhenItsRoleChangesWithinAProcess() {
+ // process1 currently runs tasks 1 and 3 and holds standbys of 0 and 2; the balanced placement gives it the
+ // opposite roles. Each task stays with the member that holds its state: tasks 2 and 0 go from standby to
+ // active on member1 and member2, and tasks 3 and 1 from active to standby on the same members.
+ final Map<String, MemberMetadataAndStateImpl> members = mkMap(
+ mkEntry("member1", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(3))), mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(2))))),
+ mkEntry("member2", memberWithTasks("process1", mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(1))), mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0))))),
+ mkEntry("member3", memberWithTasks("process2", Map.of(), Map.of())),
+ mkEntry("member4", memberWithTasks("process2", Map.of(), Map.of()))
+ );
+
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members, AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(1)),
+ statefulTopology(4, SUBTOPOLOGY_1)
+ );
+
+ assertEquals(mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(2))), result.members().get("member1").activeTasks());
+ assertEquals(mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(3))), result.members().get("member1").standbyTasks());
+ assertEquals(mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(0))), result.members().get("member2").activeTasks());
+ assertEquals(mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(1))), result.members().get("member2").standbyTasks());
+ assertAllTasksAssignedOnce(result, statefulTopology(4, SUBTOPOLOGY_1));
+ assertStandbysAssigned(result, statefulTopology(4, SUBTOPOLOGY_1), 1);
+ }
+
+ @Test
+ public void shouldAssignStandbysForStatefulTasks() {
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(
+ members("member1", "process1", "member2", "process2", "member3", "process3"),
+ AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(1)
+ ),
+ statefulTopology(3, SUBTOPOLOGY_1)
+ );
+
+ for (final String memberId : List.of("member1", "member2", "member3")) {
+ assertEquals(1, activeTaskCount(result, memberId));
+ assertEquals(1, standbyTaskCount(result, memberId));
+ }
+ assertAllTasksAssignedOnce(result, statefulTopology(3, SUBTOPOLOGY_1));
+ assertStandbysAssigned(result, statefulTopology(3, SUBTOPOLOGY_1), 1);
+ }
+
+ @Test
+ public void shouldNotAssignStandbysForStatelessTasks() {
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(
+ members("member1", "process1", "member2", "process2"),
+ AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(1)
+ ),
+ statelessTopology(4, SUBTOPOLOGY_1)
+ );
+
+ assertEquals(2, activeTaskCount(result, "member1"));
+ assertEquals(2, activeTaskCount(result, "member2"));
+ assertNoStandbyTasks(result);
+ }
+
+ @Test
+ public void shouldNotAssignStandbyToTheProcessThatOwnsTheActiveTask() {
+ // Two members per process: a standby must still land on the other process, not on the sibling member.
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(
+ members("member1_1", "process1", "member1_2", "process1", "member2_1", "process2", "member2_2", "process2"),
+ AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(1)
+ ),
+ statefulTopology(4, SUBTOPOLOGY_1)
+ );
+
+ final Set<Integer> process1Actives = mergeTasks(result, true, "member1_1", "member1_2").getOrDefault(SUBTOPOLOGY_1, Set.of());
+ final Set<Integer> process1Standbys = mergeTasks(result, false, "member1_1", "member1_2").getOrDefault(SUBTOPOLOGY_1, Set.of());
+ assertEquals(2, process1Actives.size());
+ assertEquals(2, process1Standbys.size());
+ assertTrue(process1Actives.stream().noneMatch(process1Standbys::contains),
+ "process1 holds " + process1Actives + " as active and " + process1Standbys + " as standby");
+ assertStandbysAssigned(result, statefulTopology(4, SUBTOPOLOGY_1), 1);
+ }
+
+ @Test
+ public void shouldNotAssignAnyStandbysWithInsufficientCapacityAndLogTheShortfallOnce() {
+ // A single process cannot hold a standby of its own active tasks. The assignor carries on and reports the
+ // shortfall in one INFO line for the whole assignment, not in one warning per task.
+ try (LogCaptureAppender appender = LogCaptureAppender.createAndRegister(BalancedTaskAssignor.class)) {
+ appender.setClassLogger(BalancedTaskAssignor.class, Level.INFO);
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(
+ members("member1", "process1", "member2", "process1"),
+ AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(1)
+ ),
+ statefulTopology(4, SUBTOPOLOGY_1)
+ );
+
+ assertEquals(2, activeTaskCount(result, "member1"));
+ assertEquals(2, activeTaskCount(result, "member2"));
+ assertNoStandbyTasks(result);
+ assertEquals(List.of(), appender.getMessages(Level.WARN));
+ final List<String> infoMessages = appender.getMessages(Level.INFO);
+ assertEquals(1, infoMessages.size(), "expected one INFO line, got " + infoMessages);
+ assertTrue(
+ infoMessages.get(0).startsWith("4 of 4 stateful tasks got fewer than the configured 1 standby replicas (4 replicas missing in total)"),
+ infoMessages.get(0)
+ );
+ }
+ }
+
+ @Test
+ public void shouldAssignAsManyStandbysAsCapacityAllows() {
+ // Two processes, two standby replicas requested: only one standby per task fits.
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(
+ members("member1", "process1", "member2", "process2"),
+ AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(2)
+ ),
+ statefulTopology(2, SUBTOPOLOGY_1)
+ );
+
+ assertEquals(1, activeTaskCount(result, "member1"));
+ assertEquals(1, activeTaskCount(result, "member2"));
+ assertStandbysAssigned(result, statefulTopology(2, SUBTOPOLOGY_1), 1);
+ }
+
+ @Test
+ public void shouldDistributeStatelessTasksToBalanceTotalTaskLoad() {
+ // Stateful placement leaves process1 (two members) and process2 (one member) at a load of one task per
+ // member. The stateless tasks are then dealt by active load, so process1 gets two of them and process2 one.
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members("member1_1", "process1", "member1_2", "process1", "member2_1", "process2"), AssignmentConfigsImpl.DEFAULT),
+ new MixedTopologyDescriber(3, 3)
+ );
+
+ assertEquals(2, activeTaskCount(result, "member1_1"));
+ assertEquals(2, activeTaskCount(result, "member1_2"));
+ assertEquals(2, activeTaskCount(result, "member2_1"));
+ assertAllTasksAssignedOnce(result, new MixedTopologyDescriber(3, 3));
+ assertNoStandbyTasks(result);
+ }
+
+ @Test
+ public void shouldPlaceTasksAcrossProcessesIndependentlyOfWarmupTasksAndReportedOffsets() {
+ // The placement across processes does not use lag or warm-up information, so the result must match the plain
+ // spec. Within a process a warm-up only decides which member keeps a task, and both processes here have a
+ // single member.
+ final Map<String, MemberMetadataAndStateImpl> plainMembers = members("member1", "process1", "member2", "process2");
+ final Map<String, MemberMetadataAndStateImpl> membersWithHints = mkMap(
+ mkEntry("member1", new MemberMetadataAndStateImpl(
+ Optional.of("instance1"),
+ Optional.of("rack1"),
+ "process1",
+ Map.of("zone", "a"),
+ Map.of(),
+ Map.of(),
+ mkMap(mkEntry(SUBTOPOLOGY_1, Set.of(1, 3))),
+ mkMap(mkEntry(SUBTOPOLOGY_1, mkMap(mkEntry(1, 100L), mkEntry(3, 100L)))),
+ mkMap(mkEntry(SUBTOPOLOGY_1, mkMap(mkEntry(1, 100L), mkEntry(3, 100L))))
+ )),
+ mkEntry("member2", new MemberMetadataAndStateImpl(
+ Optional.of("instance2"),
+ Optional.of("rack2"),
+ "process2",
+ Map.of("zone", "b"),
+ Map.of(),
+ Map.of(),
+ Map.of(),
+ mkMap(mkEntry(SUBTOPOLOGY_1, mkMap(mkEntry(0, 5L), mkEntry(2, 5L)))),
+ mkMap(mkEntry(SUBTOPOLOGY_1, mkMap(mkEntry(0, 500L), mkEntry(2, 500L))))
+ ))
+ );
+ final TopologyDescriber topology = statefulTopology(4, SUBTOPOLOGY_1);
+
+ final GroupAssignment plain = assignor.assign(new GroupSpecImpl(plainMembers, AssignmentConfigsImpl.DEFAULT), topology);
+ final GroupAssignment withHints = assignor.assign(new GroupSpecImpl(membersWithHints, AssignmentConfigsImpl.DEFAULT), topology);
+
+ assertEquals(plain, withHints);
+ }
+
+ @Test
+ public void shouldBeDeterministicRegardlessOfMemberOrder() {
+ final Map<String, MemberMetadataAndStateImpl> ordered = new TreeMap<>(members(
+ "member1", "process2", "member2", "process1", "member3", "process3", "member4", "process1"
+ ));
+ final Map<String, MemberMetadataAndStateImpl> reversed = new TreeMap<>((a, b) -> b.compareTo(a));
+ reversed.putAll(ordered);
+ final TopologyDescriber topology = new MixedTopologyDescriber(5, 3);
+
+ assertEquals(
+ assignor.assign(new GroupSpecImpl(ordered, AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(1)), topology),
+ assignor.assign(new GroupSpecImpl(reversed, AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(1)), topology)
+ );
+ }
+
+ @Test
+ public void shouldProduceValidAndLevelledAssignmentsForRandomGroupsWithPreviousAssignments() {
+ // Random processes of one to four members, random topologies and standby counts, and a random current
+ // assignment with stale and duplicate claims, so that the sticky fan-out within a process is exercised. The
+ // seeds are fixed and named in every failure. Consecutive seeds are mixed first, because the first draws of
+ // java.util.Random are correlated for seeds that differ only in their low bits.
+ for (int scenario = 0; scenario < 200; scenario++) {
+ final long seed = mix(2026L + scenario);
+ final Random random = new Random(seed);
+
+ final int numProcesses = 1 + random.nextInt(8);
+ final int numStatefulTasks = random.nextInt(20);
+ final int numStatelessTasks = random.nextInt(20);
+ final int numStandbyReplicas = random.nextInt(3);
+ final TopologyDescriber topology = new MixedTopologyDescriber(numStatefulTasks, numStatelessTasks);
+
+ final Map<String, MemberMetadataAndStateImpl> members = new HashMap<>();
+ final Map<String, List<String>> membersByProcess = new TreeMap<>();
+ for (int p = 0; p < numProcesses; p++) {
+ final int capacity = 1 + random.nextInt(4);
+ for (int m = 0; m < capacity; m++) {
+ final String memberId = "member" + p + "_" + m;
+ members.put(memberId, memberWithTasks("process" + p, randomClaims(random, topology), randomClaims(random, topology)));
+ membersByProcess.computeIfAbsent("process" + p, id -> new ArrayList<>()).add(memberId);
+ }
+ }
+ final String context = String.format("seed=%d processes=%d stateful=%d stateless=%d standbys=%d members=%s",
+ seed, numProcesses, numStatefulTasks, numStatelessTasks, numStandbyReplicas, membersByProcess);
+
+ final GroupAssignment result = assignor.assign(
+ new GroupSpecImpl(members, AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(numStandbyReplicas)),
+ topology
+ );
+
+ assertAllTasksAssignedOnce(result, topology);
+ // Standbys are bounded by the replicas requested and the number of other processes.
+ assertStandbysAssigned(result, topology, Math.min(numStandbyReplicas, numProcesses - 1));
+
+ // Across processes, the stateful active tasks are a fixed point of the skew loop: no process can hand
+ // one of them to another without their per-member loads crossing over.
+ final Map<String, Integer> statefulActivesByProcess = new HashMap<>();
+ membersByProcess.forEach((processId, memberIds) -> statefulActivesByProcess.put(processId,
+ memberIds.stream().mapToInt(memberId -> statefulActiveTaskCount(result, memberId)).sum()));
+ for (final String source : membersByProcess.keySet()) {
+ for (final String destination : membersByProcess.keySet()) {
+ if (!source.equals(destination)) {
+ final long loadAfterMoveAtSource = (long) (statefulActivesByProcess.get(source) - 1) * membersByProcess.get(destination).size();
+ final long loadAfterMoveAtDestination = (long) (statefulActivesByProcess.get(destination) + 1) * membersByProcess.get(source).size();
+ assertTrue(loadAfterMoveAtSource < loadAfterMoveAtDestination,
+ "a stateful task should have moved from " + source + " to " + destination + " (" + context + "): " + statefulActivesByProcess);
+ }
+ }
+ }
+
+ // Within a process, each round of the fan-out leaves every member at the floor or the ceiling of the
+ // process's per-member count: stateful actives on their own, then stateful actives and standbys, then
+ // all tasks. And no process holds a task as active and as standby.
+ membersByProcess.forEach((processId, memberIds) -> {
+ assertLevelled(memberIds.stream().map(memberId -> statefulActiveTaskCount(result, memberId)).toList(), "stateful active tasks", processId, context);
+ assertLevelled(memberIds.stream().map(memberId -> statefulActiveTaskCount(result, memberId) + standbyTaskCount(result, memberId)).toList(), "stateful active and standby tasks", processId, context);
+ assertLevelled(memberIds.stream().map(memberId -> activeTaskCount(result, memberId) + standbyTaskCount(result, memberId)).toList(), "tasks", processId, context);
+
+ final String[] memberIdArray = memberIds.toArray(new String[0]);
+ final Map<String, Set<Integer>> actives = mergeTasks(result, true, memberIdArray);
+ final Map<String, Set<Integer>> standbys = mergeTasks(result, false, memberIdArray);
+ for (final Map.Entry<String, Set<Integer>> entry : standbys.entrySet()) {
+ final Set<Integer> activePartitions = actives.getOrDefault(entry.getKey(), Set.of());
+ assertTrue(entry.getValue().stream().noneMatch(activePartitions::contains),
+ processId + " holds a task as active and standby (" + context + ")");
+ }
+ });
+ }
+ }
+
+ /** SplitMix64's finalizer, so that consecutive seeds give unrelated random sequences. */
+ private static long mix(final long seed) {
+ long mixed = (seed ^ (seed >>> 30)) * 0xbf58476d1ce4e5b9L;
+ mixed = (mixed ^ (mixed >>> 27)) * 0x94d049bb133111ebL;
+ return mixed ^ (mixed >>> 31);
+ }
+
+ /**
+ * A random quarter of the topology's tasks, plus, now and then, claims on a subtopology that no longer exists
+ * and on a partition beyond the current count, as members report after a topology change.
+ */
+ private static Map<String, Set<Integer>> randomClaims(final Random random, final TopologyDescriber topology) {
+ final Map<String, Set<Integer>> claims = new HashMap<>();
+ for (final String subtopology : topology.subtopologies()) {
+ for (int partition = 0; partition < topology.maxNumInputPartitions(subtopology); partition++) {
+ if (random.nextInt(4) == 0) {
+ claims.computeIfAbsent(subtopology, s -> new HashSet<>()).add(partition);
+ }
+ }
+ }
+ if (random.nextInt(5) == 0) {
+ claims.computeIfAbsent("removed-subtopology", s -> new HashSet<>()).add(random.nextInt(3));
+ claims.computeIfAbsent(SUBTOPOLOGY_1, s -> new HashSet<>()).add(topology.maxNumInputPartitions(SUBTOPOLOGY_1) + random.nextInt(3));
+ }
+ return claims;
+ }
+
+ /** Every count is the floor or the ceiling of the average, as one round of the fan-out guarantees. */
+ private static void assertLevelled(final List<Integer> counts, final String what, final String processId, final String context) {
+ final int total = counts.stream().mapToInt(Integer::intValue).sum();
+ final int floor = total / counts.size();
+ final int ceiling = (total + counts.size() - 1) / counts.size();
+ for (final int count : counts) {
+ assertTrue(count >= floor && count <= ceiling,
+ what + " per member of " + processId + " not levelled, expected " + floor + " or " + ceiling + " (" + context + "): " + counts);
+ }
+ }
+
+ private static int statefulActiveTaskCount(final GroupAssignment result, final String memberId) {
+ return result.members().get(memberId).activeTasks().getOrDefault(SUBTOPOLOGY_1, Set.of()).size();
+ }
+
+ private static Map<String, MemberMetadataAndStateImpl> members(final String... memberIdsAndProcessIds) {
+ final Map<String, MemberMetadataAndStateImpl> members = new HashMap<>();
+ for (int i = 0; i < memberIdsAndProcessIds.length; i += 2) {
+ members.put(memberIdsAndProcessIds[i], member(memberIdsAndProcessIds[i + 1]));
+ }
+ return members;
+ }
+
+ private static MemberMetadataAndStateImpl member(final String processId) {
+ return memberWithTasks(processId, Map.of(), Map.of());
+ }
+
+ private static MemberMetadataAndStateImpl memberWithTasks(final String processId,
+ final Map<String, Set<Integer>> activeTasks,
+ final Map<String, Set<Integer>> standbyTasks) {
+ return new MemberMetadataAndStateImpl(
+ Optional.empty(),
+ Optional.empty(),
+ processId,
+ Map.of(),
+ activeTasks,
+ standbyTasks,
+ Map.of(),
+ Map.of(),
+ Map.of()
+ );
+ }
+
+ private static TopologyDescriber statefulTopology(final int numTasks, final String... subtopologies) {
+ return new UniformTopologyDescriber(numTasks, true, List.of(subtopologies));
+ }
+
+ private static TopologyDescriber statelessTopology(final int numTasks, final String... subtopologies) {
+ return new UniformTopologyDescriber(numTasks, false, List.of(subtopologies));
+ }
+
+ private static int activeTaskCount(final GroupAssignment result, final String memberId) {
+ final MemberAssignment member = result.members().get(memberId);
+ assertNotNull(member);
+ return member.activeTasks().values().stream().mapToInt(Set::size).sum();
+ }
+
+ private static List<Integer> sortedActiveTaskCounts(final GroupAssignment result, final String... memberIds) {
+ return Stream.of(memberIds).map(memberId -> activeTaskCount(result, memberId)).sorted().toList();
+ }
+
+ /** The member holds exactly one partition of each of the given subtopologies and nothing else. */
+ private static void assertOnePartitionOfEachSubtopology(final GroupAssignment result,
+ final String memberId,
+ final String... subtopologies) {
+ final Map<String, Set<Integer>> activeTasks = result.members().get(memberId).activeTasks();
+ assertEquals(Set.of(subtopologies), activeTasks.keySet(), memberId + " holds " + activeTasks);
+ activeTasks.forEach((subtopology, partitions) ->
+ assertEquals(1, partitions.size(), memberId + " holds " + partitions + " of " + subtopology));
+ }
+
+ private static int standbyTaskCount(final GroupAssignment result, final String memberId) {
+ final MemberAssignment member = result.members().get(memberId);
+ assertNotNull(member);
+ return member.standbyTasks().values().stream().mapToInt(Set::size).sum();
+ }
+
+ private static Map<String, Set<Integer>> mergeTasks(final GroupAssignment result, final boolean active, final String... memberIds) {
+ final Map<String, Set<Integer>> merged = new HashMap<>();
+ for (final String memberId : memberIds) {
+ final MemberAssignment member = result.members().get(memberId);
+ assertNotNull(member);
+ final Map<String, Set<Integer>> tasks = active ? member.activeTasks() : member.standbyTasks();
+ tasks.forEach((subtopology, partitions) -> merged.computeIfAbsent(subtopology, s -> new HashSet<>()).addAll(partitions));
+ }
+ return merged;
+ }
+
+ private static void assertNoStandbyTasks(final GroupAssignment result) {
+ for (final Map.Entry<String, MemberAssignment> entry : result.members().entrySet()) {
+ assertTrue(entry.getValue().standbyTasks().isEmpty(), entry.getKey() + " has standby tasks " + entry.getValue().standbyTasks());
+ }
+ }
+
+ /** Every task of the topology is assigned as active task to exactly one member. */
+ private static void assertAllTasksAssignedOnce(final GroupAssignment result, final TopologyDescriber topology) {
+ final Map<TaskId, Integer> owners = new HashMap<>();
+ for (final MemberAssignment member : result.members().values()) {
+ member.activeTasks().forEach((subtopology, partitions) ->
+ partitions.forEach(partition -> owners.merge(new TaskId(subtopology, partition), 1, Integer::sum)));
+ }
+ final Set<TaskId> expected = new HashSet<>();
+ for (final String subtopology : topology.subtopologies()) {
+ for (int partition = 0; partition < topology.maxNumInputPartitions(subtopology); partition++) {
+ expected.add(new TaskId(subtopology, partition));
+ }
+ }
+ assertEquals(expected, owners.keySet());
+ owners.forEach((task, count) -> assertEquals(1, count, task + " is active on " + count + " members"));
+ }
+
+ /** Every stateful task has exactly the given number of standbys and no stateless task has any. */
+ private static void assertStandbysAssigned(final GroupAssignment result, final TopologyDescriber topology, final int expectedStandbys) {
+ final Map<TaskId, Integer> standbyCounts = new HashMap<>();
+ for (final MemberAssignment member : result.members().values()) {
+ member.standbyTasks().forEach((subtopology, partitions) ->
+ partitions.forEach(partition -> standbyCounts.merge(new TaskId(subtopology, partition), 1, Integer::sum)));
+ }
+ final Map<TaskId, Integer> expected = new HashMap<>();
+ if (expectedStandbys > 0) {
+ for (final String subtopology : topology.subtopologies()) {
+ if (topology.isStateful(subtopology)) {
+ for (int partition = 0; partition < topology.maxNumInputPartitions(subtopology); partition++) {
+ expected.put(new TaskId(subtopology, partition), expectedStandbys);
+ }
+ }
+ }
+ }
+ assertEquals(expected, standbyCounts);
+ standbyCounts.keySet().forEach(task -> assertTrue(
+ topology.isStateful(task.subtopologyId()), task + " is stateless but has standbys"));
+ }
+
+ private record UniformTopologyDescriber(int numTasks, boolean isStateful, List<String> subtopologies) implements TopologyDescriber {
+
+ @Override
+ public int maxNumInputPartitions(final String subtopologyId) throws NoSuchElementException {
+ return numTasks;
+ }
+
+ @Override
+ public boolean isStateful(final String subtopologyId) {
+ return isStateful;
+ }
+ }
+
+ /** One stateful subtopology followed by one stateless subtopology. */
+ private record MixedTopologyDescriber(int numStatefulTasks, int numStatelessTasks) implements TopologyDescriber {
+
+ @Override
+ public List<String> subtopologies() {
+ final List<String> subtopologies = new ArrayList<>();
+ if (numStatefulTasks > 0) {
+ subtopologies.add(SUBTOPOLOGY_1);
+ }
+ if (numStatelessTasks > 0) {
+ subtopologies.add(SUBTOPOLOGY_2);
+ }
+ return subtopologies;
+ }
+
+ @Override
+ public int maxNumInputPartitions(final String subtopologyId) throws NoSuchElementException {
+ return subtopologyId.equals(SUBTOPOLOGY_1) ? numStatefulTasks : numStatelessTasks;
+ }
+
+ @Override
+ public boolean isStateful(final String subtopologyId) {
+ return subtopologyId.equals(SUBTOPOLOGY_1);
+ }
+ }
+}
diff --git a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsStickyAssignorBenchmark.java b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsAssignorBenchmark.java
similarity index 87%
rename from jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsStickyAssignorBenchmark.java
rename to jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsAssignorBenchmark.java
index 5b1dc7d5c5a..f3313520c28 100644
--- a/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsStickyAssignorBenchmark.java
+++ b/jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsAssignorBenchmark.java
@@ -26,6 +26,7 @@ import org.apache.kafka.coordinator.group.api.streams.assignor.TopologyDescriber
import org.apache.kafka.coordinator.group.streams.StreamsGroupMember;
import org.apache.kafka.coordinator.group.streams.TopologyMetadata;
import org.apache.kafka.coordinator.group.streams.assignor.AssignmentConfigsImpl;
+import org.apache.kafka.coordinator.group.streams.assignor.BalancedTaskAssignor;
import org.apache.kafka.coordinator.group.streams.assignor.GroupSpecImpl;
import org.apache.kafka.coordinator.group.streams.assignor.MemberMetadataAndStateImpl;
import org.apache.kafka.coordinator.group.streams.assignor.StickyTaskAssignor;
@@ -55,13 +56,21 @@ import java.util.Optional;
import java.util.SortedMap;
import java.util.concurrent.TimeUnit;
+/**
+ * Benchmarks the built-in task assignors of the streams rebalance protocol on the same grid of groups, selected
+ * with the {@code assignorType} parameter, so that their costs can be compared configuration by configuration.
+ */
@State(Scope.Benchmark)
@Fork(value = 1)
@Warmup(iterations = 5)
@Measurement(iterations = 5)
@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
-public class StreamsStickyAssignorBenchmark {
+public class StreamsAssignorBenchmark {
+
+ public enum AssignorType {
+ STICKY, BALANCED
+ }
/**
* The assignment type is decided based on whether all the members are assigned partitions
@@ -74,7 +83,9 @@ public class StreamsStickyAssignorBenchmark {
/**
* Whether the members report offset sums for the state they hold on local disk. NONE is the behaviour of a
* group whose clients do not report offsets, OWNED_AND_DORMANT also reports state left behind by earlier
- * assignments, so several members compete as candidates for the same task.
+ * assignments, so several members compete as candidates for the same task. Only the sticky assignor reads the
+ * offsets; the balanced assignor gives the same result for both values, so run it with
+ * {@code -p reportedOffsets=NONE}.
*/
public enum ReportedOffsets {
NONE, OWNED_AND_DORMANT
@@ -85,6 +96,9 @@ public class StreamsStickyAssignorBenchmark {
*/
private static final int DORMANT_REPLICAS = 1;
+ @Param({"STICKY", "BALANCED"})
+ private AssignorType assignorType;
+
@Param({"100", "1000"})
private int memberCount;
@@ -126,7 +140,7 @@ public class StreamsStickyAssignorBenchmark {
topologyDescriber = new TopologyMetadata(metadataImage, subtopologyMap);
- taskAssignor = new StickyTaskAssignor();
+ taskAssignor = createAssignor();
Map<String, StreamsGroupMember> members = createMembers();
this.assignmentConfigs = AssignmentConfigsImpl.DEFAULT.withNumStandbyReplicas(standbyReplicas);
@@ -147,6 +161,13 @@ public class StreamsStickyAssignorBenchmark {
}
}
+ private TaskAssignor createAssignor() {
+ return switch (assignorType) {
+ case STICKY -> new StickyTaskAssignor();
+ case BALANCED -> new BalancedTaskAssignor();
+ };
+ }
+
private Map<String, StreamsGroupMember> createMembers() {
// In the rebalance case, we will add the last member as a trigger.
// This is done to keep the total members count consistent with the input.
@@ -158,8 +179,12 @@ public class StreamsStickyAssignorBenchmark {
);
}
+ /**
+ * Computes the assignment the members hold when the measured rebalance starts, with the assignor under test,
+ * and adds one member to trigger the rebalance.
+ */
private void simulateIncrementalRebalance() {
- GroupAssignment initialAssignment = new StickyTaskAssignor().assign(groupSpec, topologyDescriber);
+ GroupAssignment initialAssignment = createAssignor().assign(groupSpec, topologyDescriber);
Map<String, MemberAssignment> members = initialAssignment.members();
Map<String, MemberMetadataAndStateImpl> updatedMemberSpec = new HashMap<>();