Repository: apache/kafka
File: streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java (Java)

Code before:
                    lastFatal = new StreamsException("Encounter unexpected fatal error for task " + taskId, exception, taskId);
                }
            }

            if (lastFatal != null) {
                throw lastFatal;
            } else if (lastTaskMigrated != null) {
                throw lastTaskMigrated;
            } else {
                throw new TaskCorruptedException(aggregatedCorruptedTaskIds);
            }
        }
    }

    private void createNewTasks(final Map<TaskId, Set<TopicPartition>> activeTasksToCreate,
                                final Map<TaskId, Set<TopicPartition>> standbyTasksToCreate) {
        final Collection<Task> newActiveTasks = activeTaskCreator.createTasks(mainConsumer, activeTasksToCreate);
        final Collection<Task> newStandbyTask = standbyTaskCreator.createTasks(standbyTasksToCreate);

        if (stateUpdater == null) {
            tasks.addActiveTasks(newActiveTasks);
            tasks.addStandbyTasks(newStandbyTask);
        } else {
            final Map<TaskId, RuntimeException> taskInitExceptions = new LinkedHashMap<>();
            Stream.concat(newActiveTasks.stream(), newStandbyTask.stream())
                    .forEach(t -> addTaskToStateUpdater(t, taskInitExceptions));
            maybeThrowTaskExceptions(taskInitExceptions);
        }
    }

    private void handleTasksWithoutStateUpdater(final Map<TaskId, Set<TopicPartition>> activeTasksToCreate,
                                                final Map<TaskId, Set<TopicPartition>> standbyTasksToCreate,
                                                final Map<Task, Set<TopicPartition>> tasksToRecycle,
                                                final Set<Task> tasksToCloseClean) {
        for (final Task task : tasks.allTasks()) {
            final TaskId taskId = task.id();
            if (activeTasksToCreate.containsKey(taskId)) {
                if (task.isActive()) {
                    final Set<TopicPartition> topicPartitions = activeTasksToCreate.get(taskId);
                    if (tasks.updateActiveTaskInputPartitions(task, topicPartitions)) {
                        task.updateInputPartitions(topicPartitions, topologyMetadata.nodeToSourceTopics(task.id()));
                    }
                    task.resume();
                } else {
                    tasksToRecycle.put(task, activeTasksToCreate.get(taskId));
                }
                activeTasksToCreate.remove(taskId);
            } else if (standbyTasksToCreate.containsKey(taskId)) {
                if (!task.isActive()) {
                    updateInputPartitionsOfStandbyTaskIfTheyChanged(task, standbyTasksToCreate.get(taskId));
                    task.resume();

Review comment (suggestion):
When I started working on the state updater I had the same thought, but then I learned that task initialization might take some time due to the need of opening RocksDB instances. `handleAssignment()` is called in the  callback `ConsumerPartitionAssignor#onAssignment()` which is invoked when a group member receives its assignment from the leader. Having such time consuming operations on the rebalance path is not ideal. I do not remember exactly the reason, but I think with cooperative rebalancing the client might miss the immediate next rebalance which would lead again to more rebalances.
@ableegoldman Please keep me honest here.    

Code after:
                    lastFatal = new StreamsException("Encounter unexpected fatal error for task " + taskId, exception, taskId);
                }
            }

            if (lastFatal != null) {
                throw lastFatal;
            } else if (lastTaskMigrated != null) {
                throw lastTaskMigrated;
            } else {
                throw new TaskCorruptedException(aggregatedCorruptedTaskIds);
            }
        }
    }

    private void createNewTasks(final Map<TaskId, Set<TopicPartition>> activeTasksToCreate,
                                final Map<TaskId, Set<TopicPartition>> standbyTasksToCreate) {
        final Collection<Task> newActiveTasks = activeTaskCreator.createTasks(mainConsumer, activeTasksToCreate);
        final Collection<Task> newStandbyTask = standbyTaskCreator.createTasks(standbyTasksToCreate);

        if (stateUpdater == null) {
            tasks.addActiveTasks(newActiveTasks);
            tasks.addStandbyTasks(newStandbyTask);
        } else {
            Stream.concat(newActiveTasks.stream(), newStandbyTask.stream())
                    .forEach(stateUpdater::add);
        }
    }

    private void handleTasksWithoutStateUpdater(final Map<TaskId, Set<TopicPartition>> activeTasksToCreate,
                                                final Map<TaskId, Set<TopicPartition>> standbyTasksToCreate,
                                                final Map<Task, Set<TopicPartition>> tasksToRecycle,
                                                final Set<Task> tasksToCloseClean) {
        for (final Task task : tasks.allTasks()) {
            final TaskId taskId = task.id();
            if (activeTasksToCreate.containsKey(taskId)) {
                if (task.isActive()) {
                    final Set<TopicPartition> topicPartitions = activeTasksToCreate.get(taskId);
                    if (tasks.updateActiveTaskInputPartitions(task, topicPartitions)) {
                        task.updateInputPartitions(topicPartitions, topologyMetadata.nodeToSourceTopics(task.id()));
                    }
                    task.resume();
                } else {
                    tasksToRecycle.put(task, activeTasksToCreate.get(taskId));
                }
                activeTasksToCreate.remove(taskId);
            } else if (standbyTasksToCreate.containsKey(taskId)) {
                if (!task.isActive()) {
                    updateInputPartitionsOfStandbyTaskIfTheyChanged(task, standbyTasksToCreate.get(taskId));
                    task.resume();
                } else {
                    tasksToRecycle.put(task, standbyTasksToCreate.get(taskId));