showuon commented on a change in pull request #11493:
URL: https://github.com/apache/kafka/pull/11493#discussion_r810488636



##########
File path: 
streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
##########
@@ -1029,104 +1043,151 @@ private boolean addClientAssignments(final 
Set<TaskId> statefulTasks,
 
     /**
      * Generate an assignment that tries to preserve thread-level stickiness 
of stateful tasks without violating
-     * balance. The stateful and total task load are both balanced across 
threads. Tasks without previous owners
-     * will be interleaved by group id to spread subtopologies across threads 
and further balance the workload.
+     * balance. The tasks are balanced across threads. Tasks without previous 
owners will be interleaved by
+     * group id to spread subtopologies across threads and further balance the 
workload.
+     * threadLoad is a map that keeps track of task load per thread across 
multiple calls so actives and standbys
+     * are evenly distributed
      */
-    static Map<String, List<TaskId>> assignTasksToThreads(final 
Collection<TaskId> statefulTasksToAssign,
-                                                          final 
Collection<TaskId> statelessTasksToAssign,
-                                                          final 
SortedSet<String> consumers,
-                                                          final ClientState 
state) {
+    static Map<String, List<TaskId>> assignStatefulTasksToThreads(final 
Collection<TaskId> tasksToAssign,
+                                                                  final 
SortedSet<String> consumers,
+                                                                  final 
ClientState state,
+                                                                  final 
Map<String, Integer> threadLoad) {
         final Map<String, List<TaskId>> assignment = new HashMap<>();
         for (final String consumer : consumers) {
             assignment.put(consumer, new ArrayList<>());
         }
 
-        final List<TaskId> unassignedStatelessTasks = new 
ArrayList<>(statelessTasksToAssign);
-        Collections.sort(unassignedStatelessTasks);
-
-        final Iterator<TaskId> unassignedStatelessTasksIter = 
unassignedStatelessTasks.iterator();
+        int totalTasks = tasksToAssign.size();
+        for (final Integer threadTaskCount : threadLoad.values()) {
+            totalTasks += threadTaskCount;
+        }
 
-        final int minStatefulTasksPerThread = (int) Math.floor(((double) 
statefulTasksToAssign.size()) / consumers.size());
-        final PriorityQueue<TaskId> unassignedStatefulTasks = new 
PriorityQueue<>(statefulTasksToAssign);
+        final int minTasksPerThread = (int) Math.floor(((double) totalTasks) / 
consumers.size());
+        final PriorityQueue<TaskId> unassignedTasks = new 
PriorityQueue<>(tasksToAssign);
 
         final Queue<String> consumersToFill = new LinkedList<>();
         // keep track of tasks that we have to skip during the first pass in 
case we can reassign them later
         // using tree-map to make sure the iteration ordering over keys are 
preserved
         final Map<TaskId, String> unassignedTaskToPreviousOwner = new 
TreeMap<>();
 
-        if (!unassignedStatefulTasks.isEmpty()) {
-            // First assign stateful tasks to previous owner, up to the min 
expected tasks/thread
+        if (!unassignedTasks.isEmpty()) {
+            // First assign tasks to previous owner, up to the min expected 
tasks/thread
             for (final String consumer : consumers) {
                 final List<TaskId> threadAssignment = assignment.get(consumer);
+                // The number of tasks we have to assign here to hit 
minTasksPerThread
+                final int tasksTargetCount = minTasksPerThread - 
threadLoad.getOrDefault(consumer, 0);
 
                 for (final TaskId task : state.prevTasksByLag(consumer)) {
-                    if (unassignedStatefulTasks.contains(task)) {
-                        if (threadAssignment.size() < 
minStatefulTasksPerThread) {
+                    if (unassignedTasks.contains(task)) {
+                        final int threadTaskCount = threadAssignment.size();
+                        if (threadTaskCount < tasksTargetCount) {

Review comment:
       nit: we can use `threadAssignment.size()` to replace the  
`threadTaskCount` variable.  Same as below.

##########
File path: 
streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java
##########
@@ -1029,104 +1043,151 @@ private boolean addClientAssignments(final 
Set<TaskId> statefulTasks,
 
     /**
      * Generate an assignment that tries to preserve thread-level stickiness 
of stateful tasks without violating
-     * balance. The stateful and total task load are both balanced across 
threads. Tasks without previous owners
-     * will be interleaved by group id to spread subtopologies across threads 
and further balance the workload.
+     * balance. The tasks are balanced across threads. Tasks without previous 
owners will be interleaved by
+     * group id to spread subtopologies across threads and further balance the 
workload.
+     * threadLoad is a map that keeps track of task load per thread across 
multiple calls so actives and standbys
+     * are evenly distributed
      */
-    static Map<String, List<TaskId>> assignTasksToThreads(final 
Collection<TaskId> statefulTasksToAssign,
-                                                          final 
Collection<TaskId> statelessTasksToAssign,
-                                                          final 
SortedSet<String> consumers,
-                                                          final ClientState 
state) {
+    static Map<String, List<TaskId>> assignStatefulTasksToThreads(final 
Collection<TaskId> tasksToAssign,
+                                                                  final 
SortedSet<String> consumers,
+                                                                  final 
ClientState state,
+                                                                  final 
Map<String, Integer> threadLoad) {
         final Map<String, List<TaskId>> assignment = new HashMap<>();
         for (final String consumer : consumers) {
             assignment.put(consumer, new ArrayList<>());
         }
 
-        final List<TaskId> unassignedStatelessTasks = new 
ArrayList<>(statelessTasksToAssign);
-        Collections.sort(unassignedStatelessTasks);
-
-        final Iterator<TaskId> unassignedStatelessTasksIter = 
unassignedStatelessTasks.iterator();
+        int totalTasks = tasksToAssign.size();
+        for (final Integer threadTaskCount : threadLoad.values()) {
+            totalTasks += threadTaskCount;
+        }
 
-        final int minStatefulTasksPerThread = (int) Math.floor(((double) 
statefulTasksToAssign.size()) / consumers.size());
-        final PriorityQueue<TaskId> unassignedStatefulTasks = new 
PriorityQueue<>(statefulTasksToAssign);
+        final int minTasksPerThread = (int) Math.floor(((double) totalTasks) / 
consumers.size());
+        final PriorityQueue<TaskId> unassignedTasks = new 
PriorityQueue<>(tasksToAssign);
 
         final Queue<String> consumersToFill = new LinkedList<>();
         // keep track of tasks that we have to skip during the first pass in 
case we can reassign them later
         // using tree-map to make sure the iteration ordering over keys are 
preserved
         final Map<TaskId, String> unassignedTaskToPreviousOwner = new 
TreeMap<>();
 
-        if (!unassignedStatefulTasks.isEmpty()) {
-            // First assign stateful tasks to previous owner, up to the min 
expected tasks/thread
+        if (!unassignedTasks.isEmpty()) {
+            // First assign tasks to previous owner, up to the min expected 
tasks/thread
             for (final String consumer : consumers) {
                 final List<TaskId> threadAssignment = assignment.get(consumer);
+                // The number of tasks we have to assign here to hit 
minTasksPerThread
+                final int tasksTargetCount = minTasksPerThread - 
threadLoad.getOrDefault(consumer, 0);
 
                 for (final TaskId task : state.prevTasksByLag(consumer)) {
-                    if (unassignedStatefulTasks.contains(task)) {
-                        if (threadAssignment.size() < 
minStatefulTasksPerThread) {
+                    if (unassignedTasks.contains(task)) {
+                        final int threadTaskCount = threadAssignment.size();
+                        if (threadTaskCount < tasksTargetCount) {
                             threadAssignment.add(task);
-                            unassignedStatefulTasks.remove(task);
+                            unassignedTasks.remove(task);
                         } else {
                             unassignedTaskToPreviousOwner.put(task, consumer);
                         }
                     }
                 }
 
-                if (threadAssignment.size() < minStatefulTasksPerThread) {
+                final int threadTaskCount = threadAssignment.size();
+                if (threadTaskCount < tasksTargetCount) {
                     consumersToFill.offer(consumer);
                 }
             }
 
             // Next interleave remaining unassigned tasks amongst unfilled 
consumers
             while (!consumersToFill.isEmpty()) {
-                final TaskId task = unassignedStatefulTasks.poll();
+                final TaskId task = unassignedTasks.poll();
                 if (task != null) {
                     final String consumer = consumersToFill.poll();
                     final List<TaskId> threadAssignment = 
assignment.get(consumer);
                     threadAssignment.add(task);
-                    if (threadAssignment.size() < minStatefulTasksPerThread) {
+                    final int threadTaskCount = threadAssignment.size() + 
threadLoad.getOrDefault(consumer, 0);
+                    if (threadTaskCount < minTasksPerThread) {
                         consumersToFill.offer(consumer);
                     }
                 } else {
                     throw new TaskAssignmentException("Ran out of unassigned 
stateful tasks but some members were not at capacity");
                 }
             }
 
-            // At this point all consumers are at the min capacity, so there 
may be up to N - 1 unassigned
-            // stateful tasks still remaining that should now be distributed 
over the consumers
-            if (!unassignedStatefulTasks.isEmpty()) {
-                consumersToFill.addAll(consumers);
+            // At this point all consumers are at the min or min + 1 capacity.
+            // The min + 1 case can occur for standbys where there's fewer 
standbys than consumers and after assigning
+            // the active tasks some consumers already have min + 1 one tasks 
assigned.
+            // The tasks still remaining should now be distributed over the 
consumers that are still at min capacity

Review comment:
       Thanks for your explanation. I've got it now, because we've change the 
total tasks count from original `active tasks` into total ones including active 
tasks. Thanks.




-- 
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.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to