static Map<String, List<TaskId>> assignTasksToThreads(final Collection<TaskId> tasksToAssign,
                                                          final boolean isStateful,
                                                          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 int totalTasks = threadLoad.values().stream().reduce(tasksToAssign.size(), Integer::sum);

        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 (!unassignedTasks.isEmpty()) {
            // First assign tasks to previous owner, up to the min expected tasks/thread if these are stateful
            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);

                if (isStateful) {
                    for (final TaskId task : state.prevTasksByLag(consumer)) {
                        if (unassignedTasks.contains(task)) {
                            if (threadAssignment.size() < tasksTargetCount) {
                                threadAssignment.add(task);
                                unassignedTasks.remove(task);
                            } else {
                                unassignedTaskToPreviousOwner.put(task, consumer);
                            }
                        }
                    }
                }

                if (threadAssignment.size() < tasksTargetCount) {
                    consumersToFill.offer(consumer);
                }
            }

            // Next interleave remaining unassigned tasks amongst unfilled consumers
            while (!consumersToFill.isEmpty()) {
                final TaskId task = unassignedTasks.poll();
                if (task != null) {
                    final String consumer = consumersToFill.poll();
                    final List<TaskId> threadAssignment = assignment.get(consumer);
                    threadAssignment.add(task);
                    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 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
            if (!unassignedTasks.isEmpty()) {
                for (final String consumer : consumers) {
                    final int taskCount = assignment.get(consumer).size() + threadLoad.getOrDefault(consumer, 0);
                    if (taskCount == minTasksPerThread) {
                        consumersToFill.add(consumer);
                    }
                }

                // Go over the tasks we skipped earlier and assign them to their previous owner when possible
                for (final Map.Entry<TaskId, String> taskEntry : unassignedTaskToPreviousOwner.entrySet()) {
                    final TaskId task = taskEntry.getKey();
                    final String consumer = taskEntry.getValue();
                    if (consumersToFill.contains(consumer) && unassignedTasks.contains(task)) {
                        assignment.get(consumer).add(task);
                        unassignedTasks.remove(task);
                        // Remove this consumer since we know it is now at minCapacity + 1
                        consumersToFill.remove(consumer);
                    }
                }

                // Now just distribute the remaining unassigned stateful tasks over the consumers still at min capacity
                for (final TaskId task : unassignedTasks) {
                    final String consumer = consumersToFill.poll();
                    final List<TaskId> threadAssignment = assignment.get(consumer);
                    threadAssignment.add(task);
                }
            }
        }
        // Update threadLoad
        for (final Map.Entry<String, List<TaskId>> taskEntry : assignment.entrySet()) {
            final String consumer = taskEntry.getKey();
            final int totalCount = threadLoad.getOrDefault(consumer, 0) + taskEntry.getValue().size();
            threadLoad.put(consumer, totalCount);
        }

        return assignment;
    }
Generate an assignment that tries to preserve thread-level stickiness for stateful tasks without violating
balance. The tasks are balanced across threads. Stateful tasks without previous owners will be interleaved by
group id to spread subtopologies across threads and further balance the workload.
Stateless tasks are simply spread across threads without taking into account previous ownership.
threadLoad is a map that keeps track of task load per thread across multiple calls so active and standby
tasks are evenly distributed
streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamsPartitionAssignor.java