KAFKA-6145: Pt 2. Include offset sums in subscription

streams/src/main/java/org/apache/kafka/streams/processor/internals/ProcessorStateManager.java

@@ -180,7 +180,7 @@ public ProcessorStateManager(final TaskId taskId,
         this.storeToChangelogTopic = storeToChangelogTopic;
 
         this.baseDir = stateDirectory.directoryForTask(taskId);
-        this.checkpointFile = new OffsetCheckpoint(new File(baseDir, CHECKPOINT_FILE_NAME));
+        this.checkpointFile = new OffsetCheckpoint(stateDirectory.checkpointFileFor(taskId));
 
         log.debug("Created state store manager for task {}", taskId);
     }

streams/src/main/java/org/apache/kafka/streams/processor/internals/StateDirectory.java

@@ -107,6 +107,13 @@ public File directoryForTask(final TaskId taskId) {
         return taskDir;
     }
 
+    /**
+     * @return The File handle for the checkpoint in the given task's directory
+     */
+    File checkpointFileFor(final TaskId taskId) {
+        return new File(directoryForTask(taskId), StateManagerUtil.CHECKPOINT_FILE_NAME);
+    }
+
     /**
      * Decide if the directory of the task is empty or not
      */
@@ -285,12 +292,7 @@ public synchronized void cleanRemovedTasks(final long cleanupDelayMs) {
 
     private synchronized void cleanRemovedTasks(final long cleanupDelayMs,
                                                 final boolean manualUserCall) throws Exception {
-        final File[] taskDirs = listTaskDirectories();
-        if (taskDirs == null || taskDirs.length == 0) {
-            return; // nothing to do
-        }
-
-        for (final File taskDir : taskDirs) {
+        for (final File taskDir : listTaskDirectories()) {
             final String dirName = taskDir.getName();
             final TaskId id = TaskId.parse(dirName);
             if (!locks.containsKey(id)) {
@@ -347,8 +349,14 @@ private synchronized void cleanRemovedTasks(final long cleanupDelayMs,
      * @return The list of all the existing local directories for stream tasks
      */
     File[] listTaskDirectories() {
-        return !stateDir.exists() ? new File[0] :
-                stateDir.listFiles(pathname -> pathname.isDirectory() && PATH_NAME.matcher(pathname.getName()).matches());
+        final File[] taskDirectories =
+            stateDir.listFiles(pathname -> pathname.isDirectory() && PATH_NAME.matcher(pathname.getName()).matches());
+
+        if (!stateDir.exists() || taskDirectories == null) {
+            return new File[0];
+        } else {
+            return taskDirectories;
+        }
     }
 
     private FileChannel getOrCreateFileChannel(final TaskId taskId,

streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamTask.java

@@ -671,16 +671,16 @@ private void initializeTaskTime(final Map<TopicPartition, OffsetAndMetadata> off
     }
 
     @Override
-    public Map<TopicPartition, Long> purgableOffsets() {
-        final Map<TopicPartition, Long> purgableConsumedOffsets = new HashMap<>();
+    public Map<TopicPartition, Long> purgeableOffsets() {
+        final Map<TopicPartition, Long> purgeableConsumedOffsets = new HashMap<>();
         for (final Map.Entry<TopicPartition, Long> entry : consumedOffsets.entrySet()) {
             final TopicPartition tp = entry.getKey();
             if (topology.isRepartitionTopic(tp.topic())) {
-                purgableConsumedOffsets.put(tp, entry.getValue() + 1);
+                purgeableConsumedOffsets.put(tp, entry.getValue() + 1);
             }
         }
 
-        return purgableConsumedOffsets;
+        return purgeableConsumedOffsets;
     }
 
     private void initializeTopology() {

streams/src/main/java/org/apache/kafka/streams/processor/internals/Task.java

@@ -32,7 +32,9 @@
 import java.util.Set;
 
 public interface Task {
-    // this must be negative to distinguish a running active task from other kinds tasks which may be caught up to the same offsets
+
+    // this must be negative to distinguish a running active task from other kinds of tasks
+    // which may be caught up to the same offsets
     long LATEST_OFFSET = -2L;
 
     /*
@@ -176,7 +178,7 @@ enum TaskType {
 
     void markChangelogAsCorrupted(final Collection<TopicPartition> partitions);
 
-    default Map<TopicPartition, Long> purgableOffsets() {
+    default Map<TopicPartition, Long> purgeableOffsets() {
         return Collections.emptyMap();
     }
 
@@ -195,6 +197,4 @@ default boolean maybePunctuateStreamTime() {
     default boolean maybePunctuateSystemTime() {
         return false;
     }
-
-
 }

streams/src/main/java/org/apache/kafka/streams/processor/internals/TaskManager.java

@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.streams.processor.internals;
 
+import java.io.IOException;
+import java.util.Collections;
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.admin.DeleteRecordsResult;
 import org.apache.kafka.clients.admin.RecordsToDelete;
@@ -32,6 +34,7 @@
 import org.apache.kafka.streams.errors.TaskMigratedException;
 import org.apache.kafka.streams.processor.TaskId;
 import org.apache.kafka.streams.processor.internals.metrics.StreamsMetricsImpl;
+import org.apache.kafka.streams.state.internals.OffsetCheckpoint;
 import org.slf4j.Logger;
 
 import java.io.File;
@@ -79,6 +82,9 @@ public class TaskManager {
 
     private boolean rebalanceInProgress = false;  // if we are in the middle of a rebalance, it is not safe to commit
 
+    // includes assigned & initialized tasks and unassigned tasks we locked temporarily during rebalance
+    private Set<TaskId> lockedTaskDirectories = new HashSet<>();
+
     TaskManager(final ChangelogReader changelogReader,
                 final UUID processId,
                 final String logPrefix,
@@ -121,6 +127,8 @@ boolean isRebalanceInProgress() {
     void handleRebalanceStart(final Set<String> subscribedTopics) {
         builder.addSubscribedTopicsFromMetadata(subscribedTopics, logPrefix);
 
+        tryToLockAllTaskDirectories();
+
         rebalanceInProgress = true;
     }
 
@@ -129,6 +137,8 @@ void handleRebalanceComplete() {
         // before then the assignment has not been updated yet.
         mainConsumer.pause(mainConsumer.assignment());
 
+        releaseLockedUnassignedTaskDirectories();
+
         rebalanceInProgress = false;
     }
 
@@ -368,50 +378,105 @@ void handleLostAll() {
     }
 
     /**
+     * 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 #tryToLockAllTaskDirectories()}
+     *
      * @return Map from task id to its total offset summed across all state stores
      */
     public Map<TaskId, Long> getTaskOffsetSums() {
         final Map<TaskId, Long> taskOffsetSums = new HashMap<>();
 
-        for (final TaskId id : tasksOnLocalStorage()) {
-            if (isRunning(id)) {
-                taskOffsetSums.put(id, Task.LATEST_OFFSET);
+        for (final TaskId id : lockedTaskDirectories) {
+            final Task task = tasks.get(id);
+            if (task != null) {
+                if (task.isActive() && task.state() == RUNNING) {
+                    taskOffsetSums.put(id, Task.LATEST_OFFSET);
+                } else {
+                    taskOffsetSums.put(id, sumOfChangelogOffsets(id, task.changelogOffsets()));
+                }
             } else {
-                taskOffsetSums.put(id, 0L);
+                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;
     }
 
     /**
-     * Returns ids of tasks whose states are kept on the local storage. This includes active, standby, and previously
-     * assigned but not yet cleaned up tasks
+     * Makes a weak attempt to lock all 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.
      */
-    private Set<TaskId> tasksOnLocalStorage() {
-        // A client could contain some inactive tasks whose states are still kept on the local storage in the following scenarios:
-        // 1) the client is actively maintaining standby tasks by maintaining their states from the change log.
-        // 2) the client has just got some tasks migrated out of itself to other clients while these task states
-        //    have not been cleaned up yet (this can happen in a rolling bounce upgrade, for example).
+    private void tryToLockAllTaskDirectories() {
+        for (final File dir : stateDirectory.listTaskDirectories()) {
+            try {
+                final TaskId id = TaskId.parse(dir.getName());
+                try {
+                    if (stateDirectory.lock(id)) {
+                        lockedTaskDirectories.add(id);
+                        if (!tasks.containsKey(id)) {
+                            log.debug("Temporarily locked unassigned task {} for the upcoming rebalance", id);
+                        }
+                    }
+                } catch (final IOException e) {
+                    // if for any reason we can't lock this task dir, just move on
+                    log.warn(String.format("Exception caught while attempting to lock task %s:", id), e);
+                }
+            } catch (final TaskIdFormatException e) {
+                // ignore any unknown files that sit in the same directory
+            }
+        }
+    }
 
-        final Set<TaskId> locallyStoredTasks = new HashSet<>();
+    /**
+     * We must release the lock for any unassigned tasks that we temporarily locked in preparation for a
+     * rebalance in {@link #tryToLockAllTaskDirectories()}.
+     */
+    private void releaseLockedUnassignedTaskDirectories() {
+        final AtomicReference<RuntimeException> firstException = new AtomicReference<>(null);
 
-        final File[] stateDirs = stateDirectory.listTaskDirectories();
-        if (stateDirs != null) {
-            for (final File dir : stateDirs) {
+        final Iterator<TaskId> taskIdIterator = lockedTaskDirectories.iterator();
+        while (taskIdIterator.hasNext()) {
+            final TaskId id = taskIdIterator.next();
+            if (!tasks.containsKey(id)) {
                 try {
-                    final TaskId id = TaskId.parse(dir.getName());
-                    // if the checkpoint file exists, the state is valid.
-                    if (new File(dir, StateManagerUtil.CHECKPOINT_FILE_NAME).exists()) {
-                        locallyStoredTasks.add(id);
-                    }
-                } catch (final TaskIdFormatException e) {
-                    // there may be some unknown files that sits in the same directory,
-                    // we should ignore these files instead trying to delete them as well
+                    stateDirectory.unlock(id);
+                    taskIdIterator.remove();
+                } catch (final IOException e) {
+                    log.error(String.format("Caught the following exception while trying to unlock task %s", id), e);
+                    firstException.compareAndSet(null,
+                        new StreamsException(String.format("Failed to unlock task directory %s", id), e));
                 }
             }
         }
 
-        return locallyStoredTasks;
+        final RuntimeException fatalException = firstException.get();
+        if (fatalException != null) {
+            throw fatalException;
+        }
+    }
+
+    private long sumOfChangelogOffsets(final TaskId id, final Map<TopicPartition, Long> changelogOffsets) {
+        long offsetSum = 0L;
+        for (final Map.Entry<TopicPartition, Long> changelogEntry : changelogOffsets.entrySet()) {
+            final long offset = changelogEntry.getValue();
+
+            offsetSum += offset;
+            if (offsetSum < 0) {
+                log.warn("Sum of changelog offsets for task {} overflowed, pinning to Long.MAX_VALUE", id);
+                return Long.MAX_VALUE;
+            }
+        }
+
+        return offsetSum;
     }
 
     private void cleanupTask(final Task task) {
@@ -474,6 +539,14 @@ void shutdown(final boolean clean) {
             }
         }
 
+        try {
+            // this should be called after closing all tasks, to make sure we unlock the task dir for tasks that may
+            // have still been in CREATED at the time of shutdown, since Task#close will not do so
+            releaseLockedUnassignedTaskDirectories();
+        } catch (final RuntimeException e) {
+            firstException.compareAndSet(null, e);
+        }
+
         final RuntimeException fatalException = firstException.get();
         if (fatalException != null) {
             throw new RuntimeException("Unexpected exception while closing task", fatalException);
@@ -522,11 +595,6 @@ private Stream<Task> standbyTaskStream() {
         return tasks.values().stream().filter(t -> !t.isActive());
     }
 
-    private boolean isRunning(final TaskId id) {
-        final Task task = tasks.get(id);
-        return task != null && task.isActive() && task.state() == RUNNING;
-    }
-
     /**
      * @throws TaskMigratedException if committing offsets failed (non-EOS)
      *                               or if the task producer got fenced (EOS)
@@ -629,7 +697,7 @@ void maybePurgeCommittedRecords() {
 
             final Map<TopicPartition, RecordsToDelete> recordsToDelete = new HashMap<>();
             for (final Task task : activeTaskIterable()) {
-                for (final Map.Entry<TopicPartition, Long> entry : task.purgableOffsets().entrySet()) {
+                for (final Map.Entry<TopicPartition, Long> entry : task.purgeableOffsets().entrySet()) {
                     recordsToDelete.put(entry.getKey(), RecordsToDelete.beforeOffset(entry.getValue()));
                 }
             }
@@ -676,4 +744,8 @@ Map<MetricName, Metric> producerMetrics() {
     Set<String> producerClientIds() {
         return activeTaskCreator.producerClientIds();
     }
+
+    Set<TaskId> lockedTaskDirectories() {
+        return Collections.unmodifiableSet(lockedTaskDirectories);
+    }
 }