KAFKA-8421: Still return data during rebalance

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

@@ -28,6 +28,10 @@ public class CommitFailedException extends KafkaException {
 
     private static final long serialVersionUID = 1L;
 
+    public CommitFailedException(final String message) {
+        super(message);
+    }
+
     public CommitFailedException() {
         super("Commit cannot be completed since the group has already " +
                 "rebalanced and assigned the partitions to another member. This means that the time " +

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

@@ -18,6 +18,7 @@
 
 import org.apache.kafka.common.TopicPartition;
 
+import java.time.Duration;
 import java.util.Collection;
 import java.util.Map;
 
@@ -37,6 +38,9 @@ public interface OffsetCommitCallback {
      * @throws org.apache.kafka.clients.consumer.CommitFailedException if the commit failed and cannot be retried.
      *             This can only occur if you are using automatic group management with {@link KafkaConsumer#subscribe(Collection)},
      *             or if there is an active group with the same groupId which is using group management.
+     * @throws org.apache.kafka.clients.consumer.RetriableCommitFailedException if the commit failed but can be retried.
+     *            This can occur if, e.g. consumer instance is in the middle of a rebalance. In such cases
+     *            commit could be retried after the rebalance is completed with the {@link #poll(Duration)} call.
      * @throws org.apache.kafka.common.errors.WakeupException if {@link KafkaConsumer#wakeup()} is called before or while this
      *             function is called
      * @throws org.apache.kafka.common.errors.InterruptException if the calling thread is interrupted before or while

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

@@ -16,21 +16,17 @@
  */
 package org.apache.kafka.clients.consumer;
 
-import org.apache.kafka.common.errors.RetriableException;
+import org.apache.kafka.common.KafkaException;
 
-public class RetriableCommitFailedException extends RetriableException {
+public class RetriableCommitFailedException extends KafkaException {
 
     private static final long serialVersionUID = 1L;
 
-    public static RetriableCommitFailedException withUnderlyingMessage(String additionalMessage) {
-        return new RetriableCommitFailedException("Offset commit failed with a retriable exception. " +
-                "You should retry committing the latest consumed offsets. " +
-                "The underlying error was: " + additionalMessage);
-    }
-
     public RetriableCommitFailedException(Throwable t) {
-        super("Offset commit failed with a retriable exception. You should retry committing " +
-                "the latest consumed offsets.", t);
+        super("Offset commit failed with a retriable exception. This is usually because the consumer is not yet " +
+            "part of a group with auto partition assignment or is participating in a rebalance to reassign partitions " +
+            "right now. You should first call poll to complete the rebalance and then " +
+            "retry committing the consumed offsets for those assigned partitions.", t);
     }
 
     public RetriableCommitFailedException(String message) {

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

@@ -111,7 +111,7 @@
 public abstract class AbstractCoordinator implements Closeable {
     public static final String HEARTBEAT_THREAD_PREFIX = "kafka-coordinator-heartbeat-thread";
 
-    private enum MemberState {
+    protected enum MemberState {
         UNJOINED,    // the client is not part of a group
         REBALANCING, // the client has begun rebalancing
         STABLE,      // the client has joined and is sending heartbeats
@@ -130,11 +130,11 @@ private enum MemberState {
     private MemberState state = MemberState.UNJOINED;
     private HeartbeatThread heartbeatThread = null;
     private RequestFuture<ByteBuffer> joinFuture = null;
+    private RequestFuture<Void> findCoordinatorFuture = null;
     private Generation generation = Generation.NO_GENERATION;
     private long lastRebalanceStartMs = -1L;
     private long lastRebalanceEndMs = -1L;
 
-    private RequestFuture<Void> findCoordinatorFuture = null;
 
     /**
      * Initialize the coordination manager.

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

@@ -433,7 +433,7 @@ public boolean poll(Timer timer) {
 
         invokeCompletedOffsetCommitCallbacks();
 
-        if (subscriptions.partitionsAutoAssigned()) {
+        if (subscriptions.hasAutoAssignedPartitions()) {
             if (protocol == null) {
                 throw new IllegalStateException("User configured " + ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG +
                     " to empty while trying to subscribe for group protocol to auto assign partitions");
@@ -701,7 +701,7 @@ public void onLeavePrepare() {
         // we should reset assignment and trigger the callback before leaving group
         Set<TopicPartition> droppedPartitions = new HashSet<>(subscriptions.assignedPartitions());
 
-        if (subscriptions.partitionsAutoAssigned() && !droppedPartitions.isEmpty()) {
+        if (subscriptions.hasAutoAssignedPartitions() && !droppedPartitions.isEmpty()) {
             final Exception e;
             if (generation() != Generation.NO_GENERATION) {
                 e = invokePartitionsRevoked(droppedPartitions);
@@ -722,7 +722,7 @@ public void onLeavePrepare() {
      */
     @Override
     public boolean rejoinNeededOrPending() {
-        if (!subscriptions.partitionsAutoAssigned())
+        if (!subscriptions.hasAutoAssignedPartitions())
             return false;
 
         // we need to rejoin if we performed the assignment and metadata has changed;
@@ -761,7 +761,11 @@ public boolean refreshCommittedOffsetsIfNeeded(Timer timer) {
 
                 log.info("Setting offset for partition {} to the committed offset {}", tp, position);
                 entry.getValue().leaderEpoch().ifPresent(epoch -> this.metadata.updateLastSeenEpochIfNewer(entry.getKey(), epoch));
-                this.subscriptions.seekUnvalidated(tp, position);
+
+                // it's possible that the partition is no longer assigned when the response is received,
+                // so we need to ignore seeking if that's the case
+                if (this.subscriptions.isAssigned(tp))
+                    this.subscriptions.seekUnvalidated(tp, position);
             }
         }
         return true;
@@ -972,7 +976,7 @@ private void doAutoCommitOffsetsAsync() {
 
         commitOffsetsAsync(allConsumedOffsets, (offsets, exception) -> {
             if (exception != null) {
-                if (exception instanceof RetriableException) {
+                if (exception instanceof RetriableCommitFailedException) {
                     log.debug("Asynchronous auto-commit of offsets {} failed due to retriable error: {}", offsets,
                         exception);
                     nextAutoCommitTimer.updateAndReset(rebalanceConfig.retryBackoffMs);
@@ -1052,13 +1056,17 @@ private RequestFuture<Void> sendOffsetCommitRequest(final Map<TopicPartition, Of
         }
 
         final Generation generation;
-        if (subscriptions.partitionsAutoAssigned()) {
+        if (subscriptions.hasAutoAssignedPartitions()) {
             generation = generationIfStable();
             // if the generation is null, we are not part of an active group (and we expect to be).
-            // the only thing we can do is fail the commit and let the user rejoin the group in poll()
+            // the only thing we can do is fail the commit and let the user rejoin the group in poll();
+            // we use RetriableCommitFailedException to indicate this may not be a fatal error, and leave
+            // it to users whether they want to handle differently than CommitFailed
             if (generation == null) {
                 log.info("Failing OffsetCommit request since the consumer is not part of an active group");
-                return RequestFuture.failure(new CommitFailedException());
+                return RequestFuture.failure(new RetriableCommitFailedException("Offset commit cannot be completed since the " +
+                    "consumer is not part of an active group for auto partition assignment yet. You can try completing the rebalance " +
+                    "by calling poll() and then retry the operation"));
             }
         } else
             generation = Generation.NO_GENERATION;
@@ -1075,14 +1083,16 @@ private RequestFuture<Void> sendOffsetCommitRequest(final Map<TopicPartition, Of
         log.trace("Sending OffsetCommit request with {} to coordinator {}", offsets, coordinator);
 
         return client.send(coordinator, builder)
-                .compose(new OffsetCommitResponseHandler(offsets));
+                .compose(new OffsetCommitResponseHandler(offsets, generation));
     }
 
     private class OffsetCommitResponseHandler extends CoordinatorResponseHandler<OffsetCommitResponse, Void> {
 
+        private final Generation generation;
         private final Map<TopicPartition, OffsetAndMetadata> offsets;
 
-        private OffsetCommitResponseHandler(Map<TopicPartition, OffsetAndMetadata> offsets) {
+        private OffsetCommitResponseHandler(Map<TopicPartition, OffsetAndMetadata> offsets, Generation generation) {
+            this.generation = generation;
             this.offsets = offsets;
         }
 
@@ -1134,22 +1144,34 @@ public void handle(OffsetCommitResponse commitResponse, RequestFuture<Void> futu
                             future.raise(error);
                             return;
                         } else if (error == Errors.REBALANCE_IN_PROGRESS) {
-                            /* Consumer never tries to commit offset in between join-group and sync-group,
+                            /* Consumer should not try to commit offset in between join-group and sync-group,
                              * and hence on broker-side it is not expected to see a commit offset request
                              * during CompletingRebalance phase; if it ever happens then broker would return
-                             * this error. In this case we should just treat as a fatal CommitFailed exception.
-                             * However, we do not need to reset generations and just request re-join, such that
-                             * if the caller decides to proceed and poll, it would still try to proceed and re-join normally.
+                             * this error to indicate that we are still in the middle of a rebalance.
+                             * In this case we would throw a RetriableCommitFailedException,
+                             * request re-join but do not reset generations. If the callers decide to retry they
+                             * can go ahead and call poll to finish up the rebalance first, and then try commit again.
                              */
                             requestRejoin();
-                            future.raise(new CommitFailedException());
+                            future.raise(new RetriableCommitFailedException("Offset commit cannot be completed since the " +
+                                "consumer group is executing a rebalance at the moment. You can try completing the rebalance " +
+                                "by calling poll() and then retry commit again"));
                             return;
-                        } else if (error == Errors.UNKNOWN_MEMBER_ID
-                                || error == Errors.ILLEGAL_GENERATION) {
+                        } else if (error == Errors.UNKNOWN_MEMBER_ID) {
                             // need to reset generation and re-join group
                             resetGenerationOnResponseError(ApiKeys.OFFSET_COMMIT, error);
                             future.raise(new CommitFailedException());
                             return;
+                        } else if (error == Errors.ILLEGAL_GENERATION) {
+                            if (this.generation.equals(ConsumerCoordinator.this.generation())) {
+                                // need to reset generation and re-join group if the generation has not changed, meaning
+                                // this is indeed a illegal generation; otherwise if the generation has changed the current
+                                // generation may actually be valid so we do not reset it -- if it is not, the next request
+                                // would still fail and then we would reset generation and re-join at that time.
+                                resetGenerationOnResponseError(ApiKeys.OFFSET_COMMIT, error);
+                            }
+                            future.raise(new CommitFailedException());
+                            return;
                         } else {
                             future.raise(new KafkaException("Unexpected error in commit: " + error.message()));
                             return;

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

@@ -26,17 +26,16 @@
 import java.util.concurrent.TimeUnit;
 
 public class KafkaConsumerMetrics implements AutoCloseable {
-    private final Metrics metrics;
     private final MetricName lastPollMetricName;
     private final Sensor timeBetweenPollSensor;
     private final Sensor pollIdleSensor;
+    private final Metrics metrics;
     private long lastPollMs;
     private long pollStartMs;
     private long timeSinceLastPollMs;
 
     public KafkaConsumerMetrics(Metrics metrics, String metricGrpPrefix) {
         this.metrics = metrics;
-
         String metricGroupName = metricGrpPrefix + "-metrics";
         Measurable lastPoll = (mConfig, now) -> {
             if (lastPollMs == 0L)

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

@@ -198,7 +198,7 @@ private boolean changeSubscription(Set<String> topicsToSubscribe) {
      * @param topics The topics to add to the group subscription
      */
     synchronized boolean groupSubscribe(Collection<String> topics) {
-        if (!partitionsAutoAssigned())
+        if (!hasAutoAssignedPartitions())
             throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
         groupSubscription = new HashSet<>(groupSubscription);
         return groupSubscription.addAll(topics);
@@ -224,15 +224,18 @@ public synchronized boolean assignFromUser(Set<TopicPartition> partitions) {
 
         assignmentId++;
 
+        // update the subscribed topics
         Set<String> manualSubscribedTopics = new HashSet<>();
         Map<TopicPartition, TopicPartitionState> partitionToState = new HashMap<>();
         for (TopicPartition partition : partitions) {
             TopicPartitionState state = assignment.stateValue(partition);
             if (state == null)
                 state = new TopicPartitionState();
             partitionToState.put(partition, state);
+
             manualSubscribedTopics.add(partition.topic());
         }
+
         this.assignment.set(partitionToState);
         return changeSubscription(manualSubscribedTopics);
     }
@@ -267,7 +270,7 @@ public synchronized boolean checkAssignmentMatchedSubscription(Collection<TopicP
      * different from {@link #assignFromUser(Set)} which directly set the assignment from user inputs.
      */
     public synchronized void assignFromSubscribed(Collection<TopicPartition> assignments) {
-        if (!this.partitionsAutoAssigned())
+        if (!this.hasAutoAssignedPartitions())
             throw new IllegalArgumentException("Attempt to dynamically assign partitions while manual assignment in use");
 
         Map<TopicPartition, TopicPartitionState> assignedPartitionStates = new HashMap<>(assignments.size());
@@ -322,7 +325,7 @@ synchronized boolean matchesSubscribedPattern(String topic) {
     }
 
     public synchronized Set<String> subscription() {
-        if (partitionsAutoAssigned())
+        if (hasAutoAssignedPartitions())
             return this.subscription;
         return Collections.emptySet();
     }
@@ -416,7 +419,7 @@ synchronized List<TopicPartition> fetchablePartitions(Predicate<TopicPartition>
                 .collect(Collectors.toList());
     }
 
-    synchronized boolean partitionsAutoAssigned() {
+    public synchronized boolean hasAutoAssignedPartitions() {
         return this.subscriptionType == SubscriptionType.AUTO_TOPICS || this.subscriptionType == SubscriptionType.AUTO_PATTERN;
     }
 

clients/src/test/java/org/apache/kafka/clients/MockClient.java

@@ -302,7 +302,6 @@ private long elapsedTimeMs(long currentTimeMs, long startTimeMs) {
         return Math.max(0, currentTimeMs - startTimeMs);
     }
 
-
     private void checkTimeoutOfPendingRequests(long nowMs) {
         ClientRequest request = requests.peek();
         while (request != null && elapsedTimeMs(nowMs, request.createdTimeMs()) > request.requestTimeoutMs()) {
@@ -334,7 +333,6 @@ public void respond(RequestMatcher matcher, AbstractResponse response) {
 
     // Utility method to enable out of order responses
     public void respondToRequest(ClientRequest clientRequest, AbstractResponse response) {
-        AbstractRequest request = clientRequest.requestBuilder().build();
         requests.remove(clientRequest);
         short version = clientRequest.requestBuilder().latestAllowedVersion();
         responses.add(new ClientResponse(clientRequest.makeHeader(version), clientRequest.callback(), clientRequest.destination(),