Merge pull request #2 from leancloud/feat/whole-commit

whole commit periodically

Author: ylgrgyq

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

@@ -0,0 +1,41 @@
+package cn.leancloud.kafka.consumer;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.OffsetAndMetadata;
+import org.apache.kafka.common.TopicPartition;
+
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
+
+abstract class AbstractPartialCommitPolicy<K, V> extends AbstractCommitPolicy<K, V> {
+    private final Duration forceWholeCommitInterval;
+    private long nextWholeCommitNanos;
+
+    AbstractPartialCommitPolicy(Consumer<K, V> consumer, Duration forceWholeCommitInterval) {
+        super(consumer);
+        this.forceWholeCommitInterval = forceWholeCommitInterval;
+        this.nextWholeCommitNanos = nextForceWholeCommitTime(forceWholeCommitInterval);
+    }
+
+    Map<TopicPartition, OffsetAndMetadata> offsetsToPartialCommit() {
+        if (needWholeCommit()) {
+            final Map<TopicPartition, OffsetAndMetadata> ret = new HashMap<>(completedTopicOffsets);
+            for (TopicPartition partition : consumer.assignment()) {
+                ret.putIfAbsent(partition, consumer.committed(partition));
+            }
+            nextWholeCommitNanos = nextForceWholeCommitTime(forceWholeCommitInterval);
+            return ret;
+        } else {
+            return completedTopicOffsets;
+        }
+    }
+
+    private boolean needWholeCommit() {
+        return System.nanoTime() >= nextWholeCommitNanos;
+    }
+
+    private long nextForceWholeCommitTime(Duration forceWholeCommitInterval) {
+        return System.nanoTime() + forceWholeCommitInterval.toNanos();
+    }
+}

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

@@ -115,7 +115,7 @@ public void close() {
                 workerPool.awaitTermination(1, TimeUnit.DAYS);
             }
         } catch (InterruptedException ex) {
-            // ignore
+            Thread.currentThread().interrupt();
         }
     }
 

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

@@ -4,6 +4,8 @@
 import org.apache.kafka.clients.consumer.KafkaConsumer;
 import org.apache.kafka.clients.consumer.MockConsumer;
 import org.apache.kafka.common.serialization.Deserializer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import javax.annotation.Nullable;
 import java.time.Duration;
