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

Code before:
        return rebalanceInProgress;
    }

    void handleRebalanceStart(final Set<String> subscribedTopics) {
        topologyMetadata.addSubscribedTopicsFromMetadata(subscribedTopics, logPrefix);

        tryToLockAllNonEmptyTaskDirectories();

        rebalanceInProgress = true;
    }

    void handleRebalanceComplete() {
        // we should pause consumer only within the listener since
        // before then the assignment has not been updated yet.
        mainConsumer.pause(mainConsumer.assignment());

        releaseLockedUnassignedTaskDirectories();

        rebalanceInProgress = false;
    }

    /**
     * @throws TaskMigratedException
     */
    boolean handleCorruption(final Set<TaskId> corruptedTasks) {
        final Set<Task> corruptedActiveTasks = new HashSet<>();
        final Set<Task> corruptedStandbyTasks = new HashSet<>();

        for (final TaskId taskId : corruptedTasks) {
            final Task task = tasks.task(taskId);
            if (task.isActive()) {
                corruptedActiveTasks.add(task);
            } else {
                corruptedStandbyTasks.add(task);
            }
        }

        // Make sure to clean up any corrupted standby tasks in their entirety before committing
        // since TaskMigrated can be thrown and the resulting handleLostAll will only clean up active tasks
        closeDirtyAndRevive(corruptedStandbyTasks, true);

        // We need to commit before closing the corrupted active tasks since this will force the ongoing txn to abort
        try {
            final Collection<Task> tasksToCommit = allTasks()
                .values()
                .stream()
                .filter(t -> t.state() == Task.State.RUNNING || t.state() == Task.State.RESTORING)
                .filter(t -> !corruptedTasks.contains(t.id()))
                .collect(Collectors.toSet());
            commitTasksAndMaybeUpdateCommittableOffsets(tasksToCommit, new HashMap<>());
        } catch (final TaskCorruptedException e) {

Review comment (suggestion):
Could we do something like:

```
final Set<TopicPartition> partitionsNotToPause = tasks.allTasks().stream().flatMap(task -> task.inputPartitions().stream()).collect(Collectors.toSet());
mainConsumer.pause(mainConsumer.assignment().stream().filter(partition -> !partitionsNotToPause.contains(partition)).collect(Collectors.toSet()));
```
or

```
final Set<TopicPartition> partitionsNotToPause = tasks.allTasks().stream().flatMap(task -> task.inputPartitions().stream()).collect(Collectors.toSet());
final Set<TopicPartition> partitionsToPause = new HashSet<>(mainConsumer.assignment());
partitionsToPause.removeAll(partitionsNotToPause);
mainConsumer.pause(partitionsToPause);        
```

Code after:
        return rebalanceInProgress;
    }

    void handleRebalanceStart(final Set<String> subscribedTopics) {
        topologyMetadata.addSubscribedTopicsFromMetadata(subscribedTopics, logPrefix);

        tryToLockAllNonEmptyTaskDirectories();

        rebalanceInProgress = true;
    }

    void handleRebalanceComplete() {
        // we should pause consumer only within the listener since
        // before then the assignment has not been updated yet.
        if (stateUpdater == null) {
            mainConsumer.pause(mainConsumer.assignment());
        } else {
            // All tasks that are owned by the task manager are ready and do not need to be paused
            final Set<TopicPartition> partitionsNotToPause = tasks.allTasks()
                .stream()
                .flatMap(task -> task.inputPartitions().stream())
                .collect(Collectors.toSet());
            final Set<TopicPartition> partitionsToPause = new HashSet<>(mainConsumer.assignment());
            partitionsToPause.removeAll(partitionsNotToPause);
            mainConsumer.pause(partitionsToPause);
        }

        releaseLockedUnassignedTaskDirectories();

        rebalanceInProgress = false;
    }

    /**
     * @throws TaskMigratedException
     */
    boolean handleCorruption(final Set<TaskId> corruptedTasks) {
        final Set<Task> corruptedActiveTasks = new HashSet<>();
        final Set<Task> corruptedStandbyTasks = new HashSet<>();

        for (final TaskId taskId : corruptedTasks) {
            final Task task = tasks.task(taskId);
            if (task.isActive()) {
                corruptedActiveTasks.add(task);
            } else {
                corruptedStandbyTasks.add(task);
            }
        }

        // Make sure to clean up any corrupted standby tasks in their entirety before committing
        // since TaskMigrated can be thrown and the resulting handleLostAll will only clean up active tasks
        closeDirtyAndRevive(corruptedStandbyTasks, true);