Merge pull request #22 from leancloud/feat/no-commit-while-rebalance

no commit while partition rebalancing

Author: ylgrgyq

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

@@ -46,6 +46,7 @@ void onError(RetriableException e) {
     protected final Consumer<K, V> consumer;
     private final long syncCommitRetryIntervalMs;
     private final int maxAttemptsForEachSyncCommit;
+    private boolean commitPuased;
 
     AbstractCommitPolicy(Consumer<K, V> consumer, Duration syncCommitRetryInterval, int maxAttemptsForEachSyncCommit) {
         this.consumer = consumer;
@@ -65,6 +66,21 @@ public Set<TopicPartition> partialCommitSync(ProcessRecordsProgress progress) {
         return progress.clearCompletedPartitions(offsetsToCommit);
     }
 
+    @Override
+    public void pauseCommit() {
+        commitPuased = true;
+    }
+
+    @Override
+    public void resumeCommit() {
+        commitPuased = false;
+    }
+
+    @Override
+    public boolean commitPaused() {
+        return commitPuased;
+    }
+
     Set<TopicPartition> fullCommitSync(ProcessRecordsProgress progress) {
         commitSyncWithRetry();
 

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

@@ -5,6 +5,7 @@
 import org.apache.kafka.common.TopicPartition;
 
 import java.time.Duration;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
@@ -24,6 +25,10 @@ abstract class AbstractRecommitAwareCommitPolicy<K, V> extends AbstractCommitPol
 
     @Override
     public final Set<TopicPartition> tryCommit(boolean noPendingRecords, ProcessRecordsProgress progress) {
+        if (commitPaused()) {
+            return Collections.emptySet();
+        }
+
         if (needRecommit()) {
             commitSyncWithRetry(offsetsForRecommit());
             updateNextRecommitTime();

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

@@ -16,7 +16,7 @@ final class AutoCommitPolicy<K, V> extends AbstractCommitPolicy<K, V> {
 
     @Override
     public Set<TopicPartition> tryCommit(boolean noPendingRecords, ProcessRecordsProgress progress) {
-        if (progress.noOffsetsToCommit()) {
+        if (commitPaused() || progress.noOffsetsToCommit()) {
             return emptySet();
         }
 

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

@@ -3,7 +3,6 @@
 import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.common.TopicPartition;
 
-import java.util.Collection;
 import java.util.Set;
 
 interface CommitPolicy {
@@ -27,4 +26,25 @@ interface CommitPolicy {
      * @return those {@link TopicPartition}s which have no pending {@code ConsumerRecord}s
      */
     Set<TopicPartition> partialCommitSync(ProcessRecordsProgress progress);
+
+    /**
+     * Pause commit until {@link #resumeCommit()} is called. Next {@link #tryCommit(boolean, ProcessRecordsProgress)}
+     * will return empty Set without commit anything.
+     * <p>
+     * Please note that this method has no effect on {@link #partialCommitSync(ProcessRecordsProgress)} which will
+     * even this method was called.
+     */
+    void pauseCommit();
+
+    /**
+     * Resume commit and next {@link #tryCommit(boolean, ProcessRecordsProgress)} will do the commitment stuff.
+     */
+    void resumeCommit();
+
+    /**
+     * Check weather commit was paused.
+     *
+     * @return true when {@link #pauseCommit()} has been called
+     */
+    boolean commitPaused();
 }

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

@@ -22,4 +22,18 @@ public Set<TopicPartition> tryCommit(boolean noPendingRecords, ProcessRecordsPro
     public Set<TopicPartition> partialCommitSync(ProcessRecordsProgress progress) {
         return emptySet();
     }
+
+    @Override
+    public void pauseCommit() {
+    }
+
+    @Override
+    public void resumeCommit() {
+
+    }
+
+    @Override
+    public boolean commitPaused() {
+        return false;
+    }
 }

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

@@ -40,6 +40,7 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
             pausedPartitions = new HashSet<>(pausedPartitions);
             pausedPartitions.removeAll(policy.partialCommitSync(progress));
         }
+        policy.pauseCommit();
     }
 
     @Override
@@ -48,11 +49,14 @@ public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
             seekOnAssignedPartitions(partitions);
         }
 
+        clearProgressForRevokedPartitions();
+
+        policy.resumeCommit();
+        final Set<TopicPartition> resumedPartitions = policy.partialCommitSync(progress);
         if (!pausedPartitions.isEmpty()) {
+            pausedPartitions.removeAll(resumedPartitions);
             pausePreviousPausedPartitions(partitions);
         }
-
-        clearProgressForRevokedPartitions();
     }
 
     private void seekOnAssignedPartitions(Collection<TopicPartition> partitions) {

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

@@ -180,6 +180,15 @@ public void testInterruptOnRetry() {
         assertThat(Thread.interrupted());
     }
 
+    @Test
+    public void testPauseResume() {
+        assertThat(policy.commitPaused()).isFalse();
+        policy.pauseCommit();
+        assertThat(policy.commitPaused()).isTrue();
+        policy.resumeCommit();
+        assertThat(policy.commitPaused()).isFalse();
+    }
+
     private TopicPartition partition(int partition) {
         return new TopicPartition(testingTopic, partition);
     }

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

@@ -60,4 +60,26 @@ public void testRecommit() {
         verify(consumer, times(1)).commitSync(previousCommitOffsets);
         assertThat(policy.nextRecommitNanos()).isGreaterThan(now + Duration.ofMillis(200).toNanos());
     }
+
+    @Test
+    public void testNoRecommitWhenPaused() {
+        final Consumer<Object, Object> consumer = mock(Consumer.class);
+        final List<ConsumerRecord<Object, Object>> prevRecords = prepareConsumerRecords(toPartitions(IntStream.range(0, 10).boxed().collect(toList())), 1, 10);
+        final Map<TopicPartition, OffsetAndMetadata> previousCommitOffsets = buildCommitOffsets(prevRecords);
+        final long now = System.nanoTime();
+
+        when(consumer.assignment()).thenReturn(new HashSet<>(partitions));
+        when(consumer.committed(any())).thenAnswer(invocation ->
+                previousCommitOffsets.get(invocation.getArgument(0))
+        );
+
+        policy = new TestingPolicy(consumer, Duration.ZERO, 3, Duration.ofMillis(200));
+        policy.pauseCommit();
+
+        policy.updateNextRecommitTime(now - Duration.ofMillis(200).toNanos());
+        assertThat(policy.tryCommit(true, new ProcessRecordsProgress())).isEmpty();
+
+        verify(consumer, never()).commitSync(any());
+        assertThat(policy.nextRecommitNanos()).isEqualTo(now);
+    }
 }

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

@@ -37,6 +37,24 @@ public void tearDown() {
         consumer.close();
     }
 
+    @Test
+    public void testPauseResumeCommit() {
+        addCompleteRecordsInPolicy(progress, pendingRecords);
+
+        policy.pauseCommit();
+        assertThat(policy.tryCommit(true, progress)).isEmpty();
+        assertThat(progress.noCompletedRecords()).isFalse();
+        assertThat(progress.noPendingRecords()).isFalse();
+
+        policy.resumeCommit();
+        assertThat(policy.tryCommit(true, progress))
+                .hasSize(partitions.size())
+                .containsExactlyInAnyOrderElementsOf(partitions);
+
+        assertThat(progress.noCompletedRecords()).isTrue();
+        assertThat(progress.noPendingRecords()).isTrue();
+    }
+
     @Test
     public void testOnlyConsumedRecords() {
         assertThat(policy.tryCommit(true, progress)).isEmpty();

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

@@ -7,10 +7,7 @@
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 
-import java.util.Collection;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Set;
+import java.util.*;
 import java.util.stream.IntStream;
 
 import static cn.leancloud.kafka.consumer.TestingUtils.toPartitions;
@@ -84,6 +81,29 @@ public void testRevokePartitions() {
         verify(progress, times(1)).clearFor(revokedPartitions);
     }
 
+    @Test
+    public void testPauseResumeCommit() {
+        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, 25).boxed().collect(toList()));
+        final List<TopicPartition> partitionStillNeedsToPause = toPartitions(IntStream.range(20, 25).boxed().collect(toList()));
+
+        when(consumer.paused()).thenReturn(new HashSet<>(pausedPartitions));
+        when(policy.partialCommitSync(progress))
+                // Return empty Set on the first call
+                .thenReturn(Collections.emptySet())
+                // Return resumed partitions on the second call
+                .thenReturn(new HashSet<>(partitionToResumeAfterCommit));
+
+        listener.onPartitionsRevoked(pausedPartitions);
+        verify(policy, times(1)).pauseCommit();
+        listener.onPartitionsAssigned(assignedPartitions);
+        verify(policy, times(1)).resumeCommit();
+
+        verify(consumer, times(1)).pause(new HashSet<>(partitionStillNeedsToPause));
+    }
+
     @Test
     public void testForceSeekToBeginningForAllPartitions() {
         listener = new RebalanceListener<>(consumer, progress, policy, ConsumerSeekDestination.BEGINNING);