KAFKA-9481: Graceful handling TaskMigrated and TaskCorrupted

clients/src/main/java/org/apache/kafka/common/utils/FixedOrderMap.java

@@ -50,12 +50,6 @@ public boolean remove(final Object key, final Object value) {
         throw new UnsupportedOperationException("Removing from registeredStores is not allowed");
     }
 
-    @Deprecated
-    @Override
-    public void clear() {
-        throw new UnsupportedOperationException("Removing from registeredStores is not allowed");
-    }
-
     @Override
     public FixedOrderMap<K, V> clone() {
         throw new UnsupportedOperationException();

clients/src/test/java/org/apache/kafka/common/utils/FixedOrderMapTest.java

@@ -69,18 +69,4 @@ public void shouldForbidConditionalRemove() {
         }
         assertThat(map.get("a"), is(0));
     }
-
-    @SuppressWarnings("deprecation")
-    @Test
-    public void shouldForbidConditionalClear() {
-        final FixedOrderMap<String, Integer> map = new FixedOrderMap<>();
-        map.put("a", 0);
-        try {
-            map.clear();
-            fail("expected exception");
-        } catch (final RuntimeException e) {
-            assertThat(e, CoreMatchers.instanceOf(UnsupportedOperationException.class));
-        }
-        assertThat(map.get("a"), is(0));
-    }
 }

streams/src/main/java/org/apache/kafka/streams/errors/TaskCorruptedException.java

@@ -0,0 +1,46 @@
+/*
+ * 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.errors;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.streams.processor.TaskId;
+
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Indicates a specific task is corrupted and need to be re-initialized. It can be thrown when
+ *
+ * 1) Under EOS, if the checkpoint file does not contain offsets for corresponding store's changelogs, meaning
+ *    previously it was not close cleanly;
+ * 2) Out-of-range exception thrown during restoration, meaning that the changelog has been modified and we re-bootstrap
+ *    the store.
+ */
+public class TaskCorruptedException extends StreamsException {
+
+    private final Map<TaskId, Set<TopicPartition>> taskWithChangelogs;
+
+    public TaskCorruptedException(final Map<TaskId, Set<TopicPartition>> taskWithChangelogs) {
+        super("Tasks with changelogs " + taskWithChangelogs + " are corrupted and hence needs to be re-initialized");
+
+        this.taskWithChangelogs = taskWithChangelogs;
+    }
+
+    public Map<TaskId, Set<TopicPartition>> corruptedTaskWithChangelogs() {
+        return taskWithChangelogs;
+    }
+}

streams/src/main/java/org/apache/kafka/streams/errors/TaskMigratedException.java

@@ -17,58 +17,16 @@
 package org.apache.kafka.streams.errors;
 
 
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.streams.processor.TaskId;
-
 /**
- * Indicates that one or more tasks got migrated to another thread.
- *
- * 1) if the task field is specified, then that single task should be cleaned up and closed as "zombie" while the
- *    thread can continue as normal;
- * 2) if no tasks are specified (i.e. taskId == null), it means that the hosted thread has been fenced and all
- *    tasks are migrated, in which case the thread should rejoin the group
+ * Indicates that all tasks belongs to the thread have migrated to another thread. This exception can be thrown when
+ * the thread gets fenced (either by the consumer coordinator or by the transaction coordinator), which means it is
+ * no longer part of the group but a "zombie" already
  */
 public class TaskMigratedException extends StreamsException {
 
     private final static long serialVersionUID = 1L;
 
-    private final TaskId taskId;
-
-    public TaskMigratedException(final TaskId taskId,
-                                 final TopicPartition topicPartition,
-                                 final long endOffset,
-                                 final long pos) {
-        this(taskId, String.format("Log end offset of %s should not change while restoring: old end offset %d, current offset %d",
-            topicPartition,
-            endOffset,
-            pos), null);
-    }
-
-    public TaskMigratedException(final TaskId taskId) {
-        this(taskId, String.format("Task %s is unexpectedly closed during processing", taskId), null);
-    }
-
-    public TaskMigratedException(final TaskId taskId,
-                                 final Throwable throwable) {
-        this(taskId, String.format("Client request for task %s has been fenced due to a rebalance", taskId), throwable);
-    }
-
-    public TaskMigratedException(final TaskId taskId,
-                                 final String message,
-                                 final Throwable throwable) {
-        super(message, throwable);
-        this.taskId = taskId;
-    }
-
     public TaskMigratedException(final String message, final Throwable throwable) {
-        this(null, message + " It means all tasks belonging to this thread have been migrated", throwable);
-    }
-
-    public TaskId migratedTaskId() {
-        return taskId;
-    }
-
-    public TaskMigratedException() {
-        this(null, "A task has been migrated unexpectedly", null);
+        super(message + "; It means all tasks belonging to this thread should be migrated", throwable);
     }
 }

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

