KAFKA-9727: cleanup the state store for standby task dirty close and check null for changelogs

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

@@ -42,6 +42,7 @@ public class StandbyTask extends AbstractTask implements Task {
     private final Logger log;
     private final String logPrefix;
     private final Sensor closeTaskSensor;
+    private final boolean eosEnabled;
     private final InternalProcessorContext processorContext;
 
     private Map<TopicPartition, Long> offsetSnapshotSinceLastCommit;
@@ -71,6 +72,7 @@ public class StandbyTask extends AbstractTask implements Task {
 
         processorContext = new StandbyContextImpl(id, config, stateMgr, metrics);
         closeTaskSensor = ThreadMetrics.closeTaskSensor(Thread.currentThread().getName(), metrics);
+        this.eosEnabled = StreamsConfig.EXACTLY_ONCE.equals(config.getString(StreamsConfig.PROCESSING_GUARANTEE_CONFIG));
     }
 
     @Override
@@ -187,50 +189,41 @@ private void prepareClose(final boolean clean) {
     @Override
     public void closeClean(final Map<TopicPartition, Long> checkpoint) {
         Objects.requireNonNull(checkpoint);
-        close(true, checkpoint);
+        close(true);
 
         log.info("Closed clean");
     }
 
     @Override
     public void closeDirty() {
-        close(false, null);
+        close(false);
 
         log.info("Closed dirty");
     }
 
-    private void close(final boolean clean,
-                       final Map<TopicPartition, Long> checkpoint) {
-        if (state() == State.CREATED) {
-            // the task is created and not initialized, do nothing
-            closeTaskSensor.record();
-            transitionTo(State.CLOSED);
-            return;
-        }
-
-        if (state() == State.RUNNING) {
+    private void close(final boolean clean) {
+        if (state() == State.CREATED || state() == State.RUNNING) {
             if (clean) {
                 // since there's no written offsets we can checkpoint with empty map,
                 // and the state current offset would be used to checkpoint
                 stateMgr.checkpoint(Collections.emptyMap());
                 offsetSnapshotSinceLastCommit = new HashMap<>(stateMgr.changelogOffsets());
             }
+            final boolean wipeStateStore = !clean && eosEnabled;
+            log.info("standby task clean {}, eos enabled {}", clean, eosEnabled);
 
-            executeAndMaybeSwallow(clean, () -> {
+            executeAndMaybeSwallow(clean, () ->
                 StateManagerUtil.closeStateManager(
                     log,
                     logPrefix,
                     clean,
-                    false,
+                    wipeStateStore,
                     stateMgr,
                     stateDirectory,
-                    TaskType.STANDBY);
-                },
+                    TaskType.STANDBY),
                 "state manager close",
                 log
             );
-
-            // TODO: if EOS is enabled, we should wipe out the state stores like we did for StreamTask too
         } else {
             throw new IllegalStateException("Illegal state " + state() + " while closing standby task " + id);
         }

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

@@ -798,10 +798,16 @@ public void remove(final Collection<TopicPartition> revokedChangelogs) {
 
         for (final TopicPartition partition : revokedChangelogs) {
             final ChangelogMetadata changelogMetadata = changelogs.remove(partition);
-            if (changelogMetadata.state() != ChangelogState.REGISTERED) {
-                revokedInitializedChangelogs.add(partition);
+            if (changelogMetadata != null) {
+                if (changelogMetadata.state() != ChangelogState.REGISTERED) {
+                    revokedInitializedChangelogs.add(partition);
+                }
+                changelogMetadata.clear();
+            } else {
+                log.debug("Changelog partition {} could not be found, " +
+                    "it could be already cleaned up during the handling" +
+                    "of task corruption and never restore again", partition);
             }
-            changelogMetadata.clear();
         }
 
         removeChangelogsFromRestoreConsumer(revokedInitializedChangelogs);

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

@@ -78,7 +78,7 @@ public class StreamTask extends AbstractTask implements ProcessorNodePunctuator,
     // we want to abstract eos logic out of StreamTask, however
     // there's still an optimization that requires this info to be
     // leaked into this class, which is to checkpoint after committing if EOS is not enabled.
-    private final boolean eosDisabled;
+    private final boolean eosEnabled;
 
     private final long maxTaskIdleMs;
     private final int maxBufferedSize;
@@ -120,7 +120,7 @@ public StreamTask(final TaskId id,
 
         this.time = time;
         this.recordCollector = recordCollector;
-        eosDisabled = !StreamsConfig.EXACTLY_ONCE.equals(config.getString(StreamsConfig.PROCESSING_GUARANTEE_CONFIG));
+        eosEnabled = StreamsConfig.EXACTLY_ONCE.equals(config.getString(StreamsConfig.PROCESSING_GUARANTEE_CONFIG));
 
         final String threadId = Thread.currentThread().getName();
         closeTaskSensor = ThreadMetrics.closeTaskSensor(threadId, streamsMetrics);
@@ -326,7 +326,7 @@ public void postCommit() {
                 commitNeeded = false;
                 commitRequested = false;
 
-                if (eosDisabled) {
+                if (!eosEnabled) {
                     stateMgr.checkpoint(checkpointableOffsets());
                 }
 
@@ -487,7 +487,7 @@ private void close(final boolean clean,
             case SUSPENDED:
                 // if EOS is enabled, we wipe out the whole state store for unclean close
                 // since they are invalid to use anymore
-                final boolean wipeStateStore = !clean && !eosDisabled;
+                final boolean wipeStateStore = !clean && eosEnabled;
 
                 // first close state manager (which is idempotent) then close the record collector (which could throw),
                 // if the latter throws and we re-close dirty which would close the state manager again.

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

@@ -333,7 +333,7 @@ boolean tryToCompleteRestoration() {
                     // it is possible that if there are multiple threads within the instance that one thread
                     // trying to grab the task from the other, while the other has not released the lock since
                     // it did not participate in the rebalance. In this case we can just retry in the next iteration
-                    log.debug("Could not initialize {} due to {}; will retry", task.id(), e);
+                    log.debug("Could not initialize {} due to the following exception; will retry", task.id(), e);
                     allRunning = false;
                 }
             }

streams/src/test/java/org/apache/kafka/streams/integration/StandbyTaskEOSIntegrationTest.java

@@ -0,0 +1,153 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.integration;
+
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.serialization.IntegerSerializer;
+import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.common.utils.MockTime;
+import org.apache.kafka.streams.KafkaStreams;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.StreamsBuilder;
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.integration.utils.EmbeddedKafkaCluster;
+import org.apache.kafka.streams.integration.utils.IntegrationTestUtils;
+import org.apache.kafka.streams.kstream.Consumed;
+import org.apache.kafka.streams.processor.TaskId;
+import org.apache.kafka.streams.processor.internals.StateDirectory;
+import org.apache.kafka.streams.state.internals.OffsetCheckpoint;
+import org.apache.kafka.test.TestUtils;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Properties;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.kafka.test.TestUtils.waitForCondition;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * An integration test to verify the conversion of a dirty-closed EOS
+ * task towards a standby task is safe across restarts of the application.
+ */
+public class StandbyTaskEOSIntegrationTest {
+
+    private final String inputTopic = "input";
+
+    @ClassRule
+    public static final EmbeddedKafkaCluster CLUSTER = new EmbeddedKafkaCluster(3);
+
+    @Before
+    public void createTopics() throws Exception {
+        CLUSTER.createTopic(inputTopic, 1, 3);
+    }
+
+    @Test
+    public void surviveWithOneTaskAsStandby() throws ExecutionException, InterruptedException, IOException {
+        IntegrationTestUtils.produceKeyValuesSynchronouslyWithTimestamp(
+            inputTopic,
+            Collections.singletonList(
+                new KeyValue<>(0, 0)),
+            TestUtils.producerConfig(
+                CLUSTER.bootstrapServers(),
+                IntegerSerializer.class,
+                IntegerSerializer.class,
+                new Properties()),
+            10L);
+
+        final String appId = "eos-test-app";
+        final String stateDirPath = TestUtils.tempDirectory(appId).getPath();
+
+        final CountDownLatch instanceLatch = new CountDownLatch(1);
+
+        final String stateDirPathOne = stateDirPath + "/" + appId + "-1/";
+        final KafkaStreams streamInstanceOne =
+            buildStreamWithDirtyStateDir(appId, stateDirPathOne, instanceLatch);
+
+        final String stateDirPathTwo = stateDirPath + "/" + appId + "-2/";
+        final KafkaStreams streamInstanceTwo =
+            buildStreamWithDirtyStateDir(appId, stateDirPathTwo, instanceLatch);
+
+        streamInstanceOne.start();
+
+        streamInstanceTwo.start();
+
+        // Wait for the record to be processed
+        assertTrue(instanceLatch.await(15, TimeUnit.SECONDS));
+
+        waitForCondition(() -> streamInstanceOne.state().equals(KafkaStreams.State.RUNNING),
+            "Stream instance one should be up and running by now");
+        waitForCondition(() -> streamInstanceTwo.state().equals(KafkaStreams.State.RUNNING),
+            "Stream instance one should be up and running by now");
+
+        streamInstanceOne.close(Duration.ofSeconds(30));
+        streamInstanceTwo.close(Duration.ofSeconds(30));
+
+    }
+
+    private KafkaStreams buildStreamWithDirtyStateDir(final String appId,
+                                                      final String stateDirPath,
+                                                      final CountDownLatch recordProcessLatch) throws IOException {
+
+        final StreamsBuilder builder = new StreamsBuilder();
+        final TaskId taskId = new TaskId(0, 0);
+
+        final Properties props = props(appId, stateDirPath);
+
+        final StateDirectory stateDirectory = new StateDirectory(
+            new StreamsConfig(props), new MockTime(), true);
+
+        new OffsetCheckpoint(new File(stateDirectory.directoryForTask(taskId), ".checkpoint"))
+            .write(Collections.singletonMap(new TopicPartition("unknown-topic", 0), 5L));
+
+        assertTrue(new File(stateDirectory.directoryForTask(taskId),
+            "rocksdb/KSTREAM-AGGREGATE-STATE-STORE-0000000001").mkdirs());
+
+        builder.stream(inputTopic,
+            Consumed.with(Serdes.Integer(), Serdes.Integer()))
+            .groupByKey()
+            .count()
+            .toStream()
+            .peek((key, value) -> recordProcessLatch.countDown());
+
+        return new KafkaStreams(builder.build(), props);
+    }
+
+    private Properties props(final String appId, final String stateDirPath) {
+        final Properties streamsConfiguration = new Properties();
+        streamsConfiguration.put(StreamsConfig.APPLICATION_ID_CONFIG, appId);
+        streamsConfiguration.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, CLUSTER.bootstrapServers());
+        streamsConfiguration.put(StreamsConfig.CACHE_MAX_BYTES_BUFFERING_CONFIG, 0);
+        streamsConfiguration.put(StreamsConfig.STATE_DIR_CONFIG, stateDirPath);
+        streamsConfiguration.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1);
+        streamsConfiguration.put(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE);
+        streamsConfiguration.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.Integer().getClass());
+        streamsConfiguration.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.Integer().getClass());
+        streamsConfiguration.put(StreamsConfig.COMMIT_INTERVAL_MS_CONFIG, 1000);
+        streamsConfiguration.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+
+        return streamsConfiguration;
+    }
+}

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

@@ -62,7 +62,6 @@
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
-
 @RunWith(EasyMockRunner.class)
 public class StandbyTaskTest {
 
@@ -213,29 +212,6 @@ public void shouldReturnStateManagerChangelogOffsets() {
         EasyMock.verify(stateManager);
     }
 
-    @Test
-    public void shouldDoNothingWithCreatedStateOnClose() {
-        stateManager.close();
-        EasyMock.expectLastCall().andThrow(new AssertionError("Close should not be called")).anyTimes();
-        stateManager.flush();
-        EasyMock.expectLastCall().andThrow(new AssertionError("Flush should not be called")).anyTimes();
-        stateManager.checkpoint(EasyMock.anyObject());
-        EasyMock.expectLastCall().andThrow(new AssertionError("Checkpoint should not be called")).anyTimes();
-        EasyMock.replay(stateManager);
-        final MetricName metricName = setupCloseTaskMetric();
-
-        task = createStandbyTask();
-        final Map<TopicPartition, Long> checkpoint = task.prepareCloseClean();
-        task.closeClean(checkpoint);
-
-        assertEquals(Task.State.CLOSED, task.state());
-
-        final double expectedCloseTaskMetric = 1.0;
-        verifyCloseTaskMetric(expectedCloseTaskMetric, streamsMetrics, metricName);
-
-        EasyMock.verify(stateManager);
-    }
-
     @Test
     public void shouldNotCommitAndThrowOnCloseDirty() {
         stateManager.close();
@@ -407,12 +383,63 @@ public void shouldThrowOnCloseCleanCheckpointError() {
     public void shouldThrowIfClosingOnIllegalState() {
         task = createStandbyTask();
 
-        final Map<TopicPartition, Long> checkpoint = task.prepareCloseClean();
-        task.closeClean(checkpoint);
+        task.transitionTo(Task.State.RESTORING);
 
-        // close call are not idempotent since we are already in closed
+        // close calls are not idempotent since we are already in closed
         assertThrows(IllegalStateException.class, task::prepareCloseClean);
         assertThrows(IllegalStateException.class, task::prepareCloseDirty);
+
+        task.transitionTo(Task.State.CLOSED);
+    }
+
+    @Test
+    public void shouldCloseStateManagerOnTaskCreated() {
+        stateManager.close();
+        EasyMock.expectLastCall();
+
+        EasyMock.replay(stateManager);
+
+        final MetricName metricName = setupCloseTaskMetric();
+
+        task = createStandbyTask();
+
+        task.closeDirty();
+
+        final double expectedCloseTaskMetric = 1.0;
+        verifyCloseTaskMetric(expectedCloseTaskMetric, streamsMetrics, metricName);
+
+        EasyMock.verify(stateManager);
+
+        assertEquals(Task.State.CLOSED, task.state());
+    }
+
+    @Test
+    public void shouldDeleteStateDirOnTaskCreatedAndEOSUncleanClose() {
+        stateManager.close();
+        EasyMock.expectLastCall();
+
+        EasyMock.expect(stateManager.baseDir()).andReturn(baseDir);
+
+        EasyMock.replay(stateManager);
+
+        final MetricName metricName = setupCloseTaskMetric();
+
+        config = new StreamsConfig(mkProperties(mkMap(
+            mkEntry(StreamsConfig.APPLICATION_ID_CONFIG, applicationId),
+            mkEntry(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:2171"),
+            mkEntry(StreamsConfig.PROCESSING_GUARANTEE_CONFIG, StreamsConfig.EXACTLY_ONCE)
+        )));
+
+        task = createStandbyTask();
+
+        task.closeDirty();
+
+        final double expectedCloseTaskMetric = 1.0;
+        verifyCloseTaskMetric(expectedCloseTaskMetric, streamsMetrics, metricName);
+
+        EasyMock.verify(stateManager);
+
+        assertEquals(Task.State.CLOSED, task.state());
     }
 
     private StandbyTask createStandbyTask() {