KAFKA-20066: Assignment epochs for consumer groups

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupCoordinatorRecordHelpers.java

@@ -289,8 +289,8 @@ public static CoordinatorRecord newConsumerGroupCurrentAssignmentRecord(
                     .setMemberEpoch(member.memberEpoch())
                     .setPreviousMemberEpoch(member.previousMemberEpoch())
                     .setState(member.state().value())
-                    .setAssignedPartitions(toTopicPartitions(member.assignedPartitions()))
-                    .setPartitionsPendingRevocation(toTopicPartitions(member.partitionsPendingRevocation())),
+                    .setAssignedPartitions(toTopicPartitions(member.assignedPartitionsWithEpochs()))
+                    .setPartitionsPendingRevocation(toTopicPartitions(member.partitionsPendingRevocationWithEpochs())),
                 (short) 0
             )
         );
@@ -804,13 +804,19 @@ public static CoordinatorRecord newShareGroupStatePartitionMetadataRecord(
     }
 
     private static List<ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions> toTopicPartitions(
-        Map<Uuid, Set<Integer>> topicPartitions
-    ) {
-        List<ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions> topics = new ArrayList<>(topicPartitions.size());
-        topicPartitions.forEach((topicId, partitions) ->
+        Map<Uuid, Map<Integer, Integer>> topicPartitionsWithEpochs
+    ) {
+        List<ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions> topics = new ArrayList<>(topicPartitionsWithEpochs.size());
+        topicPartitionsWithEpochs.forEach((topicId, partitionEpochMap) -> {
+            List<Integer> partitionList = new ArrayList<>(partitionEpochMap.keySet());
+            List<Integer> epochList = partitionList.stream()
+                .map(partitionId -> partitionEpochMap.getOrDefault(partitionId, 0))
+                .toList();
             topics.add(new ConsumerGroupCurrentMemberAssignmentValue.TopicPartitions()
                 .setTopicId(topicId)
-                .setPartitions(new ArrayList<>(partitions)))
+                .setPartitions(partitionList)
+                .setAssignmentEpochs(epochList));
+            }
         );
         return topics;
     }

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/GroupMetadataManager.java

