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

Code before:
                .collect(Collectors.toSet());
        maybeLockTasks(ids);

        final Set<Task> tasksToCloseDirty = new TreeSet<>(Comparator.comparing(Task::id));
        tasksToCloseDirty.addAll(tryCloseCleanActiveTasks(activeTasks, clean, firstException));
        tasksToCloseDirty.addAll(tryCloseCleanStandbyTasks(standbyTasks, clean, firstException));

        for (final Task task : tasksToCloseDirty) {
            closeTaskDirty(task, true);
        }

        maybeUnlockTasks(ids);

        final RuntimeException exception = firstException.get();
        if (exception != null) {
            throw exception;
        }
    }

    // Returns the set of active tasks that must be closed dirty
    private Collection<Task> tryCloseCleanActiveTasks(final Collection<Task> activeTasksToClose,
                                                      final boolean clean,
                                                      final AtomicReference<RuntimeException> firstException) {
        if (!clean) {
            return activeTasksToClose;
        }
        final Comparator<Task> byId = Comparator.comparing(Task::id);
        final Set<Task> tasksToCommit = new TreeSet<>(byId);
        final Set<Task> tasksToCloseDirty = new TreeSet<>(byId);
        final Set<Task> tasksToCloseClean = new TreeSet<>(byId);
        final Map<Task, Map<TopicPartition, OffsetAndMetadata>> consumedOffsetsAndMetadataPerTask = new HashMap<>();

        // first committing all tasks and then suspend and close them clean
        for (final Task task : activeTasksToClose) {
            try {
                final Map<TopicPartition, OffsetAndMetadata> committableOffsets = task.prepareCommit(true);
                tasksToCommit.add(task);
                if (!committableOffsets.isEmpty()) {
                    consumedOffsetsAndMetadataPerTask.put(task, committableOffsets);
                }
                tasksToCloseClean.add(task);
            } catch (final TaskMigratedException e) {
                // just ignore the exception as it doesn't matter during shutdown
                tasksToCloseDirty.add(task);
            } catch (final StreamsException e) {
                e.setTaskId(task.id());
                firstException.compareAndSet(null, e);
                tasksToCloseDirty.add(task);
            } catch (final RuntimeException e) {
                firstException.compareAndSet(null, new StreamsException(e, task.id()));
                tasksToCloseDirty.add(task);

Review comment (bug):
For my own understanding. The original code here was not wrong in the sense that before this PR the `activeTasksToClose` that got passed in was the same as `activeTaskIterable()`? But with this PR, we also include pending task in `activeTasksToClose`, and that's why we need to update the code here?

Code after:
                .map(Task::id)
                .collect(Collectors.toSet());
        maybeLockTasks(ids);

        final Set<Task> tasksToCloseDirty = new TreeSet<>(Comparator.comparing(Task::id));
        tasksToCloseDirty.addAll(tryCloseCleanActiveTasks(activeTasks, clean, firstException));
        tasksToCloseDirty.addAll(tryCloseCleanStandbyTasks(standbyTasks, clean, firstException));

        for (final Task task : tasksToCloseDirty) {
            closeTaskDirty(task, true);
        }

        maybeUnlockTasks(ids);

        final RuntimeException exception = firstException.get();
        if (exception != null) {
            throw exception;
        }
    }

    // Returns the set of active tasks that must be closed dirty
    private Collection<Task> tryCloseCleanActiveTasks(final Collection<Task> activeTasksToClose,
                                                      final boolean clean,
                                                      final AtomicReference<RuntimeException> firstException) {
        if (!clean) {
            return activeTasksToClose;
        }
        final Comparator<Task> byId = Comparator.comparing(Task::id);
        final Set<Task> tasksToCommit = new TreeSet<>(byId);
        final Set<Task> tasksToCloseDirty = new TreeSet<>(byId);
        final Set<Task> tasksToCloseClean = new TreeSet<>(byId);
        final Map<Task, Map<TopicPartition, OffsetAndMetadata>> consumedOffsetsAndMetadataPerTask = new HashMap<>();

        // first committing all tasks and then suspend and close them clean
        for (final Task task : activeTasksToClose) {
            try {
                final Map<TopicPartition, OffsetAndMetadata> committableOffsets = task.prepareCommit(true);
                tasksToCommit.add(task);
                if (!committableOffsets.isEmpty()) {
                    consumedOffsetsAndMetadataPerTask.put(task, committableOffsets);
                }
                tasksToCloseClean.add(task);
            } catch (final TaskMigratedException e) {
                // just ignore the exception as it doesn't matter during shutdown
                tasksToCloseDirty.add(task);
            } catch (final StreamsException e) {
                e.setTaskId(task.id());
                firstException.compareAndSet(null, e);
                tasksToCloseDirty.add(task);
            } catch (final RuntimeException e) {
                firstException.compareAndSet(null, new StreamsException(e, task.id()));