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

Code before:
                    assignedPartitions.forEach(partitionId -> {
                        if (!partitionsOrNull.get(partitionId).remove(processIdToRemove)) {
                            log.debug("[GroupId {}] Cannot remove the process ID {} from task {}_{} because the task is " +
                                    "not owned by this process ID", groupId, processIdToRemove, subtopologyId, partitionId);
                        }
                    });
                    if (partitionsOrNull.isEmpty()) {
                        return null;
                    } else {
                        return partitionsOrNull;
                    }
                } else {
                    log.debug("[GroupId {}] Cannot remove the process ID {} from {} because it does not have any process ID",
                            groupId, processIdToRemove, subtopologyId);
                    return partitionsOrNull;
                }
            });
        });
    }

    /**
     * Adds the partition epoch based on the provided assignment.
     *
     * @param tasks     The assigned tasks.
     * @param processId The process ID.
     * @throws IllegalStateException if the existing partition has larger epoch than the new one. package-private for testing.
     */
    void addTaskProcessId(
        TasksTupleWithEpochs tasks,
        String processId
    ) {
        if (tasks != null && processId != null) {
            addTaskProcessIdFromActiveTasksWithEpochs(tasks.activeTasksWithEpochs(), processId, currentActiveTaskToProcessId);
            addTaskProcessIdToSet(tasks.standbyTasks(), processId, currentStandbyTaskToProcessIds);
            addTaskProcessIdToSet(tasks.warmupTasks(), processId, currentWarmupTaskToProcessIds);
        }
    }

    private void addTaskProcessIdFromActiveTasksWithEpochs(
        Map<String, Map<Integer, Integer>> tasksWithEpochs,
        String processId,
        TimelineHashMap<String, TimelineHashMap<Integer, String>> currentTaskProcessId
    ) {
        tasksWithEpochs.forEach((subtopologyId, assignedTaskPartitionsWithEpochs) -> {
            currentTaskProcessId.compute(subtopologyId, (__, partitionsOrNull) -> {
                if (partitionsOrNull == null) {
                    partitionsOrNull = new TimelineHashMap<>(snapshotRegistry, assignedTaskPartitionsWithEpochs.size());
                }
                for (Integer partitionId: assignedTaskPartitionsWithEpochs.keySet()) {
                    String prevValue = partitionsOrNull.put(partitionId, processId);
                    if (prevValue != null) {

Review comment (suggestion):
I think we don't compare epochs for streams now?
```suggestion
     * package-private for testing.
```

Code after:
                    assignedPartitions.forEach(partitionId -> {
                        if (!partitionsOrNull.containsKey(partitionId) || !partitionsOrNull.get(partitionId).remove(processIdToRemove)) {
                            log.debug("[GroupId {}] Cannot remove the process ID {} from task {}_{} because the task is " +
                                    "not owned by this process ID", groupId, processIdToRemove, subtopologyId, partitionId);
                        }
                    });
                    if (partitionsOrNull.isEmpty()) {
                        return null;
                    } else {
                        return partitionsOrNull;
                    }
                } else {
                    log.debug("[GroupId {}] Cannot remove the process ID {} from {} because it does not have any process ID",
                            groupId, processIdToRemove, subtopologyId);
                    return partitionsOrNull;
                }
            });
        });
    }

    /**
     * Adds the partition epoch based on the provided assignment.
     *
     * @param tasks     The assigned tasks.
     * @param processId The process ID.
     * package-private for testing.
     */
    void addTaskProcessId(
        TasksTupleWithEpochs tasks,
        String processId
    ) {
        if (tasks != null && processId != null) {
            addTaskProcessIdFromActiveTasksWithEpochs(tasks.activeTasksWithEpochs(), processId, currentActiveTaskToProcessId);
            addTaskProcessIdToSet(tasks.standbyTasks(), processId, currentStandbyTaskToProcessIds);
            addTaskProcessIdToSet(tasks.warmupTasks(), processId, currentWarmupTaskToProcessIds);
        }
    }

    private void addTaskProcessIdFromActiveTasksWithEpochs(
        Map<String, Map<Integer, Integer>> tasksWithEpochs,
        String processId,
        TimelineHashMap<String, TimelineHashMap<Integer, String>> currentTaskProcessId
    ) {
        tasksWithEpochs.forEach((subtopologyId, assignedTaskPartitionsWithEpochs) -> {
            currentTaskProcessId.compute(subtopologyId, (__, partitionsOrNull) -> {
                if (partitionsOrNull == null) {
                    partitionsOrNull = new TimelineHashMap<>(snapshotRegistry, assignedTaskPartitionsWithEpochs.size());
                }
                for (Integer partitionId : assignedTaskPartitionsWithEpochs.keySet()) {
                    String prevValue = partitionsOrNull.put(partitionId, processId);
                    if (prevValue != null) {