diff --git a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java index 529b07fd4fd74..e149bd000377e 100644 --- a/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java +++ b/group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java @@ -36,6 +36,7 @@ import java.util.Iterator; import java.util.LinkedList; import java.util.Map; +import java.util.Optional; import java.util.PriorityQueue; import java.util.Set; import java.util.stream.Collectors; @@ -212,7 +213,7 @@ private static void assignActive(final LocalState localState, final LinkedList it = activeTasks.iterator(); it.hasNext();) { final TaskId task = it.next(); final ArrayList prevMembers = localState.standbyTaskToPrevMember.get(task); - final Member prevMember = findPrevMemberWithLeastLoad(localState, prevMembers, null); + final Member prevMember = findPrevMemberWithLeastLoad(localState, prevMembers, Optional.empty()); if (prevMember != null) { final ProcessState processState = localState.processIdToState.get(prevMember.processId); if (hasUnfulfilledActiveTaskQuota(localState, processState, prevMember)) { @@ -288,44 +289,44 @@ private static boolean assignStandbyToMemberWithLeastLoad( /** * Finds the previous member with the least load for a given task. * - * @param localState The state of the assignment in progress. - * @param members The list of previous members owning the task. - * @param taskId The taskId, to check if the previous member already has the task. Can be null, if we assign it - * for the first time (e.g., during active task assignment). + * @param localState + * The state of the assignment in progress. + * @param members + * The list of previous members owning the task. + * @param standbyTaskId + * The taskId, to check if the previous member already has the task. * * @return Previous member with the least load that does not have the task, or null if no such member exists. */ private static Member findPrevMemberWithLeastLoad( final LocalState localState, final ArrayList members, - final TaskId taskId + final Optional standbyTaskId ) { if (members == null || members.isEmpty()) { return null; } - Member candidate = members.get(0); - final ProcessState candidateProcessState = localState.processIdToState.get(candidate.processId); - double candidateProcessLoad = candidateProcessState.load(); - double candidateMemberLoad = candidateProcessState.memberToTaskCounts().get(candidate.memberId); - for (int i = 1; i < members.size(); i++) { - final Member member = members.get(i); + Member candidate = null; + double candidateProcessLoad = Double.MAX_VALUE; + double candidateMemberLoad = Double.MAX_VALUE; + for (final Member member : members) { final ProcessState processState = localState.processIdToState.get(member.processId); + // A process that already owns a standby task (either as active or standby) cannot take it again + if (standbyTaskId.isPresent() && processState.hasTask(standbyTaskId.get())) { + continue; + } + final double newProcessLoad = processState.load(); - if (newProcessLoad < candidateProcessLoad && (taskId == null || !processState.hasTask(taskId))) { - final double newMemberLoad = processState.memberToTaskCounts().get(member.memberId); - if (newMemberLoad < candidateMemberLoad) { - candidateProcessLoad = newProcessLoad; - candidateMemberLoad = newMemberLoad; - candidate = member; - } + final double newMemberLoad = processState.memberToTaskCounts().get(member.memberId); + if (candidate == null || (newProcessLoad < candidateProcessLoad && newMemberLoad < candidateMemberLoad)) { + candidateProcessLoad = newProcessLoad; + candidateMemberLoad = newMemberLoad; + candidate = member; } } - if (taskId == null || !candidateProcessState.hasTask(taskId)) { - return candidate; - } - return null; + return candidate; } private static boolean hasUnfulfilledActiveTaskQuota( @@ -367,7 +368,7 @@ private static void assignStandby(final LocalState localState, final LinkedList< // prev standby tasks final ArrayList prevStandbyMembers = localState.standbyTaskToPrevMember.get(task); if (prevStandbyMembers != null && !prevStandbyMembers.isEmpty()) { - final Member prevStandbyMember = findPrevMemberWithLeastLoad(localState, prevStandbyMembers, task); + final Member prevStandbyMember = findPrevMemberWithLeastLoad(localState, prevStandbyMembers, Optional.of(task)); if (prevStandbyMember != null) { final ProcessState prevStandbyMemberProcessState = localState.processIdToState.get(prevStandbyMember.processId); if (hasUnfulfilledTaskQuota(localState, prevStandbyMemberProcessState, prevStandbyMember)) { diff --git a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorTest.java b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorTest.java index f623f196f2337..71f5646d8e8fa 100644 --- a/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorTest.java +++ b/group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorTest.java @@ -28,6 +28,7 @@ import java.util.Arrays; import java.util.HashMap; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.NoSuchElementException; @@ -1306,6 +1307,58 @@ public void shouldAssignStandbyTaskToPreviousOwnerBasedOnBelowQuotaCondition() { } + @Test + public void shouldAssignStandbyToPreviousStandbyThatDoesNotHoldTheActiveTask() { + // starting assignment [active] [standby]: + // member1/process1: [] [2] + // member2/process2: [0] [2] + // member3/process3: [1] [] + // member4/process4: [] [] + // + // active 0,1 stay on their previous owner (stickiness) + // active 2 must go to member1/process1 as it has previous standby (standby->active promotion) + // -> it should not go to member2/process2 due to load balancing of active tasks + // standby 2 should stay on member2/process2 -> no reason to move it (load of member2 stays within capacity) + // standby 0,1: one *must* go to member4/process4 due to load balancing + // -> the other one can go anywhere but member2/process2 due to load balancing + // + // expected new assignment [active] [standby] + // ("expected" here means, based on the concrete implementation -- if this test fails with a different but + // still correct result -- as laid out above --, we should just update the test) + // member1/process1: [2] [] + // member2/process2: [0] [2] + // member3/process3: [1] [0] + // member4/process4: [] [1] + final Map members = new LinkedHashMap<>(); + members.put("member1", createMemberMetadata("process1", + Map.of(), mkMap(mkEntry("test-subtopology", Set.of(2))))); + members.put("member2", createMemberMetadata("process2", + mkMap(mkEntry("test-subtopology", Set.of(0))), mkMap(mkEntry("test-subtopology", Set.of(2))))); + members.put("member3", createMemberMetadata("process3", + mkMap(mkEntry("test-subtopology", Set.of(1))), Map.of())); + members.put("member4", createMemberMetadata("process4")); + + final GroupAssignment result = assignor.assign( + new GroupSpecImpl(members, mkMap(mkEntry(NUM_STANDBY_REPLICAS_CONFIG, "1"))), + new TopologyDescriberImpl(3, true, List.of("test-subtopology")) + ); + + assertEquals(Set.of(2), getActiveTasks(result, "test-subtopology", "member1")); + // if this is not empty, but only hold standby 0 or 1, still correct + assertEquals(List.of(), getAllStandbyTaskIds(result, "member1")); + + assertEquals(Set.of(0), getActiveTasks(result, "test-subtopology", "member2")); + assertEquals(List.of(2), getAllStandbyTaskIds(result, "member2")); + + assertEquals(Set.of(1), getActiveTasks(result, "test-subtopology", "member3")); + // if this is empty, or hold standby 1 instead, still correct + assertEquals(List.of(0), getAllStandbyTaskIds(result, "member3")); + + assertEquals(Set.of(), getActiveTasks(result, "test-subtopology", "member4")); + // if this hold standby 0, or both standby 0 and 1, still correct + assertEquals(List.of(1), getAllStandbyTaskIds(result, "member4")); + } + private int getAllActiveTaskCount(GroupAssignment result, String... memberIds) { int size = 0; for (String memberId : memberIds) {