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

Code before:
                if (restoringTask.isActive()) {
                    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.
        for (final TaskId id : union(HashSet::new, lockedTaskDirectories, tasks.allTaskIds())) {
            final Task task = tasks.contains(id) ? tasks.task(id) : null;
            // Closed and uninitialized tasks don't have any offsets so we should read directly from the checkpoint
            if (task != null && 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 {}", id);
                } else {
                    taskOffsetSums.put(id, sumOfChangelogOffsets(id, changelogOffsets));
                }
            } else {
                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;
    }

    /**

Review comment (refactor):
It seems with the state updated enabled, `tasks` is actually only containing "running tasks". It seems appropriate the rename this variable to `runningTasks` (can also happen in a follow up PR).

I am actually also wondering if we still need this `Tasks` container any longer to begin with? The purpose of the `Tasks` container was to simplify `TaskManager` that manages both active and standby tasks. With the state updated (from my understanding) the `TaskManager` only manages active tasks, while standby tasks will be owned by the state-updated-thread (would it still be useful for the state-updated-thread to use `Tasks` container, given that is also own active tasks as long as they are restoring?

Code after:
                if (restoringTask.isActive()) {
                    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);
            }
        }