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

Code before:
            final Set<TaskId> statefulTasks = taskIds(topologyDescriber, false);
            assignStandby(statefulTasks);
        }

        return buildGroupAssignment(groupSpec.members().keySet());
    }

    private Set<TaskId> taskIds(final TopologyDescriber topologyDescriber, final boolean isActive) {
        final Set<TaskId> ret = new HashSet<>();
        for (final String subtopology : topologyDescriber.subtopologies()) {
            if (isActive || topologyDescriber.isStateful(subtopology)) {
                final int numberOfPartitions = topologyDescriber.maxNumInputPartitions(subtopology);
                for (int i = 0; i < numberOfPartitions; i++) {
                    ret.add(new TaskId(subtopology, i));
                }
            }
        }
        return ret;
    }

    private void initialize(final GroupSpec groupSpec, final TopologyDescriber topologyDescriber) {
        localState = new LocalState();
        localState.numStandbyReplicas =
            groupSpec.assignmentConfigs().isEmpty() ? 0
                : Integer.parseInt(groupSpec.assignmentConfigs().get("num.standby.replicas"));

        // Helpers for computing active tasks per member, and tasks per member
        localState.totalActiveTasks = 0;
        localState.totalTasks = 0;
        for (final String subtopology : topologyDescriber.subtopologies()) {
            final int numberOfPartitions = topologyDescriber.maxNumInputPartitions(subtopology);
            localState.totalTasks += numberOfPartitions;
            localState.totalActiveTasks += numberOfPartitions;
            if (topologyDescriber.isStateful(subtopology))
                localState.totalTasks += numberOfPartitions * localState.numStandbyReplicas;
        }
        localState.totalMembersWithActiveTaskCapacity = groupSpec.members().size();
        localState.totalMembersWithTaskCapacity = groupSpec.members().size();
        localState.activeTasksPerMember = computeTasksPerMember(localState.totalActiveTasks, localState.totalMembersWithActiveTaskCapacity);
        localState.tasksPerMember = computeTasksPerMember(localState.totalTasks, localState.totalMembersWithActiveTaskCapacity);

        localState.processIdToState = new HashMap<>();
        localState.activeTaskToPrevMember = new HashMap<>();
        localState.standbyTaskToPrevMember = new HashMap<>();
        for (final Map.Entry<String, AssignmentMemberSpec> memberEntry : groupSpec.members().entrySet()) {
            final String memberId = memberEntry.getKey();
            final String processId = memberEntry.getValue().processId();
            final Member member = new Member(processId, memberId);
            final AssignmentMemberSpec memberSpec = memberEntry.getValue();

            localState.processIdToState.putIfAbsent(processId, new ProcessState(processId));

Review comment (suggestion):
`allTasks` was previously just referring to active tasks. 
I renamed it to `totalActiveTasks`, and also added `totalTasks` that counts the number of tasks including standby tasks.

Code after:
            assignStandby(statefulTasks);
        }

        return buildGroupAssignment(groupSpec.members().keySet());
    }

    private LinkedList<TaskId> taskIds(final TopologyDescriber topologyDescriber, final boolean isActive) {
        final LinkedList<TaskId> ret = new LinkedList<>();
        for (final String subtopology : topologyDescriber.subtopologies()) {
            if (isActive || topologyDescriber.isStateful(subtopology)) {
                final int numberOfPartitions = topologyDescriber.maxNumInputPartitions(subtopology);
                for (int i = 0; i < numberOfPartitions; i++) {
                    ret.add(new TaskId(subtopology, i));
                }
            }
        }
        return ret;
    }

    private void initialize(final GroupSpec groupSpec, final TopologyDescriber topologyDescriber) {
        localState = new LocalState();
        localState.numStandbyReplicas =
            groupSpec.assignmentConfigs().isEmpty() ? 0
                : Integer.parseInt(groupSpec.assignmentConfigs().get("num.standby.replicas"));

        // Helpers for computing active tasks per member, and tasks per member
        localState.totalActiveTasks = 0;
        localState.totalTasks = 0;
        for (final String subtopology : topologyDescriber.subtopologies()) {
            final int numberOfPartitions = topologyDescriber.maxNumInputPartitions(subtopology);
            localState.totalTasks += numberOfPartitions;
            localState.totalActiveTasks += numberOfPartitions;
            if (topologyDescriber.isStateful(subtopology))
                localState.totalTasks += numberOfPartitions * localState.numStandbyReplicas;
        }
        localState.totalMembersWithActiveTaskCapacity = groupSpec.members().size();
        localState.totalMembersWithTaskCapacity = groupSpec.members().size();
        localState.activeTasksPerMember = computeTasksPerMember(localState.totalActiveTasks, localState.totalMembersWithActiveTaskCapacity);
        localState.tasksPerMember = computeTasksPerMember(localState.totalTasks, localState.totalMembersWithTaskCapacity);

        localState.processIdToState = new HashMap<>();
        localState.activeTaskToPrevMember = new HashMap<>();
        localState.standbyTaskToPrevMember = new HashMap<>();
        for (final Map.Entry<String, AssignmentMemberSpec> memberEntry : groupSpec.members().entrySet()) {
            final String memberId = memberEntry.getKey();
            final String processId = memberEntry.getValue().processId();
            final Member member = new Member(processId, memberId);
            final AssignmentMemberSpec memberSpec = memberEntry.getValue();

            localState.processIdToState.putIfAbsent(processId, new ProcessState(processId));
            localState.processIdToState.get(processId).addMember(memberId);