KAFKA-12523: handle TaskCorruption/TimeoutException during handleCorruption and handleRevocation

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

@@ -589,11 +589,7 @@ boolean runLoop() {
             } catch (final TaskCorruptedException e) {
                 log.warn("Detected the states of tasks " + e.corruptedTasks() + " are corrupted. " +
                         "Will close the task as dirty and re-create and bootstrap from scratch.", e);
-                try {
-                    taskManager.handleCorruption(e.corruptedTasks());
-                } catch (final TaskMigratedException taskMigrated) {
-                    handleTaskMigrated(taskMigrated);
-                }
+                taskManager.handleCorruption(e.corruptedTasks());
             } catch (final TaskMigratedException e) {
                 handleTaskMigrated(e);
             } catch (final UnsupportedVersionException e) {
@@ -611,6 +607,9 @@ boolean runLoop() {
                 if (processingMode == ProcessingMode.EXACTLY_ONCE_ALPHA || processingMode == ProcessingMode.EXACTLY_ONCE_BETA) {
                     return false;
                 }
+                // TODO: extract this catch block to encompass all of runLoop, in case exceptions slip through one of the
+                // handling methods in the other catch blocks. Note this we'll need to extract the UnsupportedVersionException
+                // as well since it invokes the exception handler, and the old-style handler relies on rethrowing exceptions
             } catch (final Throwable e) {
                 failedStreamThreadSensor.record();
                 this.streamsUncaughtExceptionHandler.accept(e);

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

@@ -167,18 +167,7 @@ void handleCorruption(final Set<TaskId> corruptedTasks) {
             }
         }
 
-        // Make sure to clean up any corrupted standby tasks in their entirety before committing
-        // since TaskMigrated can be thrown and the resulting handleLostAll will only clean up active tasks
         closeAndRevive(corruptedStandbyTasks);
-
-        commit(tasks()
-                   .values()
-                   .stream()
-                   .filter(t -> t.state() == Task.State.RUNNING || t.state() == Task.State.RESTORING)
-                   .filter(t -> !corruptedTasks.contains(t.id()))
-                   .collect(Collectors.toSet())
-        );
-
         closeAndRevive(corruptedActiveTasks);
     }
 
