KAFKA-20899: Fix "streams" StickyAssignor bug for standby tasks (#23096)

Broker side StickyTaskAssignor does not assign standby tasks correctly,
because it does not cycle through all group member correctly to find the
right candidate to host it.

Reviewers: Mingi Cho (github:ChoMinGi), Bill Bejeck <bbejeck@apache.org>

Author: mjsax

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java

@@ -27,6 +27,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;
@@ -201,7 +202,7 @@ private void assignActive(final LinkedList<TaskId> activeTasks) {
         for (final Iterator<TaskId> it = activeTasks.iterator(); it.hasNext();) {
             final TaskId task = it.next();
             final ArrayList<Member> prevMembers = localState.standbyTaskToPrevMember.get(task);
-            final Member prevMember = findPrevMemberWithLeastLoad(prevMembers, null);
+            final Member prevMember = findPrevMemberWithLeastLoad(prevMembers, Optional.empty());
             if (prevMember != null) {
                 final ProcessState processState = localState.processIdToState.get(prevMember.processId);
                 if (hasUnfulfilledActiveTaskQuota(processState, prevMember)) {
@@ -279,33 +280,31 @@ private boolean assignStandbyToMemberWithLeastLoad(PriorityQueue<ProcessState> q
      *
      * @return Previous member with the least load that does not have the task, or null if no such member exists.
      */
-    private Member findPrevMemberWithLeastLoad(final ArrayList<Member> members, final TaskId taskId) {
+    private Member findPrevMemberWithLeastLoad(final ArrayList<Member> members, final Optional<TaskId> 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 boolean hasUnfulfilledActiveTaskQuota(final ProcessState process, final Member member) {
@@ -339,7 +338,7 @@ private void assignStandby(final LinkedList<TaskId> standbyTasks) {
                 // prev standby tasks
                 final ArrayList<Member> prevStandbyMembers = localState.standbyTaskToPrevMember.get(task);
                 if (prevStandbyMembers != null && !prevStandbyMembers.isEmpty()) {
-                    final Member prevStandbyMember = findPrevMemberWithLeastLoad(prevStandbyMembers, task);
+                    final Member prevStandbyMember = findPrevMemberWithLeastLoad(prevStandbyMembers, Optional.of(task));
                     if (prevStandbyMember != null) {
                         final ProcessState prevStandbyMemberProcessState = localState.processIdToState.get(prevStandbyMember.processId);
                         if (hasUnfulfilledTaskQuota(prevStandbyMemberProcessState, prevStandbyMember)) {

group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorTest.java

@@ -24,6 +24,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;
@@ -1302,6 +1303,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<String, AssignmentMemberSpec> members = new LinkedHashMap<>();
+        members.put("member1", createAssignmentMemberSpec("process1",
+            Map.of(), mkMap(mkEntry("test-subtopology", Set.of(2)))));
+        members.put("member2", createAssignmentMemberSpec("process2",
+            mkMap(mkEntry("test-subtopology", Set.of(0))), mkMap(mkEntry("test-subtopology", Set.of(2)))));
+        members.put("member3", createAssignmentMemberSpec("process3",
+            mkMap(mkEntry("test-subtopology", Set.of(1))), Map.of()));
+        members.put("member4", createAssignmentMemberSpec("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) {