KAFKA-19128: Kafka Streams should not get offsets when close dirty (#19450)

Kafka Streams calls `prepareCommit()` in `Taskmanager#closeTaskDirty()`.
However, the dirty task must not get committed and therefore,
prepare-commit tasks such as getting offsets should not be needed as
well. The only thing needed before closing a task dirty is flushing.
Therefore, separating `flush` and `prepareCommit` could be a good fix.

Reviewers: Bill Bejeck <[email protected]>, Matthias J. Sax <[email protected]>

Author: aliehsaeedii

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

@@ -180,7 +180,7 @@ public boolean maybePunctuateSystemTime() {
     }
 
     @Override
-    public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
+    public Map<TopicPartition, OffsetAndMetadata> prepareCommit(final boolean clean) {
         throw new UnsupportedOperationException("This task is read-only");
     }
 

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

@@ -179,7 +179,7 @@ public void resume() {
      *                          or flushing state store get IO errors; such error should cause the thread to die
      */
     @Override
-    public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
+    public Map<TopicPartition, OffsetAndMetadata> prepareCommit(final boolean clean) {
         switch (state()) {
             case CREATED:
                 log.debug("Skipped preparing created task for commit");
@@ -189,15 +189,19 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
             case RUNNING:
             case SUSPENDED:
                 // do not need to flush state store caches in pre-commit since nothing would be sent for standby tasks
-                log.debug("Prepared {} task for committing", state());
+                if (!clean) {
+                    log.debug("Skipped preparing {} standby task with id {} for commit since the task is getting closed dirty.", state(), id);
+                } else {
+                    log.debug("Prepared {} task for committing", state());
+                }
 
                 break;
 
             default:
                 throw new IllegalStateException("Illegal state " + state() + " while preparing standby task " + id + " for committing ");
         }
 
-        return Collections.emptyMap();
+        return clean ? Collections.emptyMap() : null;
     }
 
     @Override

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

@@ -417,7 +417,6 @@ public void resume() {
         timeCurrentIdlingStarted = Optional.empty();
     }
 
-
     public void flush() {
         stateMgr.flushCache();
         recordCollector.flush();
@@ -429,7 +428,7 @@ public void flush() {
      * @return offsets that should be committed for this task
      */
     @Override
-    public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
+    public Map<TopicPartition, OffsetAndMetadata> prepareCommit(final boolean clean) {
         switch (state()) {
             case CREATED:
             case RESTORING:
@@ -444,6 +443,10 @@ public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
                     //
                     // TODO: this should be removed after we decouple caching with emitting
                     flush();
+                    if (!clean) {
+                        log.debug("Skipped preparing {} task with id {} for commit since the task is getting closed dirty.", state(), id);
+                        return null;
+                    }
                     hasPendingTxCommit = eosEnabled;
 
                     log.debug("Prepared {} task for committing", state());

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

@@ -201,7 +201,7 @@ default boolean maybePunctuateSystemTime() {
     /**
      * @throws StreamsException fatal error, should close the thread
      */
-    Map<TopicPartition, OffsetAndMetadata> prepareCommit();
+    Map<TopicPartition, OffsetAndMetadata> prepareCommit(final boolean clean);
 
     void postCommit(boolean enforceCheckpoint);
 

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

@@ -142,7 +142,7 @@ int commitTasksAndMaybeUpdateCommittableOffsets(final Collection<Task> tasksToCo
         for (final Task task : tasksToCommit) {
             // we need to call commitNeeded first since we need to update committable offsets
             if (task.commitNeeded()) {
-                final Map<TopicPartition, OffsetAndMetadata> offsetAndMetadata = task.prepareCommit();
+                final Map<TopicPartition, OffsetAndMetadata> offsetAndMetadata = task.prepareCommit(true);
                 if (!offsetAndMetadata.isEmpty()) {
                     consumedOffsetsAndMetadata.put(task, offsetAndMetadata);
                 }

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

@@ -278,7 +278,7 @@ private void closeDirtyAndRevive(final Collection<Task> taskWithChangelogs, fina
                     // we do not need to take the returned offsets since we are not going to commit anyways;
                     // this call is only used for active tasks to flush the cache before suspending and
                     // closing the topology
-                    task.prepareCommit();
+                    task.prepareCommit(false);
                 } catch (final RuntimeException swallow) {
                     log.warn("Error flushing cache for corrupted task {}. " +
                         "Since the task is closing dirty, the following exception is swallowed: {}",
@@ -812,7 +812,7 @@ private Map<TaskId, RuntimeException> closeAndRecycleTasks(final Map<Task, Set<T
                 //    and their changelog positions should not change at all postCommit would not write the checkpoint again.
                 // 2) for standby tasks prepareCommit should always return empty, and then in postCommit we would probably
                 //    write the checkpoint file.
-                final Map<TopicPartition, OffsetAndMetadata> offsets = task.prepareCommit();
+                final Map<TopicPartition, OffsetAndMetadata> offsets = task.prepareCommit(true);
                 if (!offsets.isEmpty()) {
                     log.error("Task {} should have been committed when it was suspended, but it reports non-empty " +
                                     "offsets {} to commit; this means it failed during last commit and hence should be closed dirty",
@@ -1264,7 +1264,7 @@ private void prepareCommitAndAddOffsetsToMap(final Set<Task> tasksToPrepare,
                                                  final Map<Task, Map<TopicPartition, OffsetAndMetadata>> consumedOffsetsPerTask) {
         for (final Task task : tasksToPrepare) {
             try {
-                final Map<TopicPartition, OffsetAndMetadata> committableOffsets = task.prepareCommit();
+                final Map<TopicPartition, OffsetAndMetadata> committableOffsets = task.prepareCommit(true);
                 if (!committableOffsets.isEmpty()) {
                     consumedOffsetsPerTask.put(task, committableOffsets);
                 }
@@ -1479,7 +1479,7 @@ private void closeTaskDirty(final Task task, final boolean removeFromTasksRegist
         try {
             // we call this function only to flush the case if necessary
             // before suspending and closing the topology
-            task.prepareCommit();
+            task.prepareCommit(false);
         } catch (final RuntimeException swallow) {
             log.warn("Error flushing cache of dirty task {}. " +
                 "Since the task is closing dirty, the following exception is swallowed: {}",
@@ -1630,7 +1630,7 @@ private Collection<Task> tryCloseCleanActiveTasks(final Collection<Task> activeT
         // first committing all tasks and then suspend and close them clean
         for (final Task task : activeTasksToClose) {
             try {
-                final Map<TopicPartition, OffsetAndMetadata> committableOffsets = task.prepareCommit();
+                final Map<TopicPartition, OffsetAndMetadata> committableOffsets = task.prepareCommit(true);
                 tasksToCommit.add(task);
                 if (!committableOffsets.isEmpty()) {
                     consumedOffsetsAndMetadataPerTask.put(task, committableOffsets);
@@ -1719,7 +1719,7 @@ private Collection<Task> tryCloseCleanStandbyTasks(final Collection<Task> standb
         // first committing and then suspend / close clean
         for (final Task task : standbyTasksToClose) {
             try {
-                task.prepareCommit();
+                task.prepareCommit(true);
                 task.postCommit(true);
                 task.suspend();
                 closeTaskClean(task);

streams/src/test/java/org/apache/kafka/streams/processor/internals/StandbyTaskTest.java

@@ -213,7 +213,7 @@ public void shouldThrowIfCommittingOnIllegalState() {
         task.suspend();
         task.closeClean();
 
-        assertThrows(IllegalStateException.class, task::prepareCommit);
+        assertThrows(IllegalStateException.class, () -> task.prepareCommit(true));
     }
 
     @Test
@@ -261,13 +261,13 @@ public void shouldFlushAndCheckpointStateManagerOnCommit() {
 
         task = createStandbyTask();
         task.initializeIfNeeded();
-        task.prepareCommit();
+        task.prepareCommit(true);
         task.postCommit(false);  // this should not checkpoint
 
-        task.prepareCommit();
+        task.prepareCommit(true);
         task.postCommit(false);  // this should checkpoint
 
-        task.prepareCommit();
+        task.prepareCommit(true);
         task.postCommit(false);  // this should not checkpoint
 
         verify(stateManager).checkpoint();
@@ -322,7 +322,7 @@ public void shouldSuspendAndCommitBeforeCloseClean() {
         task = createStandbyTask();
         task.initializeIfNeeded();
         task.suspend();
-        task.prepareCommit();
+        task.prepareCommit(true);
         task.postCommit(true);
         task.closeClean();
 
@@ -360,7 +360,7 @@ public void shouldOnlyNeedCommitWhenChangelogOffsetChanged() {
         // could commit if the offset advanced beyond threshold
         assertTrue(task.commitNeeded());
 
-        task.prepareCommit();
+        task.prepareCommit(true);
         task.postCommit(true);
     }
 
@@ -389,7 +389,7 @@ public void shouldThrowOnCloseCleanCheckpointError() {
         task = createStandbyTask();
         task.initializeIfNeeded();
 
-        task.prepareCommit();
+        task.prepareCommit(true);
         assertThrows(RuntimeException.class, () -> task.postCommit(true));
 
         assertEquals(RUNNING, task.state());