KAFKA-16758: Extend Consumer#close with an option to leave the group or not

Author: frankvicky

checkstyle/suppressions.xml

@@ -129,7 +129,7 @@
               files="(OffsetFetcher|RequestResponse)Test.java"/>
 
     <suppress checks="JavaNCSS"
-              files="RequestResponseTest.java|FetcherTest.java|FetchRequestManagerTest.java|KafkaAdminClientTest.java"/>
+              files="RequestResponseTest.java|FetcherTest.java|FetchRequestManagerTest.java|KafkaAdminClientTest.java|ConsumerMembershipManagerTest.java"/>
 
     <suppress checks="NPathComplexity"
               files="MemoryRecordsTest|MetricsTest|RequestResponseTest|TestSslUtils|AclAuthorizerBenchmark"/>

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

@@ -16,6 +16,7 @@
  */
 package org.apache.kafka.clients.consumer;
 
+import org.apache.kafka.clients.consumer.internals.ConsumerUtils;
 import org.apache.kafka.common.Metric;
 import org.apache.kafka.common.MetricName;
 import org.apache.kafka.common.PartitionInfo;
@@ -28,6 +29,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.Set;
 import java.util.regex.Pattern;
@@ -277,13 +279,91 @@ public interface Consumer<K, V> extends Closeable {
     void close();
 
     /**
+     * This method has been deprecated since Kafka 4.0 and should use {@link Consumer#close(CloseOption)} instead.
+     *
      * @see KafkaConsumer#close(Duration)
      */
+    @Deprecated
     void close(Duration timeout);
 
     /**
      * @see KafkaConsumer#wakeup()
      */
     void wakeup();
 
+    /**
+     * @see KafkaConsumer#close(CloseOption)
+     */
+    void close(final CloseOption option);
+
+    class CloseOption {
+
+        /**
+         * Specifies the group membership operation upon shutdown.
+         * By default, {@code GroupMembershipOperation.DEFAULT} will be applied, which follows the consumer's default behavior.
+         */
+        protected GroupMembershipOperation operation = GroupMembershipOperation.DEFAULT;
+
+        /**
+         * Specifies the maximum amount of time to wait for the close process to complete.
+         * This allows users to define a custom timeout for gracefully stopping the consumer.
+         * If no value is set, the default timeout {@link ConsumerUtils#DEFAULT_CLOSE_TIMEOUT_MS} will be applied.
+         */
+        protected Optional<Duration> timeout = Optional.empty();
+
+        private CloseOption() {
+        }
+
+        protected CloseOption(final CloseOption option) {
+            this.operation = option.operation;
+            this.timeout = option.timeout;
+        }
+
+        /**
+         * Static method to create a {@code CloseOption} with a custom timeout.
+         *
+         * @param timeout the maximum time to wait for the consumer to close.
+         * @return a new {@code CloseOption} instance with the specified timeout.
+         */
+        public static CloseOption timeout(final Duration timeout) {
+            CloseOption option = new CloseOption();
+            option.timeout = Optional.ofNullable(timeout);
+            return option;
+        }
+
+        /**
+         * Static method to create a {@code CloseOption} with a specified group membership operation.
+         *
+         * @param operation the group membership operation to apply. Must be one of {@code LEAVE_GROUP}, {@code REMAIN_IN_GROUP},
+         *                 or {@code DEFAULT}.
+         * @return a new {@code CloseOption} instance with the specified group membership operation.
+         */
+        public static CloseOption groupMembershipOperation(final GroupMembershipOperation operation) {
+            CloseOption option = new CloseOption();
+            option.operation = operation;
+            return option;
+        }
+
+        /**
+         * Fluent method to set the timeout for the close process.
+         *
+         * @param timeout the maximum time to wait for the consumer to close. If {@code null}, the default timeout will be used.
+         * @return this {@code CloseOption} instance.
+         */
+        public CloseOption withTimeout(final Duration timeout) {
+            this.timeout = Optional.ofNullable(timeout);
+            return this;
+        }
+
+        /**
+         * Fluent method to set the group membership operation upon shutdown.
+         *
+         * @param operation the group membership operation to apply. Must be one of {@code LEAVE_GROUP}, {@code REMAIN_IN_GROUP}, or {@code DEFAULT}.
+         * @return this {@code CloseOption} instance.
+         */
+        public CloseOption withGroupMembershipOperation(final GroupMembershipOperation operation) {
+            this.operation = operation;
+            return this;
+        }
+    }
 }

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