@@ -20,8 +20,10 @@
 import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.processor.TaskId;
 
+import java.util.Collection;
 import java.util.Set;
 
+import static org.apache.kafka.streams.processor.internals.Task.State.CLOSED;
 import static org.apache.kafka.streams.processor.internals.Task.State.CREATED;
 
 public abstract class AbstractTask implements Task {
@@ -55,6 +57,16 @@ public Set<TopicPartition> inputPartitions() {
         return partitions;
     }
 
+    @Override
+    public Collection<TopicPartition> changelogPartitions() {
+        return stateMgr.changelogPartitions();
+    }
+
+    @Override
+    public void markChangelogAsCorrupted(final Set<TopicPartition> partitions) {
+        stateMgr.markChangelogAsCorrupted(partitions);
+    }
+
     @Override
     public StateStore getStore(final String name) {
         return stateMgr.getStore(name);
@@ -70,6 +82,15 @@ public final Task.State state() {
         return state;
     }
 
+    @Override
+    public void revive() {
+        if (state == CLOSED) {
+            transitionTo(CREATED);
+        } else {
+            throw new IllegalStateException("Illegal state " + state() + " while reviving task " + id);
+        }
+    }
+
     final void transitionTo(final Task.State newState) {
         final State oldState = state();
 

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

@@ -37,6 +37,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 import static java.lang.String.format;
@@ -84,11 +85,15 @@ public static class StateStoreMetadata {
         //      update blindly with the given offset
         private Long offset;
 
+        // corrupted state store should not be included in checkpointing
+        private boolean corrupted;
+
         private StateStoreMetadata(final StateStore stateStore) {
             this.stateStore = stateStore;
             this.restoreCallback = null;
             this.recordConverter = null;
             this.changelogPartition = null;
+            this.corrupted = false;
             this.offset = null;
         }
 
@@ -282,6 +287,20 @@ Collection<TopicPartition> changelogPartitions() {
         return changelogOffsets().keySet();
     }
 
+    void markChangelogAsCorrupted(final Set<TopicPartition> partitions) {
+        for (final StateStoreMetadata storeMetadata : stores.values()) {
+            if (partitions.contains(storeMetadata.changelogPartition)) {
+                storeMetadata.corrupted = true;
+                partitions.remove(storeMetadata.changelogPartition);
+            }
+        }
+
+        if (!partitions.isEmpty()) {
+            throw new IllegalStateException("Some partitions " + partitions + " are not contained in the store list of task " +
+                taskId + " marking as corrupted, this is not expected");
+        }
+    }
+
     @Override
     public Map<TopicPartition, Long> changelogOffsets() {
         // return the current offsets for those logged stores
@@ -415,6 +434,8 @@ public void close() throws ProcessorStateException {
                     log.error("Failed to close state store {}: ", store.name(), exception);
                 }
             }
+
+            stores.clear();
         }
 
         if (firstException != null) {
@@ -439,10 +460,11 @@ public void checkpoint(final Map<TopicPartition, Long> writtenOffsets) {
 
         final Map<TopicPartition, Long> checkpointingOffsets = new HashMap<>();
         for (final StateStoreMetadata storeMetadata : stores.values()) {
-            // store is logged, persistent, and has a valid current offset
+            // store is logged, persistent, not corrupted, and has a valid current offset
             if (storeMetadata.changelogPartition != null &&
                 storeMetadata.stateStore.persistent() &&
-                storeMetadata.offset != null) {
+                storeMetadata.offset != null &&
+                !storeMetadata.corrupted) {
                 checkpointingOffsets.put(storeMetadata.changelogPartition, storeMetadata.offset);
             }
         }

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

@@ -132,7 +132,7 @@ private void maybeBeginTxn() {
             try {
                 producer.beginTransaction();
             } catch (final ProducerFencedException error) {
-                throw new TaskMigratedException(taskId, "Producer get fenced trying to begin a new transaction", error);
+                throw new TaskMigratedException("Producer get fenced trying to begin a new transaction", error);
             } catch (final KafkaException error) {
                 throw new StreamsException("Producer encounter unexpected error trying to begin a new transaction", error);
             }
@@ -169,7 +169,7 @@ public void commit(final Map<TopicPartition, OffsetAndMetadata> offsets) {
                 producer.commitTransaction();
                 transactionInFlight = false;
             } catch (final ProducerFencedException error) {
-                throw new TaskMigratedException(taskId, "Producer get fenced trying to commit a transaction", error);
+                throw new TaskMigratedException("Producer get fenced trying to commit a transaction", error);
             } catch (final TimeoutException error) {
                 // TODO KIP-447: we can consider treating it as non-fatal and retry on the thread level
                 throw new StreamsException("Timed out while committing transaction via producer for task " + taskId, error);
@@ -181,7 +181,7 @@ public void commit(final Map<TopicPartition, OffsetAndMetadata> offsets) {
             try {
                 consumer.commitSync(offsets);
             } catch (final CommitFailedException error) {
-                throw new TaskMigratedException(taskId, "Consumer committing offsets failed, " +
+                throw new TaskMigratedException("Consumer committing offsets failed, " +
                     "indicating the corresponding thread is no longer part of the group.", error);
             } catch (final TimeoutException error) {
                 // TODO KIP-447: we can consider treating it as non-fatal and retry on the thread level
@@ -216,7 +216,7 @@ private void recordSendError(final String topic, final Exception exception, fina
         } else if (exception instanceof ProducerFencedException || exception instanceof OutOfOrderSequenceException) {
             errorMessage += "\nWritten offsets would not be recorded and no more records would be sent since the producer is fenced, " +
                 "indicating the task may be migrated out.";
-            sendException = new TaskMigratedException(taskId, errorMessage, exception);
+            sendException = new TaskMigratedException(errorMessage, exception);
         } else {
             if (exception instanceof RetriableException) {
                 errorMessage += "\nThe broker is either slow or in bad state (like not having enough replicas) in responding the request, " +
@@ -314,7 +314,7 @@ public <K, V> void send(final String topic,
             if (isRecoverable(uncaughtException)) {
                 // producer.send() call may throw a KafkaException which wraps a FencedException,
                 // in this case we should throw its wrapped inner cause so that it can be captured and re-wrapped as TaskMigrationException
-                throw new TaskMigratedException(taskId, "Producer cannot send records anymore since it got fenced", uncaughtException.getCause());
+                throw new TaskMigratedException("Producer cannot send records anymore since it got fenced", uncaughtException.getCause());
             } else {
                 final String errorMessage = String.format(SEND_EXCEPTION_MESSAGE, topic, taskId, uncaughtException.toString());
                 throw new StreamsException(errorMessage, uncaughtException);

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

@@ -29,7 +29,6 @@
 import org.apache.kafka.streams.processor.internals.metrics.ThreadMetrics;
 import org.slf4j.Logger;
 
-import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
@@ -187,6 +186,16 @@ private void close(final boolean clean) {
         transitionTo(State.CLOSED);
     }
 
+    @Override
+    public boolean commitNeeded() {
+        return false;
+    }
+
+    @Override
+    public Map<TopicPartition, Long> changelogOffsets() {
+        return Collections.unmodifiableMap(stateMgr.changelogOffsets());
+    }
+
     @Override
     public void addRecords(final TopicPartition partition, final Iterable<ConsumerRecord<byte[], byte[]>> records) {
         throw new IllegalStateException("Attempted to add records to task " + id() + " for invalid input partition " + partition);
@@ -223,16 +232,4 @@ public String toString(final String indent) {
 
         return sb.toString();
     }
-
-    public boolean commitNeeded() {
-        return false;
-    }
-
-    public Collection<TopicPartition> changelogPartitions() {
-        return stateMgr.changelogPartitions();
-    }
-
-    public Map<TopicPartition, Long> changelogOffsets() {
-        return Collections.unmodifiableMap(stateMgr.changelogOffsets());
-    }
 }