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,44 @@
+/*
+ * 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.streams.processor.TaskId;
+
+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 Set<TaskId> taskIds;
+
+    public TaskCorruptedException(final Set<TaskId> taskIds) {
+        super("Tasks " + taskIds + " are corrupted and hence needs to be re-initialized");
+
+        this.taskIds = taskIds;
+    }
+
+    public Set<TaskId> corruptedTaskIds() {
+        return taskIds;
+    }
+}

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

@@ -17,9 +17,6 @@
 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.
  *
@@ -32,43 +29,7 @@ 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 have been migrated", throwable);
     }
 }

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

@@ -22,6 +22,7 @@
 
 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 {
@@ -70,6 +71,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

@@ -415,6 +415,8 @@ public void close() throws ProcessorStateException {
                     log.error("Failed to close state store {}: ", store.name(), exception);
                 }
             }
+
+            stores.clear();
         }
 
         if (firstException != null) {

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

@@ -124,7 +124,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);
             }
@@ -161,7 +161,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 K9113: currently handle timeout exception as a fatal error, should discuss whether we want to handle it
                 throw new StreamsException("Timed out while committing transaction via producer for task " + taskId, error);
@@ -173,7 +173,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 K9113: currently handle timeout exception as a fatal error
@@ -208,7 +208,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, " +
@@ -299,7 +299,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/StoreChangelogReader.java

@@ -19,6 +19,7 @@
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.consumer.InvalidOffsetException;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.errors.FencedInstanceIdException;
 import org.apache.kafka.common.errors.TimeoutException;
@@ -27,7 +28,9 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.streams.StreamsConfig;
 import org.apache.kafka.streams.errors.StreamsException;
+import org.apache.kafka.streams.errors.TaskCorruptedException;
 import org.apache.kafka.streams.errors.TaskMigratedException;
+import org.apache.kafka.streams.processor.TaskId;
 import org.apache.kafka.streams.processor.internals.ProcessorStateManager.StateStoreMetadata;
 import org.apache.kafka.streams.processor.StateRestoreListener;
 import org.slf4j.Logger;
@@ -52,7 +55,6 @@
  * The reader also maintains the source of truth for restoration state: only active tasks restoring changelog could
  * be completed, while standby tasks updating changelog would always be in restoring state after being initialized.
  */
-// TODO K9113: we need to consider how to handle InvalidOffsetException for consumer#poll / position
 public class StoreChangelogReader implements ChangelogReader {
 
     enum ChangelogState {
@@ -258,6 +260,8 @@ private boolean hasRestoredToEnd(final ChangelogMetadata metadata) {
                 // if we cannot get the position of the consumer within timeout, just return false
                 return false;
             } catch (final KafkaException e) {
+                // this also includes InvalidOffsetException, which should not happen under normal
+                // execution, hence it is also okay to wrap it as fatal StreamsException
                 throw new StreamsException("Restore consumer get unexpected error trying to get the position " +
                     " of " + partition, e);
             }
@@ -409,6 +413,15 @@ public void restore() {
             } catch (final FencedInstanceIdException e) {
                 // when the consumer gets fenced, all its tasks should be migrated
                 throw new TaskMigratedException("Restore consumer get fenced by instance-id polling records.", e);
+            } catch (final InvalidOffsetException e) {
+                log.warn("Encountered {} fetching records from restore consumer for partitions {}, " +
+                    "marking the corresponding tasks as corrupted.", e.getClass().getName(), e.partitions());
+
+                final Set<TaskId> taskIds = new HashSet<>();
+                for (final TopicPartition partition : e.partitions()) {
+                    taskIds.add(changelogs.get(partition).stateManager.taskId());
+                }
+                throw new TaskCorruptedException(taskIds);
             } catch (final KafkaException e) {
                 throw new StreamsException("Restore consumer get unexpected error polling records.", e);
             }
@@ -426,7 +439,6 @@ public void restore() {
                 restoreChangelog(changelogs.get(partition));
             }
 
-
             maybeUpdateLimitOffsetsForStandbyChangelogs();
         }
     }
@@ -679,6 +691,8 @@ private void addChangelogsToRestoreConsumer(final Set<TopicPartition> partitions
         }
         assignment.addAll(partitions);
         restoreConsumer.assign(assignment);
+
+        log.debug("Added partitions {} to the restore consumer, current assignment is {}", partitions, assignment);
     }
 
     private void pauseChangelogsFromRestoreConsumer(final Collection<TopicPartition> partitions) {
@@ -690,18 +704,18 @@ private void pauseChangelogsFromRestoreConsumer(final Collection<TopicPartition>
                 "does not contain some of the partitions " + partitions + " for pausing.");
         }
         restoreConsumer.pause(partitions);
