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/KafkaConsumer.java

@@ -1210,9 +1210,11 @@ private ConsumerRecords<K, V> poll(final Timer timer, final boolean includeMetad
                 client.maybeTriggerWakeup();
 
                 if (includeMetadataInTimeout) {
-                    if (!updateAssignmentMetadataIfNeeded(timer)) {
-                        return ConsumerRecords.empty();
-                    }
+                    // try to update assignment metadata BUT do not need to block on the timer,
+                    // since even if we are 1) in the middle of a rebalance or 2) have partitions
+                    // with unknown starting positions we may still want to return some data
+                    // as long as there are some partitions fetchable
+                    updateAssignmentMetadataIfNeeded(time.timer(1L));
                 } else {
                     while (!updateAssignmentMetadataIfNeeded(time.timer(Long.MAX_VALUE))) {
                         log.warn("Still waiting for metadata");
@@ -1285,12 +1287,6 @@ private Map<TopicPartition, List<ConsumerRecord<K, V>>> pollForFetches(Timer tim
         });
         timer.update(pollTimer.currentTimeMs());
 
-        // after the long poll, we should check whether the group needs to rebalance
-        // prior to returning data so that the group can stabilize faster
-        if (coordinator != null && coordinator.rejoinNeededOrPending()) {
-            return Collections.emptyMap();
-        }
-
         return fetcher.fetchedRecords();
     }
 
@@ -1324,6 +1320,8 @@ private Map<TopicPartition, List<ConsumerRecord<K, V>>> pollForFetches(Timer tim
      * @throws org.apache.kafka.common.errors.TimeoutException if the timeout specified by {@code default.api.timeout.ms} expires
      *            before successful completion of the offset commit
      * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
+     * @throws org.apache.kafka.common.errors.RebalanceInProgressException if this consumer instance is in the middle of a rebalance;
+     *            Commit could be retried after the rebalance is completed with the {@link #poll(Duration)} call.
      */
     @Override
     public void commitSync() {
@@ -1359,6 +1357,8 @@ public void commitSync() {
      * @throws org.apache.kafka.common.errors.TimeoutException if the timeout expires before successful completion
      *            of the offset commit
      * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
+     * @throws org.apache.kafka.common.errors.RebalanceInProgressException if this consumer instance is in the middle of a rebalance;
+     *            Commit could be retried after the rebalance is completed with the {@link #poll(Duration)} call.
      */
     @Override
     public void commitSync(Duration timeout) {
@@ -1380,7 +1380,8 @@ public void commitSync(Duration timeout) {
      * This commits offsets to Kafka. The offsets committed using this API will be used on the first fetch after every
      * rebalance and also on startup. As such, if you need to store offsets in anything other than Kafka, this API
      * should not be used. The committed offset should be the next message your application will consume,
-     * i.e. lastProcessedMessageOffset + 1.
+     * i.e. lastProcessedMessageOffset + 1. If automatic group management with {@link #subscribe(Collection)} is used,
+     * then the committed offsets must belong to the currently auto-assigned partitions.
      * <p>
      * This is a synchronous commits and will block until either the commit succeeds or an unrecoverable error is
      * encountered (in which case it is thrown to the caller), or the timeout specified by {@code default.api.timeout.ms} expires
@@ -1406,6 +1407,8 @@ public void commitSync(Duration timeout) {
      * @throws org.apache.kafka.common.errors.TimeoutException if the timeout expires before successful completion
      *            of the offset commit
      * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
+     * @throws org.apache.kafka.common.errors.RebalanceInProgressException if this consumer instance is in the middle of a rebalance;
+     *            Commit could be retried after the rebalance is completed with the {@link #poll(Duration)} call.
      */
     @Override
     public void commitSync(final Map<TopicPartition, OffsetAndMetadata> offsets) {
@@ -1418,7 +1421,8 @@ public void commitSync(final Map<TopicPartition, OffsetAndMetadata> offsets) {
     * This commits offsets to Kafka. The offsets committed using this API will be used on the first fetch after every
     * rebalance and also on startup. As such, if you need to store offsets in anything other than Kafka, this API
     * should not be used. The committed offset should be the next message your application will consume,
-    * i.e. lastProcessedMessageOffset + 1.
+    * i.e. lastProcessedMessageOffset + 1. If automatic group management with {@link #subscribe(Collection)} is used,
+    * then the committed offsets must belong to the currently auto-assigned partitions.
     * <p>
     * This is a synchronous commits and will block until either the commit succeeds, an unrecoverable error is
     * encountered (in which case it is thrown to the caller), or the timeout expires.
@@ -1444,12 +1448,15 @@ public void commitSync(final Map<TopicPartition, OffsetAndMetadata> offsets) {
     * @throws org.apache.kafka.common.errors.TimeoutException if the timeout expires before successful completion
     *            of the offset commit
     * @throws org.apache.kafka.common.errors.FencedInstanceIdException if this consumer instance gets fenced by broker.
+    * @throws org.apache.kafka.common.errors.RebalanceInProgressException if this consumer instance is in the middle of a rebalance;
+    *            Commit could be retried after the rebalance is completed with the {@link #poll(Duration)} call.
     */
     @Override
     public void commitSync(final Map<TopicPartition, OffsetAndMetadata> offsets, final Duration timeout) {
         acquireAndEnsureOpen();
         try {
             maybeThrowInvalidGroupIdException();
+            maybeThrowIfCommitOffsetsNotOwned(offsets);
             offsets.forEach(this::updateLastSeenEpochIfNewer);
             if (!coordinator.commitOffsetsSync(new HashMap<>(offsets), time.timer(timeout))) {
                 throw new TimeoutException("Timeout of " + timeout.toMillis() + "ms expired before successfully " +
@@ -1499,7 +1506,8 @@ public void commitAsync(OffsetCommitCallback callback) {
      * This commits offsets to Kafka. The offsets committed using this API will be used on the first fetch after every
      * rebalance and also on startup. As such, if you need to store offsets in anything other than Kafka, this API
      * should not be used. The committed offset should be the next message your application will consume,
-     * i.e. lastProcessedMessageOffset + 1.
+     * i.e. lastProcessedMessageOffset + 1. If automatic group management with {@link #subscribe(Collection)} is used,
+     * then the committed offsets must belong to the currently auto-assigned partitions.
      * <p>
      * This is an asynchronous call and will not block. Any errors encountered are either passed to the callback
      * (if provided) or discarded.
@@ -1519,14 +1527,29 @@ public void commitAsync(final Map<TopicPartition, OffsetAndMetadata> offsets, Of
         acquireAndEnsureOpen();
         try {
             maybeThrowInvalidGroupIdException();
+            maybeThrowIfCommitOffsetsNotOwned(offsets);
             log.debug("Committing offsets: {}", offsets);
             offsets.forEach(this::updateLastSeenEpochIfNewer);
             coordinator.commitOffsetsAsync(new HashMap<>(offsets), callback);
+        } catch (CommitFailedException e) {
+            log.error("Failed to commit offsets asynchronously because they do not belong to dynamically assigned partitions");
+            callback.onComplete(offsets, e);
         } finally {
             release();
         }
     }
 
+    private void maybeThrowIfCommitOffsetsNotOwned(final Map<TopicPartition, OffsetAndMetadata> offsets) {
+        if (subscriptions.partitionsAutoAssigned()) {
+            final Set<TopicPartition> partitions = assignment();
+            for (final TopicPartition tp : offsets.keySet()) {
+                if (!partitions.contains(tp))
+                    throw new CommitFailedException("cannot commit offsets out of owned partitions with " +
+                        "dynamic partition assigned from subscription");
+            }
+        }
+    }
+
     /**
      * Overrides the fetch offsets that the consumer will use on the next {@link #poll(Duration) poll(timeout)}. If this API
      * is invoked for the same partition more than once, the latest offset will be used on the next poll(). Note that

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;
 
@@ -45,6 +46,8 @@ public interface OffsetCommitCallback {
      *             configured groupId. See the exception for more details
      * @throws org.apache.kafka.common.KafkaException for any other unrecoverable errors (e.g. if offset metadata
      *             is too large or if the committed offset is invalid).
+     * @throws org.apache.kafka.common.errors.RebalanceInProgressException if this consumer instance is in the middle of a rebalance;
+     *             Commit could be retried after the rebalance is completed with the {@link KafkaConsumer#poll(Duration)} call.
      */
     void onComplete(Map<TopicPartition, OffsetAndMetadata> offsets, Exception exception);
 }

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

@@ -22,12 +22,6 @@ public class RetriableCommitFailedException extends RetriableException {
 
     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);

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

@@ -109,7 +109,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
@@ -125,12 +125,11 @@ private enum MemberState {
     private HeartbeatThread heartbeatThread = null;
     private boolean rejoinNeeded = true;
     private boolean needsJoinPrepare = true;
-    private MemberState state = MemberState.UNJOINED;
     private RequestFuture<ByteBuffer> joinFuture = null;
     private Node coordinator = null;
     private Generation generation = Generation.NO_GENERATION;
-
     private RequestFuture<Void> findCoordinatorFuture = null;
+    protected MemberState state = MemberState.UNJOINED;
 
     /**
      * Initialize the coordination manager.

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

@@ -36,6 +36,7 @@
 import org.apache.kafka.common.errors.FencedInstanceIdException;
 import org.apache.kafka.common.errors.GroupAuthorizationException;
 import org.apache.kafka.common.errors.InterruptException;
+import org.apache.kafka.common.errors.RebalanceInProgressException;
 import org.apache.kafka.common.errors.RetriableException;
 import org.apache.kafka.common.errors.TimeoutException;
 import org.apache.kafka.common.errors.TopicAuthorizationException;
@@ -1045,10 +1046,14 @@ private RequestFuture<Void> sendOffsetCommitRequest(final Map<TopicPartition, Of
         if (subscriptions.partitionsAutoAssigned()) {
             generation = generation();
             // 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 RebalanceInProgressException 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 RebalanceInProgressException("Offset commit cannot be completed since the " +
+                    "consumer group is not part of an active group yet. You can try completing the rebalance " +
+                    "by calling poll() and then retry the operation"));
             }
         } else
             generation = Generation.NO_GENERATION;
@@ -1127,15 +1132,19 @@ public void handle(OffsetCommitResponse commitResponse, RequestFuture<Void> futu
                             /* Consumer never tries 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. In this case we would just throw a RebalanceInProgressException
+                             * and request re-join, but we do not need to reset generations.
+                             * If the caller decides to proceed and poll, it would still try to proceed and re-join normally.
                              */
                             requestRejoin();
+                            future.raise(new RebalanceInProgressException());
+                            return;
+                        } 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.UNKNOWN_MEMBER_ID
-                                || error == Errors.ILLEGAL_GENERATION) {
+                        } else if (error == Errors.ILLEGAL_GENERATION) {
                             // need to reset generation and re-join group
                             resetGenerationOnResponseError(ApiKeys.OFFSET_COMMIT, error);
                             future.raise(new CommitFailedException());

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

@@ -222,21 +222,45 @@ public synchronized boolean assignFromUser(Set<TopicPartition> partitions) {
         if (this.assignment.partitionSet().equals(partitions))
             return false;
 
+        Map<TopicPartition, TopicPartitionState> assignedPartitionStates = partitionToStateMap(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);
+
+        this.assignment.set(assignedPartitionStates);
         return changeSubscription(manualSubscribedTopics);
     }
 
+    /**
+     * Change the assignment to the specified partitions returned from the coordinator, note this is
+     * different from {@link #assignFromUser(Set)} which directly set the assignment from user inputs.
+     */
+    public synchronized void assignFromSubscribed(Collection<TopicPartition> assignments) {
+        if (!this.partitionsAutoAssigned())
+            throw new IllegalArgumentException("Attempt to dynamically assign partitions while manual assignment in use");
+
+
+        final Map<TopicPartition, TopicPartitionState> assignedPartitionStates = partitionToStateMap(assignments);
+        assignmentId++;
+        this.assignment.set(assignedPartitionStates);
+    }
+
+    private Map<TopicPartition, TopicPartitionState> partitionToStateMap(Collection<TopicPartition> assignments) {
+        Map<TopicPartition, TopicPartitionState> map = new HashMap<>(assignments.size());
+        for (final TopicPartition tp : assignments) {
+            if (assignment.contains(tp))
+                map.put(tp, assignment.stateValue(tp));
+            else
+                map.put(tp, new TopicPartitionState());
+        }
+        return map;
+    }
+
     /**
      * @return true if assignments matches subscription, otherwise false
      */
@@ -262,20 +286,6 @@ public synchronized boolean checkAssignmentMatchedSubscription(Collection<TopicP
         return true;
     }
 
-    /**
-     * Change the assignment to the specified partitions returned from the coordinator, note this is
-     * different from {@link #assignFromUser(Set)} which directly set the assignment from user inputs.
-     */
-    public synchronized void assignFromSubscribed(Collection<TopicPartition> assignments) {
-        if (!this.partitionsAutoAssigned())
-            throw new IllegalArgumentException("Attempt to dynamically assign partitions while manual assignment in use");
-
-
-        Map<TopicPartition, TopicPartitionState> assignedPartitionStates = partitionToStateMap(assignments);
-        assignmentId++;
-        this.assignment.set(assignedPartitionStates);
-    }
-
     private void registerRebalanceListener(ConsumerRebalanceListener listener) {
         if (listener == null)
             throw new IllegalArgumentException("RebalanceListener cannot be null");
@@ -410,7 +420,7 @@ synchronized List<TopicPartition> fetchablePartitions(Predicate<TopicPartition>
                 .collect(Collectors.toList());
     }
 
-    synchronized boolean partitionsAutoAssigned() {
+    public synchronized boolean partitionsAutoAssigned() {
         return this.subscriptionType == SubscriptionType.AUTO_TOPICS || this.subscriptionType == SubscriptionType.AUTO_PATTERN;
     }
 
@@ -669,13 +679,6 @@ public synchronized ConsumerRebalanceListener rebalanceListener() {
         return rebalanceListener;
     }
 
-    private static Map<TopicPartition, TopicPartitionState> partitionToStateMap(Collection<TopicPartition> assignments) {
-        Map<TopicPartition, TopicPartitionState> map = new HashMap<>(assignments.size());
-        for (TopicPartition tp : assignments)
-            map.put(tp, new TopicPartitionState());
-        return map;
-    }
-
     private static class TopicPartitionState {
 
         private FetchState fetchState;