[GitHub] [kafka] vvcephei commented on a change in pull request #8588: KAFKA-6145: KIP-441: Improve assignment balance

[GitBox <gi...@apache.org>]

vvcephei commented on a change in pull request #8588:
URL: https://github.com/apache/kafka/pull/8588#discussion_r424718290



##########
File path: streams/src/test/java/org/apache/kafka/streams/processor/internals/assignment/AssignmentTestUtils.java
##########
@@ -75,4 +94,303 @@
     static UUID uuidForInt(final int n) {
         return new UUID(0, n);
     }
+
+    static void assertValidAssignment(final int numStandbyReplicas,
+                                      final Set<TaskId> statefulTasks,
+                                      final Set<TaskId> statelessTasks,
+                                      final Map<UUID, ClientState> assignedStates,
+                                      final StringBuilder failureContext) {
+        assertValidAssignment(
+            numStandbyReplicas,
+            0,
+            statefulTasks,
+            statelessTasks,
+            assignedStates,
+            failureContext
+        );
+    }
+
+    static void assertValidAssignment(final int numStandbyReplicas,
+                                      final int maxWarmupReplicas,
+                                      final Set<TaskId> statefulTasks,
+                                      final Set<TaskId> statelessTasks,
+                                      final Map<UUID, ClientState> assignedStates,
+                                      final StringBuilder failureContext) {
+        final Map<TaskId, Set<UUID>> assignments = new TreeMap<>();
+        for (final TaskId taskId : statefulTasks) {
+            assignments.put(taskId, new TreeSet<>());
+        }
+        for (final TaskId taskId : statelessTasks) {
+            assignments.put(taskId, new TreeSet<>());
+        }
+        for (final Map.Entry<UUID, ClientState> entry : assignedStates.entrySet()) {
+            validateAndAddActiveAssignments(statefulTasks, statelessTasks, failureContext, assignments, entry);
+            validateAndAddStandbyAssignments(statefulTasks, statelessTasks, failureContext, assignments, entry);
+        }
+
+        final AtomicInteger remainingWarmups = new AtomicInteger(maxWarmupReplicas);
+
+        final TreeMap<TaskId, Set<UUID>> misassigned =
+            assignments
+                .entrySet()
+                .stream()
+                .filter(entry -> {
+                    final int expectedActives = 1;
+                    final boolean isStateless = statelessTasks.contains(entry.getKey());
+                    final int expectedStandbys = isStateless ? 0 : numStandbyReplicas;
+                    // We'll never assign even the expected number of standbys if they don't actually fit in the cluster
+                    final int expectedAssignments = Math.min(
+                        assignedStates.size(),
+                        expectedActives + expectedStandbys
+                    );
+                    final int actualAssignments = entry.getValue().size();
+                    if (actualAssignments == expectedAssignments) {
+                        return false; // not misassigned
+                    } else {
+                        if (actualAssignments == expectedAssignments + 1 && remainingWarmups.get() > 0) {
+                            remainingWarmups.getAndDecrement();
+                            return false; // it's a warmup, so it's fine
+                        } else {
+                            return true; // misassigned
+                        }
+                    }
+                })
+                .collect(entriesToMap(TreeMap::new));
+
+        if (!misassigned.isEmpty()) {

Review comment:
       L131-158 is just gathering the information about whether each task is correctly assigned or not, based on its type and the standby configs (and maybe the warmup config). It doesn't make any assertions. So this check is actually the assertion, that no tasks are incorrectly assigned.
   
   Doing it this way is nicer, since when it fails, it tells you _all_ the incorrectly assigned tasks, not just the first one.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org