@@ -4138,9 +4138,12 @@ private CoordinatorResult<ConsumerGroupHeartbeatResponseData, CoordinatorRecord>
         ConsumerGroupMember member
     ) {
         // We will write a member epoch of -2 for this departing static member.
+        // Assignment epochs are reset to 0 so when the static member rejoins, partitions
+        // are considered assigned from epoch 0 to the new member ID.
         ConsumerGroupMember leavingStaticMember = new ConsumerGroupMember.Builder(member)
             .setMemberEpoch(LEAVE_GROUP_STATIC_MEMBER_EPOCH)
-            .setPartitionsPendingRevocation(Map.of())
+            .setPartitionsPendingRevocationWithEpochs(Map.of())
+            .resetAssignedPartitionsEpochsToZero()
             .build();
 
         return new CoordinatorResult<>(
@@ -5407,8 +5410,8 @@ public void replay(
             ConsumerGroupMember newMember = new ConsumerGroupMember.Builder(oldMember)
                 .setMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
                 .setPreviousMemberEpoch(LEAVE_GROUP_MEMBER_EPOCH)
-                .setAssignedPartitions(Map.of())
-                .setPartitionsPendingRevocation(Map.of())
+                .setAssignedPartitionsWithEpochs(Map.of())
+                .setPartitionsPendingRevocationWithEpochs(Map.of())
                 .build();
             group.updateMember(newMember);
         }

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/modern/consumer/ConsumerGroup.java

@@ -63,6 +63,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import static org.apache.kafka.coordinator.group.Utils.toOptional;
 import static org.apache.kafka.coordinator.group.Utils.toTopicPartitionMap;
@@ -672,9 +673,23 @@ public CommitPartitionValidator validateOffsetCommit(
             throw new UnsupportedVersionException("OffsetCommit version 9 or above must be used " +
                 "by members using the modern group protocol");
         }
+        // For members using the classic protocol, use strict epoch validation.
+        if (member.useClassicProtocol()) {
+            validateMemberEpoch(memberEpoch, member.memberEpoch(), true);
+            return CommitPartitionValidator.NO_OP;
+        }
 
-        validateMemberEpoch(memberEpoch, member.memberEpoch(), member.useClassicProtocol());
-        return CommitPartitionValidator.NO_OP;
+        // For member using the consumer protocol
+        // Case 1: Strict epoch match
+        if (memberEpoch == member.memberEpoch()) {
+            return CommitPartitionValidator.NO_OP;
+        }
+        // Case 2: Client epoch > broker epoch, which is an invalid request
+        if (memberEpoch > member.memberEpoch()) {
+            throw new StaleMemberEpochException(String.format("The received member epoch %d is larger than "
+                + "the expected member epoch %d.", memberEpoch, member.memberEpoch()));
+        }
+        return createAssignmentEpochValidator(member, memberEpoch);
     }
 
     /**
@@ -837,6 +852,47 @@ private void validateMemberEpoch(
         }
     }
 
+    /**
+     * Creates a validator that checks per-partition assignment epochs.
+     * A commit is rejected if the partition is not assigned to the member
+     * or if the received client-side epoch is older than the partition's assignment epoch(KIP-1251).
+     *
+     * @param member              The consumer group member.
+     * @param receivedMemberEpoch The member epoch from the offset commit request.
+     * @return A validator that checks each partition's assignment epoch.
+     */
+    private CommitPartitionValidator createAssignmentEpochValidator(
+        ConsumerGroupMember member,
+        int receivedMemberEpoch
+    ) {
+        return (topicName, topicId, partitionId) -> {
+            // Check if the partition is in the assigned partitions.
+            // If not found in assigned, check partitions pending revocation.
+            Integer assignmentEpoch = member.getAssignmentEpoch(topicId, partitionId);
+            if (assignmentEpoch == null) {
+                assignmentEpoch = member.getPendingRevocationEpoch(topicId, partitionId);
+            }
+
+            // If client-side epoch != broker-side epoch, and the partition is not assigned to this member, reject.
+            if (assignmentEpoch == null) {
+                throw new StaleMemberEpochException(String.format(
+                    "Partition %s-%d is not assigned or pending revocation for member %s. " +
+                        "Committing unassigned partitions is only allowed when member epoch matches exactly " +
+                        "(received: %d, current: %d).",
+                    topicName, partitionId, member.memberId(), receivedMemberEpoch, member.memberEpoch()));
+            }
+
+            // If the received epoch is older than when this partition was assigned,
+            // It is a zombie commit and should be rejected.
+            if (receivedMemberEpoch < assignmentEpoch) {
+                throw new StaleMemberEpochException(
+                    String.format("The received member epoch %d is older than the assignment epoch %d for partition %s-%d.",
+                        receivedMemberEpoch, assignmentEpoch, topicName, partitionId)
+                );
+            }
+        };
+    }
+
     /**
      * Computes the subscription type based on the provided information.
      *
@@ -1184,17 +1240,24 @@ public static ConsumerGroup fromClassicGroup(
             // assignment of the classic group. All the members are put in the Stable state. If the classic
             // group was in Preparing Rebalance or Completing Rebalance states, the classic members are
             // asked to rejoin the group to re-trigger a rebalance or collect their assignments.
+            int memberEpoch = classicGroup.generationId();
+            // Convert assigned partitions to epochs map
+            Map<Uuid, Map<Integer, Integer>> assignedPartitionsWithEpochs = assignedPartitions.entrySet().stream()
+                .collect(Collectors.toMap(
+                    Map.Entry::getKey,
+                    e -> e.getValue().stream().collect(Collectors.toMap(p -> p, p -> memberEpoch))
+                ));
             ConsumerGroupMember newMember = new ConsumerGroupMember.Builder(classicGroupMember.memberId())
-                .setMemberEpoch(classicGroup.generationId())
+                .setMemberEpoch(memberEpoch)
                 .setState(MemberState.STABLE)
-                .setPreviousMemberEpoch(classicGroup.generationId())
+                .setPreviousMemberEpoch(memberEpoch)
                 .setInstanceId(classicGroupMember.groupInstanceId().orElse(null))
                 .setRackId(toOptional(subscription.rackId()).orElse(null))
                 .setRebalanceTimeoutMs(classicGroupMember.rebalanceTimeoutMs())
                 .setClientId(classicGroupMember.clientId())
                 .setClientHost(classicGroupMember.clientHost())
                 .setSubscribedTopicNames(subscription.topics())
-                .setAssignedPartitions(assignedPartitions)
+                .setAssignedPartitionsWithEpochs(assignedPartitionsWithEpochs)
                 .setClassicMemberMetadata(
                     new ConsumerGroupMemberMetadataValue.ClassicMemberMetadata()
                         .setSessionTimeoutMs(classicGroupMember.sessionTimeoutMs())