lucasbru commented on code in PR #23588:
URL: https://github.com/apache/kafka/pull/23588#discussion_r4164622495


##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/ProcessState.java:
##########
@@ -44,6 +44,8 @@ public class ProcessState {
     private final Map<String, Set<TaskId>> assignedStandbyTasks;
     private final Set<TaskId> assignedTasks;
     private PriorityQueue<Map.Entry<String, Integer>> membersByLoad;
+    // The member with the fewest tasks, or null when not computed since the 
last added task.
+    private String leastLoadedMember;

Review Comment:
   This adds a second cache for "least loaded member" (a scalar recomputed by a 
full scan) alongside membersByLoad, the existing PriorityQueue that already 
answers this. Now there are two mechanisms for the same thing that have to stay 
in sync, and the new one is O(n) per miss vs O(log n) for the queue. Was 
membersByLoad not sufficient here, or is this meant to replace it?



##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java:
##########
@@ -524,6 +605,123 @@ private static void assignStandby(final LocalState 
localState, final LinkedList<
         }
     }
 
+    /**
+     * Places the standbys of {@code task} that still make it more 
rack-diverse on processes with room and returns
+     * the processes they went to. Among the equally most diverse processes 
the choice is the tag-blind picks unchanged:
+     * the previous active member, else the least-loaded previous standby 
member, each while below the quota, else
+     * the least-loaded process and its least-loaded member.
+     * <p>
+     * {@code alreadyPlaced} are the standbys of the task placed by an earlier 
call. With {@code stickyOnly},
+     * the picks stop at the first standby that would not go back to a 
previous member and add the task to
+     * {@code rackNonSticky}, leaving it and the rest to a later call.
+     */
+    private static List<ProcessState> pickRackAwareStandbys(
+        final LocalState localState,
+        final RackAwareStandbyPicker<ProcessGroup> picker,
+        final TaskId task,
+        final List<ProcessState> alreadyPlaced,
+        final boolean stickyOnly,
+        final List<TaskId> rackNonSticky
+    ) {
+        picker.startTask();
+        
picker.markUsed(localState.processIdToGroup.get(localState.statefulActiveTaskToProcess.get(task).processId()));
+        for (final ProcessState process : alreadyPlaced) {
+            
picker.markUsed(localState.processIdToGroup.get(process.processId()));
+        }
+
+        // The picker only tests groups whose value for the priority key no 
holder carries, so no process of such a group
+        // holds the task and only room is checked.
+        final Predicate<ProcessGroup> eligible = group -> 
group.hasRoom(localState);
+
+        final List<ProcessState> placed = new ArrayList<>();
+        while (alreadyPlaced.size() + placed.size() < 
localState.numStandbyReplicas) {
+            final Set<ProcessGroup> candidates = 
picker.pickCandidates(eligible);
+            if (candidates.isEmpty()) {
+                break;
+            }
+
+            // prev active or standby member on a process with room in a 
candidate group
+            final Member prevMember = findPrevMemberForStandby(
+                localState,
+                task,
+                process -> 
candidates.contains(localState.processIdToGroup.get(process.processId())) && 
leastLoadedMemberWithRoom(localState, process) != null
+            );
+            if (prevMember != null) {
+                placeRackAwareStandby(localState, picker, 
localState.processIdToState.get(prevMember.processId), prevMember.memberId, 
task, placed);
+                continue;
+            }
+
+            if (stickyOnly) {
+                rackNonSticky.add(task);
+                break;
+            }
+
+            // least loaded process of the candidate groups. It has a member 
with room, so the lookup cannot return null.
+            // Adding to that member directly keeps the heap of members by 
load out of this pass, so that it is built
+            // once, in the least-loaded pass, as without rack awareness.
+            final ProcessState processWithLeastLoad = leastLoaded(localState, 
candidates);
+            placeRackAwareStandby(localState, picker, processWithLeastLoad, 
leastLoadedMemberWithRoom(localState, processWithLeastLoad), task, placed);
+        }
+        return placed;
+    }
+
+    private static void placeRackAwareStandby(
+        final LocalState localState,
+        final RackAwareStandbyPicker<ProcessGroup> picker,
+        final ProcessState process,
+        final String memberId,
+        final TaskId task,
+        final List<ProcessState> placed
+    ) {
+        // The queue of the group catches up with the new load once the 
process reaches its head.
+        maybeUpdateTotalTasksPerMember(localState, process.addTask(memberId, 
task, false, true));
+        picker.markUsed(localState.processIdToGroup.get(process.processId()));
+        placed.add(process);
+    }
+
+    /** The least-loaded process with room of the candidate groups, which all 
have one. */
+    private static ProcessState leastLoaded(final LocalState localState, final 
Collection<ProcessGroup> candidates) {

Review Comment:
   leastLoaded() scans every candidate group linearly, comparing PQ heads one 
by one. Since each ProcessGroup already exposes an ordered 
least-loaded-with-room head, a small merge-heap over the group heads would make 
this O(log #groups) instead of O(#groups). Minor at today's group counts, but 
it's basically free given the structure already here.



##########
group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java:
##########
@@ -467,38 +496,89 @@ private static boolean hasUnfulfilledTaskQuota(
         return process.memberToTaskCounts().get(member.memberId) < 
localState.totalTasksPerMember;
     }
 
+    /**
+     * The previous member to take a standby of {@code task}: the previous 
active member, else the least-loaded previous
+     * standby member, each only while below the quota and on a process {@code 
allowed} accepts. Null if neither.
+     */
+    private static Member findPrevMemberForStandby(
+        final LocalState localState,
+        final TaskId task,
+        final Predicate<ProcessState> allowed
+    ) {
+        // prev active task
+        final Member prevActiveMember = 
localState.activeTaskToPrevMember.get(task);
+        if (prevActiveMember != null) {
+            final ProcessState prevActiveMemberProcessState = 
localState.processIdToState.get(prevActiveMember.processId);
+            if (allowed.test(prevActiveMemberProcessState) && 
!prevActiveMemberProcessState.hasTask(task) && 
hasUnfulfilledTaskQuota(localState, prevActiveMemberProcessState, 
prevActiveMember)) {
+                return prevActiveMember;
+            }
+        }
+
+        // prev standby tasks
+        final Member prevStandbyMember = 
findPrevMemberWithLeastLoad(localState, 
localState.standbyTaskToPrevMember.get(task), Optional.of(task), allowed);
+        if (prevStandbyMember != null && hasUnfulfilledTaskQuota(localState, 
localState.processIdToState.get(prevStandbyMember.processId), 
prevStandbyMember)) {
+            return prevStandbyMember;
+        }
+        return null;
+    }
+
+    /** Whether {@code task} had an active or standby member before, the only 
members a sticky pick can go back to. */
+    private static boolean hasPrevMember(final LocalState localState, final 
TaskId task) {
+        final ArrayList<Member> prevStandbyMembers = 
localState.standbyTaskToPrevMember.get(task);
+        return localState.activeTaskToPrevMember.get(task) != null || 
prevStandbyMembers != null && !prevStandbyMembers.isEmpty();
+    }
+
+    /**
+     * Assigns the standby tasks. Rack diversity ranks above stickiness: when 
{@code rack.aware.assignment.tags} is
+     * set, steps 1 and 2 place every standby that still makes its task more 
rack-diverse, using stickiness only to
+     * break ties among the most diverse processes; steps 3 and 4 are the 
tag-blind assignment of the rest. Each pair
+     * runs its sticky step for all tasks before its least-loaded step, so 
that a pick that leaves its previous members
+     * cannot take the room a later task needs to stay with one of its own. 
Steps 3 and 4 run after step 2, so a
+     * rack-aware pick may still take the room a tag-blind sticky pick needs.
+     */
     private static void assignStandby(final LocalState localState, final 
LinkedList<TaskId> standbyTasks) {

Review Comment:
   assignStandby now runs four phases (rack-aware-sticky, 
rack-aware-least-loaded, tag-blind-sticky, tag-blind-least-loaded) with three 
bookkeeping collections (currentStandbys, rackAwareStandbys, rackNonSticky) 
threaded through one long method. The per-task diversity logic got its own 
abstraction in RackAwareStandbyPicker; would it help readability to pull the 
phase ordering/precedence out into something similar, rather than following it 
inline 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: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to