KAFKA-19661 [6/N]: Use heaps also on the process-level (#20523)

In the current solution, we only use a heap to select the right process,
but resort to linear search for selecting a member within a process.
This means use cases where a lot of threads run within the same process
can yield slow assignment. The number of threads in a process shouldn’t
scale arbitrarily (our assumed case for benchmarking of 50 threads in a
single process seems quite extreme already), however, we can optimize
for this case to reduce the runtime further.

Other assignment algorithms assign directly on the member-level, but we
cannot do this in Kafka Streams, since we cannot assign tasks to
processes that already own the task. Defining a heap directly on members
would mean that we may have to skip through 10s of member before finding
one that does not belong to a process that does not yet own the member.

Instead, we can define a separate heap for each process, which keeps the
members of the process by load. We can only keep the heap as long as we
are only changing the load of the top-most member (which we usually do).
This means we keep track of a lot of heaps, but since heaps are backed
by arrays in Java, this should not result in extreme memory
inefficiencies.

In our worst-performing benchmark, this improves the runtime by ~2x on
top of the optimization above.

Also piggybacked are some minor optimizations / clean-ups:   -
initialize HashMaps and ArrayLists with the right capacity   - fix some
comments   - improve logging output

Note that this is a pure performance change, so there are no changes to
the unit tests.

Reviewers: Bill Bejeck<[email protected]>

Author: lucasbru

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

@@ -27,9 +27,9 @@
  *
  * @param instanceId   The instance ID if provided.
  * @param rackId       The rack ID if provided.
- * @param activeTasks  Reconciled active tasks
- * @param standbyTasks Reconciled standby tasks
- * @param warmupTasks  Reconciled warm-up tasks
+ * @param activeTasks  Current target active tasks
+ * @param standbyTasks Current target standby tasks
+ * @param warmupTasks  Current target warm-up tasks
  * @param processId    The process ID.
  * @param clientTags   The client tags for a rack-aware assignment.
  * @param taskOffsets  The last received cumulative task offsets of assigned tasks or dormant tasks.

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

@@ -16,9 +16,11 @@
  */
 package org.apache.kafka.coordinator.group.streams.assignor;
 
+import java.util.AbstractMap;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
+import java.util.PriorityQueue;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -36,6 +38,7 @@ public class ProcessState {
     private final Map<String, Set<TaskId>> assignedActiveTasks;
     private final Map<String, Set<TaskId>> assignedStandbyTasks;
     private final Set<TaskId> assignedTasks;
+    private PriorityQueue<Map.Entry<String, Integer>> membersByLoad;
 
     ProcessState(final String processId) {
         this.processId = processId;
@@ -45,9 +48,9 @@ public class ProcessState {
         this.assignedActiveTasks = new HashMap<>();
         this.assignedStandbyTasks = new HashMap<>();
         this.memberToTaskCounts = new HashMap<>();
+        this.membersByLoad = null;
     }
 
-
     public String processId() {
         return processId;
     }
@@ -84,7 +87,26 @@ public Map<String, Set<TaskId>> assignedStandbyTasksByMember() {
         return assignedStandbyTasks;
     }
 
-    public void addTask(final String memberId, final TaskId taskId, final boolean isActive) {
+    /**
+     * Assigns a task to a member of this process.
+     *
+     * @param memberId The member to assign to.
+     * @param taskId   The task to assign.
+     * @param isActive Whether the task is an active task (true) or a standby task (false).
+     * @return the number of tasks that `memberId` has assigned after adding the new task.
+     */
+    public int addTask(final String memberId, final TaskId taskId, final boolean isActive) {
+        int newTaskCount = addTaskInternal(memberId, taskId, isActive);
+        // We cannot efficiently add a task to a specific member and keep the memberByLoad ordered correctly.
+        // So we just drop the heap here.
+        //
+        // The order in which addTask and addTaskToLeastLoadedMember is called ensures that the heaps are built at most
+        // twice (once for active, once for standby)
+        membersByLoad = null;
+        return newTaskCount;
+    }
+
+    private int addTaskInternal(final String memberId, final TaskId taskId, final boolean isActive) {
         taskCount += 1;
         assignedTasks.add(taskId);
         if (isActive) {
@@ -94,8 +116,46 @@ public void addTask(final String memberId, final TaskId taskId, final boolean is
             assignedStandbyTasks.putIfAbsent(memberId, new HashSet<>());
             assignedStandbyTasks.get(memberId).add(taskId);
         }
-        memberToTaskCounts.put(memberId, memberToTaskCounts.get(memberId) + 1);
+        int newTaskCount = memberToTaskCounts.get(memberId) + 1;
+        memberToTaskCounts.put(memberId, newTaskCount);
         computeLoad();
+        return newTaskCount;
+    }
+
+    /**
+     * Assigns a task to the least loaded member of this process
+     *
+     * @param taskId   The task to assign.
+     * @param isActive Whether the task is an active task (true) or a standby task (false).
+     * @return the number of tasks that `memberId` has assigned after adding the new task, or -1 if the
+     *         task was not assigned to any member.
+     */
+    public int addTaskToLeastLoadedMember(final TaskId taskId, final boolean isActive) {
+        if (memberToTaskCounts.isEmpty()) {
+            return -1;
+        }
+        if (memberToTaskCounts.size() == 1) {
+            return addTaskInternal(memberToTaskCounts.keySet().iterator().next(), taskId, isActive);
+        }
+        if (membersByLoad == null) {
+            membersByLoad = new PriorityQueue<>(
+                memberToTaskCounts.size(),
+                Map.Entry.comparingByValue()
+            );
+            for (Map.Entry<String, Integer> entry : memberToTaskCounts.entrySet()) {
+                // Copy here, since map entry objects are allowed to be reused by the underlying map implementation.
+                membersByLoad.add(new AbstractMap.SimpleEntry<>(entry.getKey(), entry.getValue()));
+            }
+        }
+        Map.Entry<String, Integer> member = membersByLoad.poll();
+        if (member != null) {
+            int newTaskCount = addTaskInternal(member.getKey(), taskId, isActive);
+            member.setValue(newTaskCount);
+            membersByLoad.add(member); // Reinsert the updated member back into the priority queue
+            return newTaskCount;
+        } else {
+            throw new TaskAssignorException("No members available to assign task " + taskId);
+        }
     }
 
     private void incrementCapacity() {

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

@@ -27,7 +27,6 @@
 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;
@@ -97,11 +96,11 @@ private void initialize(final GroupSpec groupSpec, final TopologyDescriber topol
         localState.totalMembersWithActiveTaskCapacity = groupSpec.members().size();
         localState.totalMembersWithTaskCapacity = groupSpec.members().size();
         localState.activeTasksPerMember = computeTasksPerMember(localState.totalActiveTasks, localState.totalMembersWithActiveTaskCapacity);
-        localState.tasksPerMember = computeTasksPerMember(localState.totalTasks, localState.totalMembersWithTaskCapacity);
+        localState.totalTasksPerMember = computeTasksPerMember(localState.totalTasks, localState.totalMembersWithTaskCapacity);
 
-        localState.processIdToState = new HashMap<>();
-        localState.activeTaskToPrevMember = new HashMap<>();
-        localState.standbyTaskToPrevMember = new HashMap<>();
+        localState.processIdToState = new HashMap<>(localState.totalMembersWithActiveTaskCapacity);
+        localState.activeTaskToPrevMember = new HashMap<>(localState.totalActiveTasks);
+        localState.standbyTaskToPrevMember = new HashMap<>(localState.numStandbyReplicas > 0 ? (localState.totalTasks - localState.totalActiveTasks) / localState.numStandbyReplicas : 0);
         for (final Map.Entry<String, AssignmentMemberSpec> memberEntry : groupSpec.members().entrySet()) {
             final String memberId = memberEntry.getKey();
             final String processId = memberEntry.getValue().processId();
@@ -124,7 +123,7 @@ private void initialize(final GroupSpec groupSpec, final TopologyDescriber topol
                 final Set<Integer> partitionNoSet = entry.getValue();
                 for (final int partitionNo : partitionNoSet) {
                     final TaskId taskId = new TaskId(entry.getKey(), partitionNo);
-                    localState.standbyTaskToPrevMember.putIfAbsent(taskId, new ArrayList<>());
+                    localState.standbyTaskToPrevMember.putIfAbsent(taskId, new ArrayList<>(localState.numStandbyReplicas));
                     localState.standbyTaskToPrevMember.get(taskId).add(member);
                 }
             }
@@ -185,8 +184,9 @@ private void assignActive(final LinkedList<TaskId> activeTasks) {
             if (prevMember != null) {
                 final ProcessState processState = localState.processIdToState.get(prevMember.processId);
                 if (hasUnfulfilledActiveTaskQuota(processState, prevMember)) {
-                    processState.addTask(prevMember.memberId, task, true);
-                    maybeUpdateActiveTasksPerMember(processState.memberToTaskCounts().get(prevMember.memberId));
+                    int newActiveTasks = processState.addTask(prevMember.memberId, task, true);
+                    maybeUpdateActiveTasksPerMember(newActiveTasks);
+                    maybeUpdateTotalTasksPerMember(newActiveTasks);
                     it.remove();
                 }
             }
@@ -200,8 +200,9 @@ private void assignActive(final LinkedList<TaskId> activeTasks) {
             if (prevMember != null) {
                 final ProcessState processState = localState.processIdToState.get(prevMember.processId);
                 if (hasUnfulfilledActiveTaskQuota(processState, prevMember)) {
-                    processState.addTask(prevMember.memberId, task, true);
-                    maybeUpdateActiveTasksPerMember(processState.memberToTaskCounts().get(prevMember.memberId));
+                    int newActiveTasks = processState.addTask(prevMember.memberId, task, true);
+                    maybeUpdateActiveTasksPerMember(newActiveTasks);
+                    maybeUpdateTotalTasksPerMember(newActiveTasks);
                     it.remove();
                 }
             }
@@ -213,19 +214,18 @@ private void assignActive(final LinkedList<TaskId> activeTasks) {
         // 3. assign any remaining unassigned tasks
         final PriorityQueue<ProcessState> processByLoad = new PriorityQueue<>(Comparator.comparingDouble(ProcessState::load));
         processByLoad.addAll(localState.processIdToState.values());
-        for (final Iterator<TaskId> it = activeTasks.iterator(); it.hasNext();) {
-            final TaskId task = it.next();
+        for (final TaskId task: activeTasks) {
             final ProcessState processWithLeastLoad = processByLoad.poll();
             if (processWithLeastLoad == null) {
                 throw new TaskAssignorException(String.format("No process available to assign active task %s.", task));
             }
-            final String member = memberWithLeastLoad(processWithLeastLoad);
-            if (member == null) {
+            final int newTaskCount = processWithLeastLoad.addTaskToLeastLoadedMember(task, true);
+            if (newTaskCount != -1) {
+                maybeUpdateActiveTasksPerMember(newTaskCount);
+                maybeUpdateTotalTasksPerMember(newTaskCount);
+            } else {
                 throw new TaskAssignorException(String.format("No member available to assign active task %s.", task));
             }
-            processWithLeastLoad.addTask(member, task, true);
-            it.remove();
-            maybeUpdateActiveTasksPerMember(processWithLeastLoad.memberToTaskCounts().get(member));
             processByLoad.add(processWithLeastLoad); // Add it back to the queue after updating its state
         }
     }
@@ -238,11 +238,11 @@ private void maybeUpdateActiveTasksPerMember(final int activeTasksNo) {
         }
     }
 
-    private void maybeUpdateTasksPerMember(final int taskNo) {
-        if (taskNo == localState.tasksPerMember) {
+    private void maybeUpdateTotalTasksPerMember(final int taskNo) {
+        if (taskNo == localState.totalTasksPerMember) {
             localState.totalMembersWithTaskCapacity--;
             localState.totalTasks -= taskNo;
-            localState.tasksPerMember = computeTasksPerMember(localState.totalTasks, localState.totalMembersWithTaskCapacity);
+            localState.totalTasksPerMember = computeTasksPerMember(localState.totalTasks, localState.totalMembersWithTaskCapacity);
         }
     }
 
@@ -253,10 +253,10 @@ private boolean assignStandbyToMemberWithLeastLoad(PriorityQueue<ProcessState> q
         }
         boolean found = false;
         if (!processWithLeastLoad.hasTask(taskId)) {
-            final String memberId = memberWithLeastLoad(processWithLeastLoad);
-            if (memberId != null) {
-                processWithLeastLoad.addTask(memberId, taskId, false);
+            final int newTaskCount = processWithLeastLoad.addTaskToLeastLoadedMember(taskId, false);
+            if (newTaskCount != -1) {
                 found = true;
+                maybeUpdateTotalTasksPerMember(newTaskCount);
             }
         } else if (!queue.isEmpty()) {
             found = assignStandbyToMemberWithLeastLoad(queue, taskId);
@@ -303,26 +303,12 @@ private Member findPrevMemberWithLeastLoad(final ArrayList<Member> members, fina
         return null;
     }
 
-    private String memberWithLeastLoad(final ProcessState processWithLeastLoad) {
-        final Map<String, Integer> members = processWithLeastLoad.memberToTaskCounts();
-        if (members.isEmpty()) {
-            return null;
-        }
-        if (members.size() == 1) {
-            return members.keySet().iterator().next();
-        }
-        final Optional<String> memberWithLeastLoad = processWithLeastLoad.memberToTaskCounts().entrySet().stream()
-            .min(Map.Entry.comparingByValue())
-            .map(Map.Entry::getKey);
-        return memberWithLeastLoad.orElse(null);
-    }
-
     private boolean hasUnfulfilledActiveTaskQuota(final ProcessState process, final Member member) {
         return process.memberToTaskCounts().get(member.memberId) < localState.activeTasksPerMember;
     }
 
     private boolean hasUnfulfilledTaskQuota(final ProcessState process, final Member member) {
-        return process.memberToTaskCounts().get(member.memberId) < localState.tasksPerMember;
+        return process.memberToTaskCounts().get(member.memberId) < localState.totalTasksPerMember;
     }
 
     private void assignStandby(final LinkedList<TaskId> standbyTasks) {
@@ -339,8 +325,8 @@ private void assignStandby(final LinkedList<TaskId> standbyTasks) {
                 if (prevActiveMember != null) {
                     final ProcessState prevActiveMemberProcessState = localState.processIdToState.get(prevActiveMember.processId);
                     if (!prevActiveMemberProcessState.hasTask(task) && hasUnfulfilledTaskQuota(prevActiveMemberProcessState, prevActiveMember)) {
-                        prevActiveMemberProcessState.addTask(prevActiveMember.memberId, task, false);
-                        maybeUpdateTasksPerMember(prevActiveMemberProcessState.memberToTaskCounts().get(prevActiveMember.memberId));
+                        int newTaskCount = prevActiveMemberProcessState.addTask(prevActiveMember.memberId, task, false);
+                        maybeUpdateTotalTasksPerMember(newTaskCount);
                         continue;
                     }
                 }
@@ -352,8 +338,8 @@ private void assignStandby(final LinkedList<TaskId> standbyTasks) {
                     if (prevStandbyMember != null) {
                         final ProcessState prevStandbyMemberProcessState = localState.processIdToState.get(prevStandbyMember.processId);
                         if (hasUnfulfilledTaskQuota(prevStandbyMemberProcessState, prevStandbyMember)) {
-                            prevStandbyMemberProcessState.addTask(prevStandbyMember.memberId, task, false);
-                            maybeUpdateTasksPerMember(prevStandbyMemberProcessState.memberToTaskCounts().get(prevStandbyMember.memberId));
+                            int newTaskCount = prevStandbyMemberProcessState.addTask(prevStandbyMember.memberId, task, false);
+                            maybeUpdateTotalTasksPerMember(newTaskCount);
                             continue;
                         }
                     }
@@ -430,6 +416,6 @@ private static class LocalState {
         int totalMembersWithActiveTaskCapacity;
         int totalMembersWithTaskCapacity;
         int activeTasksPerMember;
-        int tasksPerMember;
+        int totalTasksPerMember;
     }
 }

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

@@ -38,4 +38,9 @@ public int compareTo(final TaskId other) {
             .compare(this, other);
     }
 
+    @Override
+    public String toString() {
+        return subtopologyId + '_' + partition;
+    }
+
 }