KAFKA-15276: Implement event plumbing for ConsumerRebalanceListener callbacks

checkstyle/suppressions.xml

@@ -127,6 +127,9 @@
     <suppress checks="NPathComplexity"
               files="MemoryRecordsTest|MetricsTest|RequestResponseTest|TestSslUtils|AclAuthorizerBenchmark"/>
 
+    <suppress checks="MethodLength"
+              files="ConsumerTestBuilder.java"/>
+
     <suppress checks="(WhitespaceAround|LocalVariableName|ImportControl|AvoidStarImport)"
               files="Murmur3Test.java"/>
 

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

@@ -131,8 +131,9 @@ void initializeResources() {
      * </ol>
      */
     void runOnce() {
-        // If there are errors processing any events, the error will be thrown immediately. This will have
-        // the effect of closing the background thread.
+        // Process the events—if any—that were produced by the application thread. It is possible that when processing
+        // an event generates an error. In such cases, the processor will log an exception, but we do not want those
+        // errors to be propagated to the caller.
         applicationEventProcessor.process();
 
         final long currentTimeMs = time.milliseconds();

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

@@ -0,0 +1,46 @@
+/*
+ * 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.internals;
+
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+
+/**
+ * This class just provides a static name for the methods in the {@link ConsumerRebalanceListener} interface
+ * for a bit more compile time assurance.
+ */
+public enum ConsumerRebalanceListenerMethodName {
+
+    ON_PARTITIONS_REVOKED("onPartitionsRevoked"),
+    ON_PARTITIONS_ASSIGNED("onPartitionsAssigned"),
+    ON_PARTITIONS_LOST("onPartitionsLost");
+
+    private final String fullyQualifiedMethodName;
+
+    ConsumerRebalanceListenerMethodName(String methodName) {
+        this.fullyQualifiedMethodName = String.format("%s.%s", ConsumerRebalanceListener.class.getSimpleName(), methodName);
+    }
+
+    /**
+     * Provides the fully-qualified method name, e.g. {@code ConsumerRebalanceListener.onPartitionsRevoked}. This
+     * is used for log messages.
+     *
+     * @return Fully-qualified method name
+     */
+    public String fullyQualifiedMethodName() {
+        return fullyQualifiedMethodName;
+    }
+}

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

@@ -31,7 +31,6 @@
 import org.apache.kafka.common.KafkaException;
 import org.apache.kafka.common.errors.InterruptException;
 import org.apache.kafka.common.errors.TimeoutException;
-import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.metrics.KafkaMetricsContext;
 import org.apache.kafka.common.metrics.MetricConfig;
 import org.apache.kafka.common.metrics.Metrics;
@@ -51,8 +50,8 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
 public final class ConsumerUtils {
@@ -206,22 +205,39 @@ public static void refreshCommittedOffsets(final Map<TopicPartition, OffsetAndMe
         }
     }
 
-    public static <T> T getResult(CompletableFuture<T> future, Timer timer) {
+    public static <T> T getResult(Future<T> future, Timer timer) {
         try {
             return future.get(timer.remainingMs(), TimeUnit.MILLISECONDS);
         } catch (ExecutionException e) {
-            Throwable t = e.getCause();
-
-            if (t instanceof WakeupException)
-                throw new WakeupException();
-            else if (t instanceof KafkaException)
-                throw (KafkaException) t;
-            else
-                throw new KafkaException(t);
+            throw maybeWrapAsKafkaException(e.getCause());
         } catch (InterruptedException e) {
             throw new InterruptException(e);
         } catch (java.util.concurrent.TimeoutException e) {
             throw new TimeoutException(e);
         }
     }
+
+    public static <T> T getResult(Future<T> future) {
+        try {
+            return future.get();
+        } catch (ExecutionException e) {
+            throw maybeWrapAsKafkaException(e.getCause());
+        } catch (InterruptedException e) {
+            throw new InterruptException(e);
+        }
+    }
+
+    public static KafkaException maybeWrapAsKafkaException(Throwable t) {
+        if (t instanceof KafkaException)
+            return (KafkaException) t;
+        else
+            return new KafkaException(t);
+    }
+
+    public static KafkaException maybeWrapAsKafkaException(Throwable t, String message) {
+        if (t instanceof KafkaException)
+            return (KafkaException) t;
+        else
+            return new KafkaException(message, t);
+    }
 }

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

@@ -19,6 +19,7 @@
 import org.apache.kafka.clients.CommonClientConfigs;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 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;
 import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
 import org.apache.kafka.clients.consumer.internals.events.GroupMetadataUpdateEvent;
@@ -206,6 +207,14 @@ public NetworkClientDelegate.PollResult poll(long currentTimeMs) {
         return new NetworkClientDelegate.PollResult(heartbeatRequestState.heartbeatIntervalMs, Collections.singletonList(request));
     }
 
+    /**
+     * Returns the {@link MembershipManager} that this request manager is using to track the state of the group.
+     * This is provided so that the {@link ApplicationEventProcessor} can access the state for querying or updating.
+     */
+    public MembershipManager membershipManager() {
+        return membershipManager;
+    }
+
     /**
      * Returns the delay for which the application thread can safely wait before it should be responsive
      * to results from the request managers. For example, the subscription state can change when heartbeats

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

@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.clients.consumer.internals;
 
+import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
+import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
 import org.apache.kafka.common.Uuid;
 import org.apache.kafka.common.message.ConsumerGroupHeartbeatResponseData;
 
@@ -145,6 +147,17 @@ public interface MembershipManager {
      */
     void onSubscriptionUpdated();
 
+    /**
+     * Signals that a {@link ConsumerRebalanceListener} callback has completed. This is invoked when the
+     * application thread has completed the callback and has submitted a
+     * {@link ConsumerRebalanceListenerCallbackCompletedEvent} to the network I/O thread. At this point, we
+     * notify the state machine that it's complete so that it can move to the next appropriate step of the
+     * rebalance process.
+     *
+     * @param event Event with details about the callback that was executed
+     */
+    void consumerRebalanceListenerCallbackCompleted(ConsumerRebalanceListenerCallbackCompletedEvent event);
+
     /**
      * Transition to the {@link MemberState#JOINING} state to attempt joining a group.
      */

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

@@ -21,6 +21,11 @@
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
 import org.apache.kafka.clients.consumer.internals.Utils.TopicIdPartitionComparator;
 import org.apache.kafka.clients.consumer.internals.Utils.TopicPartitionComparator;
+import org.apache.kafka.clients.consumer.internals.events.BackgroundEventHandler;
+import org.apache.kafka.clients.consumer.internals.events.CompletableBackgroundEvent;
+import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackCompletedEvent;
+import org.apache.kafka.clients.consumer.internals.events.ConsumerRebalanceListenerCallbackNeededEvent;
+import org.apache.kafka.clients.consumer.internals.events.ErrorBackgroundEvent;
 import org.apache.kafka.common.ClusterResource;
 import org.apache.kafka.common.ClusterResourceListener;
 import org.apache.kafka.common.KafkaException;
@@ -48,6 +53,10 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 
+import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_ASSIGNED;
+import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_LOST;
+import static org.apache.kafka.clients.consumer.internals.ConsumerRebalanceListenerMethodName.ON_PARTITIONS_REVOKED;
+
 /**
  * Group manager for a single consumer that has a group id defined in the config
  * {@link ConsumerConfig#GROUP_ID_CONFIG}, to use the Kafka-based offset management capability,
@@ -245,14 +254,22 @@ public class MembershipManagerImpl implements MembershipManager, ClusterResource
      */
     private final Optional<ClientTelemetryReporter> clientTelemetryReporter;
 
+    /**
+     * Serves as the conduit by which we can report events to the application thread. This is needed as we send
+     * {@link ConsumerRebalanceListenerCallbackNeededEvent callbacks} and, if needed,
+     * {@link ErrorBackgroundEvent errors} to the application thread.
+     */
+    private final BackgroundEventHandler backgroundEventHandler;
+
     public MembershipManagerImpl(String groupId,
                                  Optional<String> groupInstanceId,
                                  Optional<String> serverAssignor,
                                  SubscriptionState subscriptions,
                                  CommitRequestManager commitRequestManager,
                                  ConsumerMetadata metadata,
                                  LogContext logContext,
-                                 Optional<ClientTelemetryReporter> clientTelemetryReporter) {
+                                 Optional<ClientTelemetryReporter> clientTelemetryReporter,
+                                 BackgroundEventHandler backgroundEventHandler) {
         this.groupId = groupId;
         this.state = MemberState.UNSUBSCRIBED;
         this.serverAssignor = serverAssignor;
@@ -266,6 +283,7 @@ public MembershipManagerImpl(String groupId,
         this.currentAssignment = new HashMap<>();
         this.log = logContext.logger(MembershipManagerImpl.class);
         this.clientTelemetryReporter = clientTelemetryReporter;
+        this.backgroundEventHandler = backgroundEventHandler;
     }
 
     /**
@@ -1089,7 +1107,7 @@ private CompletableFuture<Void> invokeOnPartitionsRevokedCallback(Set<TopicParti
         // current behaviour.
         Optional<ConsumerRebalanceListener> listener = subscriptions.rebalanceListener();
         if (!partitionsRevoked.isEmpty() && listener.isPresent()) {
-            throw new UnsupportedOperationException("User-defined callbacks not supported yet");
+            return enqueueConsumerRebalanceListenerCallback(ON_PARTITIONS_REVOKED, partitionsRevoked);
         } else {
             return CompletableFuture.completedFuture(null);
         }
@@ -1100,7 +1118,7 @@ private CompletableFuture<Void> invokeOnPartitionsAssignedCallback(Set<TopicPart
         // the current behaviour.
         Optional<ConsumerRebalanceListener> listener = subscriptions.rebalanceListener();
         if (listener.isPresent()) {
-            throw new UnsupportedOperationException("User-defined callbacks not supported yet");
+            return enqueueConsumerRebalanceListenerCallback(ON_PARTITIONS_ASSIGNED, partitionsAssigned);
         } else {
             return CompletableFuture.completedFuture(null);
         }
@@ -1111,12 +1129,61 @@ private CompletableFuture<Void> invokeOnPartitionsLostCallback(Set<TopicPartitio
         // behaviour.
         Optional<ConsumerRebalanceListener> listener = subscriptions.rebalanceListener();
         if (!partitionsLost.isEmpty() && listener.isPresent()) {
-            throw new UnsupportedOperationException("User-defined callbacks not supported yet");
+            return enqueueConsumerRebalanceListenerCallback(ON_PARTITIONS_LOST, partitionsLost);
         } else {
             return CompletableFuture.completedFuture(null);
         }
     }
 
+    /**
+     * Enqueue a {@link ConsumerRebalanceListenerCallbackNeededEvent} to trigger the execution of the
+     * appropriate {@link ConsumerRebalanceListener} {@link ConsumerRebalanceListenerMethodName method} on the
+     * application thread.
+     *
+     * <p/>
+     *
+     * Because the reconciliation process (run in the background thread) will be blocked by the application thread
+     * until it completes this, we need to provide a {@link CompletableFuture} by which to remember where we left off.
+     *
+     * @param methodName Callback method that needs to be executed on the application thread
+     * @param partitions Partitions to supply to the callback method
+     * @return Future that will be chained within the rest of the reconciliation logic
+     */
+    private CompletableFuture<Void> enqueueConsumerRebalanceListenerCallback(ConsumerRebalanceListenerMethodName methodName,
+                                                                             Set<TopicPartition> partitions) {
+        SortedSet<TopicPartition> sortedPartitions = new TreeSet<>(TOPIC_PARTITION_COMPARATOR);
+        sortedPartitions.addAll(partitions);
+        CompletableBackgroundEvent<Void> event = new ConsumerRebalanceListenerCallbackNeededEvent(methodName, sortedPartitions);
+        backgroundEventHandler.add(event);
+        log.debug("The event to trigger the {} method execution was enqueued successfully", methodName.fullyQualifiedMethodName());
+        return event.future();
+    }
+
+    @Override
+    public void consumerRebalanceListenerCallbackCompleted(ConsumerRebalanceListenerCallbackCompletedEvent event) {
+        ConsumerRebalanceListenerMethodName methodName = event.methodName();
+        Optional<KafkaException> error = event.error();
+        CompletableFuture<Void> future = event.future();
+
+        if (error.isPresent()) {
+            Exception e = error.get();
+            log.warn(
+                "The {} method completed with an error ({}); signaling to continue to the next phase of rebalance",
+                methodName.fullyQualifiedMethodName(),
+                e.getMessage()
+            );
+
+            future.completeExceptionally(e);
+        } else {
+            log.debug(
+                "The {} method completed successfully; signaling to continue to the next phase of rebalance",
+                methodName.fullyQualifiedMethodName()
+            );
+
+            future.complete(null);
+        }
+    }
+
     /**
      * Log partitions being revoked that were already paused, since the pause flag will be
      * effectively lost.