(feat) clear progress based on assigned partitions in consumer

Author: ylgrgyq

src/main/java/cn/leancloud/kafka/consumer/AbstractRecommitAwareCommitPolicy.java

@@ -58,7 +58,7 @@ private Map<TopicPartition, OffsetAndMetadata> offsetsForRecommit() {
         for (TopicPartition partition : consumer.assignment()) {
             final OffsetAndMetadata offset = consumer.committed(partition);
             if (offset != null) {
-                ret.putIfAbsent(partition, offset);
+                ret.put(partition, offset);
             }
         }
 

src/main/java/cn/leancloud/kafka/consumer/RebalanceListener.java

@@ -22,7 +22,6 @@ final class RebalanceListener<K, V> implements ConsumerRebalanceListener {
     private final ProcessRecordsProgress progress;
     private final Set<TopicPartition> knownPartitions;
     private final ConsumerSeekDestination forceSeekTo;
-    private Set<TopicPartition> lastRevokedPartitions;
     private Set<TopicPartition> pausedPartitions;
 
     RebalanceListener(Consumer<K, V> consumer, ProcessRecordsProgress progress, CommitPolicy policy, ConsumerSeekDestination forceSeekTo) {
@@ -32,12 +31,10 @@ final class RebalanceListener<K, V> implements ConsumerRebalanceListener {
         this.pausedPartitions = Collections.emptySet();
         this.knownPartitions = new HashSet<>();
         this.forceSeekTo = forceSeekTo;
-        this.lastRevokedPartitions = new HashSet<>();
     }
 
     @Override
     public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
-        lastRevokedPartitions = new HashSet<>(partitions);
         pausedPartitions = consumer.paused();
         if (!pausedPartitions.isEmpty()) {
             pausedPartitions = new HashSet<>(pausedPartitions);
@@ -55,14 +52,7 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
             pausePreviousPausedPartitions(partitions);
         }
 
-        // We can't use removeAll because of the bad performance of removeAll(#Collection) on HashSet
-        for (TopicPartition p : partitions) {
-            lastRevokedPartitions.remove(p);
-        }
-
-        clearProgressForRevokedPartitions(lastRevokedPartitions);
-
-        lastRevokedPartitions.clear();
+        clearProgressForRevokedPartitions();
     }
 
     private void seekOnAssignedPartitions(Collection<TopicPartition> partitions) {
@@ -95,10 +85,12 @@ private void pausePreviousPausedPartitions(Collection<TopicPartition> partitions
         pausedPartitions = Collections.emptySet();
     }
 
-    private void clearProgressForRevokedPartitions(Collection<TopicPartition> revokedPartitions) {
-        // revoke those partitions which was revoked and not reassigned
-        if (!revokedPartitions.isEmpty()) {
-            progress.clearFor(revokedPartitions);
+    private void clearProgressForRevokedPartitions() {
+        final Set<TopicPartition> partitionsWithProgress = progress.allPartitions();
+        partitionsWithProgress.removeAll(consumer.assignment());
+
+        if (!partitionsWithProgress.isEmpty()) {
+            progress.clearFor(partitionsWithProgress);
         }
     }
 }

src/test/java/cn/leancloud/kafka/consumer/RebalanceListenerTest.java

@@ -53,23 +53,35 @@ public void testPauseNotFinishedPartitionsOnPartitionAssign() {
         verify(consumer, times(1)).pause(new HashSet<>(partitionStillNeedsToPause));
     }
 
+    @Test
+    public void testPausedPartitionsWasAllRevoked() {
+        listener = new RebalanceListener<>(consumer, progress, policy, ConsumerSeekDestination.NONE);
+        final List<TopicPartition> pausedPartitions = toPartitions(IntStream.range(0, 30).boxed().collect(toList()));
+        final List<TopicPartition> partitionToResumeAfterCommit = toPartitions(IntStream.range(0, 20).boxed().collect(toList()));
+        final List<TopicPartition> assignedPartitions = toPartitions(IntStream.range(10, 20).boxed().collect(toList()));
+
+        when(consumer.paused()).thenReturn(new HashSet<>(pausedPartitions));
+        when(policy.partialCommitSync(progress)).thenReturn(new HashSet<>(partitionToResumeAfterCommit));
+
+        listener.onPartitionsRevoked(pausedPartitions);
+        listener.onPartitionsAssigned(assignedPartitions);
+
+        verify(consumer, never()).pause(any());
+    }
+
     @Test
     public void testRevokePartitions() {
+        final Set<TopicPartition> allPartitions = new HashSet<>(toPartitions(IntStream.range(0, 30).boxed().collect(toList())));
+        final Set<TopicPartition> reAssignedPartitions = new HashSet<>(toPartitions(IntStream.range(0, 20).boxed().collect(toList())));
+        final Set<TopicPartition> revokedPartitions = new HashSet<>(toPartitions(IntStream.range(20, 30).boxed().collect(toList())));
         final ProcessRecordsProgress progress = mock(ProcessRecordsProgress.class);
+        when(progress.allPartitions()).thenReturn(allPartitions);
+        when(consumer.assignment()).thenReturn(reAssignedPartitions);
         listener = new RebalanceListener<>(consumer, progress, policy, ConsumerSeekDestination.NONE);
-        final List<TopicPartition> allPartitions = toPartitions(IntStream.range(0, 30).boxed().collect(toList()));
-        final List<TopicPartition> reAssignedPartitions = toPartitions(IntStream.range(0, 20).boxed().collect(toList()));
-        final Set<TopicPartition> revokedPartitions = new HashSet<>(toPartitions(IntStream.range(20, 30).boxed().collect(toList())));
-
-        doAnswer(invocation -> {
-            assertThat(revokedPartitions).containsExactlyInAnyOrderElementsOf(invocation.getArgument(0));
-            return null;
-        }).when(progress).clearFor(anySet());
 
-        listener.onPartitionsRevoked(allPartitions);
         listener.onPartitionsAssigned(reAssignedPartitions);
 
-        verify(progress, times(1)).clearFor(anyCollection());
+        verify(progress, times(1)).clearFor(revokedPartitions);
     }
 
     @Test

src/test/java/cn/leancloud/kafka/consumer/integration/Bootstrap.java

@@ -318,7 +318,7 @@ public LcKafkaConsumer<Integer, String> buildConsumer(String consumerName, TestS
                     new StringDeserializer())
                     .workerPool(workerPool, false)
                     .recommitInterval(Duration.ofSeconds(1))
-                    .buildAsync();
+                    .buildPartialAsync();
         }
     }
 
@@ -341,7 +341,7 @@ public LcKafkaConsumer<Integer, String> buildConsumer(String consumerName, TestS
                     new IntegerDeserializer(),
                     new StringDeserializer())
                     .recommitInterval(Duration.ofSeconds(1))
-                    .buildAsync();
+                    .buildPartialAsync();
         }
     }
 

src/test/java/cn/leancloud/kafka/consumer/integration/JoinLeaveGroupTest.java

@@ -26,7 +26,7 @@ public void runTest(TestContext cxt, TestStatistics statistics) throws Exception
         lingerConsumer.subscribe(Collections.singletonList(cxt.topic()));
 
         final List<LcKafkaConsumer<Integer, String>> postJoinConsumers = new ArrayList<>();
-        for (int i = 0; i < 5; ++i) {
+        for (int i = 0; i < 8; ++i) {
             final LcKafkaConsumer<Integer, String> consumer = cxt.factory().buildConsumer(name() + "-" + i, statistics);
             consumer.subscribe(Collections.singletonList(cxt.topic()));
             postJoinConsumers.add(consumer);

src/test/java/cn/leancloud/kafka/consumer/integration/TestRunner.java

@@ -16,7 +16,7 @@ public class TestRunner implements Closeable {
 
     TestRunner(String topic, ConsumerFactory factory) {
         this.factory = factory;
-        this.producer = new TestingProducer(Duration.ofMillis(100), 2);
+        this.producer = new TestingProducer(Duration.ofMillis(100), 4);
         this.context = new TestContext(topic, producer, factory);
     }