Merge pull request #12 from leancloud/feat/force-seek-to-beginning

support force seek offsets to beginning or end

Author: ylgrgyq

.codecov.yml

@@ -2,4 +2,4 @@ coverage:
   status:
     project:
       default:
-        threshold: 0.25
+        threshold: 0.1%

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

@@ -0,0 +1,47 @@
+package cn.leancloud.kafka.consumer;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.TopicPartition;
+
+import java.util.Collection;
+
+/**
+ * The destination for a consumer reset offset operation.
+ */
+public enum ConsumerSeekDestination {
+    /**
+     * Do nothing. Do not reset offset of the given consumer to anywhere.
+     */
+    NONE {
+        @Override
+        public void seek(Consumer<?, ?> consumer, Collection<TopicPartition> partitions) {
+
+        }
+    },
+    /**
+     * Reset offsets for all the given partitions of a given consumer to beginning.
+     */
+    BEGINNING {
+        @Override
+        public void seek(Consumer<?, ?> consumer, Collection<TopicPartition> partitions) {
+            consumer.seekToBeginning(partitions);
+        }
+    },
+    /**
+     * Reset offsets for all the given partitions of a given consumer to end.
+     */
+    END {
+        @Override
+        public void seek(Consumer<?, ?> consumer, Collection<TopicPartition> partitions) {
+            consumer.seekToEnd(partitions);
+        }
+    };
+
+    /**
+     * Reset offsets for all the given partitions of a given consumer.
+     *
+     * @param consumer   the consumer to reset offset
+     * @param partitions the partitions to reset offsets
+     */
+    abstract public void seek(Consumer<?, ?> consumer, Collection<TopicPartition> partitions);
+}

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

@@ -93,8 +93,8 @@ boolean timeout() {
 
     Fetcher(LcKafkaConsumerBuilder<K, V> consumerBuilder) {
         this(consumerBuilder.getConsumer(), consumerBuilder.getPollTimeout(), consumerBuilder.getConsumerRecordHandler(),
-                consumerBuilder.getWorkerPool(), consumerBuilder.getPolicy(), consumerBuilder.gracefulShutdownTimeout(),
-                consumerBuilder.handleRecordTimeout());
+                consumerBuilder.getWorkerPool(), consumerBuilder.getPolicy(), consumerBuilder.getGracefulShutdownTimeout(),
+                consumerBuilder.getHandleRecordTimeout());
     }
 
     Fetcher(Consumer<K, V> consumer,
@@ -153,7 +153,7 @@ public void run() {
 
                 unsubscribedStatus = UnsubscribedStatus.ERROR;
                 markClosed();
-                logger.error("Fetcher quit with unexpected exception. Will rebalance after poll timeout.", ex);
+                logger.error("Fetcher quit with unexpected exception.", ex);
                 break;
             }
         }

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

