From c27a0ffb26322a78f14ad4fd205b26abfea5d3f6 Mon Sep 17 00:00:00 2001
From: ableegoldman <sophie@confluent.io>
Date: Wed, 24 Jun 2020 15:42:56 -0700
Subject: [PATCH 1/3] validateClean

---
 .../processor/internals/AbstractTask.java     |  1 +
 .../internals/ActiveTaskCreator.java          |  2 +-
 .../processor/internals/StandbyTask.java      |  2 +-
 .../internals/StandbyTaskCreator.java         |  2 +-
 .../processor/internals/StreamTask.java       | 21 ++++++++++++-------
 .../streams/processor/internals/Task.java     |  2 +-
 .../processor/internals/StandbyTaskTest.java  |  6 +++---
 .../processor/internals/StreamTaskTest.java   | 10 ++++-----
 .../processor/internals/TaskManagerTest.java  |  2 +-
 9 files changed, 27 insertions(+), 21 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/AbstractTask.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/AbstractTask.java
index 044825a21f607..d7c7bd5b8a26a 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/AbstractTask.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/AbstractTask.java
@@ -108,4 +108,5 @@ public void update(final Set<TopicPartition> topicPartitions, final Map<String,
         this.inputPartitions = topicPartitions;
         topology.updateSourceTopics(nodeToSourceTopics);
     }
