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

Code before:
    }

    public boolean hasStartupTasks() {
        return !tasksForLocalState.isEmpty();
    }

    public Task removeStartupTask(final TaskId taskId) {
        final Task task = tasksForLocalState.remove(taskId);
        if (task != null) {
            lockedTasksToOwner.replace(taskId, Thread.currentThread());
        }
        return task;
    }

    public void closeStartupTasks() {
        closeStartupTasks(t -> true);
    }

    private void closeStartupTasks(final Predicate<Task> predicate) {
        if (!tasksForLocalState.isEmpty()) {
            // "drain" Tasks first to ensure that we don't try to close Tasks that another thread is attempting to close
            final Set<Task> drainedTasks = new HashSet<>(tasksForLocalState.size());
            for (final Map.Entry<TaskId, Task> entry : tasksForLocalState.entrySet()) {
                if (predicate.test(entry.getValue()) && removeStartupTask(entry.getKey()) != null) {
                    // only add to our list of drained Tasks if we exclusively "claimed" a Task from tasksForLocalState
                    // to ensure we don't accidentally try to drain the same Task multiple times from concurrent threads
                    drainedTasks.add(entry.getValue());
                }
            }

            // now that we have exclusive ownership of the drained tasks, close them
            for (final Task task : drainedTasks) {
                task.suspend();
                task.closeClean();
            }
        }
    }

    public Map<TaskId, Long> taskOffsetSums(final Set<TaskId> tasks) {
        return taskOffsetSums.entrySet().stream()
                .filter(e -> tasks.contains(e.getKey()))
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    }

    public void updateTaskOffsets(final TaskId taskId, final Map<TopicPartition, Long> changelogOffsets) {
        if (changelogOffsets.isEmpty()) {
            return;
        }

        taskOffsetSums.put(taskId, sumOfChangelogOffsets(taskId, changelogOffsets));
    }

Review comment (bug):
It would be nice to add a unit test that verifies the output of this method after updating the task offsets. 

Also, I think this method should be considered the source of truth for the current lag of each taskId, so I would suggest modifying this method to return an unmodifiable map instead of a modifiable one. However, this will break the TaskManager when it tries to put a `Task.LATEST_OFFSET` for the tasks that didn't have a cached offset, but we can do that on the scope of this method. 

In other words, it is making this method to return an offset (either a valid one or Task.LATEST_OFFSET) for every single TaskId received in the arguments. 

Code after:
                    tasksInLocalState.add(task);
                }
            }
        }
    }

    public boolean hasStartupTasks() {
        return !tasksInLocalState.isEmpty();
    }

    public synchronized boolean removeStartupState(final TaskId taskId) {
        final boolean removed = tasksInLocalState.remove(taskId);
        if (removed) {
            lockedTasksToOwner.put(taskId, Thread.currentThread());
        }
        return removed;
    }


    private void unlockStartupStores() {
        for (final TaskId task : tasksInLocalState) {
            unlock(task);
        }
    }

    public Map<TaskId, Long> taskOffsetSums(final Set<TaskId> tasks) {
        return taskOffsetSums.entrySet().stream()
                .filter(e -> tasks.contains(e.getKey()))
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    }

    public void updateTaskOffsets(final TaskId taskId, final Map<TopicPartition, Long> changelogOffsets) {
        if (!changelogOffsets.isEmpty()) {
            taskOffsetSums.put(taskId, sumOfChangelogOffsets(taskId, changelogOffsets));
        }
    }

    public void removeTaskOffsets(final TaskId taskId) {
        taskOffsetSums.remove(taskId);
    }

    private long sumOfChangelogOffsets(final TaskId taskId, final Map<TopicPartition, Long> changelogOffsets) {
        long offsetSum = 0L;
        for (final Map.Entry<TopicPartition, Long> changelogEntry : changelogOffsets.entrySet()) {
            final long offset = changelogEntry.getValue();

            if (offset != OffsetCheckpoint.OFFSET_UNKNOWN) {
                if (offset < 0) {
                    throw new StreamsException(
                            new IllegalStateException("Expected not to get a sentinel offset, but got: " + changelogEntry),
                            taskId);