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

Code before:
        if (stateUpdater != null) {
            // If there are failed tasks handling them first
            for (final StateUpdater.ExceptionAndTask exceptionAndTask : stateUpdater.drainExceptionsAndFailedTasks()) {
                final Task failedTask = exceptionAndTask.task();
                closeTaskDirty(failedTask, false);
            }

            final Map<TaskId, CompletableFuture<StateUpdater.RemovedTaskResult>> futures = new LinkedHashMap<>();
            for (final Task task : stateUpdater.tasks()) {
                final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(task.id());
                futures.put(task.id(), future);
            }
            final Set<Task> tasksToCloseClean = new HashSet<>();
            final Set<Task> tasksToCloseDirty = new HashSet<>();
            addToTasksToClose(futures, tasksToCloseClean, tasksToCloseDirty);
            stateUpdater.shutdown(Duration.ofMillis(Long.MAX_VALUE));

            for (final Task task : tasksToCloseClean) {
                tasks.addTask(task);
            }
            for (final Task task : tasksToCloseDirty) {
                closeTaskDirty(task, false);
            }
            // Handling all failures that occurred during the remove process
            for (final StateUpdater.ExceptionAndTask exceptionAndTask : stateUpdater.drainExceptionsAndFailedTasks()) {
                final Task failedTask = exceptionAndTask.task();
                closeTaskDirty(failedTask, false);
            }

            // If there is anything left unhandled due to timeouts, handling now
            for (final Task task : stateUpdater.tasks()) {
                closeTaskDirty(task, false);
            }
        }
    }

    private void shutdownSchedulingTaskManager() {
        if (schedulingTaskManager != null) {
            schedulingTaskManager.shutdown(Duration.ofMillis(Long.MAX_VALUE));
        }
    }

    /**
     * Closes and cleans up after the provided tasks, including closing their corresponding task producers
     */
    void closeAndCleanUpTasks(final Collection<Task> activeTasks, final Collection<Task> standbyTasks, final boolean clean) {
        final AtomicReference<RuntimeException> firstException = new AtomicReference<>(null);

        final Set<TaskId> ids =
            activeTasks.stream()
                .map(Task::id)

Review comment (bug):
Given that we do remove failed tasks above now, do we still need this step? Could we have added newly failed tasks in between? Or it this redundant now, and part of the bug is, that we did execute this cleanup logic in the incorrect order?

Code after:
        if (stateUpdater != null) {
            // If there are failed tasks handling them first
            for (final StateUpdater.ExceptionAndTask exceptionAndTask : stateUpdater.drainExceptionsAndFailedTasks()) {
                final Task failedTask = exceptionAndTask.task();
                closeTaskDirty(failedTask, false);
            }

            final Map<TaskId, CompletableFuture<StateUpdater.RemovedTaskResult>> futures = new LinkedHashMap<>();
            for (final Task task : stateUpdater.tasks()) {
                final CompletableFuture<StateUpdater.RemovedTaskResult> future = stateUpdater.remove(task.id());
                futures.put(task.id(), future);
            }
            final Set<Task> tasksToCloseClean = new HashSet<>();
            final Set<Task> tasksToCloseDirty = new HashSet<>();
            addToTasksToClose(futures, tasksToCloseClean, tasksToCloseDirty);
            // at this point we removed all tasks, so the shutdown should not take a lot of time
            stateUpdater.shutdown(Duration.ofMinutes(1L));

            for (final Task task : tasksToCloseClean) {
                tasks.addTask(task);
            }
            for (final Task task : tasksToCloseDirty) {
                closeTaskDirty(task, false);
            }
            // Handling all failures that occurred during the remove process
            for (final StateUpdater.ExceptionAndTask exceptionAndTask : stateUpdater.drainExceptionsAndFailedTasks()) {
                final Task failedTask = exceptionAndTask.task();
                closeTaskDirty(failedTask, false);
            }

            // If there is anything left unhandled due to timeouts, handling now
            for (final Task task : stateUpdater.tasks()) {
                closeTaskDirty(task, false);
            }
        }
    }

    private void shutdownSchedulingTaskManager() {
        if (schedulingTaskManager != null) {
            schedulingTaskManager.shutdown(Duration.ofMinutes(5L));
        }
    }

    /**
     * Closes and cleans up after the provided tasks, including closing their corresponding task producers
     */
    void closeAndCleanUpTasks(final Collection<Task> activeTasks, final Collection<Task> standbyTasks, final boolean clean) {
        final AtomicReference<RuntimeException> firstException = new AtomicReference<>(null);

        final Set<TaskId> ids =
            activeTasks.stream()