+
 }
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/ActiveTaskCreator.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/ActiveTaskCreator.java
index 0a1f47e139e6c..012ff20f8b980 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/ActiveTaskCreator.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/ActiveTaskCreator.java
@@ -184,7 +184,7 @@ StreamTask createActiveTaskFromStandby(final StandbyTask standbyTask,
         final ProcessorStateManager stateManager = standbyTask.stateMgr;
         final LogContext logContext = getLogContext(standbyTask.id);
 
-        standbyTask.closeAndRecycleState();
+        standbyTask.closeCleanAndRecycleState();
         stateManager.transitionTaskType(TaskType.ACTIVE, logContext);
 
         return createActiveTask(
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTask.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTask.java
index ffd09f14059c8..1aa68bc721f2e 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTask.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTask.java
@@ -189,7 +189,7 @@ public void closeDirty() {
     }
 
     @Override
-    public void closeAndRecycleState() {
+    public void closeCleanAndRecycleState() {
         streamsMetrics.removeAllTaskLevelSensors(Thread.currentThread().getName(), id.toString());
         if (state() == State.SUSPENDED) {
             stateMgr.recycle();
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTaskCreator.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTaskCreator.java
index 443db8e8dd94a..b5f2b74f3f453 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTaskCreator.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StandbyTaskCreator.java
@@ -112,7 +112,7 @@ StandbyTask createStandbyTaskFromActive(final StreamTask streamTask,
         final InternalProcessorContext context = streamTask.processorContext();
         final ProcessorStateManager stateManager = streamTask.stateMgr;
 
-        streamTask.closeAndRecycleState();
+        streamTask.closeCleanAndRecycleState();
         stateManager.transitionTaskType(TaskType.STANDBY, getLogContext(streamTask.id()));
 
         return createStandbyTask(
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamTask.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamTask.java
index 6e8bf40ab79c6..4b27436023b4b 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamTask.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/StreamTask.java
@@ -463,6 +463,7 @@ private Map<TopicPartition, Long> extractPartitionTimes() {
 
     @Override
     public void closeClean() {
+        validateClean();
         streamsMetrics.removeAllTaskLevelSensors(Thread.currentThread().getName(), id.toString());
         close(true);
         log.info("Closed clean");
@@ -482,7 +483,8 @@ public void update(final Set<TopicPartition> topicPartitions, final Map<String,
     }
 
     @Override
-    public void closeAndRecycleState() {
+    public void closeCleanAndRecycleState() {
+        validateClean();
         streamsMetrics.removeAllTaskLevelSensors(Thread.currentThread().getName(), id.toString());
         switch (state()) {
             case SUSPENDED:
@@ -515,17 +517,20 @@ private void writeCheckpoint() {
         stateMgr.checkpoint(checkpointableOffsets());
     }
 
-    /**
-     * You must commit a task and checkpoint the state manager before closing as this will release the state dir lock
-     */
-    private void close(final boolean clean) {
-        if (clean && commitNeeded) {
-            // It may be that we failed to commit a task during handleRevocation, but "forgot" this and tried to
-            // closeClean in handleAssignment. We should throw if we detect this to force the TaskManager to closeDirty
+    private void validateClean() {
+        // It may be that we failed to commit a task during handleRevocation, but "forgot" this and tried to
+        // closeClean in handleAssignment. We should throw if we detect this to force the TaskManager to closeDirty
+        if (commitNeeded) {
             log.debug("Tried to close clean but there was pending uncommitted data, this means we failed to"
                           + " commit and should close as dirty instead");
             throw new TaskMigratedException("Tried to close dirty task as clean");
         }
+    }
+
+    /**
+     * You must commit a task and checkpoint the state manager before closing as this will release the state dir lock
+     */
+    private void close(final boolean clean) {
         switch (state()) {
             case SUSPENDED:
                 // first close state manager (which is idempotent) then close the record collector
diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/Task.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/Task.java
index 0200870b7aa9a..103c2317b639b 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/Task.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/Task.java
@@ -153,7 +153,7 @@ enum TaskType {
     /**
      * Attempt a clean close but do not close the underlying state
      */
-    void closeAndRecycleState();
+    void closeCleanAndRecycleState();
 
     /**
      * Revive a closed task to a created one; should never throw an exception
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StandbyTaskTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StandbyTaskTest.java
index 82f33c4697b9a..f98d6304e5842 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StandbyTaskTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StandbyTaskTest.java
@@ -502,13 +502,13 @@ public void shouldRecycleTask() {
         EasyMock.replay(stateManager);
 
         task = createStandbyTask();
-        assertThrows(IllegalStateException.class, () -> task.closeAndRecycleState()); // CREATED
+        assertThrows(IllegalStateException.class, () -> task.closeCleanAndRecycleState()); // CREATED
 
         task.initializeIfNeeded();
-        assertThrows(IllegalStateException.class, () -> task.closeAndRecycleState()); // RUNNING
+        assertThrows(IllegalStateException.class, () -> task.closeCleanAndRecycleState()); // RUNNING
 
         task.suspend();
-        task.closeAndRecycleState(); // SUSPENDED
+        task.closeCleanAndRecycleState(); // SUSPENDED
 
         // Currently, there are no metrics registered for standby tasks.
         // This is a regression test so that, if we add some, we will be sure to deregister them.
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
index c6ffe74e6b8ce..01fb2bbafd168 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
@@ -1761,7 +1761,7 @@ public void shouldUnregisterMetricsInCloseAndRecycle() {
 
         task.suspend();
         assertThat(getTaskMetrics(), not(empty()));
-        task.closeAndRecycleState();
+        task.closeCleanAndRecycleState();
         assertThat(getTaskMetrics(), empty());
     }
 
@@ -1805,16 +1805,16 @@ public void shouldOnlyRecycleSuspendedTasks() {
         EasyMock.replay(stateManager, recordCollector);
 
         task = createStatefulTask(createConfig(false, "100"), true);
-        assertThrows(IllegalStateException.class, () -> task.closeAndRecycleState()); // CREATED
+        assertThrows(IllegalStateException.class, () -> task.closeCleanAndRecycleState()); // CREATED
 
         task.initializeIfNeeded();
-        assertThrows(IllegalStateException.class, () -> task.closeAndRecycleState()); // RESTORING
+        assertThrows(IllegalStateException.class, () -> task.closeCleanAndRecycleState()); // RESTORING
 
         task.completeRestoration();
-        assertThrows(IllegalStateException.class, () -> task.closeAndRecycleState()); // RUNNING
+        assertThrows(IllegalStateException.class, () -> task.closeCleanAndRecycleState()); // RUNNING
 
         task.suspend();
-        task.closeAndRecycleState(); // SUSPENDED
+        task.closeCleanAndRecycleState(); // SUSPENDED
 
         EasyMock.verify(stateManager, recordCollector);
     }
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
index fcfbb1f455672..23166d14d084b 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/TaskManagerTest.java
@@ -2723,7 +2723,7 @@ public void closeDirty() {
         }
 
         @Override
-        public void closeAndRecycleState() {
+        public void closeCleanAndRecycleState() {
             transitionTo(State.CLOSED);
         }
 

From 3c8967f8f336c653a56575d2c29c15db4b16d33a Mon Sep 17 00:00:00 2001
From: ableegoldman <sophie@confluent.io>
Date: Wed, 24 Jun 2020 15:51:03 -0700
Subject: [PATCH 2/3] add unit tests

---
 .../processor/internals/AbstractTask.java     |  1 -
 .../processor/internals/StreamTaskTest.java   | 25 ++++++++++++++++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/streams/src/main/java/org/apache/kafka/streams/processor/internals/AbstractTask.java b/streams/src/main/java/org/apache/kafka/streams/processor/internals/AbstractTask.java
index d7c7bd5b8a26a..044825a21f607 100644
--- a/streams/src/main/java/org/apache/kafka/streams/processor/internals/AbstractTask.java
+++ b/streams/src/main/java/org/apache/kafka/streams/processor/internals/AbstractTask.java
@@ -108,5 +108,4 @@ public void update(final Set<TopicPartition> topicPartitions, final Map<String,
         this.inputPartitions = topicPartitions;
         topology.updateSourceTopics(nodeToSourceTopics);
     }
-
 }
diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
index 01fb2bbafd168..1e55afd799e8f 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
@@ -44,6 +44,7 @@
 import org.apache.kafka.streams.errors.LockException;
 import org.apache.kafka.streams.errors.ProcessorStateException;
 import org.apache.kafka.streams.errors.StreamsException;
+import org.apache.kafka.streams.errors.TaskMigratedException;
 import org.apache.kafka.streams.processor.PunctuationType;
 import org.apache.kafka.streams.processor.Punctuator;
 import org.apache.kafka.streams.processor.StateStore;
@@ -1752,7 +1753,7 @@ public void shouldUnregisterMetricsInCloseDirty() {
     }
 
     @Test
-    public void shouldUnregisterMetricsInCloseAndRecycle() {
+    public void shouldUnregisterMetricsInCloseCleanAndRecycleState() {
         EasyMock.expect(stateManager.changelogPartitions()).andReturn(Collections.emptySet()).anyTimes();
         EasyMock.expect(recordCollector.offsets()).andReturn(Collections.emptyMap()).anyTimes();
         EasyMock.replay(stateManager, recordCollector);
@@ -1798,6 +1799,28 @@ public void shouldUpdatePartitions() {
         assertThat(task.inputPartitions(), equalTo(newPartitions));
     }
 
+    @Test
+    public void shouldThrowIfCleanClosingDirtyTask() {
+        task = createStatefulTask(createConfig(false, "100"), true);
+
+        task.addRecords(partition1, singletonList(getConsumerRecord(partition1, 0)));
+        task.process(0L);
+        assertTrue(task.commitNeeded());
+
+        assertThrows(TaskMigratedException.class, () -> task.closeClean());
+    }
+
+    @Test
+    public void shouldThrowIfRecyclingDirtyTask() {
+        task = createStatefulTask(createConfig(false, "100"), true);
+
+        task.addRecords(partition1, singletonList(getConsumerRecord(partition1, 0)));
+        task.process(0L);
+        assertTrue(task.commitNeeded());
+
+        assertThrows(TaskMigratedException.class, () -> task.closeCleanAndRecycleState());
+    }
+
     @Test
     public void shouldOnlyRecycleSuspendedTasks() {
         stateManager.recycle();

From fc71861dfcb093d5d430f8e7a3fa8786c1d12fb1 Mon Sep 17 00:00:00 2001
From: ableegoldman <sophie@confluent.io>
Date: Wed, 24 Jun 2020 15:55:21 -0700
Subject: [PATCH 3/3] fix test setup

---
 .../kafka/streams/processor/internals/StreamTaskTest.java | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
index 1e55afd799e8f..960747047cfc4 100644
--- a/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
+++ b/streams/src/test/java/org/apache/kafka/streams/processor/internals/StreamTaskTest.java
@@ -1801,7 +1801,9 @@ public void shouldUpdatePartitions() {
 
     @Test
     public void shouldThrowIfCleanClosingDirtyTask() {
-        task = createStatefulTask(createConfig(false, "100"), true);
+        task = createStatelessTask(createConfig(false, "0"), StreamsConfig.METRICS_LATEST);
+        task.initializeIfNeeded();
+        task.completeRestoration();
 
         task.addRecords(partition1, singletonList(getConsumerRecord(partition1, 0)));
         task.process(0L);
@@ -1812,7 +1814,9 @@ public void shouldThrowIfCleanClosingDirtyTask() {
 
     @Test
     public void shouldThrowIfRecyclingDirtyTask() {
-        task = createStatefulTask(createConfig(false, "100"), true);
+        task = createStatelessTask(createConfig(false, "0"), StreamsConfig.METRICS_LATEST);
+        task.initializeIfNeeded();
+        task.completeRestoration();
 
         task.addRecords(partition1, singletonList(getConsumerRecord(partition1, 0)));
         task.process(0L);