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

Code before:
            );
            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() {
        capacity++;
        computeLoad();
    }

    public void computeLoad() {
        if (capacity <= 0) {
            this.load = -1;
        } else {
            this.load = (double) taskCount / capacity;
        }
    }

    public void addMember(final String member) {
        this.memberToTaskCounts.put(member, 0);
        incrementCapacity();
    }

    public boolean hasCapacity() {
        return this.load < 1.0;
    }

    public int compareTo(final ProcessState other) {
        int loadCompare = Double.compare(this.load, other.load());
        if (loadCompare == 0) {
            return Integer.compare(other.capacity, this.capacity);
        }
        return loadCompare;
    }

    public boolean hasTask(final TaskId taskId) {
        return assignedTasks.contains(taskId);
    }


Review comment (performance):
Creating new AbstractMap.SimpleEntry objects for each entry is inefficient. Consider using a custom wrapper class or storing member IDs directly with their task counts in a separate data structure.

Code after:
        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() {
        capacity++;
        computeLoad();
    }

    public void computeLoad() {
        if (capacity <= 0) {
            this.load = -1;
        } else {
            this.load = (double) taskCount / capacity;
        }
    }