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

Code before:
    private void maybeRemoveTaskProcessId(
        StreamsGroupMember oldMember
    ) {
        if (oldMember != null) {
            removeTaskProcessIds(oldMember.assignedTasks(), oldMember.processId());
            removeTaskProcessIds(oldMember.tasksPendingRevocation(), oldMember.processId());
        }
    }

    void removeTaskProcessIds(
        TasksTuple tasks,
        String processId
    ) {
        if (tasks != null) {
            removeTaskProcessIds(tasks.activeTasks(), currentActiveTaskToProcessId, processId);
            removeTaskProcessIdsFromSet(tasks.standbyTasks(), currentStandbyTaskToProcessIds, processId);
            removeTaskProcessIdsFromSet(tasks.warmupTasks(), currentWarmupTaskToProcessIds, processId);
        }
    }

    /**
     * Removes the task process IDs based on the provided assignment.
     *
     * @param assignment    The assignment.
     * @param expectedProcessId The expected process ID.
     * @throws IllegalStateException if the process ID does not exist. package-private for testing.
     */
    private void removeTaskProcessIds(
        Map<String, Set<Integer>> assignment,
        TimelineHashMap<String, TimelineHashMap<Integer, String>> currentTasksProcessId,
        String expectedProcessId
    ) {
        assignment.forEach((subtopologyId, assignedPartitions) -> {
            currentTasksProcessId.compute(subtopologyId, (__, partitionsOrNull) -> {
                if (partitionsOrNull != null) {
                    assignedPartitions.forEach(partitionId -> {
                        String prevValue = partitionsOrNull.get(partitionId);
                        if (!Objects.equals(prevValue, expectedProcessId)) {
                            log.warn(
                                String.format("Cannot remove the process ID %s from task %s_%s because the partition is " +
                                    "still owned at a different process ID %s", expectedProcessId, subtopologyId, partitionId, prevValue));
                        } else {
                            partitionsOrNull.remove(partitionId);
                        }
                    });
                    if (partitionsOrNull.isEmpty()) {
                        return null;
                    } else {
                        return partitionsOrNull;
                    }
                } else {

Review comment (suggestion):
```suggestion
     * package-private for testing.
```

Code after:
            newMember.processId()
        );
        addTaskProcessId(
            newMember.tasksPendingRevocation(),
            newMember.processId()
        );
    }

    /**
     * Removes the task process IDs for the provided member.
     *
     * @param oldMember The old member.
     */
    private void maybeRemoveTaskProcessId(
        StreamsGroupMember oldMember
    ) {
        if (oldMember != null) {
            removeTaskProcessIds(oldMember.assignedTasks(), oldMember.processId());
            removeTaskProcessIds(oldMember.tasksPendingRevocation(), oldMember.processId());
        }
    }

    void removeTaskProcessIds(
        TasksTupleWithEpochs tasks,
        String processId
    ) {
        if (tasks != null) {
            removeTaskProcessIds(tasks.activeTasksWithEpochs(), currentActiveTaskToProcessId, processId);
            removeTaskProcessIdsFromSet(tasks.standbyTasks(), currentStandbyTaskToProcessIds, processId);
            removeTaskProcessIdsFromSet(tasks.warmupTasks(), currentWarmupTaskToProcessIds, processId);
        }
    }

    /**
     * Removes the task process IDs based on the provided assignment.
     *
     * @param assignment    The assignment.
     * @param expectedProcessId The expected process ID.
     * package-private for testing.
     */
    private void removeTaskProcessIds(
        Map<String, Map<Integer, Integer>> assignment,
        TimelineHashMap<String, TimelineHashMap<Integer, String>> currentTasksProcessId,
        String expectedProcessId
    ) {
        assignment.forEach((subtopologyId, assignedPartitions) -> {
            currentTasksProcessId.compute(subtopologyId, (__, partitionsOrNull) -> {
                if (partitionsOrNull != null) {
                    assignedPartitions.keySet().forEach(partitionId -> {
                        String prevValue = partitionsOrNull.get(partitionId);
                        if (Objects.equals(prevValue, expectedProcessId)) {