KAFKA-2686: Reset needsPartitionAssignment in SubscriptionState.assign()

Author: Guozhang Wang <[email protected]>

Reviewers: Jason Gustafson, Jun Rao

Closes apache#352 from guozhangwang/K2686

Author: guozhangwang

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

@@ -629,6 +629,9 @@ public Set<String> subscription() {
      * assign partitions. Topic subscriptions are not incremental. This list will replace the current
      * assignment (if there is one). Note that it is not possible to combine topic subscription with group management
      * with manual partition assignment through {@link #assign(List)}.
+     *
+     * If the given list of topics is empty, it is treated the same as {@link #unsubscribe()}.
+     *
      * <p>
      * As part of group management, the consumer will keep track of the list of consumers that belong to a particular
      * group and will trigger a rebalance operation if one of the following events trigger -
@@ -653,9 +656,14 @@ public Set<String> subscription() {
     public void subscribe(List<String> topics, ConsumerRebalanceListener listener) {
         acquire();
         try {
-            log.debug("Subscribed to topic(s): {}", Utils.join(topics, ", "));
-            this.subscriptions.subscribe(topics, listener);
-            metadata.setTopics(subscriptions.groupSubscription());
+            if (topics.isEmpty()) {
+                // treat subscribing to empty topic list as the same as unsubscribing
+                this.unsubscribe();
+            } else {
+                log.debug("Subscribed to topic(s): {}", Utils.join(topics, ", "));
+                this.subscriptions.subscribe(topics, listener);
+                metadata.setTopics(subscriptions.groupSubscription());
+            }
         } finally {
             release();
         }
@@ -666,6 +674,9 @@ public void subscribe(List<String> topics, ConsumerRebalanceListener listener) {
      * assign partitions. Topic subscriptions are not incremental. This list will replace the current
      * assignment (if there is one). It is not possible to combine topic subscription with group management
      * with manual partition assignment through {@link #assign(List)}.
+     *
+     * If the given list of topics is empty, it is treated the same as {@link #unsubscribe()}.
+     *
      * <p>
      * This is a short-hand for {@link #subscribe(List, ConsumerRebalanceListener)}, which
      * uses a noop listener. If you need the ability to either seek to particular offsets, you should prefer
@@ -715,6 +726,7 @@ public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
     public void unsubscribe() {
         acquire();
         try {
+            log.debug("Unsubscribed all topics or patterns and assigned partitions");
             this.subscriptions.unsubscribe();
             this.coordinator.resetGeneration();
             this.metadata.needMetadataForAllTopics(false);
@@ -739,7 +751,7 @@ public void assign(List<TopicPartition> partitions) {
         acquire();
         try {
             log.debug("Subscribed to partition(s): {}", Utils.join(partitions, ", "));
-            this.subscriptions.assign(partitions);
+            this.subscriptions.assignFromUser(partitions);
             Set<String> topics = new HashSet<>();
             for (TopicPartition tp : partitions)
                 topics.add(tp.topic());

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

@@ -76,7 +76,7 @@ public Set<TopicPartition> assignment() {
     public void rebalance(Collection<TopicPartition> newAssignment) {
         // TODO: Rebalance callbacks
         this.records.clear();
-        this.subscriptions.changePartitionAssignment(newAssignment);
+        this.subscriptions.assignFromSubscribed(newAssignment);
     }
 
     @Override
@@ -112,7 +112,7 @@ public void subscribe(List<String> topics, final ConsumerRebalanceListener liste
     @Override
     public void assign(List<TopicPartition> partitions) {
         ensureNotClosed();
-        this.subscriptions.assign(partitions);
+        this.subscriptions.assignFromUser(partitions);
     }
 
     @Override

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

@@ -169,7 +169,7 @@ protected void onJoin(int generation,
         subscriptions.needRefreshCommits();
 
         // update partition assignment
-        subscriptions.changePartitionAssignment(assignment.partitions());
+        subscriptions.assignFromSubscribed(assignment.partitions());
 
         // give the assignor a chance to update internal state based on the received assignment
         assignor.onAssignment(assignment);

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

@@ -28,8 +28,8 @@
 
 /**
  * A class for tracking the topics, partitions, and offsets for the consumer. A partition
- * is "assigned" either directly with {@link #assign(List)} (manual assignment)
- * or with {@link #changePartitionAssignment(Collection)} (automatic assignment).
+ * is "assigned" either directly with {@link #assignFromUser(Collection)} (manual assignment)
+ * or with {@link #assignFromSubscribed(Collection)} (automatic assignment from subscription).
  *
  * Once assigned, the partition is not considered "fetchable" until its initial position has
  * been set with {@link #seek(TopicPartition, long)}. Fetchable partitions track a fetch
@@ -129,12 +129,16 @@ public void groupSubscribe(Collection<String> topics) {
     }
 
     public void needReassignment() {
-        //
         this.groupSubscription.retainAll(subscription);
         this.needsPartitionAssignment = true;
     }
 
-    public void assign(List<TopicPartition> partitions) {
+    /**
+     * Change the assignment to the specified partitions provided by the user,
+     * note this is different from {@link #assignFromSubscribed(Collection)}
+     * whose input partitions are provided from the subscribed topics.
+     */
+    public void assignFromUser(Collection<TopicPartition> partitions) {
         if (!this.subscription.isEmpty() || this.subscribedPattern != null)
             throw new IllegalStateException(SUBSCRIPTION_EXCEPTION_MESSAGE);
 
@@ -146,6 +150,22 @@ public void assign(List<TopicPartition> partitions) {
                 addAssignedPartition(partition);
 
         this.assignment.keySet().retainAll(this.userAssignment);
+
+        this.needsPartitionAssignment = false;
+    }
+
+    /**
+     * Change the assignment to the specified partitions returned from the coordinator,
+     * note this is different from {@link #assignFromUser(Collection)} which directly set the assignment from user inputs
+     */
+    public void assignFromSubscribed(Collection<TopicPartition> assignments) {
+        for (TopicPartition tp : assignments)
+            if (!this.subscription.contains(tp.topic()))
+                throw new IllegalArgumentException("Assigned partition " + tp + " for non-subscribed topic.");
+        this.assignment.clear();
+        for (TopicPartition tp: assignments)
+            addAssignedPartition(tp);
+        this.needsPartitionAssignment = false;
     }
 
     public void subscribe(Pattern pattern, ConsumerRebalanceListener listener) {
@@ -306,16 +326,6 @@ public boolean partitionAssignmentNeeded() {
         return this.needsPartitionAssignment;
     }
 
-    public void changePartitionAssignment(Collection<TopicPartition> assignments) {
-        for (TopicPartition tp : assignments)
-            if (!this.subscription.contains(tp.topic()))
-                throw new IllegalArgumentException("Assigned partition " + tp + " for non-subscribed topic.");
-        this.assignment.clear();
-        for (TopicPartition tp: assignments)
-            addAssignedPartition(tp);
-        this.needsPartitionAssignment = false;
-    }
-
     public boolean isAssigned(TopicPartition tp) {
         return assignment.containsKey(tp);
     }

clients/src/test/java/org/apache/kafka/clients/consumer/KafkaConsumerTest.java

@@ -17,15 +17,20 @@
 package org.apache.kafka.clients.consumer;
 
 import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicPartition;
 import org.apache.kafka.common.serialization.ByteArrayDeserializer;
 import org.apache.kafka.test.MockMetricsReporter;
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.Collections;
 import java.util.Properties;
 
 public class KafkaConsumerTest {
 
+    private final String topic = "test";
+    private final TopicPartition tp0 = new TopicPartition("test", 0);
+
     @Test
     public void testConstructorClose() throws Exception {
         Properties props = new Properties();
@@ -46,4 +51,31 @@ public void testConstructorClose() throws Exception {
         }
         Assert.fail("should have caught an exception and returned");
     }
+
+    @Test
+    public void testSubscription() {
+        Properties props = new Properties();
+        props.setProperty(ConsumerConfig.CLIENT_ID_CONFIG, "testSubscription");
+        props.setProperty(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9999");
+        props.setProperty(ConsumerConfig.METRIC_REPORTER_CLASSES_CONFIG, MockMetricsReporter.class.getName());
+
+        KafkaConsumer<byte[], byte[]> consumer = new KafkaConsumer<byte[], byte[]>(
+            props, new ByteArrayDeserializer(), new ByteArrayDeserializer());
+
+        consumer.subscribe(Collections.singletonList(topic));
+        Assert.assertEquals(Collections.singleton(topic), consumer.subscription());
+        Assert.assertTrue(consumer.assignment().isEmpty());
+
+        consumer.subscribe(Collections.<String>emptyList());
+        Assert.assertTrue(consumer.subscription().isEmpty());
+        Assert.assertTrue(consumer.assignment().isEmpty());
+
+        consumer.assign(Collections.singletonList(tp0));
+        Assert.assertTrue(consumer.subscription().isEmpty());
+        Assert.assertEquals(Collections.singleton(tp0), consumer.assignment());
+
+        consumer.unsubscribe();
+        Assert.assertTrue(consumer.subscription().isEmpty());
+        Assert.assertTrue(consumer.assignment().isEmpty());
+    }
 }

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

@@ -193,7 +193,7 @@ public void testIllegalGeneration() {
 
         // illegal_generation will cause re-partition
         subscriptions.subscribe(Arrays.asList(topicName), rebalanceListener);
-        subscriptions.changePartitionAssignment(Collections.singletonList(tp));
+        subscriptions.assignFromSubscribed(Collections.singletonList(tp));
 
         time.sleep(sessionTimeoutMs);
         RequestFuture<Void> future = coordinator.sendHeartbeatRequest(); // should send out the heartbeat
@@ -217,7 +217,7 @@ public void testUnknownConsumerId() {
 
         // illegal_generation will cause re-partition
         subscriptions.subscribe(Arrays.asList(topicName), rebalanceListener);
-        subscriptions.changePartitionAssignment(Collections.singletonList(tp));
+        subscriptions.assignFromSubscribed(Collections.singletonList(tp));
 
         time.sleep(sessionTimeoutMs);
         RequestFuture<Void> future = coordinator.sendHeartbeatRequest(); // should send out the heartbeat
@@ -413,7 +413,7 @@ public void testInvalidSessionTimeout() {
 
     @Test
     public void testCommitOffsetOnly() {
-        subscriptions.assign(Arrays.asList(tp));
+        subscriptions.assignFromUser(Arrays.asList(tp));
 
         client.prepareResponse(consumerMetadataResponse(node, Errors.NONE.code()));
         coordinator.ensureCoordinatorKnown();
@@ -430,7 +430,7 @@ public void testCommitOffsetOnly() {
 
     @Test
     public void testCommitOffsetMetadata() {
-        subscriptions.assign(Arrays.asList(tp));
+        subscriptions.assignFromUser(Arrays.asList(tp));
 
         client.prepareResponse(consumerMetadataResponse(node, Errors.NONE.code()));
         coordinator.ensureCoordinatorKnown();
@@ -473,7 +473,7 @@ public void testResetGeneration() {
         // now switch to manual assignment
         subscriptions.unsubscribe();
         coordinator.resetGeneration();
-        subscriptions.assign(Arrays.asList(tp));
+        subscriptions.assignFromUser(Arrays.asList(tp));
 
         // the client should not reuse generation/memberId from auto-subscribed generation
         client.prepareResponse(new MockClient.RequestMatcher() {
@@ -612,7 +612,7 @@ public void testRefreshOffset() {
         client.prepareResponse(consumerMetadataResponse(node, Errors.NONE.code()));
         coordinator.ensureCoordinatorKnown();
 
-        subscriptions.assign(Arrays.asList(tp));
+        subscriptions.assignFromUser(Arrays.asList(tp));
         subscriptions.needRefreshCommits();
         client.prepareResponse(offsetFetchResponse(tp, Errors.NONE.code(), "", 100L));
         coordinator.refreshCommittedOffsetsIfNeeded();
@@ -625,7 +625,7 @@ public void testRefreshOffsetLoadInProgress() {
         client.prepareResponse(consumerMetadataResponse(node, Errors.NONE.code()));
         coordinator.ensureCoordinatorKnown();
 
-        subscriptions.assign(Arrays.asList(tp));
+        subscriptions.assignFromUser(Arrays.asList(tp));
         subscriptions.needRefreshCommits();
         client.prepareResponse(offsetFetchResponse(tp, Errors.OFFSET_LOAD_IN_PROGRESS.code(), "", 100L));
         client.prepareResponse(offsetFetchResponse(tp, Errors.NONE.code(), "", 100L));
@@ -639,7 +639,7 @@ public void testRefreshOffsetNotCoordinatorForConsumer() {
         client.prepareResponse(consumerMetadataResponse(node, Errors.NONE.code()));
         coordinator.ensureCoordinatorKnown();
 
-        subscriptions.assign(Arrays.asList(tp));
+        subscriptions.assignFromUser(Arrays.asList(tp));
         subscriptions.needRefreshCommits();
         client.prepareResponse(offsetFetchResponse(tp, Errors.NOT_COORDINATOR_FOR_GROUP.code(), "", 100L));
         client.prepareResponse(consumerMetadataResponse(node, Errors.NONE.code()));
@@ -654,7 +654,7 @@ public void testRefreshOffsetWithNoFetchableOffsets() {
         client.prepareResponse(consumerMetadataResponse(node, Errors.NONE.code()));
         coordinator.ensureCoordinatorKnown();
 
-        subscriptions.assign(Arrays.asList(tp));
+        subscriptions.assignFromUser(Arrays.asList(tp));
         subscriptions.needRefreshCommits();
         client.prepareResponse(offsetFetchResponse(tp, Errors.NONE.code(), "", -1L));
         coordinator.refreshCommittedOffsetsIfNeeded();