@@ -332,8 +321,8 @@ private void handleCloseAndRecycle(final Set<Task> tasksToRecycle,
                 //    write the checkpoint file.
                 final Map<TopicPartition, OffsetAndMetadata> offsets = task.prepareCommit();
                 if (!offsets.isEmpty()) {
-                    log.error("Task {} should has been committed when it was suspended, but it reports non-empty " +
-                                    "offsets {} to commit; it means it fails during last commit and hence should be closed dirty",
+                    log.error("Task {} should have been committed when it was suspended, but it reports non-empty " +
+                                    "offsets {} to commit; this means it failed during last commit and hence should be closed dirty",
                             task.id(), offsets);
 
                     tasksToCloseDirty.add(task);
@@ -515,7 +504,43 @@ void handleRevocation(final Collection<TopicPartition> revokedPartitions) {
             commitOffsetsOrTransaction(consumedOffsetsPerTask);
         } catch (final RuntimeException e) {
             log.error("Exception caught while committing those revoked tasks " + revokedActiveTasks, e);
-            firstException.compareAndSet(null, e);
+
+            // if we hit a TaskCorruptedException, we have to peel off any revoked tasks since they will be gone by
+            // the time the exception is bubbled up through poll()
+            if (e instanceof TaskCorruptedException) {
+                final Set<TaskId> corruptedTasks = ((TaskCorruptedException) e).corruptedTasks();
+                final Set<TaskId> stillAssignedCorruptedTasks = new HashSet<>();
+                final Set<TaskId> revokedCorruptedTasks = new HashSet<>();
+
+                for (final TaskId taskId : corruptedTasks) {
+                    final Task task = tasks.task(taskId);
+                    if (!revokedActiveTasks.contains(task)) {
+                        stillAssignedCorruptedTasks.add(taskId);
+                    } else {
+                        revokedCorruptedTasks.add(taskId);
+                    }
+                }
+
+                // If any of the corrupted tasks are being revoked, we need to filter them from the
+                // TaskCorruptedException otherwise we will try to revive tasks we no longer own
+                // Any revoked tasks will be detected and closed dirty in handleAssignment and wipe their state
+                if (!stillAssignedCorruptedTasks.isEmpty()) {
+                    log.info("Removing revoked tasks {} from the corrupted list, the remaining tasks will be revived " +
+                        "after the rebalance has completed.", revokedCorruptedTasks);
+
+                    // TODO: usually we would always save the first seen exception to rethrow, but TaskCorruptedException
+                    // must be bubbled up to the StreamThread to ensure corrupted state is cleared and revived
+                    firstException.set(new TaskCorruptedException(stillAssignedCorruptedTasks));
+                } else {
+                    // If all corrupted tasks are being revoked, we can just close them dirty in handleAssignment
+                    // and can swallow the exception
+                    log.info("All the corrupted tasks are being revoked so the exception can be swallowed: {}",
+                             revokedCorruptedTasks);
+                }
+            } else {
+                // TODO: KIP-572 need to handle TimeoutException, may be rethrown from committing offsets under ALOS
+                firstException.compareAndSet(null, e);
+            }
         }
 
         // only try to complete post-commit if committing succeeded;

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

@@ -2333,74 +2333,6 @@ void runOnce() {
         verify(taskManager);
     }
 
-    @Test
-    @SuppressWarnings("unchecked")
-    public void shouldCatchTaskMigratedExceptionOnOnTaskCorruptedExceptionPath() {
-        final TaskManager taskManager = EasyMock.createNiceMock(TaskManager.class);
-        final Consumer<byte[], byte[]> consumer = mock(Consumer.class);
-        final ConsumerGroupMetadata consumerGroupMetadata = mock(ConsumerGroupMetadata.class);
-        expect(consumer.groupMetadata()).andStubReturn(consumerGroupMetadata);
-        expect(consumerGroupMetadata.groupInstanceId()).andReturn(Optional.empty());
-        consumer.subscribe((Collection<String>) anyObject(), anyObject());
-        EasyMock.expectLastCall().anyTimes();
-        consumer.unsubscribe();
-        EasyMock.expectLastCall().anyTimes();
-        EasyMock.replay(consumerGroupMetadata);
-        final Task task1 = mock(Task.class);
-        final Task task2 = mock(Task.class);
-        final TaskId taskId1 = new TaskId(0, 0);
-        final TaskId taskId2 = new TaskId(0, 2);
-
-        final Set<TaskId> corruptedTasks = singleton(taskId1);
-
-        expect(task1.state()).andReturn(Task.State.RUNNING).anyTimes();
-        expect(task1.id()).andReturn(taskId1).anyTimes();
-        expect(task2.state()).andReturn(Task.State.RUNNING).anyTimes();
-        expect(task2.id()).andReturn(taskId2).anyTimes();
-
-        taskManager.handleCorruption(corruptedTasks);
-        expectLastCall().andThrow(new TaskMigratedException("Task migrated",
-                                                            new RuntimeException("non-corrupted task migrated")));
-
-        taskManager.handleLostAll();
-        expectLastCall();
-
-        EasyMock.replay(task1, task2, taskManager, consumer);
-
-        final StreamsMetricsImpl streamsMetrics =
-            new StreamsMetricsImpl(metrics, CLIENT_ID, StreamsConfig.METRICS_LATEST, mockTime);
-        final StreamThread thread = new StreamThread(
-            mockTime,
-            config,
-            null,
-            consumer,
-            consumer,
-            null,
-            null,
-            taskManager,
-            streamsMetrics,
-            internalTopologyBuilder,
-            CLIENT_ID,
-            new LogContext(""),
-            new AtomicInteger(),
-            new AtomicLong(Long.MAX_VALUE),
-            null,
-            HANDLER,
-            null
-        ) {
-            @Override
-            void runOnce() {
-                setState(State.PENDING_SHUTDOWN);
-                throw new TaskCorruptedException(corruptedTasks);
-            }
-        }.updateThreadMetadata(getSharedAdminClientId(CLIENT_ID));
-
-        thread.setState(StreamThread.State.STARTING);
-        thread.runLoop();
-
-        verify(taskManager);
-    }
-
     @Test
     public void shouldNotCommitNonRunningNonRestoringTasks() {
         final TaskManager taskManager = EasyMock.createNiceMock(TaskManager.class);

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

@@ -690,7 +690,6 @@ public void suspend() {
 
         task00.setChangelogOffsets(singletonMap(t1p0, 0L));
         taskManager.handleCorruption(singleton(taskId00));
-        assertThat(task00.commitPrepared, is(true));
         assertThat(task00.state(), is(Task.State.CREATED));
         assertThat(task00.partitionsForOffsetReset, equalTo(taskId00Partitions));
         assertThat(taskManager.activeTaskMap(), is(singletonMap(taskId00, task00)));
@@ -718,7 +717,6 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() {
         topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString());
         expectLastCall().anyTimes();
         expectRestoreToBeCompleted(consumer, changeLogReader);
-        consumer.commitSync(eq(emptyMap()));
         expect(consumer.assignment()).andReturn(taskId00Partitions);
         replay(activeTaskCreator, topologyBuilder, consumer, changeLogReader);
 
@@ -731,7 +729,6 @@ public void shouldCommitNonCorruptedTasksOnTaskCorruptedException() {
         corruptedTask.setChangelogOffsets(singletonMap(t1p0, 0L));
         taskManager.handleCorruption(singleton(taskId00));
 
-        assertTrue(nonCorruptedTask.commitPrepared);
         assertThat(nonCorruptedTask.partitionsForOffsetReset, equalTo(Collections.emptySet()));
         assertThat(corruptedTask.partitionsForOffsetReset, equalTo(taskId00Partitions));
 
@@ -774,48 +771,6 @@ public void shouldNotCommitNonRunningNonCorruptedTasks() {
         verify(consumer);
     }
 
-    @Test
-    public void shouldCleanAndReviveCorruptedStandbyTasksBeforeCommittingNonCorruptedTasks() {
-        final ProcessorStateManager stateManager = EasyMock.createStrictMock(ProcessorStateManager.class);
-        stateManager.markChangelogAsCorrupted(taskId00Partitions);
-        replay(stateManager);
-
-        final StateMachineTask corruptedStandby = new StateMachineTask(taskId00, taskId00Partitions, false, stateManager);
-        final StateMachineTask runningNonCorruptedActive = new StateMachineTask(taskId01, taskId01Partitions, true, stateManager) {
-            @Override
-            public Map<TopicPartition, OffsetAndMetadata> prepareCommit() {
-                throw new TaskMigratedException("You dropped out of the group!", new RuntimeException());
-            }
-        };
-
-        // handleAssignment
-        expect(standbyTaskCreator.createTasks(eq(taskId00Assignment))).andStubReturn(singleton(corruptedStandby));
-        expect(activeTaskCreator.createTasks(anyObject(), eq(taskId01Assignment))).andStubReturn(singleton(runningNonCorruptedActive));
-        topologyBuilder.addSubscribedTopicsFromAssignment(anyObject(), anyString());
-        expectLastCall().anyTimes();
-
-        expectRestoreToBeCompleted(consumer, changeLogReader);
-
-        replay(activeTaskCreator, standbyTaskCreator, topologyBuilder, consumer, changeLogReader);
-
-        taskManager.handleAssignment(taskId01Assignment, taskId00Assignment);
-        assertThat(taskManager.tryToCompleteRestoration(time.milliseconds(), null), is(true));
-
-        // make sure this will be committed and throw
-        assertThat(runningNonCorruptedActive.state(), is(Task.State.RUNNING));
-        assertThat(corruptedStandby.state(), is(Task.State.RUNNING));
-
-        runningNonCorruptedActive.setCommitNeeded();
-
-        corruptedStandby.setChangelogOffsets(singletonMap(t1p0, 0L));
-        assertThrows(TaskMigratedException.class, () -> taskManager.handleCorruption(singleton(taskId00)));
-
-
-        assertThat(corruptedStandby.commitPrepared, is(true));
-        assertThat(corruptedStandby.state(), is(Task.State.CREATED));
-        verify(consumer);
-    }
-
     @Test
     public void shouldCloseStandbyUnassignedTasksWhenCreatingNewTasks() {
         final Task task00 = new StateMachineTask(taskId00, taskId00Partitions, false);