+
+        log.debug("Paused partitions {} from the restore consumer", partitions);
     }
 
     private void removeChangelogsFromRestoreConsumer(final Collection<TopicPartition> partitions) {
         final Set<TopicPartition> assignment = new HashSet<>(restoreConsumer.assignment());
 
-        // the current assignment should contain the all partitions to remove
-        if (!assignment.containsAll(partitions)) {
-            throw new IllegalStateException("The current assignment " + assignment + " " +
-                "does not contain some of the partitions " + partitions + " for removing.");
+        if (assignment.removeAll(partitions)) {
+            restoreConsumer.assign(assignment);
+
+            log.debug("Removed partitions {} from the restore consumer, current assignment is {}", partitions, assignment);
         }
-        assignment.removeAll(partitions);
-        restoreConsumer.assign(assignment);
     }
 
     private void resumeChangelogsFromRestoreConsumer(final Collection<TopicPartition> partitions) {
@@ -713,6 +727,8 @@ private void resumeChangelogsFromRestoreConsumer(final Collection<TopicPartition
                 "does not contain some of the partitions " + partitions + " for resuming.");
         }
         restoreConsumer.resume(partitions);
+
+        log.debug("Resumed partitions {} from the restore consumer", partitions);
     }
 
     private void prepareChangelogs(final Set<ChangelogMetadata> newPartitionsToRestore) {
@@ -759,6 +775,8 @@ private void prepareChangelogs(final Set<ChangelogMetadata> newPartitionsToResto
                 } catch (final TimeoutException e) {
                     // if we cannot find the starting position at the beginning, just use the default 0L
                 } catch (final KafkaException e) {
+                    // this also includes InvalidOffsetException, which should not happen under normal
+                    // execution, hence it is also okay to wrap it as fatal StreamsException
                     throw new StreamsException("Restore consumer get unexpected error trying to get the position " +
                         " of " + partition, e);
                 }

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

@@ -69,10 +69,9 @@ public class StreamTask extends AbstractTask implements ProcessorNodePunctuator,
     // visible for testing
     static final byte LATEST_MAGIC_BYTE = 1;
 
+    private final Time time;
     private final Logger log;
     private final String logPrefix;
-    private final Time time;
-    private final String threadId;
     private final Consumer<byte[], byte[]> consumer;
 
     // we want to abstract eos logic out of StreamTask, however
@@ -82,7 +81,6 @@ public class StreamTask extends AbstractTask implements ProcessorNodePunctuator,
 
     private final long maxTaskIdleMs;
     private final int maxBufferedSize;
-    private final StreamsMetricsImpl streamsMetrics;
     private final PartitionGroup partitionGroup;
     private final RecordCollector recordCollector;
     private final PartitionGroup.RecordInfo recordInfo;
@@ -121,11 +119,10 @@ public StreamTask(final TaskId id,
         log = logContext.logger(getClass());
 
         this.time = time;
-        this.streamsMetrics = streamsMetrics;
         this.recordCollector = recordCollector;
         eosDisabled = !StreamsConfig.EXACTLY_ONCE.equals(config.getString(StreamsConfig.PROCESSING_GUARANTEE_CONFIG));
 
-        threadId = Thread.currentThread().getName();
+        final String threadId = Thread.currentThread().getName();
         closeTaskSensor = ThreadMetrics.closeTaskSensor(threadId, streamsMetrics);
         final String taskId = id.toString();
         if (streamsMetrics.version() == Version.FROM_0100_TO_24) {
@@ -435,7 +432,6 @@ private void close(final boolean clean) {
 
         partitionGroup.close();
         closeTaskSensor.record();
-        streamsMetrics.removeAllTaskLevelSensors(threadId, id.toString());
 
         transitionTo(State.CLOSED);
     }
@@ -692,9 +688,10 @@ private void closeRecordCollector(final boolean clean) {
     @Override
     public void addRecords(final TopicPartition partition, final Iterable<ConsumerRecord<byte[], byte[]>> records) {
         if (state() == State.CLOSED || state() == State.CLOSING) {
-            log.info("Stream task {} is already closed, probably because it got unexpectedly migrated to another thread already. " +
-                         "Notifying the thread to trigger a new rebalance immediately.", id());
-            throw new TaskMigratedException(id());
+            // a task is only closing / closed when 1) task manager is closing, 2) a rebalance is undergoing;
+            // in either case we can just log it and move on without notifying the thread since the consumer
+            // would soon be updated to not return any records for this task anymore.
+            log.info("Stream task {} is already in {} state, skip adding records to it.", id(), state());
         }
 
         final int newQueueSize = partitionGroup.addRawRecords(partition, records);