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

Code before:
                    tasks.addPendingTaskToCloseDirty(restoringTask.id());
                    stateUpdater.remove(restoringTask.id());
                }
            }
        }
    }

    public void signalResume() {
        if (stateUpdater != null) {
            stateUpdater.signalResume();
        }
    }

    /**
     * Compute the offset total summed across all stores in a task. Includes offset sum for any tasks we own the
     * lock for, which includes assigned and unassigned tasks we locked in {@link #tryToLockAllNonEmptyTaskDirectories()}.
     * Does not include stateless or non-logged tasks.
     */
    public Map<TaskId, Long> getTaskOffsetSums() {
        final Map<TaskId, Long> taskOffsetSums = new HashMap<>();

        final Map<TaskId, Task> tasks = allTasks();
        final Set<TaskId> lockedTaskDirectoriesOfNonOwnedTasksAndClosedAndCreatedTasks =
            union(HashSet::new, lockedTaskDirectories, tasks.keySet());
        for (final Task task : tasks.values()) {
            if (task.state() != State.CREATED && task.state() != State.CLOSED) {
                final Map<TopicPartition, Long> changelogOffsets = task.changelogOffsets();
                if (changelogOffsets.isEmpty()) {
                    log.debug("Skipping to encode apparently stateless (or non-logged) offset sum for task {}",
                        task.id());
                } else {
                    taskOffsetSums.put(task.id(), sumOfChangelogOffsets(task.id(), changelogOffsets));
                }
                lockedTaskDirectoriesOfNonOwnedTasksAndClosedAndCreatedTasks.remove(task.id());
            }
        }

        // Not all tasks will create directories, and there may be directories for tasks we don't currently own,
        // so we consider all tasks that are either owned or on disk. This includes stateless tasks, which should
        // just have an empty changelogOffsets map.
        for (final TaskId id : lockedTaskDirectoriesOfNonOwnedTasksAndClosedAndCreatedTasks) {
            final File checkpointFile = stateDirectory.checkpointFileFor(id);
            try {
                if (checkpointFile.exists()) {
                    taskOffsetSums.put(id, sumOfChangelogOffsets(id, new OffsetCheckpoint(checkpointFile).read()));
                }
            } catch (final IOException e) {
                log.warn(String.format("Exception caught while trying to read checkpoint for task %s:", id), e);
            }
        }


Review comment (suggestion):
Ah, I recommended this change thinking that `lockedTaskDirectories` always includes all `ClosedAndCreatedTasks` -- I think it does right? So it should be enough to assign this to `lockedTaskDirectories`.

Code after:
                    tasks.addPendingTaskToCloseDirty(restoringTask.id());
                    stateUpdater.remove(restoringTask.id());
                }
            }
        }
    }

    public void signalResume() {
        if (stateUpdater != null) {
            stateUpdater.signalResume();
        }
    }

    /**
     * Compute the offset total summed across all stores in a task. Includes offset sum for any tasks we own the
     * lock for, which includes assigned and unassigned tasks we locked in {@link #tryToLockAllNonEmptyTaskDirectories()}.
     * Does not include stateless or non-logged tasks.
     */
    public Map<TaskId, Long> getTaskOffsetSums() {
        final Map<TaskId, Long> taskOffsetSums = new HashMap<>();

        // Not all tasks will create directories, and there may be directories for tasks we don't currently own,
        // so we consider all tasks that are either owned or on disk. This includes stateless tasks, which should
        // just have an empty changelogOffsets map.
        final Map<TaskId, Task> tasks = allTasks();
        final Set<TaskId> lockedTaskDirectoriesOfNonOwnedTasksAndClosedAndCreatedTasks =
            union(HashSet::new, lockedTaskDirectories, tasks.keySet());
        for (final Task task : tasks.values()) {
            if (task.state() != State.CREATED && task.state() != State.CLOSED) {
                final Map<TopicPartition, Long> changelogOffsets = task.changelogOffsets();
                if (changelogOffsets.isEmpty()) {
                    log.debug("Skipping to encode apparently stateless (or non-logged) offset sum for task {}",
                        task.id());
                } else {
                    taskOffsetSums.put(task.id(), sumOfChangelogOffsets(task.id(), changelogOffsets));
                }
                lockedTaskDirectoriesOfNonOwnedTasksAndClosedAndCreatedTasks.remove(task.id());
            }
        }

        for (final TaskId id : lockedTaskDirectoriesOfNonOwnedTasksAndClosedAndCreatedTasks) {
            final File checkpointFile = stateDirectory.checkpointFileFor(id);
            try {
                if (checkpointFile.exists()) {
                    taskOffsetSums.put(id, sumOfChangelogOffsets(id, new OffsetCheckpoint(checkpointFile).read()));
                }
            } catch (final IOException e) {
                log.warn(String.format("Exception caught while trying to read checkpoint for task %s:", id), e);
            }
        }