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

Code before:
        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> createdAndClosedTasks = new HashSet<>();
        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));
                }
            } else {
                createdAndClosedTasks.add(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.
        final Set<TaskId> lockedTaskDirectoriesOfNonOwnedTasks = new HashSet<>(lockedTaskDirectories);
        lockedTaskDirectoriesOfNonOwnedTasks.removeAll(tasks.keySet());
        for (final TaskId id : union(HashSet::new, lockedTaskDirectoriesOfNonOwnedTasks, createdAndClosedTasks)) {
            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);
            }
        }

        return taskOffsetSums;
    }

    /**
     * Makes a weak attempt to lock all non-empty task directories in the state dir. We are responsible for computing and
     * reporting the offset sum for any unassigned tasks we obtain the lock for in the upcoming rebalance. Tasks

Review comment (suggestion):
nit: if you want to do it with fewer collections, you could inititialize `lockedTaskDirectoriesOfNonOwnedTasks` earlier, and just remove directly from that set in the `if` branch, instead of adding to `createdAndClosedTasks` in the `else` branch.

Code after:
        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);
            }
        }

        return taskOffsetSums;
    }

    /**
     * Makes a weak attempt to lock all non-empty task directories in the state dir. We are responsible for computing and
     * reporting the offset sum for any unassigned tasks we obtain the lock for in the upcoming rebalance. Tasks
     * that we locked but didn't own will be released at the end of the rebalance (unless of course we were
     * assigned the task as a result of the rebalance). This method should be idempotent.