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

Code before:
     * @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) {
                        log.debug("[GroupId {}] Cannot set the process ID of {}-{} to {} because the partition is " +
                            "still owned by process ID {}", groupId, subtopologyId, partitionId, processId, prevValue);
                    }
                }
                return partitionsOrNull;
            });
        });
    }

    private void addTaskProcessIdToSet(
        Map<String, Set<Integer>> tasks,
        String processId,
        TimelineHashMap<String, TimelineHashMap<Integer, Set<String>>> currentTaskProcessId
    ) {
        tasks.forEach((subtopologyId, assignedTaskPartitions) -> {
            currentTaskProcessId.compute(subtopologyId, (__, partitionsOrNull) -> {
                if (partitionsOrNull == null) {
                    partitionsOrNull = new TimelineHashMap<>(snapshotRegistry, assignedTaskPartitions.size());
                }
                for (Integer partitionId : assignedTaskPartitions) {
                    partitionsOrNull.computeIfAbsent(partitionId, ___ -> new HashSet<>()).add(processId);
                }
                return partitionsOrNull;

Review comment (suggestion):
nit: missing space
```suggestion
                for (Integer partitionId : assignedTaskPartitionsWithEpochs.keySet()) {
```

Code after:
     * @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) {
                        log.debug("[GroupId {}] Setting the process ID of {}-{} to {} even though the partition is " +
                            "still owned by process ID {}", groupId, subtopologyId, partitionId, processId, prevValue);
                    }
                }
                return partitionsOrNull;
            });
        });
    }

    private void addTaskProcessIdToSet(
        Map<String, Set<Integer>> tasks,
        String processId,
        TimelineHashMap<String, TimelineHashMap<Integer, Set<String>>> currentTaskProcessId
    ) {
        tasks.forEach((subtopologyId, assignedTaskPartitions) -> {
            currentTaskProcessId.compute(subtopologyId, (__, partitionsOrNull) -> {
                if (partitionsOrNull == null) {
                    partitionsOrNull = new TimelineHashMap<>(snapshotRegistry, assignedTaskPartitions.size());
                }
                for (Integer partitionId : assignedTaskPartitions) {
                    partitionsOrNull.computeIfAbsent(partitionId, ___ -> new HashSet<>()).add(processId);
                }
                return partitionsOrNull;