@@ -51,6 +51,7 @@ int code() {
     private final ExecutorService workerPool;
     private final CommitPolicy<K, V> policy;
     private final boolean shutdownWorkerPoolOnStop;
+    private final ConsumerSeekDestination forceSeekTo;
     private volatile State state;
 
     LcKafkaConsumer(LcKafkaConsumerBuilder<K, V> builder) {
@@ -60,6 +61,7 @@ int code() {
         this.shutdownWorkerPoolOnStop = builder.isShutdownWorkerPoolOnStop();
         this.policy = builder.getPolicy();
         this.fetcher = new Fetcher<>(builder);
+        this.forceSeekTo = builder.getForceSeekTo();
         this.fetcherThread = new Thread(fetcher);
     }
 
@@ -89,7 +91,7 @@ public synchronized CompletableFuture<UnsubscribedStatus> subscribe(Collection<S
 
         ensureInInit();
 
-        consumer.subscribe(topics, new RebalanceListener<>(consumer, policy));
+        consumer.subscribe(topics, new RebalanceListener<>(consumer, policy, forceSeekTo));
 
         fetcherThread.setName(fetcherThreadName(topics));
         fetcherThread.start();
@@ -118,7 +120,7 @@ public synchronized CompletableFuture<UnsubscribedStatus> subscribe(Pattern patt
 
         ensureInInit();
 
-        consumer.subscribe(pattern, new RebalanceListener<>(consumer, policy));
+        consumer.subscribe(pattern, new RebalanceListener<>(consumer, policy, forceSeekTo));
 
         fetcherThread.setName(fetcherThreadName(pattern));
         fetcherThread.start();
@@ -189,8 +191,11 @@ private CompletableFuture<UnsubscribedStatus> setupUnsubscribedFuture() {
         // returned CompletableFuture
         final CompletableFuture<UnsubscribedStatus> ret = new CompletableFuture<>();
         fetcher.unsubscribeStatusFuture().thenAccept(status -> {
-            close();
-            ret.complete(status);
+            try {
+                close();
+            } finally {
+                ret.complete(status);
+            }
         });
         return ret;
     }

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

@@ -84,6 +84,7 @@ private static void requireArgument(boolean expression, String template, Object.
     private Duration handleRecordTimeout = Duration.ZERO;
     private Map<String, Object> configs;
     private ConsumerRecordHandler<K, V> consumerRecordHandler;
+    private ConsumerSeekDestination forceSeekTo = ConsumerSeekDestination.NONE;
     @Nullable
     private Consumer<K, V> consumer;
     @Nullable
@@ -362,6 +363,25 @@ public LcKafkaConsumerBuilder<K, V> workerPool(ExecutorService workerPool, boole
         return this;
     }
 
+    /**
+     * After subscribed to any topic, force seek offsets for every assigned partitions to beginning or end no matter
+     * where the recorded offsets for the group of this consumer are. The difference between this configuration with
+     * "auto.offset.reset" in kafka configurations is that "auto.offset.reset" only reset offset to beginning or end
+     * when there's no recorded offsets or the previous recorded offsets was removed, while {@code forceSeekTo()} will
+     * reset offsets even when there are offsets records.
+     * <p>
+     * Please note that the reset offset operation will take place only once when a partition is assigned on the
+     * first time. If a partition was revoked then assigned again, the offset of this partition will not reset again.
+     *
+     * @param destination where to seek to
+     * @return this
+     */
+    public LcKafkaConsumerBuilder<K, V> forceSeekTo(ConsumerSeekDestination destination) {
+        requireNonNull(destination, "destination");
+        this.forceSeekTo = destination;
+        return this;
+    }
+
     /**
      * Build a consumer which commits offset automatically at fixed interval. It is both OK for with or without a
      * worker thread pool. But without a worker pool, please tune the {@code max.poll.interval.ms} in
@@ -530,11 +550,11 @@ Duration getPollTimeout() {
         return pollTimeout;
     }
 
-    Duration gracefulShutdownTimeout() {
+    Duration getGracefulShutdownTimeout() {
         return gracefulShutdownTimeout;
     }
 
-    Duration handleRecordTimeout() {
+    Duration getHandleRecordTimeout() {
         return handleRecordTimeout;
     }
 
@@ -543,6 +563,10 @@ CommitPolicy<K, V> getPolicy() {
         return policy;
     }
 
+    ConsumerSeekDestination getForceSeekTo() {
+        return forceSeekTo;
+    }
+
     private Consumer<K, V> buildConsumer(boolean autoCommit) {
         checkConfigs(BasicConsumerConfigs.values());
         ENABLE_AUTO_COMMIT.set(configs, Boolean.toString(autoCommit));

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

@@ -11,19 +11,24 @@
 import java.util.HashSet;
 import java.util.Set;
 
+import static cn.leancloud.kafka.consumer.ConsumerSeekDestination.NONE;
 import static java.util.stream.Collectors.toSet;
 
 final class RebalanceListener<K, V> implements ConsumerRebalanceListener {
     private static final Logger logger = LoggerFactory.getLogger(RebalanceListener.class);
 
     private final CommitPolicy<K, V> policy;
     private final Consumer<K, V> consumer;
+    private final Set<TopicPartition> knownPartitions;
+    private final ConsumerSeekDestination forceSeekTo;
     private Set<TopicPartition> pausedPartitions;
 
-    RebalanceListener(Consumer<K, V> consumer, CommitPolicy<K, V> policy) {
+    RebalanceListener(Consumer<K, V> consumer, CommitPolicy<K, V> policy, ConsumerSeekDestination forceSeekTo) {
         this.policy = policy;
         this.consumer = consumer;
         this.pausedPartitions = Collections.emptySet();
+        this.knownPartitions = new HashSet<>();
+        this.forceSeekTo = forceSeekTo;
     }
 
     @Override
@@ -37,6 +42,21 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
 
     @Override
     public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
+        if (forceSeekTo != NONE) {
+            final Set<TopicPartition> newPartitions = new HashSet<>();
+            for (TopicPartition p : partitions) {
+                if (!knownPartitions.contains(p)) {
+                    newPartitions.add(p);
+                    logger.info("Assigned new partition: {}, force seeking it's offset to {}", p, forceSeekTo);
+                }
+            }
+
+            if (!newPartitions.isEmpty()) {
+                forceSeekTo.seek(consumer, newPartitions);
+                knownPartitions.addAll(newPartitions);
+            }
+        }
+
         if (!pausedPartitions.isEmpty()) {
             final Set<TopicPartition> partitionToPause = partitions
                     .stream()

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

@@ -115,6 +115,14 @@ public void testNegativeHandleRecordTimeout() {
                 .hasMessage("handleRecordTimeout: PT-1S (expect positive or zero duration)");
     }
 
+    @Test
+    public void testNullForceSeekDestination() {
+        assertThatThrownBy(() -> LcKafkaConsumerBuilder.newBuilder(configs, testingHandler, keyDeserializer, valueDeserializer)
+                .forceSeekTo(null))
+                .isInstanceOf(NullPointerException.class)
+                .hasMessage("destination");
+    }
+
     @Test
     public void testNegativeShutdownTimeoutMs() {
         assertThatThrownBy(() -> LcKafkaConsumerBuilder.newBuilder(configs, testingHandler, keyDeserializer, valueDeserializer)
@@ -231,6 +239,7 @@ public void testSyncConsumer() {
         final LcKafkaConsumer<Object, Object> consumer = LcKafkaConsumerBuilder.newBuilder(configs, testingHandler)
                 .mockKafkaConsumer(new MockConsumer<>(OffsetResetStrategy.LATEST))
                 .pollTimeout(Duration.ofMillis(1000))
+                .forceSeekTo(ConsumerSeekDestination.END)
                 .maxPendingAsyncCommits(100)
                 .workerPool(workerPool, false)
                 .recommitInterval(Duration.ofMinutes(20))

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

@@ -5,25 +5,27 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 
+import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.stream.IntStream;
 
 import static cn.leancloud.kafka.consumer.TestingUtils.toPartitions;
 import static java.util.stream.Collectors.toList;
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.*;
 
 public class RebalanceListenerTest {
     private CommitPolicy<Object, Object> policy;
     private Consumer<Object, Object> consumer;
-    private RebalanceListener listener;
+    private RebalanceListener<Object, Object> listener;
 
     @Before
     public void setUp() {
         policy = mock(CommitPolicy.class);
         consumer = mock(Consumer.class);
-        listener = new RebalanceListener(consumer, policy);
     }
 
     @After
@@ -33,6 +35,7 @@ public void tearDown() {
 
     @Test
     public void testPauseNotFinishedPartitionsOnPartitionAssign() {
+        listener = new RebalanceListener<>(consumer, 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()));
@@ -46,4 +49,82 @@ public void testPauseNotFinishedPartitionsOnPartitionAssign() {
 
         verify(consumer, times(1)).pause(new HashSet<>(partitionStillNeedsToPause));
     }
+
+    @Test
+    public void testForceSeekToBeginningForAllPartitions() {
+        listener = new RebalanceListener<>(consumer, policy, ConsumerSeekDestination.BEGINNING);
+
+        final List<TopicPartition> assignedPartitions = toPartitions(IntStream.range(0, 30).boxed().collect(toList()));
+        listener.onPartitionsAssigned(assignedPartitions);
+
+        verify(consumer, times(1)).seekToBeginning(argThat(partitions ->
+                new HashSet<>(assignedPartitions).equals(new HashSet<>(partitions))
+        ));
+    }
+
+    @Test
+    public void testForceSeekToBeginningForNewPartitions() {
+        listener = new RebalanceListener<>(consumer, policy, ConsumerSeekDestination.BEGINNING);
+
+        final List<TopicPartition> initialPartitions = toPartitions(IntStream.range(0, 15).boxed().collect(toList()));
+        final List<TopicPartition> newAssignedPartitions = toPartitions(IntStream.range(10, 30).boxed().collect(toList()));
+
+        listener.onPartitionsAssigned(initialPartitions);
+        listener.onPartitionsAssigned(newAssignedPartitions);
+
+        ArgumentCaptor<Collection<TopicPartition>> seekToBeginningArgs = ArgumentCaptor.forClass(Collection.class);
+
+        verify(consumer, times(2)).seekToBeginning(seekToBeginningArgs.capture());
+
+        assertThat(seekToBeginningArgs.getAllValues().get(0)).containsExactlyInAnyOrderElementsOf(initialPartitions);
+
+        newAssignedPartitions.removeAll(initialPartitions);
+        assertThat(seekToBeginningArgs.getAllValues().get(1)).containsExactlyInAnyOrderElementsOf(newAssignedPartitions);
+    }
+
+    @Test
+    public void testForceSeekToBeginningWithoutNewPartitions() {
+        listener = new RebalanceListener<>(consumer, policy, ConsumerSeekDestination.BEGINNING);
+
+        final List<TopicPartition> assignedPartitions = toPartitions(IntStream.range(0, 30).boxed().collect(toList()));
+        listener.onPartitionsAssigned(assignedPartitions);
+        listener.onPartitionsAssigned(assignedPartitions);
+        listener.onPartitionsAssigned(assignedPartitions);
+
+        verify(consumer, times(1)).seekToBeginning(argThat(partitions ->
+                new HashSet<>(assignedPartitions).equals(new HashSet<>(partitions))
+        ));
+    }
+
+    @Test
+    public void testForceSeekToEndForAllPartitions() {
+        listener = new RebalanceListener<>(consumer, policy, ConsumerSeekDestination.END);
+
+        final List<TopicPartition> assignedPartitions = toPartitions(IntStream.range(0, 30).boxed().collect(toList()));
+        listener.onPartitionsAssigned(assignedPartitions);
+
+        verify(consumer, times(1)).seekToEnd(argThat(partitions ->
+                new HashSet<>(assignedPartitions).equals(new HashSet<>(partitions))
+        ));
+    }
+
+    @Test
+    public void testForceSeekToEndForNewPartitions() {
+        listener = new RebalanceListener<>(consumer, policy, ConsumerSeekDestination.END);
+
+        final List<TopicPartition> initialPartitions = toPartitions(IntStream.range(0, 15).boxed().collect(toList()));
+        final List<TopicPartition> newAssignedPartitions = toPartitions(IntStream.range(10, 30).boxed().collect(toList()));
+
+        listener.onPartitionsAssigned(initialPartitions);
+        listener.onPartitionsAssigned(newAssignedPartitions);
+
+        ArgumentCaptor<Collection<TopicPartition>> seekToEndArgs = ArgumentCaptor.forClass(Collection.class);
+
+        verify(consumer, times(2)).seekToEnd(seekToEndArgs.capture());
+
+        assertThat(seekToEndArgs.getAllValues().get(0)).containsExactlyInAnyOrderElementsOf(initialPartitions);
+
+        newAssignedPartitions.removeAll(initialPartitions);
+        assertThat(seekToEndArgs.getAllValues().get(1)).containsExactlyInAnyOrderElementsOf(newAssignedPartitions);
+    }
 }

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

@@ -1,9 +1,6 @@
 package cn.leancloud.kafka.consumer.integration;
 
-import cn.leancloud.kafka.consumer.ConsumerRecordHandler;
-import cn.leancloud.kafka.consumer.LcKafkaConsumer;
-import cn.leancloud.kafka.consumer.LcKafkaConsumerBuilder;
-import cn.leancloud.kafka.consumer.NamedThreadFactory;
+import cn.leancloud.kafka.consumer.*;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;