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

Code before:
        }

        private void removeUpdatingAndPausedTasks() {
            changelogReader.clear();
            updatingTasks.forEach((id, task) -> {
                task.maybeCheckpoint(true);
                removedTasks.add(task);
            });
            updatingTasks.clear();
            pausedTasks.forEach((id, task) -> {
                removedTasks.add(task);
            });
            pausedTasks.clear();
        }

        private void initializeTask(final Task task) {
            boolean initializationFailed = false;
            if (task.state() == Task.State.CREATED) {
                try {
                    task.initializeIfNeeded();
                } catch (final StreamsException streamsException) {
                    addToExceptionsAndFailedTasksThenRemoveFromUpdatingTasks(new ExceptionAndTasks(Collections.singleton(task), streamsException));
                    initializationFailed = true;
                }
            }
            if (!initializationFailed) {
                if (isStateless(task)) {
                    addToRestoredTasks((StreamTask) task);
                    updatingTasks.remove(task.id());
                    log.info("Stateless active task " + task.id() + " was added to the restored tasks of the state updater");
                } else {
                    changelogReader.register(task.changelogPartitions(), task.stateManager());
                    if (task.isActive()) {
                        log.info("Stateful active task " + task.id() + " was added to the state updater");
                        changelogReader.enforceRestoreActive();
                    } else {
                        log.info("Standby task " + task.id() + " was added to the state updater");
                        if (updatingTasks.size() == 1) {
                            changelogReader.transitToUpdateStandby();
                        }
                    }
                }
            }
        }

        private void removeTask(final TaskId taskId) {
            final Task task;
            if (updatingTasks.containsKey(taskId)) {
                task = updatingTasks.get(taskId);
                task.maybeCheckpoint(true);
                final Collection<TopicPartition> changelogPartitions = task.changelogPartitions();

Review comment (refactor):
Why do we need this flag? Adding the code in the `if`-branch to the `try`-block after `task.initializeIfNeeded();` should have the same effect without needing a flag.
For readability, I would extract the code that starts with `if (isStateless(task))` to a separate method named `handleInitializedTask()` or `postInitializeTask()` or similar. 

Code after:
                task.maybeCheckpoint(true);
                removedTasks.add(task);
            });
            updatingTasks.clear();
            pausedTasks.forEach((id, task) -> {
                removedTasks.add(task);
            });
            pausedTasks.clear();
        }

        private void addTask(final Task task) {
            final Task existingTask = updatingTasks.putIfAbsent(task.id(), task);
            if (existingTask != null) {
                throw new IllegalStateException((existingTask.isActive() ? "Active" : "Standby") + " task " + task.id() + " already exist, " +
                    "should not try to add another " + (task.isActive() ? "active" : "standby") + " task with the same id. " + BUG_ERROR_MESSAGE);
            }
        }

        private void initializeTask(final Task task) {
            if (task.state() == Task.State.CREATED) {
                try {
                    task.initializeIfNeeded();
                } catch (final StreamsException streamsException) {
                    addToExceptionsAndFailedTasksThenRemoveFromUpdatingTasks(new ExceptionAndTasks(Collections.singleton(task), streamsException));
                    return;
                }
            }
            postInitializeTask(task);
        }

        private void postInitializeTask(final Task task) {
            if (isStateless(task)) {
                addToRestoredTasks((StreamTask) task);
                updatingTasks.remove(task.id());
                log.info("Stateless active task " + task.id() + " was added to the restored tasks of the state updater");
            } else {
                changelogReader.register(task.changelogPartitions(), task.stateManager());
                if (task.isActive()) {
                    log.info("Stateful active task " + task.id() + " was added to the state updater");
                    changelogReader.enforceRestoreActive();
                } else {
                    log.info("Standby task " + task.id() + " was added to the state updater");
                    if (updatingTasks.size() == 1) {
                        changelogReader.transitToUpdateStandby();
                    }
                }
            }
        }

        private void removeTask(final TaskId taskId) {
            final Task task;