@@ -0,0 +1,37 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.clients.consumer;
+
+/**
+ * Enum to specify the group membership operation upon leaving group.
+ *
+ * <ul>
+ *   <li><b>{@code LEAVE_GROUP}</b>:  means the consumer will leave the group.</li>
+ *   <li><b>{@code REMAIN_IN_GROUP}</b>: means the consumer will remain in the group.</li>
+ *   <li><b>{@code DEFAULT}</b>: Applies the default behavior:
+ *     <ul>
+ *       <li>For <b>static members</b>: The consumer will remain in the group.</li>
+ *       <li>For <b>dynamic members</b>: The consumer will leave the group.</li>
+ *     </ul>
+ *   </li>
+ * </ul>
+ */
+public enum GroupMembershipOperation {
+    LEAVE_GROUP,
+    REMAIN_IN_GROUP,
+    DEFAULT
+}

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

@@ -1776,6 +1776,7 @@ public void close() {
      * @throws org.apache.kafka.common.KafkaException for any other error during close
      */
     @Override
+    @SuppressWarnings("deprecation")
     public void close(Duration timeout) {
         delegate.close(timeout);
     }
@@ -1790,6 +1791,17 @@ public void wakeup() {
         delegate.wakeup();
     }
 
+    /**
+     * This method allows the caller to specify shutdown behavior using the {@link CloseOption} class.
+     * If {@code null} is provided, the default behavior will be applied, equivalent to providing a new {@link CloseOption} instance.
+     *
+     * @param option see {@link CloseOption}
+     */
+    @Override
+    public void close(CloseOption option) {
+        delegate.close(option);
+    }
+
     // Functions below are for testing only
     String clientId() {
         return delegate.clientId();

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

@@ -557,6 +557,7 @@ public void close() {
     }
 
     @Override
+    @SuppressWarnings("deprecation")
     public synchronized void close(Duration timeout) {
         this.closed = true;
     }
@@ -570,6 +571,11 @@ public synchronized void wakeup() {
         wakeup.set(true);
     }
 
+    @Override
+    public void close(CloseOption option) {
+        this.closed = true;
+    }
+
     /**
      * Schedule a task to be executed during a poll(). One enqueued task will be executed per {@link #poll(Duration)}
      * invocation. You can use this repeatedly to mock out multiple responses to poll invocations.

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

@@ -19,6 +19,7 @@
 import org.apache.kafka.clients.ClientResponse;
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.clients.GroupRebalanceConfig;
+import org.apache.kafka.clients.consumer.GroupMembershipOperation;
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.Node;
 import org.apache.kafka.common.errors.AuthenticationException;
@@ -1114,22 +1115,26 @@ private boolean isProtocolTypeInconsistent(String protocolType) {
      */
     @Override
     public final void close() {
-        close(time.timer(0));
+        close(time.timer(0), GroupMembershipOperation.DEFAULT);
     }
 
     /**
      * @throws KafkaException if the rebalance callback throws exception
      */
-    protected void close(Timer timer) {
+    protected void close(Timer timer, GroupMembershipOperation membershipOperation) {
         try {
             closeHeartbeatThread();
         } finally {
             // Synchronize after closing the heartbeat thread since heartbeat thread
             // needs this lock to complete and terminate after close flag is set.
             synchronized (this) {
-                if (rebalanceConfig.leaveGroupOnClose) {
+                // If membershipOperation is REMAIN_IN_GROUP, never send leave group request.
+                // If membershipOperation is DEFAULT, leave group based on rebalanceConfig.leaveGroupOnClose.
+                // Otherwise, leave group only if membershipOperation is LEAVE_GROUP.
+                if (GroupMembershipOperation.REMAIN_IN_GROUP != membershipOperation &&
+                    (GroupMembershipOperation.LEAVE_GROUP == membershipOperation || rebalanceConfig.leaveGroupOnClose)) {
                     onLeavePrepare();
-                    maybeLeaveGroup("the consumer is being closed");
+                    maybeLeaveGroup(membershipOperation, "the consumer is being closed");
                 }
 
                 // At this point, there may be pending commits (async commits or sync commits that were
@@ -1151,31 +1156,28 @@ protected void handlePollTimeoutExpiry() {
             "either by increasing max.poll.interval.ms or by reducing the maximum size of batches " +
             "returned in poll() with max.poll.records.");
 
-        maybeLeaveGroup("consumer poll timeout has expired.");
+        maybeLeaveGroup(GroupMembershipOperation.DEFAULT, "consumer poll timeout has expired.");
     }
 
     /**
-     * Sends LeaveGroupRequest and logs the {@code leaveReason}, unless this member is using static membership or is already
-     * not part of the group (ie does not have a valid member id, is in the UNJOINED state, or the coordinator is unknown).
+     * Sends LeaveGroupRequest and logs the {@code leaveReason}, unless this member is using static membership
+     * with the default consumer group membership operation, or is already not part of the group (i.e., does not have a
+     * valid member ID, is in the UNJOINED state, or the coordinator is unknown).
      *
+     * @param membershipOperation the operation on consumer group membership that the consumer will perform when closing
      * @param leaveReason the reason to leave the group for logging
      * @throws KafkaException if the rebalance callback throws exception
      */
-    public synchronized RequestFuture<Void> maybeLeaveGroup(String leaveReason) {
+    public synchronized RequestFuture<Void> maybeLeaveGroup(GroupMembershipOperation membershipOperation, 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()) {
-            // this is a minimal effort attempt to leave the group. we do not
-            // attempt any resending if the request fails or times out.
+        if (shouldSendLeaveGroupRequest(membershipOperation)) {
             log.info("Member {} sending LeaveGroup request to coordinator {} due to {}",
                 generation.memberId, coordinator, leaveReason);
+
             LeaveGroupRequest.Builder request = new LeaveGroupRequest.Builder(
                 rebalanceConfig.groupId,
-                Collections.singletonList(new MemberIdentity().setMemberId(generation.memberId).setReason(JoinGroupRequest.maybeTruncateReason(leaveReason)))
+                List.of(new MemberIdentity().setMemberId(generation.memberId).setReason(JoinGroupRequest.maybeTruncateReason(leaveReason)))
             );
 
             future = client.send(coordinator, request).compose(new LeaveGroupResponseHandler(generation));
@@ -1187,6 +1189,17 @@ public synchronized RequestFuture<Void> maybeLeaveGroup(String leaveReason) {
         return future;
     }
 
+    private boolean shouldSendLeaveGroupRequest(GroupMembershipOperation membershipOperation) {
+        // According to KIP-1092, static members can leave the group if the consumer group membership operation is LEAVE_GROUP.
+        // If the operation is REMAIN_IN_GROUP, this method "maybeLeaveGroup" will not be invoked.
+        return membershipOperation == GroupMembershipOperation.LEAVE_GROUP ||
+            // Below is the default behavior for consumer group membership:
+            // 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.
+            (isDynamicMember() && !coordinatorUnknown() && state != MemberState.UNJOINED && generation.hasMemberId());
+    }
+
     protected boolean isDynamicMember() {
         return rebalanceConfig.groupInstanceId.isEmpty();
     }

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

@@ -18,6 +18,7 @@
 
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.GroupMembershipOperation;
 import org.apache.kafka.clients.consumer.internals.NetworkClientDelegate.PollResult;
 import org.apache.kafka.clients.consumer.internals.events.ApplicationEventProcessor;
 import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
@@ -221,7 +222,17 @@ public NetworkClientDelegate.PollResult poll(long currentTimeMs) {
      */
     @Override
     public PollResult pollOnClose(long currentTimeMs) {
-        if (membershipManager().isLeavingGroup()) {
+        AbstractMembershipManager<R> membershipManager = membershipManager();
+        GroupMembershipOperation leaveGroupOperation = membershipManager.leaveGroupOperation();
+
+        if (membershipManager.isLeavingGroup() &&
+            // Default operation: both static and dynamic consumers will send a leave heartbeat
+            (GroupMembershipOperation.DEFAULT == leaveGroupOperation ||
+                // Leave operation: both static and dynamic consumers will send a leave heartbeat
+                GroupMembershipOperation.LEAVE_GROUP == leaveGroupOperation ||
+                // Remain in group: only static consumers will send a leave heartbeat, while dynamic members will not
+                membershipManager.groupInstanceId().isPresent())
+        ) {
             NetworkClientDelegate.UnsentRequest request = makeHeartbeatRequest(currentTimeMs, true);
             return new NetworkClientDelegate.PollResult(heartbeatRequestState.heartbeatIntervalMs, Collections.singletonList(request));
         }