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

Code before:
            currentTaskProcessId.compute(subtopologyId, (__, partitionsOrNull) -> {
                if (partitionsOrNull == null) {
                    partitionsOrNull = new TimelineHashMap<>(snapshotRegistry, assignedTaskPartitionsWithEpochs.size());
                }
                for (Map.Entry<Integer, Integer> partitionEntry : assignedTaskPartitionsWithEpochs.entrySet()) {
                    Integer partitionId = partitionEntry.getKey();
                    String prevValue = partitionsOrNull.get(partitionId);
                    
                    if (prevValue == null) {
                        partitionsOrNull.put(partitionId, processId);
                    } else {
                        String memberId = null;
                        for (Map.Entry<String, StreamsGroupMember> memberEntry : members.entrySet()) {
                            if (memberEntry.getValue().processId().equals(prevValue)) {
                                memberId = memberEntry.getKey();
                                break;
                            }
                        }
                        if (memberId != null && 
                            members.get(memberId).assignedTasks().activeTasksWithEpochs().get(subtopologyId).get(partitionId) <= partitionEntry.getValue()) {
                            partitionsOrNull.put(partitionId, processId);
                        } else {
                            throw new IllegalStateException(
                                String.format("[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):
Applied the same version of the fix as in ConsumerGroup#addPartitionEpochs. The way epochs are implemented for the Streams groups is not very conducive to this use case which is why it's a bit convoluted.

Code after:
            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;
            });
        });
    }

    public StreamsGroupDescribeResponseData.DescribedGroup asDescribedGroup(
        long committedOffset
    ) {
        StreamsGroupDescribeResponseData.DescribedGroup describedGroup = new StreamsGroupDescribeResponseData.DescribedGroup()
            .setGroupId(groupId)
            .setGroupEpoch(groupEpoch.get(committedOffset))
            .setGroupState(state.get(committedOffset).toString())
            .setAssignmentEpoch(targetAssignmentEpoch.get(committedOffset))
            .setTopology(
                configuredTopology.get(committedOffset)
                    .filter(ConfiguredTopology::isReady)
                    .map(ConfiguredTopology::asStreamsGroupDescribeTopology)
                    .orElse(
                        topology.get(committedOffset)
                            .map(StreamsTopology::asStreamsGroupDescribeTopology)
                            .orElseThrow(() -> new IllegalStateException("There should always be a topology for a streams group."))
                    )