simplification

Author: lucasbru

group-coordinator/src/main/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilder.java

@@ -165,8 +165,7 @@ public StreamsGroupMember build() {
                 if (member.memberEpoch() != targetAssignmentEpoch) {
                     return computeNextAssignment(
                         member.memberEpoch(),
-                        member.assignedTasks(),
-                        member.tasksPendingRevocation()
+                        member.assignedTasks()
                     );
                 } else {
                     return member;
@@ -193,8 +192,7 @@ public StreamsGroupMember build() {
                 // its state towards the latest target assignment.
                 return computeNextAssignment(
                     member.memberEpoch() + 1,
-                    member.assignedTasks(),
-                    member.tasksPendingRevocation()
+                    member.assignedTasks()
                 );
 
             case UNRELEASED_TASKS:
@@ -203,8 +201,7 @@ public StreamsGroupMember build() {
                 // of the unreleased tasks when they become available.
                 return computeNextAssignment(
                     member.memberEpoch(),
-                    member.assignedTasks(),
-                    member.tasksPendingRevocation()
+                    member.assignedTasks()
                 );
 
             case UNKNOWN:
@@ -220,8 +217,7 @@ public StreamsGroupMember build() {
 
                 return computeNextAssignment(
                     targetAssignmentEpoch,
-                    member.assignedTasks(),
-                    member.tasksPendingRevocation()
+                    member.assignedTasks()
                 );
         }
 
@@ -321,16 +317,12 @@ private static boolean computeAssignmentDifferenceForOneSubtopology(final String
      * another member, as defined by the `isUnreleasedTask` predicate).
      *
      * Epoch Handling:
-     * - For tasks in resultAssignedTasks, the epoch from currentAssignment is preserved.
-     * - For tasks in resultTasksPendingAssignment, the epoch is determined as follows:
-     *   - If the task is present in memberTasksPendingRevocation, its epoch from that map is used
-     *     (to preserve the revocation epoch for tasks being revoked and re-assigned).
-     *   - Otherwise, the targetAssignmentEpoch is used.
+     * - For tasks in resultAssignedTasks and resultTasksPendingRevocation, the epoch from currentAssignment is preserved.
+     * - For tasks in resultTasksPendingAssignment, the targetAssignmentEpoch is used.
      */
     private boolean computeAssignmentDifferenceWithEpoch(Map<String, Map<Integer, Integer>> currentAssignment,
                                                          Map<String, Set<Integer>> targetAssignment,
                                                          int targetAssignmentEpoch,
-                                                         Map<String, Map<Integer, Integer>> memberTasksPendingRevocation,
                                                          Map<String, Map<Integer, Integer>> resultAssignedTasks,
                                                          Map<String, Map<Integer, Integer>> resultTasksPendingRevocation,
                                                          Map<String, Map<Integer, Integer>> resultTasksPendingAssignment,
@@ -346,7 +338,6 @@ private boolean computeAssignmentDifferenceWithEpoch(Map<String, Map<Integer, In
                 currentAssignment.getOrDefault(subtopologyId, Map.of()),
                 targetAssignment.getOrDefault(subtopologyId, Set.of()),
                 targetAssignmentEpoch,
-                memberTasksPendingRevocation.getOrDefault(subtopologyId, Map.of()),
                 resultAssignedTasks,
                 resultTasksPendingRevocation,
                 resultTasksPendingAssignment,
@@ -360,7 +351,6 @@ private static boolean computeAssignmentDifferenceForOneSubtopologyWithEpoch(fin
                                                                                  final Map<Integer, Integer> currentTasksForThisSubtopology,
                                                                                  final Set<Integer> targetTasksForThisSubtopology,
                                                                                  final int targetAssignmentEpoch,
-                                                                                 final Map<Integer, Integer> currentTasksPendingRevocationForThisSubtopology,
                                                                                  final Map<String, Map<Integer, Integer>> resultAssignedTasks,
                                                                                  final Map<String, Map<Integer, Integer>> resultTasksPendingRevocation,
                                                                                  final Map<String, Map<Integer, Integer>> resultTasksPendingAssignment,
@@ -387,9 +377,7 @@ private static boolean computeAssignmentDifferenceForOneSubtopologyWithEpoch(fin
         Map<Integer, Integer> resultTasksPendingAssignmentForThisSubtopology = new HashMap<>();
         for (Integer taskId : targetTasksForThisSubtopology) {
             if (!resultAssignedTasksForThisSubtopology.containsKey(taskId)) {
-                // Use the epoch from pending revocation if the task is there, otherwise use targetAssignmentEpoch
-                Integer epoch = currentTasksPendingRevocationForThisSubtopology.getOrDefault(taskId, targetAssignmentEpoch);
-                resultTasksPendingAssignmentForThisSubtopology.put(taskId, epoch);
+                resultTasksPendingAssignmentForThisSubtopology.put(taskId, targetAssignmentEpoch);
             }
         }
         boolean hasUnreleasedTasks = resultTasksPendingAssignmentForThisSubtopology.keySet().removeIf(taskId ->
@@ -414,15 +402,13 @@ private static boolean computeAssignmentDifferenceForOneSubtopologyWithEpoch(fin
     /**
      * Computes the next assignment.
      *
-     * @param memberEpoch                  The epoch of the member to use. This may be different from
-     *                                     the epoch in {@link CurrentAssignmentBuilder#member}.
-     * @param memberAssignedTasks          The assigned tasks of the member to use.
-     * @param memberTasksPendingRevocation The tasks pending revocation of the member.
+     * @param memberEpoch         The epoch of the member to use. This may be different from
+     *                            the epoch in {@link CurrentAssignmentBuilder#member}.
+     * @param memberAssignedTasks The assigned tasks of the member to use.
      * @return A new StreamsGroupMember.
      */
     private StreamsGroupMember computeNextAssignment(int memberEpoch,
-                                                     TasksTupleWithEpochs memberAssignedTasks,
-                                                     TasksTupleWithEpochs memberTasksPendingRevocation) {
+                                                     TasksTupleWithEpochs memberAssignedTasks) {
         Map<String, Map<Integer, Integer>> newActiveAssignedTasks = new HashMap<>();
         Map<String, Map<Integer, Integer>> newActiveTasksPendingRevocation = new HashMap<>();
         Map<String, Map<Integer, Integer>> newActiveTasksPendingAssignment = new HashMap<>();
@@ -437,7 +423,6 @@ private StreamsGroupMember computeNextAssignment(int memberEpoch,
             memberAssignedTasks.activeTasksWithEpochs(),
             targetAssignment.activeTasks(),
             targetAssignmentEpoch,
-            memberTasksPendingRevocation.activeTasksWithEpochs(),
             newActiveAssignedTasks,
             newActiveTasksPendingRevocation,
             newActiveTasksPendingAssignment,

group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/CurrentAssignmentBuilderTest.java

@@ -834,12 +834,8 @@ public void testUnknownState(TaskRole taskRole) {
         );
     }
 
-    /**
-     * Tests epoch assignment from assigned tasks.
-     * This tests the first lookup path in addEpochsToTasks where epochs are found in assigned tasks.
-     */
     @Test
-    public void testEpochLookupFromAssignedTasks() {
+    public void testAssignmentEpochsShouldBePreservedFromPreviousAssignment() {
         final int memberEpoch = 10;
 
         // Create a member with tasks that have specific epochs in assigned tasks
@@ -880,73 +876,10 @@ public void testEpochLookupFromAssignedTasks() {
         );
     }
 
-    /**
-     * Tests epoch assignment from tasks pending revocation.
-     * This tests the second lookup path in addEpochsToTasks where epochs are found in pending revocation.
-     * When tasks pending revocation are brought back into the target assignment after the member revokes them,
-     * the epochs should be preserved from the pending revocation.
-     */
-    @Test
-    public void testEpochLookupFromPendingRevocation() {
-        final int memberEpoch = 10;
-
-        // Create a member in UNREVOKED_TASKS state with:
-        // - Some tasks assigned (5, 6) with epochs (9, 10)
-        // - Some tasks pending revocation (1, 2, 3, 4) with specific epochs (5, 6, 7, 8)
-        StreamsGroupMember member = new StreamsGroupMember.Builder(MEMBER_NAME)
-            .setState(MemberState.UNREVOKED_TASKS)
-            .setProcessId(PROCESS_ID)
-            .setMemberEpoch(memberEpoch)
-            .setPreviousMemberEpoch(memberEpoch)
-            .setAssignedTasks(mkTasksTupleWithEpochs(TaskRole.ACTIVE,
-                mkTasksWithEpochs(SUBTOPOLOGY_ID1, Map.of(5, 9)),
-                mkTasksWithEpochs(SUBTOPOLOGY_ID2, Map.of(6, 10))))
-            .setTasksPendingRevocation(mkTasksTupleWithEpochs(TaskRole.ACTIVE,
-                mkTasksWithEpochs(SUBTOPOLOGY_ID1, Map.of(1, 5, 2, 6)),
-                mkTasksWithEpochs(SUBTOPOLOGY_ID2, Map.of(3, 7, 4, 8))))
-            .build();
-
-        // Target assignment brings back the pending tasks (1, 2, 3, 4) along with current assigned (5, 6)
-        // The member has revoked the pending tasks (they're not in owned assignment)
-        // So it should transition to next epoch and assign the tasks
-        StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(memberEpoch + 1, mkTasksTuple(TaskRole.ACTIVE,
-                mkTasks(SUBTOPOLOGY_ID1, 1, 2, 5),
-                mkTasks(SUBTOPOLOGY_ID2, 3, 4, 6)))
-            .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
-            .withCurrentStandbyTaskProcessIds((subtopologyId, partitionId) -> Set.of())
-            .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Set.of())
-            .withOwnedAssignment(mkTasksTuple(TaskRole.ACTIVE,
-                mkTasks(SUBTOPOLOGY_ID1, 5),  // Only owns the assigned tasks, not the pending revocation
-                mkTasks(SUBTOPOLOGY_ID2, 6)))
-            .build();
-
-        // Verify that:
-        // - Tasks from pending revocation (1@5, 2@6, 3@7, 4@8) retain their epochs when re-assigned
-        // - Tasks from assigned (5@9, 6@10) retain their epochs
-        assertEquals(
-            new StreamsGroupMember.Builder(MEMBER_NAME)
-                .setState(MemberState.STABLE)
-                .setProcessId(PROCESS_ID)
-                .setMemberEpoch(memberEpoch + 1)
-                .setPreviousMemberEpoch(memberEpoch)
-                .setAssignedTasks(mkTasksTupleWithEpochs(TaskRole.ACTIVE,
-                    mkTasksWithEpochs(SUBTOPOLOGY_ID1, Map.of(1, 5, 2, 6, 5, 9)),
-                    mkTasksWithEpochs(SUBTOPOLOGY_ID2, Map.of(3, 7, 4, 8, 6, 10))))
-                .setTasksPendingRevocation(TasksTupleWithEpochs.EMPTY)
-                .build(),
-            updatedMember
-        );
-    }
-
-    /**
-     * Tests fallback to default epoch.
-     * This tests the third path in addEpochsToTasks where epochs are not found in assigned or pending,
-     * so the default epoch (target assignment epoch) is used.
-     */
     @Test
-    public void testEpochFallbackToDefault() {
+    public void testNewlyAssignedTasksGetTargetAssignmentEpoch() {
         final int memberEpoch = 10;
+        final int targetAssignmentEpoch = 11;
 
         // Create a member with empty assignments
         StreamsGroupMember member = new StreamsGroupMember.Builder(MEMBER_NAME)
@@ -960,22 +893,22 @@ public void testEpochFallbackToDefault() {
 
         // New tasks are assigned
         StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(memberEpoch + 1, mkTasksTuple(TaskRole.ACTIVE,
+            .withTargetAssignment(targetAssignmentEpoch, mkTasksTuple(TaskRole.ACTIVE,
                 mkTasks(SUBTOPOLOGY_ID1, 1, 2),
                 mkTasks(SUBTOPOLOGY_ID2, 3, 4)))
             .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
             .withCurrentStandbyTaskProcessIds((subtopologyId, partitionId) -> Set.of())
             .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Set.of())
             .build();
 
-        // Verify that all tasks use the default epoch (memberEpoch + 1)
+        // Verify that all tasks use the target assignment epoch
         assertEquals(
             new StreamsGroupMember.Builder(MEMBER_NAME)
                 .setState(MemberState.STABLE)
                 .setProcessId(PROCESS_ID)
-                .setMemberEpoch(memberEpoch + 1)
+                .setMemberEpoch(targetAssignmentEpoch)
                 .setPreviousMemberEpoch(memberEpoch)
-                .setAssignedTasks(mkTasksTupleWithCommonEpoch(TaskRole.ACTIVE, memberEpoch + 1,
+                .setAssignedTasks(mkTasksTupleWithCommonEpoch(TaskRole.ACTIVE, targetAssignmentEpoch,
                     mkTasks(SUBTOPOLOGY_ID1, 1, 2),
                     mkTasks(SUBTOPOLOGY_ID2, 3, 4)))
                 .setTasksPendingRevocation(TasksTupleWithEpochs.EMPTY)
@@ -986,14 +919,16 @@ public void testEpochFallbackToDefault() {
 
     /**
      * Tests mixed epoch assignment scenarios.
-     * This tests all three paths in addEpochsToTasks:
-     * - Some epochs from assigned tasks
-     * - Some epochs from pending revocation
-     * - Some epochs fallback to default
+     * - Some epochs from previously assigned tasks (Tasks 1, 2).
+     *   This happens regardless of whether the assigned task is reconciled (owned) by the client (Task 1) or not (Task 2)
+     * - Some newly assigned task (Task 5) which should get the target assignment epoch.
+     * - Some tasks are revoked by the member (Task 3, 4). One is immediately reassigned, which also gets
+     *   the target assignment epoch (Task 3).
      */
     @Test
-    public void testMixedEpochLookup() {
+    public void testMixedPreservedAndNewAssignmentEpochs() {
         final int memberEpoch = 10;
+        final int targetAssignmentEpoch = 11;
 
         // Create a member with:
         // - Tasks 1, 2 in assigned with epochs 5, 6
@@ -1009,80 +944,31 @@ public void testMixedEpochLookup() {
                 mkTasksWithEpochs(SUBTOPOLOGY_ID2, Map.of(3, 7, 4, 8))))
             .build();
 
-        // Target assignment includes:
-        // - Task 1 from assigned (should use epoch 5)
-        // - Task 3 from pending revocation (should use epoch 7)
-        // - Task 5 (new task, should use default epoch memberEpoch + 1 = 11)
-        // The member revokes tasks 2, 3, 4 (not in owned), transitions to next epoch
+        // The member revokes tasks 3, 4 (not in owned), transitions to next epoch
         StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(memberEpoch + 1, mkTasksTuple(TaskRole.ACTIVE,
-                mkTasks(SUBTOPOLOGY_ID1, 1),
+            .withTargetAssignment(targetAssignmentEpoch, mkTasksTuple(TaskRole.ACTIVE,
+                mkTasks(SUBTOPOLOGY_ID1, 1, 2),
                 mkTasks(SUBTOPOLOGY_ID2, 3, 5)))
             .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> null)
             .withCurrentStandbyTaskProcessIds((subtopologyId, partitionId) -> Set.of())
             .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Set.of())
             .withOwnedAssignment(mkTasksTuple(TaskRole.ACTIVE,
-                mkTasks(SUBTOPOLOGY_ID1, 1)))  // Only owns task 1 (task 2 is being revoked, tasks 3,4 already revoked)
+                mkTasks(SUBTOPOLOGY_ID1, 1)))  // Only owns task 1 (task 2 is not yet reconciled, tasks 3,4 already revoked)
             .build();
 
         // Verify mixed epoch assignment:
-        // - Task 1 from SUBTOPOLOGY_ID1 should have epoch 5 (from assigned)
-        // - Task 3 from SUBTOPOLOGY_ID2 should have epoch 7 (from pending revocation)
-        // - Task 5 from SUBTOPOLOGY_ID2 should have epoch 11 (default)
-        assertEquals(
-            new StreamsGroupMember.Builder(MEMBER_NAME)
-                .setState(MemberState.STABLE)
-                .setProcessId(PROCESS_ID)
-                .setMemberEpoch(memberEpoch + 1)
-                .setPreviousMemberEpoch(memberEpoch)
-                .setAssignedTasks(mkTasksTupleWithEpochs(TaskRole.ACTIVE,
-                    mkTasksWithEpochs(SUBTOPOLOGY_ID1, Map.of(1, 5)),
-                    mkTasksWithEpochs(SUBTOPOLOGY_ID2, Map.of(3, 7, 5, memberEpoch + 1))))
-                .setTasksPendingRevocation(TasksTupleWithEpochs.EMPTY)
-                .build(),
-            updatedMember
-        );
-    }
-
-    /**
-     * Tests that assigned tasks take precedence over pending revocation.
-     * If a task exists in both assigned and pending revocation, the epoch from assigned should be used.
-     */
-    @Test
-    public void testAssignedTasksHavePrecedenceOverPendingRevocation() {
-        final int memberEpoch = 10;
-
-        // Create a member with task 1 in both assigned (epoch 5) and pending revocation (epoch 7)
-        // This shouldn't normally happen, but tests the precedence logic
-        StreamsGroupMember member = new StreamsGroupMember.Builder(MEMBER_NAME)
-            .setState(MemberState.STABLE)
-            .setProcessId(PROCESS_ID)
-            .setMemberEpoch(memberEpoch)
-            .setPreviousMemberEpoch(memberEpoch)
-            .setAssignedTasks(mkTasksTupleWithEpochs(TaskRole.ACTIVE,
-                mkTasksWithEpochs(SUBTOPOLOGY_ID1, Map.of(1, 5, 2, 6))))
-            .setTasksPendingRevocation(mkTasksTupleWithEpochs(TaskRole.ACTIVE,
-                mkTasksWithEpochs(SUBTOPOLOGY_ID1, Map.of(1, 7))))
-            .build();
-
-        // Target assignment includes task 1
-        StreamsGroupMember updatedMember = new CurrentAssignmentBuilder(member)
-            .withTargetAssignment(memberEpoch + 1, mkTasksTuple(TaskRole.ACTIVE,
-                mkTasks(SUBTOPOLOGY_ID1, 1, 2)))
-            .withCurrentActiveTaskProcessId((subtopologyId, partitionId) -> PROCESS_ID)
-            .withCurrentStandbyTaskProcessIds((subtopologyId, partitionId) -> Set.of())
-            .withCurrentWarmupTaskProcessIds((subtopologyId, partitionId) -> Set.of())
-            .build();
-
-        // Verify that task 1 uses epoch 5 from assigned, not epoch 7 from pending revocation
+        // - Task 1 from SUBTOPOLOGY_ID1 should have epoch 5 (previous assignment epoch)
+        // - Task 3 from SUBTOPOLOGY_ID2 should have epoch 11 (target assignment epoch)
+        // - Task 5 from SUBTOPOLOGY_ID2 should have epoch 11 (target assignment epoch)
         assertEquals(
             new StreamsGroupMember.Builder(MEMBER_NAME)
                 .setState(MemberState.STABLE)
                 .setProcessId(PROCESS_ID)
-                .setMemberEpoch(memberEpoch + 1)
+                .setMemberEpoch(targetAssignmentEpoch)
                 .setPreviousMemberEpoch(memberEpoch)
                 .setAssignedTasks(mkTasksTupleWithEpochs(TaskRole.ACTIVE,
-                    mkTasksWithEpochs(SUBTOPOLOGY_ID1, Map.of(1, 5, 2, 6))))
+                    mkTasksWithEpochs(SUBTOPOLOGY_ID1, Map.of(1, 5, 2, 6)),
+                    mkTasksWithEpochs(SUBTOPOLOGY_ID2, Map.of(3, targetAssignmentEpoch, 5, targetAssignmentEpoch))))
                 .setTasksPendingRevocation(TasksTupleWithEpochs.EMPTY)
                 .build(),
             updatedMember