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

Code before:
            // "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));
    }

    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);
                }
                offsetSum += offset;
                if (offsetSum < 0) {
                    log.warn("Sum of changelog offsets for task {} overflowed, pinning to Long.MAX_VALUE", taskId);
                    return Long.MAX_VALUE;

Review comment (suggestion):
maybe flip this to `!changelogOffsets.isEmpty()` and move the `taskOffsetSums.put` in the body of the `if` statement - IMHO it's easier to read

Code after:
        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);
                }
                offsetSum += offset;
                if (offsetSum < 0) {
                    log.warn("Sum of changelog offsets for task {} overflowed, pinning to Long.MAX_VALUE", taskId);
                    return Long.MAX_VALUE;
                }
            }
        }

        return offsetSum;
    }

    public UUID initializeProcessId() {
        if (!hasPersistentStores) {
            final UUID processId = UUID.randomUUID();
            log.info("Created new process id: {}", processId);
            return processId;
        }

        if (!lockStateDirectory()) {