Repository: apache/kafka
File: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignor.java (Java)

Code before:
                    processState.addTask(prevMember.memberId, task, true);
                    maybeUpdateActiveTasksPerMember(processState.memberToTaskCounts().get(prevMember.memberId));
                    it.remove();
                }
            }
        }

        // To achieve an initially range-based assignment, sort by subtopology
        activeTasks.sort(Comparator.comparing(TaskId::subtopologyId).thenComparing(TaskId::partition));

        // 3. assign any remaining unassigned tasks
        final PriorityQueue<ProcessState> processByLoad = new PriorityQueue<>(Comparator.comparingDouble(ProcessState::load));
        processByLoad.addAll(localState.processIdToState.values());
        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 int newTaskCount = processWithLeastLoad.addTaskToLeastLoadedMember(task, true);
            if (newTaskCount != -1) {
                maybeUpdateActiveTasksPerMember(newTaskCount);
            } else {
                throw new TaskAssignorException(String.format("No member available to assign active task %s.", task));
            }
            processByLoad.add(processWithLeastLoad); // Add it back to the queue after updating its state
        }
    }

    private void maybeUpdateActiveTasksPerMember(final int activeTasksNo) {
        if (activeTasksNo == localState.activeTasksPerMember) {
            localState.totalMembersWithActiveTaskCapacity--;
            localState.totalActiveTasks -= activeTasksNo;
            localState.activeTasksPerMember = computeTasksPerMember(localState.totalActiveTasks, localState.totalMembersWithActiveTaskCapacity);
        }
    }

    private void maybeUpdateTasksPerMember(final int taskNo) {
        if (taskNo == localState.tasksPerMember) {
            localState.totalMembersWithTaskCapacity--;
            localState.totalTasks -= taskNo;
            localState.tasksPerMember = computeTasksPerMember(localState.totalTasks, localState.totalMembersWithTaskCapacity);
        }
    }

    private boolean assignStandbyToMemberWithLeastLoad(PriorityQueue<ProcessState> queue, TaskId taskId) {
        final ProcessState processWithLeastLoad = queue.poll();
        if (processWithLeastLoad == null) {
            return false;
        }
        boolean found = false;
        if (!processWithLeastLoad.hasTask(taskId)) {

Review comment (suggestion):
I removed the iterator here, since we don't really need to remove from `activeTasks` - we will not use it afterwards. So a foreach loop is simpler.

Code after:
                if (hasUnfulfilledActiveTaskQuota(processState, prevMember)) {
                    int newActiveTasks = processState.addTask(prevMember.memberId, task, true);
                    maybeUpdateActiveTasksPerMember(newActiveTasks);
                    maybeUpdateTotalTasksPerMember(newActiveTasks);
                    it.remove();
                }
            }
        }

        // To achieve an initially range-based assignment, sort by subtopology
        activeTasks.sort(Comparator.comparing(TaskId::subtopologyId).thenComparing(TaskId::partition));

        // 3. assign any remaining unassigned tasks
        final PriorityQueue<ProcessState> processByLoad = new PriorityQueue<>(Comparator.comparingDouble(ProcessState::load));
        processByLoad.addAll(localState.processIdToState.values());
        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 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));
            }
            processByLoad.add(processWithLeastLoad); // Add it back to the queue after updating its state
        }
    }

    private void maybeUpdateActiveTasksPerMember(final int activeTasksNo) {
        if (activeTasksNo == localState.activeTasksPerMember) {
            localState.totalMembersWithActiveTaskCapacity--;
            localState.totalActiveTasks -= activeTasksNo;
            localState.activeTasksPerMember = computeTasksPerMember(localState.totalActiveTasks, localState.totalMembersWithActiveTaskCapacity);
        }
    }

    private void maybeUpdateTotalTasksPerMember(final int taskNo) {
        if (taskNo == localState.totalTasksPerMember) {
            localState.totalMembersWithTaskCapacity--;
            localState.totalTasks -= taskNo;
            localState.totalTasksPerMember = computeTasksPerMember(localState.totalTasks, localState.totalMembersWithTaskCapacity);
        }
    }

    private boolean assignStandbyToMemberWithLeastLoad(PriorityQueue<ProcessState> queue, TaskId taskId) {
        final ProcessState processWithLeastLoad = queue.poll();
        if (processWithLeastLoad == null) {
            return false;