@@ -22,6 +24,8 @@
  * @param <V> the type of value for records consumed from Kafka
  */
 public final class LcKafkaConsumerBuilder<K, V> {
+    private static final Logger logger = LoggerFactory.getLogger(LcKafkaConsumerBuilder.class);
+
     /**
      * Create a {@code LcKafkaConsumerBuilder} used to build {@link LcKafkaConsumer}.
      *
@@ -84,6 +88,8 @@ private static void requireArgument(boolean expression, String template, Object.
     private Deserializer<V> valueDeserializer;
     @Nullable
     private CommitPolicy<K, V> policy;
+    @Nullable
+    private Duration forceWholeCommitInterval;
 
     private LcKafkaConsumerBuilder(Map<String, Object> kafkaConsumerConfigs,
                                    ConsumerRecordHandler<K, V> consumerRecordHandler) {
@@ -109,14 +115,14 @@ private LcKafkaConsumerBuilder(Map<String, Object> kafkaConsumerConfigs,
      * If 0, poll operation will return immediately with any records that are available currently in the buffer,
      * else returns empty.
      * <p>
-     * Must not be negative.
+     * Must not be negative. And the default {@code pollTimeoutMillis} is 100.
      *
-     * @param pollTimeoutMs the poll timeout in milliseconds
+     * @param pollTimeoutMillis the poll timeout in milliseconds
      * @return this
      */
-    public LcKafkaConsumerBuilder<K, V> pollTimeoutMillis(long pollTimeoutMs) {
-        requireArgument(pollTimeoutMs >= 0, "pollTimeoutMillis: %s (expect >= 0)", pollTimeoutMs);
-        this.pollTimeout = pollTimeoutMs;
+    public LcKafkaConsumerBuilder<K, V> pollTimeoutMillis(long pollTimeoutMillis) {
+        requireArgument(pollTimeoutMillis >= 0, "pollTimeoutMillis: %s (expect >= 0)", pollTimeoutMillis);
+        this.pollTimeout = pollTimeoutMillis;
         return this;
     }
 
@@ -127,6 +133,7 @@ public LcKafkaConsumerBuilder<K, V> pollTimeoutMillis(long pollTimeoutMs) {
      * If 0, poll operation will return immediately with any records that are available currently in the buffer,
      * else returns empty.
      * <p>
+     * The default {@code pollTimeout} is 100 millis seconds.
      *
      * @param pollTimeout the poll timeout duration
      * @return this
@@ -140,6 +147,8 @@ public LcKafkaConsumerBuilder<K, V> pollTimeout(Duration pollTimeout) {
     /**
      * Sets the amount of time to wait after calling {@link LcKafkaConsumer#close()} for
      * consumed records to handle before actually shutting down.
+     * <p>
+     * The default {@code gracefulShutdownTimeoutMillis} is 10_000.
      *
      * @param gracefulShutdownTimeoutMillis the graceful shutdown timeout in milliseconds
      * @return this
@@ -154,6 +163,8 @@ public LcKafkaConsumerBuilder<K, V> gracefulShutdownTimeoutMillis(long gracefulS
     /**
      * Sets the amount of time to wait after calling {@link LcKafkaConsumer#close()} for
      * consumed records to handle before actually shutting down.
+     * <p>
+     * The default {@code gracefulShutdownTimeout} is 10 seconds.
      *
      * @param gracefulShutdownTimeout the graceful shutdown timeout duration
      * @return this
@@ -168,6 +179,8 @@ public LcKafkaConsumerBuilder<K, V> gracefulShutdownTimeout(Duration gracefulShu
      * When using async consumer to commit offset asynchronously, this argument can force consumer to do a synchronous
      * commit after there's already this ({@code maxPendingAsyncCommits}) many async commits on the fly without
      * response from broker.
+     * <p>
+     * The default {@code maxPendingAsyncCommits} is 10.
      *
      * @param maxPendingAsyncCommits do a synchronous commit when pending async commits beyond this limit
      * @return this
@@ -179,6 +192,56 @@ public LcKafkaConsumerBuilder<K, V> maxPendingAsyncCommits(int maxPendingAsyncCo
         return this;
     }
 
+    /**
+     * The interval to commit all partitions and it's completed offsets to broker on a partial commit consumer.
+     * <p>
+     * This configuration is only valid and is required on partial commit consumer build with
+     * {@link LcKafkaConsumerBuilder#buildPartialSync()} or {@link LcKafkaConsumerBuilder#buildPartialAsync()}.
+     * For these kind of consumers, usually they only commit offsets of a partition when there was records consumed from
+     * that partition and all these consumed records was handled successfully. But we must periodically commit those
+     * subscribed partitions who have not have any records too. Otherwise, after commit offset log retention timeout,
+     * Kafka broker may forget where the current commit offset of these partition for the consumer are. Then, when the
+     * consumer crashed and recovered, if the consumer set "auto.offset.reset" configuration to "earliest", it may
+     * consume a already consumed record again. So please make sure that {@code forceWholeCommitIntervalInMillis}
+     * is within log retention time set on Kafka broker.
+     * <p>
+     * The default {@code forceWholeCommitInterval} is 1 hour.
+     *
+     * @param forceWholeCommitIntervalInMillis the interval in millis seconds to do a whole commit
+     * @return this
+     */
+    public LcKafkaConsumerBuilder<K, V> forceWholeCommitIntervalInMillis(long forceWholeCommitIntervalInMillis) {
+        requireArgument(forceWholeCommitIntervalInMillis > 0,
+                "forceWholeCommitIntervalInMillis: %s (expected > 0)", forceWholeCommitIntervalInMillis);
+
+        this.forceWholeCommitInterval = Duration.ofMillis(forceWholeCommitIntervalInMillis);
+        return this;
+    }
+
+    /**
+     * The interval to commit all partitions and it's completed offsets to broker on a partial commit consumer.
+     * <p>
+     * This configuration is only valid on partial commit consumer build with
+     * {@link LcKafkaConsumerBuilder#buildPartialSync()} or {@link LcKafkaConsumerBuilder#buildPartialAsync()}.
+     * For these kind of consumers, usually they only commit offsets of a partition when there was records consumed from
+     * that partition and all these consumed records was handled successfully. But we must periodically commit those
+     * subscribed partitions who have not have any records too. Otherwise, after commit offset log retention timeout,
+     * Kafka broker may forget where the current commit offset of these partition for the consumer are. Then, when the
+     * consumer crashed and recovered, if the consumer set "auto.offset.reset" configuration to "earliest", it may
+     * consume a already consumed record again. So please make sure that {@code forceWholeCommitInterval}
+     * is within log retention time set on Kafka broker.
+     * <p>
+     * The default {@code forceWholeCommitInterval} is 1 hour.
+     *
+     * @param forceWholeCommitInterval the interval to do a whole commit
+     * @return this
+     */
+    public LcKafkaConsumerBuilder<K, V> forceWholeCommitInterval(Duration forceWholeCommitInterval) {
+        requireNonNull(forceWholeCommitInterval, "forceWholeCommitInterval");
+        this.forceWholeCommitInterval = forceWholeCommitInterval;
+        return this;
+    }
+
     /**
      * Internal testing usage only.
      * <p>
@@ -304,8 +367,15 @@ public <K1 extends K, V1 extends V> LcKafkaConsumer<K1, V1> buildSync() {
      * @return this
      */
     public <K1 extends K, V1 extends V> LcKafkaConsumer<K1, V1> buildPartialSync() {
+        if (forceWholeCommitInterval == null) {
+            logger.warn("Force whole commit interval is not set for a partial commit consumer, the default " +
+                    "interval of 1 hour will be used.");
+            forceWholeCommitInterval = Duration.ofHours(1);
+        }
+        assert forceWholeCommitInterval != null;
+
         consumer = buildConsumer(false);
-        policy = new PartialSyncCommitPolicy<>(consumer);
+        policy = new PartialSyncCommitPolicy<>(consumer, forceWholeCommitInterval);
         return doBuild();
     }
 
@@ -365,8 +435,15 @@ public <K1 extends K, V1 extends V> LcKafkaConsumer<K1, V1> buildAsync() {
      * @return this
      */
     public <K1 extends K, V1 extends V> LcKafkaConsumer<K1, V1> buildPartialAsync() {
+        if (forceWholeCommitInterval == null) {
+            logger.warn("Force whole commit interval is not set for a partial commit consumer, the default " +
+                    "interval of 30 seconds will be used.");
+            forceWholeCommitInterval = Duration.ofSeconds(30);
+        }
+        assert forceWholeCommitInterval != null;
+
         consumer = buildConsumer(false);
-        policy = new PartialAsyncCommitPolicy<>(consumer, maxPendingAsyncCommits);
+        policy = new PartialAsyncCommitPolicy<>(consumer, forceWholeCommitInterval, maxPendingAsyncCommits);
         return doBuild();
     }
 

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

@@ -6,17 +6,18 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.time.Duration;
 import java.util.*;
 
-final class PartialAsyncCommitPolicy<K, V> extends AbstractCommitPolicy<K, V> {
+final class PartialAsyncCommitPolicy<K, V> extends AbstractPartialCommitPolicy<K, V> {
     private static final Logger logger = LoggerFactory.getLogger(PartialAsyncCommitPolicy.class);
 
     private final int maxPendingAsyncCommits;
     private int pendingAsyncCommitCounter;
     private boolean forceSync;
 
-    PartialAsyncCommitPolicy(Consumer<K, V> consumer, int maxPendingAsyncCommits) {
-        super(consumer);
+    PartialAsyncCommitPolicy(Consumer<K, V> consumer, Duration forceWholeCommitInterval, int maxPendingAsyncCommits) {
+        super(consumer, forceWholeCommitInterval);
         this.maxPendingAsyncCommits = maxPendingAsyncCommits;
     }
 
@@ -28,13 +29,13 @@ public Set<TopicPartition> tryCommit(boolean noPendingRecords) {
 
         final Set<TopicPartition> partitions = getCompletedPartitions(noPendingRecords);
         if (forceSync || pendingAsyncCommitCounter >= maxPendingAsyncCommits) {
-            consumer.commitSync(completedTopicOffsets);
+            consumer.commitSync(offsetsToPartialCommit());
             pendingAsyncCommitCounter = 0;
             forceSync = false;
             clearCachedCompletedPartitionsRecords(partitions, noPendingRecords);
         } else {
             ++pendingAsyncCommitCounter;
-            consumer.commitAsync(completedTopicOffsets, (offsets, exception) -> {
+            consumer.commitAsync(offsetsToPartialCommit(), (offsets, exception) -> {
                 --pendingAsyncCommitCounter;
                 assert pendingAsyncCommitCounter >= 0 : "actual: " + pendingAsyncCommitCounter;
                 if (exception != null) {

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

@@ -3,12 +3,13 @@
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.common.TopicPartition;
 
+import java.time.Duration;
 import java.util.Collections;
 import java.util.Set;
 
-final class PartialSyncCommitPolicy<K, V> extends AbstractCommitPolicy<K, V> {
-    PartialSyncCommitPolicy(Consumer<K, V> consumer) {
-        super(consumer);
+final class PartialSyncCommitPolicy<K, V> extends AbstractPartialCommitPolicy<K, V> {
+    PartialSyncCommitPolicy(Consumer<K, V> consumer, Duration forceWholeCommitInterval) {
+        super(consumer, forceWholeCommitInterval);
     }
 
     @Override
@@ -17,7 +18,7 @@ public Set<TopicPartition> tryCommit(boolean noPendingRecords) {
             return Collections.emptySet();
         }
 
-        consumer.commitSync(completedTopicOffsets);
+        consumer.commitSync(offsetsToPartialCommit());
 
         final Set<TopicPartition> partitions = getCompletedPartitions(noPendingRecords);
         clearCachedCompletedPartitionsRecords(partitions, noPendingRecords);

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

@@ -0,0 +1,103 @@
+package cn.leancloud.kafka.consumer;
+
+import org.apache.kafka.clients.consumer.*;
+import org.apache.kafka.common.TopicPartition;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.time.Duration;
+import java.util.*;
+
+import static cn.leancloud.kafka.consumer.TestingUtils.*;
+import static java.util.Comparator.comparing;
+import static java.util.stream.Collectors.toList;
+import static java.util.function.BinaryOperator.maxBy;
+import static java.util.stream.IntStream.range;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
+
+public class AbstractPartialCommitPolicyTest {
+    private static class TestingPartialCommitPolicy extends AbstractPartialCommitPolicy<Object, Object> {
+        TestingPartialCommitPolicy(Consumer<Object, Object> consumer, Duration forceWholeCommitInterval) {
+            super(consumer, forceWholeCommitInterval);
+        }
+
+        @Override
+        public Set<TopicPartition> tryCommit(boolean noPendingRecords) {
+            throw new UnsupportedOperationException();
+        }
+    }
+
+    private MockConsumer<Object, Object> consumer;
+    private AbstractPartialCommitPolicy<Object, Object> policy;
+
+    @Before
+    public void setUp() throws Exception {
+        consumer = new MockConsumer<>(OffsetResetStrategy.LATEST);
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        consumer.close();
+    }
+
+    @Test
+    public void testPartialCommit() {
+        policy = new TestingPartialCommitPolicy(consumer, Duration.ofHours(1));
+        final List<TopicPartition> partitions = toPartitions(range(0, 30).boxed().collect(toList()));
+        final List<ConsumerRecord<Object, Object>> pendingRecords = prepareConsumerRecords(partitions, 1, 10);
+        final List<ConsumerRecord<Object, Object>> completedRecords =
+                completeRecords(pendingRecords.subList(0, pendingRecords.size() / 2));
+
+        final Map<TopicPartition, OffsetAndMetadata> completeOffsets = buildCommitOffsets(completedRecords);
+
+        assertThat(policy.offsetsToPartialCommit()).isEqualTo(completeOffsets);
+    }
+
+    @Test
+    public void testWholeCommit() throws Exception {
+        policy = new TestingPartialCommitPolicy(consumer, Duration.ofMillis(200));
+        final Map<TopicPartition, OffsetAndMetadata> previousCommitOffsets =
+                commitRecords(completeRecords(prepareConsumerRecords(toPartitions(range(0, 10).boxed().collect(toList())), 1, 10)));
+
+        policy.offsetsToPartialCommit().clear();
+        policy.topicOffsetHighWaterMark().clear();
+        final Map<TopicPartition, OffsetAndMetadata> newOffsetsToCommit =
+                buildCommitOffsets(completeRecords(prepareConsumerRecords(toPartitions(range(10, 20).boxed().collect(toList())), 1, 10)));
+
+
+        Thread.sleep(200);
+        assertThat(policy.offsetsToPartialCommit())
+                .hasSize(previousCommitOffsets.size() + newOffsetsToCommit.size())
+                .containsAllEntriesOf(previousCommitOffsets)
+                .containsAllEntriesOf(newOffsetsToCommit);
+    }
+
+    private List<ConsumerRecord<Object, Object>> completeRecords(List<ConsumerRecord<Object, Object>> records) {
+        for (ConsumerRecord<Object, Object> record : records) {
+            policy.addPendingRecord(record);
+            policy.completeRecord(record);
+        }
+        return records;
+    }
+
+    private Map<TopicPartition, OffsetAndMetadata> commitRecords(List<ConsumerRecord<Object, Object>> records) {
+        final Map<TopicPartition, OffsetAndMetadata> offsets = buildCommitOffsets(records);
+        assignPartitions(consumer, offsets.keySet(), 0L);
+        fireConsumerRecords(consumer, records);
+        consumer.poll(0);
+        consumer.commitSync(offsets);
+        return offsets;
+    }
+
+    private Map<TopicPartition, OffsetAndMetadata> buildCommitOffsets(List<ConsumerRecord<Object, Object>> records) {
+        final Map<TopicPartition, OffsetAndMetadata> completeOffsets = new HashMap<>();
+        for (ConsumerRecord<Object, Object> record : records) {
+            completeOffsets.merge(new TopicPartition(testingTopic, record.partition()),
+                    new OffsetAndMetadata(record.offset() + 1),
+                    maxBy(comparing(OffsetAndMetadata::offset)));
+        }
+        return completeOffsets;
+    }
+}