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

Code before:
            .filter(Objects::nonNull)
            .forEach(removedTaskResult -> {
                if (removedTaskResult.exception().isPresent()) {
                    tasksToCloseDirtyFromStateUpdater.add(removedTaskResult.task());
                } else {
                    tasksToCloseCleanFromStateUpdater.add(removedTaskResult.task());
                }
            });
    }

    private Task checkIfTaskFailed(final StateUpdater.RemovedTaskResult removedTaskResult,
                                   final Map<TaskId, RuntimeException> failedTasks) {
        final Task task = removedTaskResult.task();
        if (removedTaskResult.exception().isPresent()) {
            failedTasks.put(task.id(), removedTaskResult.exception().get());
            tasks.addFailedTask(task);
            return null;
        }
        return task;
    }

    private StateUpdater.RemovedTaskResult waitForFuture(final TaskId taskId,
                                                         final CompletableFuture<StateUpdater.RemovedTaskResult> future) {
        final StateUpdater.RemovedTaskResult removedTaskResult;
        try {
            removedTaskResult = future.get(1, TimeUnit.MINUTES);
            if (removedTaskResult == null) {
                throw new IllegalStateException("Task " + taskId + " was not found in the state updater. "
                    + BUG_ERROR_MESSAGE);
            }
            return removedTaskResult;
        } catch (final ExecutionException executionException) {
            log.warn("An exception happened when removing task {} from the state updater. The task was added to the " +
                    "failed task in the state updater: ",
                taskId, executionException);
            return null;
        } catch (final InterruptedException shouldNotHappen) {
            Thread.currentThread().interrupt();
            log.error(INTERRUPTED_ERROR_MESSAGE, shouldNotHappen);
            throw new IllegalStateException(INTERRUPTED_ERROR_MESSAGE, shouldNotHappen);
        } catch (final java.util.concurrent.TimeoutException timeoutException) {
            log.warn("The state updater wasn't able to remove task {} in time. The state updater thread may be dead. "
                    + BUG_ERROR_MESSAGE, taskId, timeoutException);
            return null;
        }
    }

    private Map<TaskId, Set<TopicPartition>> pendingTasksToCreate(final Map<TaskId, Set<TopicPartition>> tasksToCreate) {
        final Map<TaskId, Set<TopicPartition>> pendingTasks = new HashMap<>();
        final Iterator<Map.Entry<TaskId, Set<TopicPartition>>> iter = tasksToCreate.entrySet().iterator();
        while (iter.hasNext()) {

Review comment (suggestion):
Do I understand correctly, that this is the core fix, to avoid blocking forever?

(Plus the additional cleanup code add to `shutdownStateUpdater` in case we really hit the timeout here?)

Code after:
            .filter(Objects::nonNull)
            .forEach(removedTaskResult -> {
                if (removedTaskResult.exception().isPresent()) {
                    tasksToCloseDirtyFromStateUpdater.add(removedTaskResult.task());
                } else {
                    tasksToCloseCleanFromStateUpdater.add(removedTaskResult.task());
                }
            });
    }

    private Task checkIfTaskFailed(final StateUpdater.RemovedTaskResult removedTaskResult,
                                   final Map<TaskId, RuntimeException> failedTasks) {
        final Task task = removedTaskResult.task();
        if (removedTaskResult.exception().isPresent()) {
            failedTasks.put(task.id(), removedTaskResult.exception().get());
            tasks.addFailedTask(task);
            return null;
        }
        return task;
    }

    private StateUpdater.RemovedTaskResult waitForFuture(final TaskId taskId,
                                                         final CompletableFuture<StateUpdater.RemovedTaskResult> future) {
        final StateUpdater.RemovedTaskResult removedTaskResult;
        try {
            removedTaskResult = future.get(5, TimeUnit.MINUTES);
            if (removedTaskResult == null) {
                throw new IllegalStateException("Task " + taskId + " was not found in the state updater. "
                    + BUG_ERROR_MESSAGE);
            }
            return removedTaskResult;
        } catch (final ExecutionException executionException) {
            log.warn("An exception happened when removing task {} from the state updater. The task was added to the " +
                    "failed task in the state updater: ",
                taskId, executionException);
            return null;
        } catch (final InterruptedException shouldNotHappen) {
            Thread.currentThread().interrupt();
            log.error(INTERRUPTED_ERROR_MESSAGE, shouldNotHappen);
            throw new IllegalStateException(INTERRUPTED_ERROR_MESSAGE, shouldNotHappen);
        } catch (final java.util.concurrent.TimeoutException timeoutException) {
            log.warn("The state updater wasn't able to remove task {} in time. The state updater thread may be dead. "
                    + BUG_ERROR_MESSAGE, taskId, timeoutException);
            return null;
        }
    }

    private Map<TaskId, Set<TopicPartition>> pendingTasksToCreate(final Map<TaskId, Set<TopicPartition>> tasksToCreate) {
        final Map<TaskId, Set<TopicPartition>> pendingTasks = new HashMap<>();
        final Iterator<Map.Entry<TaskId, Set<TopicPartition>>> iter = tasksToCreate.entrySet().iterator();
        while (iter.hasNext()) {