KAFKA-9231: Streams Threads may die from recoverable errors with EOS enabled

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

@@ -23,6 +23,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.ProcessorContext;
 import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.processor.TaskId;
@@ -172,7 +173,13 @@ public String toString(final String indent) {
      * Flush all state stores owned by this task
      */
     void flushState() {
-        stateMgr.flush();
+        try {
+            stateMgr.flush();
+        } catch (final ProcessorStateException e) {
+            if (e.getCause() instanceof RecoverableClientException) {
+                throw new TaskMigratedException(this, e);
+            }
+        }
     }
 
     /**

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

@@ -148,7 +148,7 @@ private RuntimeException suspendRunningTasks(final Set<TaskId> runningTasksToSus
                 // swallow and move on since we are rebalancing
                 log.info("Failed to suspend stream task {} since it got migrated to another thread already. " +
                     "Closing it as zombie and moving on.", id);
-                firstException.compareAndSet(null, closeZombieTask(task));
+                tryCloseZombieTask(task);
                 prevActiveTasks.remove(id);
             } catch (final RuntimeException e) {
                 log.error("Suspending stream task {} failed due to the following error:", id, e);

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

@@ -97,14 +97,12 @@ Collection<T> running() {
         return running.values();
     }
 
-    RuntimeException closeZombieTask(final T task) {
+    void tryCloseZombieTask(final T task) {
         try {
             task.close(false, true);
         } catch (final RuntimeException e) {
             log.warn("Failed to close zombie {} {} due to {}; ignore and proceed.", taskTypeName, task.id(), e.toString());
-            return e;
         }
-        return null;
     }
 
     boolean hasRunningTasks() {
@@ -259,7 +257,7 @@ void shutdown(final boolean clean) {
             } catch (final TaskMigratedException e) {
                 log.info("Failed to close {} {} since it got migrated to another thread already. " +
                     "Closing it as zombie and move on.", taskTypeName, task.id());
-                firstException.compareAndSet(null, closeZombieTask(task));
+                tryCloseZombieTask(task);
             } catch (final RuntimeException t) {
                 log.error("Failed while closing {} {} due to the following error:",
                     task.getClass().getSimpleName(),

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

@@ -101,6 +101,9 @@ public void update(final ConsumerRecord<byte[], byte[]> record) {
     }
 
     public void flushState() {
+        // this could theoretically throw a ProcessorStateException caused by a ProducerFencedException,
+        // but in practice this shouldn't happen for global state update tasks, since the stores are not
+        // logged and there are no downstream operators after global stores.
         stateMgr.flush();
         stateMgr.checkpoint(offsets);
     }

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

@@ -33,6 +33,7 @@
 import org.apache.kafka.common.errors.SecurityDisabledException;
 import org.apache.kafka.common.errors.SerializationException;
 import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.errors.UnknownProducerIdException;
 import org.apache.kafka.common.errors.UnknownServerException;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.header.Headers;
@@ -179,21 +180,22 @@ public void onCompletion(final RecordMetadata metadata,
                         offsets.put(tp, metadata.offset());
                     } else {
                         if (sendException == null) {
-                            if (exception instanceof ProducerFencedException) {
+                            if (exception instanceof ProducerFencedException || exception instanceof UnknownProducerIdException) {
                                 log.warn(LOG_MESSAGE, topic, exception.getMessage(), exception);
 
                                 // KAFKA-7510 put message key and value in TRACE level log so we don't leak data by default
                                 log.trace("Failed message: (key {} value {} timestamp {}) topic=[{}] partition=[{}]", key, value, timestamp, topic, partition);
 
-                                sendException = new ProducerFencedException(
+                                sendException = new RecoverableClientException(
                                     String.format(
                                         EXCEPTION_MESSAGE,
                                         logPrefix,
                                         "producer got fenced",
                                         timestamp,
                                         topic,
                                         exception.toString()
-                                    )
+                                    ),
+                                    exception
                                 );
                             } else {
                                 if (productionExceptionIsFatal(exception)) {
@@ -233,13 +235,13 @@ public void onCompletion(final RecordMetadata metadata,
                 String.format("%sFailed to send record to topic %s due to timeout.", logPrefix, topic),
                 e
             );
-        } catch (final Exception uncaughtException) {
+        } catch (final RuntimeException uncaughtException) {
             if (uncaughtException instanceof KafkaException &&
                 uncaughtException.getCause() instanceof ProducerFencedException) {
                 final KafkaException kafkaException = (KafkaException) 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 (ProducerFencedException) kafkaException.getCause();
+                throw new RecoverableClientException("Caught a wrapped ProducerFencedException", kafkaException);
             } else {
                 throw new StreamsException(
                     String.format(
@@ -264,15 +266,23 @@ private void checkForException() {
     @Override
     public void flush() {
         log.debug("Flushing producer");
-        producer.flush();
+        try {
+            producer.flush();
+        } catch (final ProducerFencedException | UnknownProducerIdException e) {
+            throw new RecoverableClientException("Caught a recoverable exception while flushing", e);
+        }
         checkForException();
     }
 
     @Override
     public void close() {
         log.debug("Closing producer");
         if (producer != null) {
-            producer.close();
+            try {
+                producer.close();
+            } catch (final ProducerFencedException | UnknownProducerIdException e) {
+                throw new RecoverableClientException("Caught a recoverable exception while closing", e);
+            }
             producer = null;
         }
         checkForException();

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

@@ -0,0 +1,33 @@
+/*
+ * 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.processor.internals;
+
+import org.apache.kafka.common.KafkaException;
+
+/**
+ * Denotes an exception that is recoverable by re-creating the client (ie, the client is no longer in a valid state),
+ * as opposed to retriable (the failure was transient, so the same client can be used again later),
+ * or fatal (the request was actually invalid, so retrying or recovering would not help)
+ *
+ * This class also serves the dual purpose of capturing the stack trace as early as possible,
+ * at the site of the Producer call, since the exeptions that cause this don't record stack traces.
+ */
+public class RecoverableClientException extends KafkaException {
+    public RecoverableClientException(final String message, final Throwable cause) {
+        super(message, cause);
+    }
+}

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

@@ -122,6 +122,9 @@ public void commit() {
     }
 
     private void flushAndCheckpointState() {
+        // this could theoretically throw a ProcessorStateException caused by a ProducerFencedException,
+        // but in practice this shouldn't happen for standby tasks, since they don't produce to changelog topics
+        // or downstream topics.
         stateMgr.flush();
         stateMgr.checkpoint(Collections.emptyMap());
     }

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

@@ -27,6 +27,7 @@
 import org.apache.kafka.common.errors.AuthorizationException;
 import org.apache.kafka.common.errors.ProducerFencedException;
 import org.apache.kafka.common.errors.TimeoutException;
+import org.apache.kafka.common.errors.UnknownProducerIdException;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.metrics.Sensor;
 import org.apache.kafka.common.metrics.stats.Avg;
@@ -337,8 +338,8 @@ public void initializeTopology() {
         if (eosEnabled) {
             try {
                 this.producer.beginTransaction();
-            } catch (final ProducerFencedException fatal) {
-                throw new TaskMigratedException(this, fatal);
+            } catch (final ProducerFencedException | UnknownProducerIdException e) {
+                throw new TaskMigratedException(this, e);
             }
             transactionInFlight = true;
         }
@@ -438,8 +439,8 @@ public boolean process() {
             if (recordInfo.queue().size() == maxBufferedSize) {
                 consumer.resume(singleton(partition));
             }
-        } catch (final ProducerFencedException fatal) {
-            throw new TaskMigratedException(this, fatal);
+        } catch (final RecoverableClientException e) {
+            throw new TaskMigratedException(this, e);
         } catch (final KafkaException e) {
             final String stackTrace = getStacktraceString(e);
             throw new StreamsException(format("Exception caught in process. taskId=%s, " +
@@ -488,8 +489,8 @@ public void punctuate(final ProcessorNode node, final long timestamp, final Punc
 
         try {
             node.punctuate(timestamp, punctuator);
-        } catch (final ProducerFencedException fatal) {
-            throw new TaskMigratedException(this, fatal);
+        } catch (final RecoverableClientException e) {
+            throw new TaskMigratedException(this, e);
         } catch (final KafkaException e) {
             throw new StreamsException(String.format("%sException caught while punctuating processor '%s'", logPrefix, node.name()), e);
         } finally {
@@ -558,7 +559,7 @@ void commit(final boolean startNewTransaction, final Map<TopicPartition, Long> p
             } else {
                 consumer.commitSync(consumedOffsetsAndMetadata);
             }
-        } catch (final CommitFailedException | ProducerFencedException error) {
+        } catch (final CommitFailedException | ProducerFencedException | UnknownProducerIdException error) {
             throw new TaskMigratedException(this, error);
         }
 
@@ -581,8 +582,8 @@ protected void flushState() {
         super.flushState();
         try {
             recordCollector.flush();
-        } catch (final ProducerFencedException fatal) {
-            throw new TaskMigratedException(this, fatal);
+        } catch (final RecoverableClientException e) {
+            throw new TaskMigratedException(this, e);
         }
     }
 
@@ -664,7 +665,7 @@ void suspend(final boolean clean,
 
                     try {
                         recordCollector.close();
-                    } catch (final ProducerFencedException e) {
+                    } catch (final RecoverableClientException e) {
                         taskMigratedException = new TaskMigratedException(this, e);
                     } finally {
                         producer = null;

streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java

@@ -188,10 +188,11 @@ void openDB(final ProcessorContext context) {
             throw new ProcessorStateException(fatal);
         }
 
-        maybeSetUpMetricsRecorder(context, configs);
-
         openRocksDB(dbOptions, columnFamilyOptions);
         open = true;
+
+        // Do this last because the prior operations could throw exceptions.
+        maybeSetUpMetricsRecorder(context, configs);
     }
 
     private void maybeSetUpMetricsRecorder(final ProcessorContext context, final Map<String, Object> configs) {

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

@@ -37,15 +37,16 @@
 import org.easymock.EasyMock;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.function.ThrowingRunnable;
 
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
-import org.junit.function.ThrowingRunnable;
 
+import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.CoreMatchers.nullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
@@ -235,14 +236,45 @@ public void shouldCloseTaskOnSuspendIfTaskMigratedException() {
         t1.suspend();
         EasyMock.expectLastCall().andThrow(new TaskMigratedException());
         t1.close(false, true);
-        EasyMock.expectLastCall();
+        EasyMock.expectLastCall().andThrow(new RuntimeException("any exception"));
         EasyMock.replay(t1);
 
         assertThat(suspendTask(), nullValue());
         assertTrue(assignedTasks.runningTaskIds().isEmpty());
         EasyMock.verify(t1);
     }
 
+    @Test
+    public void shouldCloseUncleanAndThenRethrowOnShutdownIfRuntimeException() {
+        mockTaskInitialization();
+
+        t1.close(true, false);
+        EasyMock.expectLastCall().andThrow(new RuntimeException("any first exception"));
+        t1.close(false, false);
+        EasyMock.expectLastCall().andThrow(new RuntimeException("any second exception"));
+        EasyMock.replay(t1);
+        addAndInitTask();
+        try {
+            assignedTasks.shutdown(true);
+            fail("expected a runtime exception");
+        } catch (final RuntimeException e) {
+            assertThat(e.getMessage(), is("any first exception"));
+        }
+    }
+
+    @Test
+    public void shouldCloseWithoutExceptionOnShutdownIfTaskMigratedException() {
+        mockTaskInitialization();
+
+        t1.close(true, false);
+        EasyMock.expectLastCall().andThrow(new TaskMigratedException());
+        t1.close(false, true);
+        EasyMock.expectLastCall().andThrow(new RuntimeException("any second exception"));
+        EasyMock.replay(t1);
+        addAndInitTask();
+        assignedTasks.shutdown(true);
+    }
+
     @Test
     public void shouldResumeMatchingSuspendedTasks() {
         mockRunningTaskSuspension();
@@ -603,7 +635,7 @@ public void action(final StreamTask task) {
             public Set<TaskId> taskIds() {
                 return assignedTasks.suspendedTaskIds();
             }
-            
+
         }.createTaskAndClear();
     }
 

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

@@ -27,6 +27,7 @@
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.PartitionInfo;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.errors.ProducerFencedException;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.Headers;
 import org.apache.kafka.common.header.internals.RecordHeader;
@@ -191,6 +192,26 @@ public synchronized Future<RecordMetadata> send(final ProducerRecord record, fin
         collector.send("topic1", "3", "0", null, null, stringSerializer, stringSerializer, streamPartitioner);
     }
 
+    @SuppressWarnings("unchecked")
+    @Test(expected = RecoverableClientException.class)
+    public void shouldThrowRecoverableExceptionWhenProducerFencedInCallback() {
+        final RecordCollector collector = new RecordCollectorImpl(
+            "test",
+            logContext,
+            new DefaultProductionExceptionHandler(),
+            new Metrics().sensor("skipped-records"));
+        collector.init(new MockProducer(cluster, true, new DefaultPartitioner(), byteArraySerializer, byteArraySerializer) {
+            @Override
+            public synchronized Future<RecordMetadata> send(final ProducerRecord record, final Callback callback) {
+                callback.onCompletion(null, new ProducerFencedException("asdf"));
+                return null;
+            }
+        });
+
+        collector.send("topic1", "3", "0", null, null, stringSerializer, stringSerializer, streamPartitioner);
+        collector.send("topic1", "3", "0", null, null, stringSerializer, stringSerializer, streamPartitioner);
+    }
+
     @SuppressWarnings("unchecked")
     @Test
     public void shouldThrowStreamsExceptionOnSubsequentCallIfASendFailsWithDefaultExceptionHandler() {