KAFKA-9620: Do not throw in the middle of consumer user callbacks

clients/src/main/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinator.java

@@ -333,7 +333,7 @@ protected void onJoinComplete(int generation,
                                   ByteBuffer assignmentBuffer) {
         log.debug("Executing onJoinComplete with generation {} and memberId {}", generation, memberId);
 
-        // only the leader is responsible for monitoring for metadata changes (i.e. partition changes)
+        // Only the leader is responsible for monitoring for metadata changes (i.e. partition changes)
         if (!isLeader)
             assignmentSnapshot = null;
 
@@ -377,12 +377,12 @@ protected void onJoinComplete(int generation,
             );
 
             if (!revokedPartitions.isEmpty()) {
-                // revoke partitions that were previously owned but no longer assigned;
+                // Revoke partitions that were previously owned but no longer assigned;
                 // note that we should only change the assignment (or update the assignor's state)
                 // AFTER we've triggered  the revoke callback
                 firstException.compareAndSet(null, invokePartitionsRevoked(revokedPartitions));
 
-                // if revoked any partitions, need to re-join the group afterwards
+                // If revoked any partitions, need to re-join the group afterwards
                 log.debug("Need to revoke partitions {} and re-join the group", revokedPartitions);
                 requestRejoin();
             }
@@ -392,21 +392,32 @@ protected void onJoinComplete(int generation,
         // were not explicitly requested, so we update the joined subscription here.
         maybeUpdateJoinedSubscription(assignedPartitions);
 
-        // give the assignor a chance to update internal state based on the received assignment
+        // Give the assignor a chance to update internal state based on the received assignment
         groupMetadata = new ConsumerGroupMetadata(rebalanceConfig.groupId, generation, memberId, rebalanceConfig.groupInstanceId);
-        assignor.onAssignment(assignment, groupMetadata);
 
-        // reschedule the auto commit starting from now
+        // Catch any exception here to make sure we could complete the user callback.
+        try {
+            assignor.onAssignment(assignment, groupMetadata);
+        } catch (Exception e) {
+            firstException.compareAndSet(null, e);
+        }
+
+        // Reschedule the auto commit starting from now
         if (autoCommitEnabled)
             this.nextAutoCommitTimer.updateAndReset(autoCommitIntervalMs);
 
         subscriptions.assignFromSubscribed(assignedPartitions);
 
-        // add partitions that were not previously owned but are now assigned
+        // Add partitions that were not previously owned but are now assigned
         firstException.compareAndSet(null, invokePartitionsAssigned(addedPartitions));
 
-        if (firstException.get() != null)
-            throw new KafkaException("User rebalance callback throws an error", firstException.get());
+        if (firstException.get() != null) {
+            if (firstException.get() instanceof KafkaException) {
+                throw (KafkaException) firstException.get();
+            } else {
+                throw new KafkaException("User rebalance callback throws an error", firstException.get());
+            }
+        }
     }
 
     void maybeUpdateSubscriptionMetadata() {

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

@@ -84,7 +84,7 @@ static void wipeStateStores(final Logger log, final ProcessorStateManager stateM
             Utils.delete(stateMgr.baseDir());
         } catch (final IOException fatalException) {
             // since it is only called under dirty close, we always swallow the exception
-            log.warn("Failed to wiping state stores for task {}", stateMgr.taskId());
+            log.warn("Failed to wiping state stores for task {} due to {}", stateMgr.taskId(), fatalException);
         }
     }
 

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

@@ -198,12 +198,19 @@ public void handleAssignment(final Map<TaskId, Set<TopicPartition>> activeTasks,
         }
 
         if (!taskCloseExceptions.isEmpty()) {
+            for (final Map.Entry<TaskId, RuntimeException> entry : taskCloseExceptions.entrySet()) {
+                if (!(entry.getValue() instanceof TaskMigratedException)) {
+                    throw new RuntimeException(
+                        "Unexpected failure to close " + taskCloseExceptions.size() +
+                            " task(s) [" + taskCloseExceptions.keySet() + "]. " +
+                            "First unexpected exception (for task " + entry.getKey() + ") follows.", entry.getValue()
+                    );
+                }
+            }
+
             final Map.Entry<TaskId, RuntimeException> first = taskCloseExceptions.entrySet().iterator().next();
-            throw new RuntimeException(
-                "Unexpected failure to close " + taskCloseExceptions.size() +
-                    " task(s) [" + taskCloseExceptions.keySet() + "]. " +
-                    "First exception (for task " + first.getKey() + ") follows.", first.getValue()
-            );
+            // If all exceptions are task-migrated, we would just throw the first one.
+            throw first.getValue();
         }
 
         if (!activeTasksToCreate.isEmpty()) {
@@ -293,12 +300,11 @@ boolean tryToCompleteRestoration() {
      * @throws TaskMigratedException if the task producer got fenced (EOS only)
      */
     void handleRevocation(final Collection<TopicPartition> revokedPartitions) {
-        final Set<TaskId> revokedTasks = new HashSet<>();
         final Set<TopicPartition> remainingPartitions = new HashSet<>(revokedPartitions);
 
         for (final Task task : tasks.values()) {
             if (remainingPartitions.containsAll(task.inputPartitions())) {
-                revokedTasks.add(task.id());
+                task.suspend();
             }
             remainingPartitions.removeAll(task.inputPartitions());
         }
@@ -308,11 +314,6 @@ void handleRevocation(final Collection<TopicPartition> revokedPartitions) {
                 "due to race condition of consumer detecting the heartbeat failure, or the tasks " +
                 "have been cleaned up by the handleAssignment callback.", remainingPartitions);
         }
-
-        for (final TaskId taskId : revokedTasks) {
-            final Task task = tasks.get(taskId);
-            task.suspend();
-        }
     }
 
     /**