(fix) partial async commit

Author: ylgrgyq

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

@@ -26,23 +26,23 @@ abstract class AbstractCommitPolicy<K, V> implements CommitPolicy<K, V> {
     }
 
     @Override
-    public void addPendingRecord(ConsumerRecord<K, V> record) {
+    public void markPendingRecord(ConsumerRecord<K, V> record) {
         topicOffsetHighWaterMark.merge(
                 new TopicPartition(record.topic(), record.partition()),
                 record.offset() + 1,
                 Math::max);
     }
 
     @Override
-    public void completeRecord(ConsumerRecord<K, V> record) {
+    public void markCompletedRecord(ConsumerRecord<K, V> record) {
         completedTopicOffsets.merge(
                 new TopicPartition(record.topic(), record.partition()),
                 new OffsetAndMetadata(record.offset() + 1L),
                 maxBy(comparing(OffsetAndMetadata::offset)));
     }
 
     @Override
-    public Set<TopicPartition> partialCommit() {
+    public Set<TopicPartition> syncPartialCommit() {
         consumer.commitSync(completedTopicOffsets);
         final Set<TopicPartition> partitions = checkCompletedPartitions();
         completedTopicOffsets.clear();

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

@@ -73,7 +73,10 @@ private void commit(Map<TopicPartition, OffsetAndMetadata> offsets) {
         } else {
             asyncCommit(offsets);
         }
-        // update next recommit time even if async commit failed, we tolerate this situation
+
+        // for our commit policy, no matter syncCommit or asyncCommit we use, we always
+        // commit all assigned offsets, so we can update recommit time here safely. And
+        // we don't mind that if the async commit request failed, we tolerate this situation
         updateNextRecommitTime();
     }
 

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

@@ -6,11 +6,42 @@
 import java.util.Set;
 
 interface CommitPolicy<K, V> {
-    void addPendingRecord(ConsumerRecord<K, V> record);
+    /**
+     * Mark an {@link ConsumerRecord} as pending before processing it. So {@link CommitPolicy} can know which and
+     * how many records we need to process. It is called by {@link Fetcher} when {@code Fetcher} fetched any
+     * {@link ConsumerRecord}s from Broker.
+     *
+     * @param record the {@link ConsumerRecord} need to process
+     */
+    void markPendingRecord(ConsumerRecord<K, V> record);
 
-    void completeRecord(ConsumerRecord<K, V> record);
+    /**
+     * Mark an {@link ConsumerRecord} as completed after processing it. So {@link CommitPolicy} can know which and
+     * how many records we have processed. It is called by {@link Fetcher} when {@code Fetcher} make sure that
+     * a {@code ConsumerRecord} was processed successfully.
+     *
+     * @param record the {@link ConsumerRecord} processed
+     */
+    void markCompletedRecord(ConsumerRecord<K, V> record);
 
+    /**
+     * Try commit offset for any {@link TopicPartition}s which has processed {@link ConsumerRecord}s based on the
+     * intrinsic policy of this {@link CommitPolicy}. This method is called whenever there're any
+     * {@link ConsumerRecord} processed.
+     *
+     * @param noPendingRecords is there any pending records which have not been processed. Though {@link CommitPolicy}
+     *                         can calculate this value by itself, we still pass this value as {@link Fetcher} can
+     *                         calculate this value much quicker
+     * @return those {@link TopicPartition}s which have no pending {@code ConsumerRecord}s
+     */
     Set<TopicPartition> tryCommit(boolean noPendingRecords);
 
-    Set<TopicPartition> partialCommit();
+    /**
+     * Do a dedicated partition commit synchronously which only commit those {@link ConsumerRecord}s that have
+     * processed but have not been committed yet. Usually it is called when {@link LcKafkaConsumer} is about to
+     * shutdown or when some partitions was revoked.
+     *
+     * @return those {@link TopicPartition}s which have no pending {@code ConsumerRecord}s
+     */
+    Set<TopicPartition> syncPartialCommit();
 }

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

@@ -25,10 +25,6 @@ static class TimeoutFuture<K, V> implements Future<ConsumerRecord<K, V>> {
         private final long timeoutAtNanos;
         private final Time time;
 
-        TimeoutFuture(Future<ConsumerRecord<K, V>> wrappedFuture) {
-            this(wrappedFuture, Long.MAX_VALUE);
-        }
-
         TimeoutFuture(Future<ConsumerRecord<K, V>> wrappedFuture, long timeoutInNanos) {
             this(wrappedFuture, timeoutInNanos, Time.SYSTEM);
         }
@@ -193,13 +189,13 @@ private void dispatchFetchedRecords(ConsumerRecords<K, V> records) {
                 return record;
             });
             pendingFutures.put(record, timeoutAwareFuture(future));
-            policy.addPendingRecord(record);
+            policy.markPendingRecord(record);
         }
     }
 
-    private TimeoutFuture<K, V> timeoutAwareFuture(Future<ConsumerRecord<K, V>> future) {
+    private Future<ConsumerRecord<K, V>> timeoutAwareFuture(Future<ConsumerRecord<K, V>> future) {
         if (unlimitedHandleRecordTime()) {
-            return new TimeoutFuture<>(future);
+            return future;
         } else {
             return new TimeoutFuture<>(future, handleRecordTimeoutNanos);
         }
@@ -219,7 +215,7 @@ assert record != null;
         assert !future.isCancelled();
         final Future<ConsumerRecord<K, V>> v = pendingFutures.remove(record);
         assert v != null;
-        policy.completeRecord(record);
+        policy.markCompletedRecord(record);
     }
 
     private void processTimeoutRecords() throws TimeoutException {
@@ -228,6 +224,7 @@ private void processTimeoutRecords() throws TimeoutException {
         }
 
         for (Map.Entry<ConsumerRecord<K, V>, Future<ConsumerRecord<K, V>>> entry : pendingFutures.entrySet()) {
+            // we can sure that this conversion must be success
             final TimeoutFuture<K, V> future = (TimeoutFuture<K, V>) entry.getValue();
             if (future.timeout()) {
                 future.cancel(false);
@@ -262,7 +259,7 @@ private void gracefulShutdown(UnsubscribedStatus unsubscribedStatus) {
         long shutdownTimeout = 0L;
         try {
             shutdownTimeout = waitPendingFuturesDone();
-            policy.partialCommit();
+            policy.syncPartialCommit();
             pendingFutures.clear();
         } catch (Exception ex) {
             logger.error("Graceful shutdown got unexpected exception", ex);
@@ -285,7 +282,7 @@ private long waitPendingFuturesDone() {
                 assert remain >= 0;
                 final ConsumerRecord<K, V> record = future.get(remain, TimeUnit.MILLISECONDS);
                 assert record != null;
-                policy.completeRecord(record);
+                policy.markCompletedRecord(record);
             } catch (TimeoutException ex) {
                 future.cancel(false);
             } catch (InterruptedException ex) {

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

@@ -15,12 +15,12 @@ static <K, V> NoOpCommitPolicy<K, V> getInstance() {
     }
 
     @Override
-    public void addPendingRecord(ConsumerRecord<K, V> record) {
+    public void markPendingRecord(ConsumerRecord<K, V> record) {
 
     }
 
     @Override
-    public void completeRecord(ConsumerRecord<K, V> record) {
+    public void markCompletedRecord(ConsumerRecord<K, V> record) {
 
     }
 
@@ -30,7 +30,7 @@ public Set<TopicPartition> tryCommit(boolean noPendingRecords) {
     }
 
     @Override
-    public Set<TopicPartition> partialCommit() {
+    public Set<TopicPartition> syncPartialCommit() {
         return Collections.emptySet();
     }
 }

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

@@ -10,39 +10,42 @@
 import java.time.Duration;
 import java.util.*;
 
-final class PartialAsyncCommitPolicy<K, V> extends AbstractPartialCommitPolicy<K, V> {
+final class PartialAsyncCommitPolicy<K, V> extends AbstractRecommitAwareCommitPolicy<K, V> {
     private static final Logger logger = LoggerFactory.getLogger(PartialAsyncCommitPolicy.class);
 
     private final int maxPendingAsyncCommits;
     private final OffsetCommitCallback callback;
+    private final Map<TopicPartition, OffsetAndMetadata> pendingAsyncCommitOffset;
     private int pendingAsyncCommitCounter;
     private boolean forceSync;
 
     PartialAsyncCommitPolicy(Consumer<K, V> consumer, Duration forceWholeCommitInterval, int maxPendingAsyncCommits) {
         super(consumer, forceWholeCommitInterval);
         this.maxPendingAsyncCommits = maxPendingAsyncCommits;
         this.callback = new AsyncCommitCallback();
+        this.pendingAsyncCommitOffset = new HashMap<>();
     }
 
     @Override
     public Set<TopicPartition> tryCommit(boolean noPendingRecords) {
-        final Map<TopicPartition, OffsetAndMetadata> offsets = offsetsForPartialCommit();
+        final boolean syncCommit = forceSync || pendingAsyncCommitCounter >= maxPendingAsyncCommits;
+        final Map<TopicPartition, OffsetAndMetadata> offsets = offsetsForPartialCommit(syncCommit);
+
         if (offsets.isEmpty()) {
             return Collections.emptySet();
         } else {
             final Set<TopicPartition> partitions = getCompletedPartitions(noPendingRecords);
-            if (forceSync || pendingAsyncCommitCounter >= maxPendingAsyncCommits) {
+            if (syncCommit) {
                 consumer.commitSync(offsets);
+                pendingAsyncCommitOffset.clear();
                 pendingAsyncCommitCounter = 0;
                 forceSync = false;
                 clearCachedCompletedPartitionsRecords(partitions, noPendingRecords);
             } else {
                 ++pendingAsyncCommitCounter;
                 consumer.commitAsync(offsets, callback);
+                pendingAsyncCommitOffset.putAll(offsets);
             }
-
-            // update next recommit time even if async commit failed, we tolerate this situation
-            updateNextRecommitTime();
             return partitions;
         }
     }
@@ -57,6 +60,23 @@ boolean forceSync() {
         return forceSync;
     }
 
+    private Map<TopicPartition, OffsetAndMetadata> offsetsForPartialCommit(boolean syncCommit) {
+        final Map<TopicPartition, OffsetAndMetadata> offsets;
+        if (needRecommit()) {
+            offsets = offsetsForRecommit();
+            // we tolerate the commit failure when using async commit
+            updateNextRecommitTime();
+        } else if (syncCommit) {
+            offsets = completedTopicOffsets;
+        } else {
+            offsets = new HashMap<>(completedTopicOffsets);
+            for (Map.Entry<TopicPartition, OffsetAndMetadata> entry : pendingAsyncCommitOffset.entrySet()) {
+                offsets.remove(entry.getKey(), entry.getValue());
+            }
+        }
+        return offsets;
+    }
+
     private class AsyncCommitCallback implements OffsetCommitCallback {
         @Override
         public void onComplete(Map<TopicPartition, OffsetAndMetadata> offsets, Exception exception) {
@@ -68,11 +88,12 @@ public void onComplete(Map<TopicPartition, OffsetAndMetadata> offsets, Exception
                 logger.warn("Failed to commit offset: " + offsets + " asynchronously", exception);
                 forceSync = true;
             } else {
-                final Map<TopicPartition, OffsetAndMetadata> completeOffsets =
-                        offsets == completedTopicOffsets ? new HashMap<>(offsets) : offsets;
-                for (Map.Entry<TopicPartition, OffsetAndMetadata> entry : completeOffsets.entrySet()) {
-                    completedTopicOffsets.remove(entry.getKey(), entry.getValue());
+                for (Map.Entry<TopicPartition, OffsetAndMetadata> entry : offsets.entrySet()) {
                     topicOffsetHighWaterMark.remove(entry.getKey(), entry.getValue().offset());
+                    // we don't clear pendingAsyncCommitOffset here
+                    // because they are cleared on sync commit which occurs every maxPendingAsyncCommits async commits
+                    // todo: maybe we don't need to clean completedTopicOffsets too
+                    completedTopicOffsets.remove(entry.getKey(), entry.getValue());
                 }
             }
         }

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

@@ -9,7 +9,7 @@
 import java.util.Map;
 import java.util.Set;
 
-final class PartialSyncCommitPolicy<K, V> extends AbstractPartialCommitPolicy<K, V> {
+final class PartialSyncCommitPolicy<K, V> extends AbstractRecommitAwareCommitPolicy<K, V> {
     PartialSyncCommitPolicy(Consumer<K, V> consumer, Duration forceWholeCommitInterval) {
         super(consumer, forceWholeCommitInterval);
     }
@@ -19,7 +19,6 @@ public Set<TopicPartition> tryCommit(boolean noPendingRecords) {
         final Map<TopicPartition, OffsetAndMetadata> offsets = offsetsForPartialCommit();
         if (!offsets.isEmpty()) {
             consumer.commitSync(offsets);
-            updateNextRecommitTime();
         }
 
         if (completedTopicOffsets.isEmpty()) {
@@ -30,4 +29,14 @@ public Set<TopicPartition> tryCommit(boolean noPendingRecords) {
             return partitions;
         }
     }
+
+    private Map<TopicPartition, OffsetAndMetadata> offsetsForPartialCommit() {
+        if (needRecommit()) {
+            final Map<TopicPartition, OffsetAndMetadata> offsets = offsetsForRecommit();
+            updateNextRecommitTime();
+            return offsets;
+        } else {
+            return completedTopicOffsets;
+        }
+    }
 }

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

@@ -36,7 +36,7 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
         pausedPartitions = consumer.paused();
         if (!pausedPartitions.isEmpty()) {
             pausedPartitions = new HashSet<>(pausedPartitions);
-            pausedPartitions.removeAll(policy.partialCommit());
+            pausedPartitions.removeAll(policy.syncPartialCommit());
         }
     }