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

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

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import java.time.Duration;
 import java.util.Collection;
 
 import org.apache.kafka.common.TopicPartition;
@@ -118,7 +119,9 @@ public interface ConsumerRebalanceListener {
     /**
      * A callback method the user can implement to provide handling of offset commits to a customized store on the start
      * of a rebalance operation. This method will be called before a rebalance operation starts and after the consumer
-     * stops fetching data. It is recommended that offsets should be committed in this callback to either Kafka or a
+     * stops fetching data. It can also be called when consumer is being closed ({@link KafkaConsumer#close(Duration)})
+     * or is unsubscribing ({@link KafkaConsumer#unsubscribe()}).
+     * It is recommended that offsets should be committed in this callback to either Kafka or a
      * custom offset store to prevent duplicate data.
      * <p>
      * For examples on usage of this API, see Usage Examples section of {@link KafkaConsumer KafkaConsumer}
@@ -170,7 +173,8 @@ public interface ConsumerRebalanceListener {
      * to other consumers during a rebalance event. However, during exceptional scenarios when the consumer realized that it
      * does not own this partition any longer, i.e. not revoked via a normal rebalance event, then this method would be invoked.
      * <p>
-     * For example, this function is called if a consumer's session timeout has expired.
+     * For example, this function is called if a consumer's session timeout has expired, or if a fatal error has been
+     * received indicating the consumer is no longer part of the group.
      * <p>
      * By default it will just trigger {@link ConsumerRebalanceListener#onPartitionsRevoked}; for users who want to distinguish
      * the handling logic of revoked partitions v.s. lost partitions, they can override the default implementation.

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

@@ -1054,8 +1054,10 @@ public void unsubscribe() {
         try {
             fetcher.clearBufferedDataForUnassignedPartitions(Collections.emptySet());
             this.subscriptions.unsubscribe();
-            if (this.coordinator != null)
+            if (this.coordinator != null) {
+                this.coordinator.onLeaveGroup();
                 this.coordinator.maybeLeaveGroup("the consumer unsubscribed from all topics");
+            }
             log.info("Unsubscribed all topics or patterns and assigned partitions");
         } finally {
             release();

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

@@ -202,6 +202,11 @@ protected abstract void onJoinComplete(int generation,
                                            String protocol,
                                            ByteBuffer memberAssignment);
 
+    /**
+     * Invoked prior to each leave group event. This is typically used to cleanup assigned partitions
+     */
+    protected void onLeaveGroup() {}
+
     /**
      * Visible for testing.
      *
@@ -839,6 +844,7 @@ protected void close(Timer timer) {
             // needs this lock to complete and terminate after close flag is set.
             synchronized (this) {
                 if (rebalanceConfig.leaveGroupOnClose) {
+                    onLeaveGroup();
                     maybeLeaveGroup("the consumer is being closed");
                 }
 
@@ -854,51 +860,33 @@ protected void close(Timer timer) {
         }
     }
 
-    public synchronized RequestFuture<Void> maybeLeaveGroup(String leaveReason) {
-        return maybeLeaveGroup(leaveReason, false);
-    }
-
     /**
-     * Leave the current group and reset local generation/memberId.
-     * @param leaveReason reason to attempt leaving the group
-     * @throws KafkaException if the callback throws exception
+     * @throws KafkaException if the rebalance callback throws exception
      */
-    public synchronized RequestFuture<Void> maybeLeaveGroup(String leaveReason, boolean dueToHeartbeatExpiration) {
-        LeaveGroupRequest.Builder request = null;
+    public synchronized RequestFuture<Void> maybeLeaveGroup(String leaveReason) {
+        RequestFuture<Void> future = null;
+
         // Starting from 2.3, only dynamic members will send LeaveGroupRequest to the broker,
         // consumer with valid group.instance.id is viewed as static member that never sends LeaveGroup,
         // and the membership expiration is only controlled by session timeout.
         if (isDynamicMember() && !coordinatorUnknown() &&
-                state != MemberState.UNJOINED && generation.hasMemberId()) {
+            state != MemberState.UNJOINED && generation.hasMemberId()) {
             // this is a minimal effort attempt to leave the group. we do not
             // attempt any resending if the request fails or times out.
             log.info("Member {} sending LeaveGroup request to coordinator {} due to {}",
-                     generation.memberId, coordinator, leaveReason);
-            request = new LeaveGroupRequest.Builder(
+                generation.memberId, coordinator, leaveReason);
+            LeaveGroupRequest.Builder request = new LeaveGroupRequest.Builder(
                 rebalanceConfig.groupId,
                 Collections.singletonList(new MemberIdentity().setMemberId(generation.memberId))
             );
-        }
 
-        // we need to reset generation first in order to trigger the rebalance callback if necessary, before sending
-        // the leave group which may trigger the rebalance; if leave-group is caused by heartbeat poll expiration
-        // at the heartbeat thread, we should not trigger the rebalance callback and will wait until the caller thread
-        // calling consumer.poll to trigger it.
-        if (dueToHeartbeatExpiration) {
-            resetGeneration();
-        } else {
-            resetGenerationOnLeaveGroup();
-        }
-
-        if (request != null) {
-            RequestFuture<Void> future = client.send(coordinator, request)
-                .compose(new LeaveGroupResponseHandler());
+            future = client.send(coordinator, request).compose(new LeaveGroupResponseHandler());
             client.pollNoWakeup();
-
-            return future;
         }
 
-        return null;
+        resetGenerationOnLeaveGroup();
+
+        return future;
     }
 
     protected boolean isDynamicMember() {
@@ -1141,7 +1129,7 @@ public void run() {
                                                     "the poll loop is spending too much time processing messages. " +
                                                     "You can address this either by increasing max.poll.interval.ms or by reducing " +
                                                     "the maximum size of batches returned in poll() with max.poll.records.";
-                            maybeLeaveGroup(leaveReason, true);
+                            maybeLeaveGroup(leaveReason);
                         } else if (!heartbeat.shouldHeartbeat(now)) {
                             // poll again after waiting for the retry backoff in case the heartbeat failed or the
                             // coordinator disconnected

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

@@ -297,8 +297,8 @@ private Exception invokePartitionsRevoked(final Set<TopicPartition> revokedParti
         return null;
     }
 
-    private Exception invokePartitionsLost(final String rootCause, final Set<TopicPartition> lostPartitions) {
-        log.info("Lost previously assigned partitions {} due to {}", Utils.join(lostPartitions, ", "), rootCause);
+    private Exception invokePartitionsLost(final Set<TopicPartition> lostPartitions) {
+        log.info("Lost previously assigned partitions {}", Utils.join(lostPartitions, ", "));
 
         ConsumerRebalanceListener listener = subscriptions.rebalanceListener();
         try {
@@ -650,9 +650,9 @@ protected void onJoinPrepare(int generation, String memberId) {
             revokedPartitions = new HashSet<>(subscriptions.assignedPartitions());
 
             if (!revokedPartitions.isEmpty()) {
-                exception = invokePartitionsLost(
-                    "generation has been reset since consumer is no longer part of the group",
-                    revokedPartitions);
+                log.info("Giving away all assigned partitions as lost since generation has been reset," +
+                    "indicating that consumer is no longer part of the group");
+                exception = invokePartitionsLost(revokedPartitions);
 
                 subscriptions.assignFromSubscribed(Collections.emptySet());
             }
@@ -694,16 +694,12 @@ protected void onJoinPrepare(int generation, String memberId) {
         }
     }
 
-    /**
-     * @throws KafkaException if the rebalance callback throws exception
-     */
     @Override
-    synchronized void resetGenerationOnLeaveGroup() {
-        // we can reset assignment upon leaving group and trigger the callback as well
+    public void onLeaveGroup() {
+        // we should reset assignment and trigger the callback before leaving group
         Set<TopicPartition> droppedPartitions = new HashSet<>(subscriptions.assignedPartitions());
 
         if (subscriptions.partitionsAutoAssigned() && !droppedPartitions.isEmpty()) {
-
             final Exception e = invokePartitionsRevoked(droppedPartitions);
 
             subscriptions.assignFromSubscribed(Collections.emptySet());
@@ -712,8 +708,6 @@ synchronized void resetGenerationOnLeaveGroup() {
                 throw new KafkaException("User rebalance callback throws an error", e);
             }
         }
-
-        super.resetGenerationOnLeaveGroup();
     }
 
     /**
@@ -726,29 +720,8 @@ public boolean rejoinNeededOrPending() {
 
         // we need to rejoin if we performed the assignment and metadata has changed;
         // also for those owned-but-no-longer-existed partitions we should drop them as lost
-        if (assignmentSnapshot != null && !assignmentSnapshot.matches(metadataSnapshot)) {
-            Set<TopicPartition> ownedPartitions = new HashSet<>(subscriptions.assignedPartitions());
-            Set<TopicPartition> noLongerExistingPartitions = ownedPartitions.stream()
-                .filter(tp -> !metadataSnapshot.partitionsPerTopic().containsKey(tp.topic()))
-                .collect(Collectors.toSet());
-
-            if (!noLongerExistingPartitions.isEmpty()) {
-                ownedPartitions.removeAll(noLongerExistingPartitions);
-
-                Exception e = invokePartitionsLost("topic metadata has changed and partitions " +
-                        noLongerExistingPartitions + " do not exist any more, owned partitions left are " +
-                        ownedPartitions,
-                    noLongerExistingPartitions);
-
-                subscriptions.assignFromSubscribed(ownedPartitions);
-
-                if (e != null) {
-                    throw new KafkaException("User rebalance callback throws an error", e);
-                }
-            }
-
+        if (assignmentSnapshot != null && !assignmentSnapshot.matches(metadataSnapshot))
             return true;
-        }
 
         // we need to join if our subscription has changed since the last join
         if (joinedSubscription != null && !joinedSubscription.equals(subscriptions.subscription())) {

clients/src/test/java/org/apache/kafka/clients/consumer/internals/ConsumerCoordinatorTest.java

@@ -705,11 +705,9 @@ public boolean matches(AbstractRequest body) {
 
         coordinator.poll(time.timer(Long.MAX_VALUE));
 
-        Collection<TopicPartition> lost = Collections.singleton(t2p);
         revoked = getRevoked(newAssigned, oldAssigned);
-        revoked.removeAll(lost);
-        assertEquals(protocol == COOPERATIVE, revoked.isEmpty());
-        revokedCount += revoked.isEmpty() ? 0 : 1;
+        assertFalse(revoked.isEmpty());
+        revokedCount += 1;
         Collection<TopicPartition> added = getAdded(newAssigned, oldAssigned);
 
         assertFalse(coordinator.rejoinNeededOrPending());
@@ -719,8 +717,7 @@ public boolean matches(AbstractRequest body) {
         assertEquals(revoked.isEmpty() ? null : revoked, rebalanceListener.revoked);
         assertEquals(3, rebalanceListener.assignedCount);
         assertEquals(added, rebalanceListener.assigned);
-        assertEquals(1, rebalanceListener.lostCount);
-        assertEquals(lost, rebalanceListener.lost);
+        assertEquals(0, rebalanceListener.lostCount);
     }
 
     @Test