root
/
kafka
mirror of https://github.com/apache/kafka.git
1
0
Fork 0
Code Issues Actions 2 Packages Projects Releases Wiki Activity

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<bbejeck@apache.org>
This commit is contained in:
Lucas Brutschy 2025-09-15 17:19:53 +02:00 committed by GitHub
parent 749c2d91d5
commit 8628d74c49
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 100 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

@ -16,9 +16,11 @@
*/ */
package org.apache.kafka.coordinator.group.streams.assignor; package org.apache.kafka.coordinator.group.streams.assignor;
import java.util.AbstractMap;
import java.util.HashMap; import java.util.HashMap;
import java.util.HashSet; import java.util.HashSet;
import java.util.Map; import java.util.Map;
import java.util.PriorityQueue;
import java.util.Set; import java.util.Set;
import java.util.stream.Collectors; 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>> assignedActiveTasks;
private final Map<String, Set<TaskId>> assignedStandbyTasks; private final Map<String, Set<TaskId>> assignedStandbyTasks;
private final Set<TaskId> assignedTasks; private final Set<TaskId> assignedTasks;
private PriorityQueue<Map.Entry<String, Integer>> membersByLoad;
ProcessState(final String processId) { ProcessState(final String processId) {
this.processId = processId; this.processId = processId;
@ -45,9 +48,9 @@ public class ProcessState {
this.assignedActiveTasks = new HashMap<>(); this.assignedActiveTasks = new HashMap<>();
this.assignedStandbyTasks = new HashMap<>(); this.assignedStandbyTasks = new HashMap<>();
this.memberToTaskCounts = new HashMap<>(); this.memberToTaskCounts = new HashMap<>();
this.membersByLoad = null;
} }
public String processId() { public String processId() {
return processId; return processId;
} }
@ -84,7 +87,26 @@ public class ProcessState {
return assignedStandbyTasks; 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; taskCount += 1;
assignedTasks.add(taskId); assignedTasks.add(taskId);
if (isActive) { if (isActive) {
@ -94,8 +116,46 @@ public class ProcessState {
assignedStandbyTasks.putIfAbsent(memberId, new HashSet<>()); assignedStandbyTasks.putIfAbsent(memberId, new HashSet<>());
assignedStandbyTasks.get(memberId).add(taskId); assignedStandbyTasks.get(memberId).add(taskId);
} }
memberToTaskCounts.put(memberId, memberToTaskCounts.get(memberId) + 1); int newTaskCount = memberToTaskCounts.get(memberId) + 1;
memberToTaskCounts.put(memberId, newTaskCount);
computeLoad(); 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() { private void incrementCapacity() {

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

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

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

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