You are viewing a plain text version of this content. The canonical link for it is here.
Posted to jira@kafka.apache.org by "suzhiking (via GitHub)" <gi...@apache.org> on 2026/09/10 07:43:17 UTC
[PR] KAFKA-20998: Add balanced task assignor for streams groups [kafka]
suzhiking opened a new pull request, #23419:
URL: https://github.com/apache/kafka/pull/23419
The streams rebalance protocol (KIP-1071) ships a single built-in task
assignor, "sticky", which minimizes task movement between rebalances. The
classic protocol also offers the HighAvailabilityTaskAssignor, whose
placement spreads the tasks of each subtopology over the instances and
evens out the per-thread task load, at the price of moving more tasks
when the membership changes. This change adds that placement strategy to
the streams protocol as a second built-in assignor named "balanced".
The new BalancedTaskAssignor is a translation of the placement half of
the classic HighAvailabilityTaskAssignor: stateful active tasks are dealt
round-robin over the processes in sorted order and then moved between
processes while a move reduces the skew of the per-member load; standby
tasks go to the least loaded process that does not hold the task and are
evened out the same way; stateless active tasks fill in the gaps by
active task load. Since the broker assigns tasks to members rather than
to processes, the assignor also spreads each process's tasks over its
members, keeping a task on the member that currently owns it where the
quota allows, which is what the classic client does when it distributes
tasks over its stream threads. The cross-process placement does not
depend on the previous assignment.
The parts of the classic assignor that are not assignor concerns in the
streams protocol are deliberately left out: warm-up tasks and probing
rebalances are handled by the group coordinator's assignment refiner
(KAFKA-20665), and rack-aware standby task assignment is tracked in
KAFKA-20999 for both built-in assignors. Consequently the assignor does
not read the members' warm-up tasks or reported offsets.
The assignor is appended to the built-in list, so group.streams.assignors
defaults to "sticky,balanced" and a group can select it with
streams.assignor.name=balanced; "sticky" remains the default assignor.
The KIP names this assignor "highly_available"; the name "balanced"
reflects that high availability is delivered by the refiner rather than
by the assignor, and the KIP is to be amended accordingly.
Testing: BalancedTaskAssignorTest covers even distribution across
processes and members, subtopology interleaving, the contrast with the
sticky assignor on the same previous assignment, member stickiness within
a process, standby placement and capacity shortfall, stateless task
placement, independence from warm-up and offset inputs, determinism, and
a randomized validity check. GroupCoordinatorConfigTest is updated for
the new default, and StreamsGroupHeartbeatRequestTest selects the
balanced assignor through the group config. StreamsBalancedAssignorBenchmark
mirrors the sticky assignor benchmark; on the largest grid points the
balanced assignor takes 1-5x the sticky assignor's time, except for
groups of 1000 single-member processes where the all-pairs skew loop
dominates (about 6-18 ms per assignment).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-20998: Add balanced task assignor for streams groups [kafka]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #23419:
URL: https://github.com/apache/kafka/pull/23419#issuecomment-5724640732
A label of 'needs-attention' was automatically added to this PR in order to raise the
attention of the committers. Once this issue has been triaged, the `triage` label
should be removed to prevent this automation from happening again.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-20998: Add balanced task assignor for streams groups [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru commented on code in PR #23419:
URL: https://github.com/apache/kafka/pull/23419#discussion_r4131409373
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignorTest.java:
##########
@@ -0,0 +1,703 @@
+/*
+ * 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() {
Review Comment:
Nice. Could we also add a case for standbys within a process, i.e. a standby staying on its previous owner, and a task that changes role (standby -> active) within the same process?
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java:
##########
@@ -0,0 +1,535 @@
+/*
+ * 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 stateless active tasks, then standby
+ * tasks. Each round levels the members' total task counts, keeping a task on the member that currently owns it
+ * where that does not leave another member short.</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#warmupTasks()}, {@link MemberAssignmentState#taskOffsets()}
+ * and {@link MemberAssignmentState#taskEndOffsets()} are not read here), 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 owns a task, used to keep the task on that member if possible.
+ private final Map<TaskId, String> currentActiveOwner = new HashMap<>();
+ private final Map<TaskId, String> currentStandbyOwner = 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(currentActiveOwner, memberId, currentAssignment.activeTasks());
+ recordCurrentOwner(currentStandbyOwner, memberId, currentAssignment.standbyTasks());
+ }
+
+ private static void recordCurrentOwner(final Map<TaskId, String> currentOwner,
+ 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 own 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 three rounds, as the classic client spreads a process's
+ * tasks over its stream threads: stateful active tasks first, then stateless active tasks, then standby
+ * tasks. Each round levels the members' total task counts on top of what the previous rounds placed, so the
+ * stateful tasks are even on their own, the active tasks are even after the second round, and the standbys
+ * fill up the members with fewer active tasks. The classic client levels the standbys before the stateless
+ * tasks; the active tasks go first here, as in the {@link StickyTaskAssignor}, so that a member never
+ * processes two tasks while another only maintains standbys.
+ *
+ * @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, currentActiveOwner, taskCountByMember, activeTasksByMember);
+ distributeTasksOverMembers(statelessActiveTasks, currentActiveOwner, taskCountByMember, activeTasksByMember);
Review Comment:
The classic client (`StreamsPartitionAssignor`, around line 1025) distributes stateful actives, then standbys, then stateless actives. Since the goal of this assignor is to reproduce the classic outcome, I'd rather use the same order here, even if the sticky assignor does it differently. We can still improve on it in a follow-up.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java:
##########
@@ -0,0 +1,535 @@
+/*
+ * 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 stateless active tasks, then standby
+ * tasks. Each round levels the members' total task counts, keeping a task on the member that currently owns it
+ * where that does not leave another member short.</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#warmupTasks()}, {@link MemberAssignmentState#taskOffsets()}
+ * and {@link MemberAssignmentState#taskEndOffsets()} are not read here), 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 owns a task, used to keep the task on that member if possible.
+ private final Map<TaskId, String> currentActiveOwner = new HashMap<>();
+ private final Map<TaskId, String> currentStandbyOwner = 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(currentActiveOwner, memberId, currentAssignment.activeTasks());
Review Comment:
The classic client decides the previous owner from `prevTasksByLag`, i.e. all stateful tasks the thread reported offsets for, regardless of whether it was active or standby. Here an active only sticks to the previous active owner and a standby only to the previous standby owner, so e.g. a standby that gets promoted to active on the same process doesn't stay on its member. Could we use the union of active and standby tasks for stateful stickiness, to get closer to classic?
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignorTest.java:
##########
@@ -0,0 +1,703 @@
+/*
+ * 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 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 shouldIgnoreWarmupTasksAndReportedOffsets() {
+ // The balanced strategy does not use lag or warm-up information; the result must match the plain spec.
+ 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 all actives, 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 -> activeTaskCount(result, memberId)).toList(), "active tasks", processId, context);
Review Comment:
This asserts that actives are levelled after the stateless round, which only holds with the current round order. If we switch to the classic order (see my other comment), I think this needs to check the classic properties instead (stateful actives levelled, then total levelled).
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java:
##########
@@ -0,0 +1,535 @@
+/*
+ * 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 stateless active tasks, then standby
+ * tasks. Each round levels the members' total task counts, keeping a task on the member that currently owns it
+ * where that does not leave another member short.</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#warmupTasks()}, {@link MemberAssignmentState#taskOffsets()}
+ * and {@link MemberAssignmentState#taskEndOffsets()} are not read here), 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 owns a task, used to keep the task on that member if possible.
+ private final Map<TaskId, String> currentActiveOwner = new HashMap<>();
+ private final Map<TaskId, String> currentStandbyOwner = 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(currentActiveOwner, memberId, currentAssignment.activeTasks());
+ recordCurrentOwner(currentStandbyOwner, memberId, currentAssignment.standbyTasks());
+ }
+
+ private static void recordCurrentOwner(final Map<TaskId, String> currentOwner,
+ 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 own 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 three rounds, as the classic client spreads a process's
+ * tasks over its stream threads: stateful active tasks first, then stateless active tasks, then standby
+ * tasks. Each round levels the members' total task counts on top of what the previous rounds placed, so the
+ * stateful tasks are even on their own, the active tasks are even after the second round, and the standbys
+ * fill up the members with fewer active tasks. The classic client levels the standbys before the stateless
+ * tasks; the active tasks go first here, as in the {@link StickyTaskAssignor}, so that a member never
+ * processes two tasks while another only maintains standbys.
+ *
+ * @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, currentActiveOwner, taskCountByMember, activeTasksByMember);
+ distributeTasksOverMembers(statelessActiveTasks, currentActiveOwner, taskCountByMember, activeTasksByMember);
+ distributeTasksOverMembers(standbyTasks, currentStandbyOwner, taskCountByMember, standbyTasksByMember);
Review Comment:
Also, in the classic `assignTasksToThreads` only stateful tasks stick to their previous owner (`if (isStateful)`), stateless tasks are just dealt out without looking at previous ownership. Here the stateless round passes `currentActiveOwner`, so they stick as well. Can we match classic here?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-20998: Add balanced task assignor for streams groups [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru merged PR #23419:
URL: https://github.com/apache/kafka/pull/23419
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] KAFKA-20998: Add balanced task assignor for streams groups [kafka]
Posted by "lucasbru (via GitHub)" <gi...@apache.org>.
lucasbru commented on code in PR #23419:
URL: https://github.com/apache/kafka/pull/23419#discussion_r4124360939
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java:
##########
@@ -0,0 +1,444 @@
+/*
+ * 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.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+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 evenly over its members, keeping a task on the member that currently
+ * owns it where the quota allows.</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#warmupTasks()}, {@link MemberAssignmentState#taskOffsets()}
+ * and {@link MemberAssignmentState#taskEndOffsets()} are not read here), and rack-aware placement is tracked
+ * separately.
+ * <p>
+ * The assignor is not registered as a built-in assignor yet, so a group cannot select it through the group
+ * configuration {@code streams.assignor.name}. Appending it to the built-in list in {@code GroupCoordinatorConfig},
+ * which also adds its name to the default of {@code group.streams.assignors}, is deferred to a follow-up together
+ * with the optimization of the skew loop, so that the name becomes selectable only once the assignor is ready for
+ * production-sized groups.
+ */
+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().activeTasks.add(task);
+ }
+
+ balanceTasksOverProcesses(processes, process -> process.activeTasks);
+ }
+
+ 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);
+
+ 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) {
+ log.warn("Unable to assign {} of {} standby tasks for task [{}]. " +
+ "There is not enough available capacity. You should increase the number of " +
+ "application instances to maintain the requested number of standby replicas.",
+ remainingReplicas, numStandbyReplicas, task);
+ }
+ }
+
+ 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.activeTasks.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> sourceTasks = tasksToBalance.apply(source);
+ final SortedSet<TaskId> destinationTasks = tasksToBalance.apply(destination);
+ // Iterate over a copy, since the moves below modify the source's tasks.
+ final Iterator<TaskId> sourceIterator = new ArrayList<>(sourceTasks).iterator();
+ while (shouldMoveATask(source, destination) && sourceIterator.hasNext()) {
Review Comment:
Since we always take the lowest TaskIds from the source first, the moves tend to take whole subtopologies. E.g. P1 with 3 members, P2 with 1, two subtopologies with 4 partitions each: round-robin gives P2 A1,A3,B1,B3, and balancing moves A1 and A3 to P1, so all of A ends up on P1. Does not need to be addressed in this PR, because the classic assignor also doesn't do it, but keep in mind for further improvements.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java:
##########
@@ -0,0 +1,444 @@
+/*
+ * 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.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+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 evenly over its members, keeping a task on the member that currently
+ * owns it where the quota allows.</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#warmupTasks()}, {@link MemberAssignmentState#taskOffsets()}
+ * and {@link MemberAssignmentState#taskEndOffsets()} are not read here), and rack-aware placement is tracked
+ * separately.
+ * <p>
+ * The assignor is not registered as a built-in assignor yet, so a group cannot select it through the group
+ * configuration {@code streams.assignor.name}. Appending it to the built-in list in {@code GroupCoordinatorConfig},
+ * which also adds its name to the default of {@code group.streams.assignors}, is deferred to a follow-up together
+ * with the optimization of the skew loop, so that the name becomes selectable only once the assignor is ready for
+ * production-sized groups.
+ */
+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().activeTasks.add(task);
+ }
+
+ balanceTasksOverProcesses(processes, process -> process.activeTasks);
+ }
+
+ 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);
+
+ 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) {
+ log.warn("Unable to assign {} of {} standby tasks for task [{}]. " +
+ "There is not enough available capacity. You should increase the number of " +
+ "application instances to maintain the requested number of standby replicas.",
+ remainingReplicas, numStandbyReplicas, task);
+ }
+ }
+
+ 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.activeTasks.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> sourceTasks = tasksToBalance.apply(source);
+ final SortedSet<TaskId> destinationTasks = tasksToBalance.apply(destination);
+ // Iterate over a copy, since the moves below modify the source's tasks.
+ final Iterator<TaskId> sourceIterator = new ArrayList<>(sourceTasks).iterator();
+ while (shouldMoveATask(source, destination) && sourceIterator.hasNext()) {
+ final TaskId taskToMove = sourceIterator.next();
+ if (!destination.hasTask(taskToMove)) {
+ sourceTasks.remove(taskToMove);
+ 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;
+ }
+
+ private static int computeTasksPerMember(final int numberOfTasks, final int numberOfMembers) {
+ int tasksPerMember = numberOfTasks / numberOfMembers;
+ if (numberOfTasks % numberOfMembers > 0) {
+ tasksPerMember++;
+ }
+ return tasksPerMember;
+ }
+
+ /**
+ * 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> activeTasks = new TreeSet<>();
+ private final SortedSet<TaskId> standbyTasks = new TreeSet<>();
+ // The member of this process that currently owns a task, used to keep the task on that member if possible.
+ private final Map<TaskId, String> currentActiveOwner = new HashMap<>();
+ private final Map<TaskId, String> currentStandbyOwner = 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(currentActiveOwner, memberId, currentAssignment.activeTasks());
+ recordCurrentOwner(currentStandbyOwner, memberId, currentAssignment.standbyTasks());
+ }
+
+ private static void recordCurrentOwner(final Map<TaskId, String> currentOwner,
+ 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 own 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 assignedTaskCount() {
+ return activeTasks.size() + standbyTasks.size();
+ }
+
+ private double assignedTaskLoad() {
+ return ((double) assignedTaskCount()) / capacity();
+ }
+
+ private double activeTaskLoad() {
+ return ((double) activeTasks.size()) / capacity();
+ }
+
+ private boolean hasTask(final TaskId task) {
+ return activeTasks.contains(task) || standbyTasks.contains(task);
+ }
+
+ /**
+ * Spreads the process's tasks evenly over its members: first the active tasks, then the standby tasks on
+ * top of them. A task stays on the member that currently owns it as long as that member is below its quota;
+ * the remaining tasks go to the least loaded member.
+ *
+ * @return The assignment of every member of this process, including members that received no task.
+ */
+ private Map<String, MemberAssignment> distributeTasksOverMembers() {
+ final Map<String, Integer> totalTaskCountByMember = 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) {
+ totalTaskCountByMember.put(memberId, 0);
+ activeTasksByMember.put(memberId, new HashSet<>());
+ standbyTasksByMember.put(memberId, new HashSet<>());
+ }
+
+ distributeTasksOverMembers(activeTasks, currentActiveOwner, totalTaskCountByMember,
+ computeTasksPerMember(activeTasks.size(), capacity()), activeTasksByMember);
Review Comment:
Using ceil here as the stickiness quota can leave a member of the same process short by more than one. E.g. one process, members A and B own t0,t1 and t2,t3, and a third member C joins: quota is ceil(4/3)=2, so A and B both keep their two tasks and C gets nothing (2/2/0 instead of 2/1/1). The classic `assignTasksToThreads` uses floor for the sticky pass and only then hands out the remainder, which avoids this. Same applies to the standby pass below with the total quota.
##########
group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignorTest.java:
##########
@@ -0,0 +1,585 @@
+/*
+ * 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.MemberAssignment;
+import org.apache.kafka.coordinator.group.api.streams.assignor.TaskAssignorException;
+import org.apache.kafka.coordinator.group.api.streams.assignor.TopologyDescriber;
+
+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 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 shouldNotAssignAnyStandbysWithInsufficientCapacity() {
+ // A single process cannot hold a standby of its own active tasks; the assignor warns and carries on.
+ 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);
+ }
+
+ @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 shouldIgnoreWarmupTasksAndReportedOffsets() {
+ // The balanced strategy does not use lag or warm-up information; the result must match the plain spec.
+ 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 shouldProduceValidAndEvenAssignmentsForRandomInput() {
+ final long seed = System.nanoTime();
Review Comment:
This test never sets a previous assignment and always uses the same number of members per process, so it doesn't exercise the stickiness path within a process (it wouldn't catch the 2/2/0 case above). Could we add previous owners and varying member counts? Also, I'd prefer a fixed seed, or at least logging the seed up front, so failures are easy to reproduce.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java:
##########
@@ -0,0 +1,444 @@
+/*
+ * 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.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+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 evenly over its members, keeping a task on the member that currently
+ * owns it where the quota allows.</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#warmupTasks()}, {@link MemberAssignmentState#taskOffsets()}
+ * and {@link MemberAssignmentState#taskEndOffsets()} are not read here), and rack-aware placement is tracked
+ * separately.
+ * <p>
+ * The assignor is not registered as a built-in assignor yet, so a group cannot select it through the group
+ * configuration {@code streams.assignor.name}. Appending it to the built-in list in {@code GroupCoordinatorConfig},
+ * which also adds its name to the default of {@code group.streams.assignors}, is deferred to a follow-up together
Review Comment:
I'd keep the follow-up plan in the PR description / JIRA rather than the Javadoc, since it will go stale as soon as the assignor is registered.
##########
jmh-benchmarks/src/main/java/org/apache/kafka/jmh/assignor/StreamsBalancedAssignorBenchmark.java:
##########
@@ -0,0 +1,212 @@
+/*
+ * 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.jmh.assignor;
+
+import org.apache.kafka.coordinator.common.runtime.CoordinatorMetadataImage;
+import org.apache.kafka.coordinator.group.api.streams.assignor.AssignmentConfigs;
+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.TaskAssignor;
+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.topics.ConfiguredSubtopology;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Threads;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.SortedMap;
+import java.util.concurrent.TimeUnit;
+
+@State(Scope.Benchmark)
+@Fork(value = 1)
+@Warmup(iterations = 5)
+@Measurement(iterations = 5)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+public class StreamsBalancedAssignorBenchmark {
Review Comment:
This is almost identical to `StreamsStickyAssignorBenchmark`. Could we add an `@Param` for the assignor there instead? That would also avoid running the `ReportedOffsets` dimension, which has no effect on this assignor and doubles the number of runs.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java:
##########
@@ -0,0 +1,444 @@
+/*
+ * 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.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+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 evenly over its members, keeping a task on the member that currently
+ * owns it where the quota allows.</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#warmupTasks()}, {@link MemberAssignmentState#taskOffsets()}
+ * and {@link MemberAssignmentState#taskEndOffsets()} are not read here), and rack-aware placement is tracked
+ * separately.
+ * <p>
+ * The assignor is not registered as a built-in assignor yet, so a group cannot select it through the group
+ * configuration {@code streams.assignor.name}. Appending it to the built-in list in {@code GroupCoordinatorConfig},
+ * which also adds its name to the default of {@code group.streams.assignors}, is deferred to a follow-up together
+ * with the optimization of the skew loop, so that the name becomes selectable only once the assignor is ready for
+ * production-sized groups.
+ */
+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().activeTasks.add(task);
+ }
+
+ balanceTasksOverProcesses(processes, process -> process.activeTasks);
+ }
+
+ 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);
+
+ 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) {
+ log.warn("Unable to assign {} of {} standby tasks for task [{}]. " +
Review Comment:
This logs one warning per task on every assignment, on the broker. With a single process and 1000 stateful tasks that's 1000 WARN lines per rebalance. Could we collect the tasks and log once?
Also, I think we could drop severity to INFO at most. This is kind of expected to happen, and on the broker we might get assignor runs from many many groups and frequently, so this will be pretty noisy.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java:
##########
@@ -0,0 +1,444 @@
+/*
+ * 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.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+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 evenly over its members, keeping a task on the member that currently
+ * owns it where the quota allows.</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#warmupTasks()}, {@link MemberAssignmentState#taskOffsets()}
+ * and {@link MemberAssignmentState#taskEndOffsets()} are not read here), and rack-aware placement is tracked
+ * separately.
+ * <p>
+ * The assignor is not registered as a built-in assignor yet, so a group cannot select it through the group
+ * configuration {@code streams.assignor.name}. Appending it to the built-in list in {@code GroupCoordinatorConfig},
+ * which also adds its name to the default of {@code group.streams.assignors}, is deferred to a follow-up together
+ * with the optimization of the skew loop, so that the name becomes selectable only once the assignor is ready for
+ * production-sized groups.
+ */
+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().activeTasks.add(task);
+ }
+
+ balanceTasksOverProcesses(processes, process -> process.activeTasks);
+ }
+
+ 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);
+
+ 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) {
+ log.warn("Unable to assign {} of {} standby tasks for task [{}]. " +
+ "There is not enough available capacity. You should increase the number of " +
+ "application instances to maintain the requested number of standby replicas.",
+ remainingReplicas, numStandbyReplicas, task);
+ }
+ }
+
+ 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.activeTasks.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> sourceTasks = tasksToBalance.apply(source);
+ final SortedSet<TaskId> destinationTasks = tasksToBalance.apply(destination);
+ // Iterate over a copy, since the moves below modify the source's tasks.
+ final Iterator<TaskId> sourceIterator = new ArrayList<>(sourceTasks).iterator();
Review Comment:
Nit: I think we don't need the copy here. Only the source set is modified, so iterating `sourceTasks` directly and using `iterator.remove()` should work and saves an allocation per process pair per pass.
##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/BalancedTaskAssignor.java:
##########
@@ -0,0 +1,444 @@
+/*
+ * 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.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+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 evenly over its members, keeping a task on the member that currently
+ * owns it where the quota allows.</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#warmupTasks()}, {@link MemberAssignmentState#taskOffsets()}
+ * and {@link MemberAssignmentState#taskEndOffsets()} are not read here), and rack-aware placement is tracked
+ * separately.
+ * <p>
+ * The assignor is not registered as a built-in assignor yet, so a group cannot select it through the group
+ * configuration {@code streams.assignor.name}. Appending it to the built-in list in {@code GroupCoordinatorConfig},
+ * which also adds its name to the default of {@code group.streams.assignors}, is deferred to a follow-up together
+ * with the optimization of the skew loop, so that the name becomes selectable only once the assignor is ready for
+ * production-sized groups.
+ */
+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().activeTasks.add(task);
+ }
+
+ balanceTasksOverProcesses(processes, process -> process.activeTasks);
+ }
+
+ 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);
+
+ 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) {
+ log.warn("Unable to assign {} of {} standby tasks for task [{}]. " +
+ "There is not enough available capacity. You should increase the number of " +
+ "application instances to maintain the requested number of standby replicas.",
+ remainingReplicas, numStandbyReplicas, task);
+ }
+ }
+
+ 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.activeTasks.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> sourceTasks = tasksToBalance.apply(source);
+ final SortedSet<TaskId> destinationTasks = tasksToBalance.apply(destination);
+ // Iterate over a copy, since the moves below modify the source's tasks.
+ final Iterator<TaskId> sourceIterator = new ArrayList<>(sourceTasks).iterator();
+ while (shouldMoveATask(source, destination) && sourceIterator.hasNext()) {
+ final TaskId taskToMove = sourceIterator.next();
+ if (!destination.hasTask(taskToMove)) {
+ sourceTasks.remove(taskToMove);
+ 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;
+ }
+
+ private static int computeTasksPerMember(final int numberOfTasks, final int numberOfMembers) {
+ int tasksPerMember = numberOfTasks / numberOfMembers;
+ if (numberOfTasks % numberOfMembers > 0) {
+ tasksPerMember++;
+ }
+ return tasksPerMember;
+ }
+
+ /**
+ * 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> activeTasks = new TreeSet<>();
+ private final SortedSet<TaskId> standbyTasks = new TreeSet<>();
+ // The member of this process that currently owns a task, used to keep the task on that member if possible.
+ private final Map<TaskId, String> currentActiveOwner = new HashMap<>();
+ private final Map<TaskId, String> currentStandbyOwner = 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(currentActiveOwner, memberId, currentAssignment.activeTasks());
+ recordCurrentOwner(currentStandbyOwner, memberId, currentAssignment.standbyTasks());
+ }
+
+ private static void recordCurrentOwner(final Map<TaskId, String> currentOwner,
+ 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 own 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 assignedTaskCount() {
+ return activeTasks.size() + standbyTasks.size();
+ }
+
+ private double assignedTaskLoad() {
+ return ((double) assignedTaskCount()) / capacity();
+ }
+
+ private double activeTaskLoad() {
+ return ((double) activeTasks.size()) / capacity();
+ }
+
+ private boolean hasTask(final TaskId task) {
+ return activeTasks.contains(task) || standbyTasks.contains(task);
+ }
+
+ /**
+ * Spreads the process's tasks evenly over its members: first the active tasks, then the standby tasks on
+ * top of them. A task stays on the member that currently owns it as long as that member is below its quota;
+ * the remaining tasks go to the least loaded member.
+ *
+ * @return The assignment of every member of this process, including members that received no task.
+ */
+ private Map<String, MemberAssignment> distributeTasksOverMembers() {
+ final Map<String, Integer> totalTaskCountByMember = 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) {
+ totalTaskCountByMember.put(memberId, 0);
+ activeTasksByMember.put(memberId, new HashSet<>());
+ standbyTasksByMember.put(memberId, new HashSet<>());
+ }
+
+ distributeTasksOverMembers(activeTasks, currentActiveOwner, totalTaskCountByMember,
+ computeTasksPerMember(activeTasks.size(), capacity()), activeTasksByMember);
+ distributeTasksOverMembers(standbyTasks, currentStandbyOwner, totalTaskCountByMember,
+ computeTasksPerMember(assignedTaskCount(), capacity()), standbyTasksByMember);
+
+ return memberIds.stream().collect(Collectors.toMap(
+ Function.identity(),
+ memberId -> new MemberAssignment(
+ toCompactedTaskIds(activeTasksByMember.get(memberId)),
+ toCompactedTaskIds(standbyTasksByMember.get(memberId))
+ )
+ ));
+ }
+
+ /**
+ * Distributes the tasks of one role (active or standby) over the members.
+ *
+ * @param tasksToDistribute The tasks of this role placed on the process.
+ * @param currentOwner The member of this process that currently owns a task, if any.
+ * @param totalTaskCountByMember The number of tasks of <em>both</em> roles each member holds so far; the
+ * quota and the least-loaded choice are based on this total, and it is
+ * updated as tasks are distributed.
+ * @param quota The number of tasks a member may hold in total before it stops keeping
+ * its current tasks.
+ * @param distributedTasksByMember The tasks of this 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> totalTaskCountByMember,
+ final int quota,
+ final Map<String, Set<TaskId>> distributedTasksByMember) {
+ final List<TaskId> unassignedTasks = new ArrayList<>();
+ for (final TaskId task : tasksToDistribute) {
Review Comment:
Stateful and stateless actives share one quota here, so stickiness can put all stateful tasks on one member. E.g. members A and B, stateful s0,s1 and stateless l0,l1, A previously owned s0,s1: A keeps both stateful tasks and B only gets the stateless ones. The sticky assignor and the classic thread assignment balance stateful tasks separately. Is that intentional?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: jira-unsubscribe@kafka.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org