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

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

@@ -171,13 +171,23 @@ void handleCorruption(final Set<TaskId> corruptedTasks) {
         // 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())
-        );
+        // We need to commit before closing the corrupted active tasks since this will force the ongoing txn to abort
+        try {
+            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())
+            );
+        } catch (final TaskCorruptedException e) {
+            log.info("Some additional tasks were found corrupted while trying to commit, these will be added to the " +
+                         "tasks to clean and revive: {}", e.corruptedTasks());
+            for (final TaskId taskId : e.corruptedTasks()) {
+                final Task task = tasks.task(taskId);
+                corruptedActiveTasks.put(task, task.changelogPartitions());
+            }
+        }
 
         closeAndRevive(corruptedActiveTasks);
     }
@@ -332,8 +342,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);
@@ -509,41 +519,62 @@ void handleRevocation(final Collection<TopicPartition> revokedPartitions) {
             prepareCommitAndAddOffsetsToMap(commitNeededActiveTasks, consumedOffsetsPerTask);
         }
 
-        // even if commit failed, we should still continue and complete suspending those tasks,
-        // so we would capture any exception and throw
+        // even if commit failed, we should still continue and complete suspending those tasks, so we would capture
+        // any exception and rethrow it at the end
+        final Set<TaskId> corruptedTasks = new HashSet<>();
         try {
             commitOffsetsOrTransaction(consumedOffsetsPerTask);
         } catch (final RuntimeException e) {
             log.error("Exception caught while committing those revoked tasks " + revokedActiveTasks, e);
-            firstException.compareAndSet(null, e);
-        }
 
-        // only try to complete post-commit if committing succeeded;
-        // we enforce checkpointing upon suspending a task: if it is resumed later we just
-        // proceed normally, if it is going to be closed we would checkpoint by then
-        if (firstException.get() == null) {
-            for (final Task task : revokedActiveTasks) {
-                try {
-                    task.postCommit(true);
-                } catch (final RuntimeException e) {
-                    log.error("Exception caught while post-committing task " + task.id(), e);
-                    firstException.compareAndSet(null, e);
+            // If we hit a TaskCorruptedException, we should just handle the cleanup for those corrupted tasks right here
+            if (e instanceof TaskCorruptedException) {
+                corruptedTasks.addAll(((TaskCorruptedException) e).corruptedTasks());
+                final Map<Task, Collection<TopicPartition>> corruptedTasksWithChangelogs = new HashMap<>();
+                for (final TaskId taskId : corruptedTasks) {
+                    final Task task = tasks.task(taskId);
+                    task.markChangelogAsCorrupted(task.changelogPartitions());
+                    corruptedTasksWithChangelogs.put(task, task.changelogPartitions());
                 }
+                closeAndRevive(corruptedTasksWithChangelogs);
+            } else {
+                // TODO: KIP-572 need to handle TimeoutException, may be rethrown from committing offsets under ALOS
+                // currently we just let the thread die but we can probably just close those tasks as dirty and proceed
+                firstException.compareAndSet(null, e);
             }
+        }
 
-            if (shouldCommitAdditionalTasks) {
-                for (final Task task : commitNeededActiveTasks) {
+        // only try to complete post-commit if committing succeeded, or if we hit a TaskCorruptedException then we
+        // can still checkpoint the uncorrupted tasks (if any)
+        // we enforce checkpointing upon suspending a task: if it is resumed later we just proceed normally, if it is
+        // going to be closed we would checkpoint by then
+        if (firstException.get() == null || !corruptedTasks.isEmpty()) {
+            for (final Task task : revokedActiveTasks) {
+                if (!corruptedTasks.contains(task.id())) {
                     try {
-                        // for non-revoking active tasks, we should not enforce checkpoint
-                        // since if it is EOS enabled, no checkpoint should be written while
-                        // the task is in RUNNING tate
-                        task.postCommit(false);
+                        task.postCommit(true);
                     } catch (final RuntimeException e) {
                         log.error("Exception caught while post-committing task " + task.id(), e);
                         firstException.compareAndSet(null, e);
                     }
                 }
             }
+
+            if (shouldCommitAdditionalTasks) {
+                for (final Task task : commitNeededActiveTasks) {
+                    if (!corruptedTasks.contains(task.id())) {
+                        try {
+                            // for non-revoking active tasks, we should not enforce checkpoint
+                            // since if it is EOS enabled, no checkpoint should be written while
+                            // the task is in RUNNING tate
+                            task.postCommit(false);
+                        } catch (final RuntimeException e) {
+                            log.error("Exception caught while post-committing task " + task.id(), e);
+                            firstException.compareAndSet(null, e);
+                        }
+                    }
+                }
+            }
         }
 
         for (final Task task : revokedActiveTasks) {

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));