Repository: apache/kafka
File: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/consumer/ConsumerGroup.java (Java)

Code before:
        if (oldMember != null) {
            removePartitionEpochs(oldMember.assignedPartitions(), oldMember.memberEpoch());
            removePartitionEpochs(oldMember.partitionsPendingRevocation(), oldMember.memberEpoch());
        }
    }

    /**
     * Removes the partition epochs based on the provided assignment.
     *
     * @param assignment    The assignment.
     * @param expectedEpoch The expected epoch.
     * package-private for testing.
     */
    void removePartitionEpochs(
        Map<Uuid, Set<Integer>> assignment,
        int expectedEpoch
    ) {
        assignment.forEach((topicId, assignedPartitions) -> {
            currentPartitionEpoch.compute(topicId, (__, partitionsOrNull) -> {
                if (partitionsOrNull != null) {
                    assignedPartitions.forEach(partitionId -> {
                        Integer prevValue = partitionsOrNull.get(partitionId);
                        if (prevValue != null && prevValue == expectedEpoch) {
                            partitionsOrNull.remove(partitionId);
                        } else {
                            // GroupId added for context. 
                            log.debug(
                                String.format("[Group %s]: Cannot remove the epoch %d from %s-%s because the partition is " +
                                    "still owned at a different epoch %d", groupId, expectedEpoch, topicId, partitionId, prevValue));
                        }
                    });
                    if (partitionsOrNull.isEmpty()) {
                        return null;
                    } else {
                        return partitionsOrNull;
                    }
                } else {
                    log.debug(
                        String.format("[Group %s]: Cannot remove the epoch %d from %s because it does not have any epoch",
                            groupId, expectedEpoch, topicId));
                    return partitionsOrNull;
                }
            });
        });
    }

    /**
     * Adds the partitions epoch based on the provided assignment.
     *
     * @param assignment    The assignment.
     * @param epoch         The new epoch.

Review comment (suggestion):
```suggestion
```
Can we drop this comment? The same for streams.

Code after:
        if (oldMember != null) {
            removePartitionEpochs(oldMember.assignedPartitions(), oldMember.memberEpoch());
            removePartitionEpochs(oldMember.partitionsPendingRevocation(), oldMember.memberEpoch());
        }
    }

    /**
     * Removes the partition epochs based on the provided assignment.
     *
     * @param assignment    The assignment.
     * @param expectedEpoch The expected epoch.
     * package-private for testing.
     */
    void removePartitionEpochs(
        Map<Uuid, Set<Integer>> assignment,
        int expectedEpoch
    ) {
        assignment.forEach((topicId, assignedPartitions) -> {
            currentPartitionEpoch.compute(topicId, (__, partitionsOrNull) -> {
                if (partitionsOrNull != null) {
                    assignedPartitions.forEach(partitionId -> {
                        Integer prevValue = partitionsOrNull.get(partitionId);
                        if (prevValue != null && prevValue == expectedEpoch) {
                            partitionsOrNull.remove(partitionId);
                        } else {
                            log.debug("[GroupId {}] Cannot remove the epoch {} from {}-{} because the partition is " +
                                    "still owned at a different epoch {}", groupId, expectedEpoch, topicId, partitionId, prevValue);
                        }
                    });
                    if (partitionsOrNull.isEmpty()) {
                        return null;
                    } else {
                        return partitionsOrNull;
                    }
                } else {
                    log.debug("[GroupId {}] Cannot remove the epoch {} from {} because it does not have any epoch",
                            groupId, expectedEpoch, topicId);
                    return partitionsOrNull;
                }
            });
        });
    }

    /**
     * Adds the partitions epoch based on the provided assignment.
     *
     * @param assignment    The assignment.
     * @param epoch         The new epoch.
     * @throws IllegalStateException if updating a partition with a smaller or equal epoch.
     * package-private for testing.
     */