KAFKA-8179: Part 3, Add PartitionsLost API for resetGenerations and metadata/subscription change

clients/src/main/java/org/apache/kafka/clients/consumer/ConsumerRebalanceListener.java

@@ -95,7 +95,7 @@ public interface ConsumerRebalanceListener {
      * invocation of {@link KafkaConsumer#poll(java.time.Duration)} in which this callback is being executed. This means it is not
      * necessary to catch these exceptions and re-attempt to wakeup or interrupt the consumer thread.
      *
-     * @param partitions The list of partitions that were assigned to the consumer on the last rebalance
+     * @param partitions The list of partitions that were assigned to the consumer and now need to be revoked
      * @throws org.apache.kafka.common.errors.WakeupException If raised from a nested call to {@link KafkaConsumer}
      * @throws org.apache.kafka.common.errors.InterruptException If raised from a nested call to {@link KafkaConsumer}
      */
@@ -122,4 +122,30 @@ public interface ConsumerRebalanceListener {
      * @throws org.apache.kafka.common.errors.InterruptException If raised from a nested call to {@link KafkaConsumer}
      */
     void onPartitionsAssigned(Collection<TopicPartition> partitions);
+
+    /**
+     * A callback method the user can implement to provide handling of cleaning up resources for partitions that have already
+     * been re-assigned to other consumers. This method will not be called during normal execution as the owned partitions would
+     * first be revoked by calling the {@link ConsumerRebalanceListener#onPartitionsRevoked} first, before being re-assigned
+     * to other consumers. However, when the consumer is being kicked out of the group and hence were not aware when the new
+     * group is formed without itself, this function will then be called when the consumer finally be notified about the
+     * partitions that have already been emigrated to other consumers.
+     *
+     * By default it will just trigger {@link ConsumerRebalanceListener#onPartitionsRevoked}; for advanced users who want to distinguish
+     * the handling logic of revoked partitions v.s. emigrated partitions, they can override the default implementation.
+     *
+     * It is possible
+     * for a {@link org.apache.kafka.common.errors.WakeupException} or {@link org.apache.kafka.common.errors.InterruptException}
+     * to be raised from one these nested invocations. In this case, the exception will be propagated to the current
+     * invocation of {@link KafkaConsumer#poll(java.time.Duration)} in which this callback is being executed. This means it is not
+     * necessary to catch these exceptions and re-attempt to wakeup or interrupt the consumer thread.
+     *
+     * @param partitions The list of partitions that were assigned to the consumer and now have been re-assigned
+     *                  to other consumers
+     * @throws org.apache.kafka.common.errors.WakeupException If raised from a nested call to {@link KafkaConsumer}
+     * @throws org.apache.kafka.common.errors.InterruptException If raised from a nested call to {@link KafkaConsumer}
+     */
+    default void onPartitionsLost(Collection<TopicPartition> partitions) {
+        onPartitionsRevoked(partitions);
+    }
 }

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

@@ -521,7 +521,7 @@ public void handle(JoinGroupResponse joinResponse, RequestFuture<ByteBuffer> fut
                 future.raise(error);
             } else if (error == Errors.UNKNOWN_MEMBER_ID) {
                 // reset the member id and retry immediately
-                resetGeneration();
+                resetGeneration(this.getClass().getName() + " encountering " + error);
                 log.debug("Attempt to join group failed due to unknown member id.");
                 future.raise(error);
             } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
@@ -643,7 +643,7 @@ public void handle(SyncGroupResponse syncResponse,
                 } else if (error == Errors.UNKNOWN_MEMBER_ID
                         || error == Errors.ILLEGAL_GENERATION) {
                     log.debug("SyncGroup failed: {}", error.message());
-                    resetGeneration();
+                    resetGeneration(this.getClass().getName() + " encountering " + error);
                     future.raise(error);
                 } else if (error == Errors.COORDINATOR_NOT_AVAILABLE
                         || error == Errors.NOT_COORDINATOR) {
@@ -797,7 +797,9 @@ final synchronized boolean hasValidMemberId() {
     /**
      * Reset the generation and memberId because we have fallen out of the group.
      */
-    protected synchronized void resetGeneration() {
+    protected synchronized void resetGeneration(String errorMessage) {
+        log.debug("Resetting generation and requesting re-join group due to {}", errorMessage);
+
         this.generation = Generation.NO_GENERATION;
         this.rejoinNeeded = true;
         this.state = MemberState.UNJOINED;
@@ -859,7 +861,7 @@ public synchronized void maybeLeaveGroup(String leaveReason) {
             client.pollNoWakeup();
         }
 
-        resetGeneration();
+        resetGeneration("consumer proactively leaving group");
     }
 
     protected boolean isDynamicMember() {
@@ -913,14 +915,14 @@ public void handle(HeartbeatResponse heartbeatResponse, RequestFuture<Void> futu
                 future.raise(error);
             } else if (error == Errors.ILLEGAL_GENERATION) {
                 log.info("Attempt to heartbeat failed since generation {} is not current", generation.generationId);
-                resetGeneration();
+                resetGeneration(this.getClass().getName() + " encountering " + error);
                 future.raise(error);
             } else if (error == Errors.FENCED_INSTANCE_ID) {
                 log.error("Received fatal exception: group.instance.id gets fenced");
                 future.raise(error);
             } else if (error == Errors.UNKNOWN_MEMBER_ID) {
                 log.info("Attempt to heartbeat failed for since member id {} is not valid.", generation.memberId);
-                resetGeneration();
+                resetGeneration(this.getClass().getName() + " encountering " + error);
                 future.raise(error);
             } else if (error == Errors.GROUP_AUTHORIZATION_FAILED) {
                 future.raise(new GroupAuthorizationException(rebalanceConfig.groupId));

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

@@ -70,6 +70,7 @@
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
 /**
@@ -234,7 +235,7 @@ private void maybeUpdateJoinedSubscription(Set<TopicPartition> assignedPartition
         // into the subscriptions as long as they still match the subscribed pattern
 
         Set<String> addedTopics = new HashSet<>();
-        //this is a copy because its handed to listener below
+        // this is a copy because its handed to listener below
         for (TopicPartition tp : assignedPartitions) {
             if (!joinedSubscription.contains(tp.topic()))
                 addedTopics.add(tp.topic());
@@ -252,6 +253,72 @@ private void maybeUpdateJoinedSubscription(Set<TopicPartition> assignedPartition
         }
     }
 
+    private void maybeAssignPartitions(final Set<TopicPartition> assignedPartitions) {
+        if (!assignedPartitions.isEmpty()) {
+            log.info("Setting newly assigned partitions: {}", Utils.join(assignedPartitions, ", "));
+
+            ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
+            try {
+                listener.onPartitionsAssigned(assignedPartitions);
+            } catch (WakeupException | InterruptException e) {
+                throw e;
+            } catch (Exception e) {
+                log.error("User provided listener {} failed on partition assignment", listener.getClass().getName(), e);
+            }
+        }
+    }
+
+    private void maybeRevokePartitions(final Predicate<TopicPartition> predicate) {
+        Set<TopicPartition> ownedPartitions = new HashSet<>(subscriptions.assignedPartitions());
+
+        Set<TopicPartition> revokedPartitions = ownedPartitions.stream().filter(predicate).collect(Collectors.toSet());
+
+        if (!revokedPartitions.isEmpty()) {
+            log.info("Revoke previously assigned partitions {}", revokedPartitions);
+
+            ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
+            try {
+                listener.onPartitionsRevoked(revokedPartitions);
+            } catch (WakeupException | InterruptException e) {
+                throw e;
+            } catch (Exception e) {
+                log.error("User provided listener {} failed on partition being revocation", listener.getClass().getName(), e);
+            }
+
+            Set<TopicPartition> leftPartitions = new HashSet<>(ownedPartitions);
+            leftPartitions.removeAll(revokedPartitions);
+
+            subscriptions.assignFromSubscribed(leftPartitions);
+        }
+    }
+
+    private void maybeLosePartitions(final String rootCause, final Predicate<TopicPartition> predicate) {
+        if (subscriptions.partitionsAutoAssigned()) {
+            Set<TopicPartition> ownedPartitions = new HashSet<>(subscriptions.assignedPartitions());
+
+            Set<TopicPartition> lostPartitions = metadataSnapshot == null ? ownedPartitions :
+                ownedPartitions.stream().filter(predicate).collect(Collectors.toSet());
+
+            if (!lostPartitions.isEmpty()) {
+                log.info("Lost previously assigned partitions {} due to {}", lostPartitions, rootCause);
+
+                ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
+                try {
+                    listener.onPartitionsLost(lostPartitions);
+                } catch (WakeupException | InterruptException e) {
+                    throw e;
+                } catch (Exception e) {
+                    log.error("User provided listener {} failed on partition being lost", listener.getClass().getName(), e);
+                }
+
+                Set<TopicPartition> leftPartitions = new HashSet<>(ownedPartitions);
+                leftPartitions.removeAll(lostPartitions);
+
+                subscriptions.assignFromSubscribed(leftPartitions);
+            }
+        }
+    }
+
     @Override
     protected void onJoinComplete(int generation,
                                   String memberId,
@@ -292,28 +359,38 @@ protected void onJoinComplete(int generation,
             this.nextAutoCommitTimer.updateAndReset(autoCommitIntervalMs);
 
         // execute the user's callback after rebalance
-        ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
-
         switch (protocol) {
             case EAGER:
                 if (!ownedPartitions.isEmpty()) {
                     log.info("Coordinator has owned partitions {} that are not revoked with {} protocol, " +
                         "it is likely client is woken up before a previous pending rebalance completes its callback", ownedPartitions, protocol);
                 }
 
-                log.info("Setting newly assigned partitions: {}", Utils.join(assignedPartitions, ", "));
-                try {
-                    listener.onPartitionsAssigned(assignedPartitions);
-                } catch (WakeupException | InterruptException e) {
-                    throw e;
-                } catch (Exception e) {
-                    log.error("User provided listener {} failed on partition assignment", listener.getClass().getName(), e);
-                }
+                // TODO: after refactored the error handling, we need to exclude owned partitions from assigned partitions
+
+                maybeAssignPartitions(assignedPartitions);
                 break;
 
             case COOPERATIVE:
-                assignAndRevoke(listener, assignedPartitions, ownedPartitions);
+                Set<TopicPartition> addedPartitions = new HashSet<>(assignedPartitions);
+                Set<TopicPartition> revokedPartitions = new HashSet<>(ownedPartitions);
+                addedPartitions.removeAll(ownedPartitions);
+                revokedPartitions.removeAll(assignedPartitions);
 
+                log.info("Updating with newly assigned partitions: {}, compare with already owned partitions: {}, " +
+                        "newly added partitions: {}, revoking partitions: {}",
+                    Utils.join(assignedPartitions, ", "),
+                    Utils.join(ownedPartitions, ", "),
+                    Utils.join(addedPartitions, ", "),
+                    Utils.join(revokedPartitions, ", "));
+
+                // add partitions that was not previously owned but are now assigned
+                maybeAssignPartitions(addedPartitions);
+
+                // revoked partitions that was previously owned but no longer assigned
+                maybeRevokePartitions(tp -> !assignedPartitions.contains(tp));
+
+                // request re-join based on leader communicated error code
                 if (assignment.error() == ConsumerProtocol.AssignmentError.NEED_REJOIN) {
                     requestRejoin();
                 }
@@ -323,38 +400,6 @@ protected void onJoinComplete(int generation,
 
     }
 
-    private void assignAndRevoke(final ConsumerRebalanceListener listener,
-                                 final Set<TopicPartition> assignedPartitions,
-                                 final Set<TopicPartition> ownedPartitions) {
-        Set<TopicPartition> addedPartitions = new HashSet<>(assignedPartitions);
-        Set<TopicPartition> revokedPartitions = new HashSet<>(ownedPartitions);
-        addedPartitions.removeAll(ownedPartitions);
-        revokedPartitions.removeAll(assignedPartitions);
-
-        log.info("Updating with newly assigned partitions: {}, compare with already owned partitions: {}, " +
-                "newly added partitions: {}, revoking partitions: {}",
-            Utils.join(assignedPartitions, ", "),
-            Utils.join(ownedPartitions, ", "),
-            Utils.join(addedPartitions, ", "),
-            Utils.join(revokedPartitions, ", "));
-
-        try {
-            listener.onPartitionsAssigned(addedPartitions);
-        } catch (WakeupException | InterruptException e) {
-            throw e;
-        } catch (Exception e) {
-            log.error("User provided listener {} failed on partition assignment", listener.getClass().getName(), e);
-        }
-
-        try {
-            listener.onPartitionsRevoked(revokedPartitions);
-        } catch (WakeupException | InterruptException e) {
-            throw e;
-        } catch (Exception e) {
-            log.error("User provided listener {} failed on partition revocation", listener.getClass().getName(), e);
-        }
-    }
-
     void maybeUpdateSubscriptionMetadata() {
         int version = metadata.updateVersion();
         if (version > metadataSnapshot.version) {
@@ -587,45 +632,45 @@ protected void onJoinPrepare(int generation, String memberId) {
         maybeAutoCommitOffsetsSync(time.timer(rebalanceConfig.rebalanceTimeoutMs));
 
         // execute the user's callback before rebalance
-        ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
-
         switch (protocol) {
             case EAGER:
-                Set<TopicPartition> revokedPartitions = new HashSet<>(subscriptions.assignedPartitions());
-                log.info("Revoking previously assigned partitions {}", revokedPartitions);
-                try {
-                    listener.onPartitionsRevoked(revokedPartitions);
-                } catch (WakeupException | InterruptException e) {
-                    throw e;
-                } catch (Exception e) {
-                    log.error("User provided listener {} failed on partition revocation", listener.getClass().getName(), e);
-                }
-
-                // also clear the assigned partitions since all have been revoked
-                subscriptions.assignFromSubscribed(Collections.emptySet());
-
+                // revoke all partitions
+                maybeRevokePartitions(tp -> true);
                 break;
 
             case COOPERATIVE:
+                // only revoke those partitions that are not in the subscription any more.
+                maybeRevokePartitions(tp -> !subscriptions.subscription().contains(tp.topic()));
                 break;
         }
 
         isLeader = false;
         subscriptions.resetGroupSubscription();
     }
 
+    @Override
+    public void resetGeneration(String errorMessage) {
+        // revoke all partitions
+        maybeLosePartitions(errorMessage, tp -> true);
+        super.resetGeneration(errorMessage);
+    }
+
     @Override
     public boolean rejoinNeededOrPending() {
         if (!subscriptions.partitionsAutoAssigned())
             return false;
 
         // we need to rejoin if we performed the assignment and metadata has changed
-        if (assignmentSnapshot != null && !assignmentSnapshot.matches(metadataSnapshot))
+        if (assignmentSnapshot != null && !assignmentSnapshot.matches(metadataSnapshot)) {
+            maybeLosePartitions("topic metadata has changed and therefore some topics may not exist any more",
+                tp -> !metadataSnapshot.partitionsPerTopic().containsKey(tp.topic()));
             return true;
+        }
 
         // we need to join if our subscription has changed since the last join
-        if (joinedSubscription != null && !joinedSubscription.equals(subscriptions.subscription()))
+        if (joinedSubscription != null && !joinedSubscription.equals(subscriptions.subscription())) {
             return true;
+        }
 
         return super.rejoinNeededOrPending();
     }
@@ -1033,7 +1078,7 @@ public void handle(OffsetCommitResponse commitResponse, RequestFuture<Void> futu
                         } else if (error == Errors.UNKNOWN_MEMBER_ID
                                 || error == Errors.ILLEGAL_GENERATION) {
                             // need to reset generation and re-join group
-                            resetGeneration();
+                            resetGeneration(this.getClass().getName() + " encountering " + error);
                             future.raise(new CommitFailedException());
                             return;
                         } else {
@@ -1173,6 +1218,10 @@ private MetadataSnapshot(SubscriptionState subscription, Cluster cluster, int ve
         boolean matches(MetadataSnapshot other) {
             return version == other.version || partitionsPerTopic.equals(other.partitionsPerTopic);
         }
+
+        Map<String, Integer> partitionsPerTopic() {
+            return partitionsPerTopic;
+        }
     }
 
     private static class OffsetCommitCompletion {