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

Code before:
        return memberToTaskCounts;
    }

    public Set<TaskId> assignedActiveTasks() {
        return assignedActiveTasks.values().stream()
            .flatMap(Set::stream)
            .collect(Collectors.toSet());
    }

    public Map<String, Set<TaskId>> assignedActiveTasksByMember() {
        return assignedActiveTasks;
    }

    public Set<TaskId> assignedStandbyTasks() {
        return assignedStandbyTasks.values().stream()
            .flatMap(Set::stream)
            .collect(Collectors.toSet());
    }

    public Map<String, Set<TaskId>> assignedStandbyTasksByMember() {
        return assignedStandbyTasks;
    }

    public void addTask(final String memberId, final TaskId taskId, final boolean isActive) {
        addTaskInternal(memberId, taskId, isActive);
        membersByLoad = null; // reset, since it may not be sorted anymore
    }

    private int addTaskInternal(final String memberId, final TaskId taskId, final boolean isActive) {
        taskCount += 1;
        assignedTasks.add(taskId);
        if (isActive) {
            assignedActiveTasks.putIfAbsent(memberId, new HashSet<>());
            assignedActiveTasks.get(memberId).add(taskId);
        } else {
            assignedStandbyTasks.putIfAbsent(memberId, new HashSet<>());
            assignedStandbyTasks.get(memberId).add(taskId);
        }
        int newTaskCount = memberToTaskCounts.get(memberId) + 1;
        memberToTaskCounts.put(memberId, newTaskCount);
        computeLoad();
        return newTaskCount;
    }

    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);
        }

Review comment (performance):
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)

Code after:
        return memberToTaskCounts;
    }

    public Set<TaskId> assignedActiveTasks() {
        return assignedActiveTasks.values().stream()
            .flatMap(Set::stream)
            .collect(Collectors.toSet());
    }

    public Map<String, Set<TaskId>> assignedActiveTasksByMember() {
        return assignedActiveTasks;
    }

    public Set<TaskId> assignedStandbyTasks() {
        return assignedStandbyTasks.values().stream()
            .flatMap(Set::stream)
            .collect(Collectors.toSet());
    }

    public Map<String, Set<TaskId>> assignedStandbyTasksByMember() {
        return assignedStandbyTasks;
    }

    /**
     * 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) {
            assignedActiveTasks.putIfAbsent(memberId, new HashSet<>());
            assignedActiveTasks.get(memberId).add(taskId);
        } else {
            assignedStandbyTasks.putIfAbsent(memberId, new HashSet<>());
            assignedStandbyTasks.get(memberId).add(taskId);