Merge pull request #16 from leancloud/fix/partial-commit

fix partial commit

Author: ylgrgyq

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

@@ -7,13 +7,8 @@
 import org.apache.kafka.common.errors.RetriableException;
 
 import java.time.Duration;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 
-import static java.util.Comparator.comparing;
-import static java.util.function.BinaryOperator.maxBy;
 import static java.util.stream.Collectors.toSet;
 
 abstract class AbstractCommitPolicy<K, V> implements CommitPolicy<K, V> {
@@ -49,81 +44,118 @@ void onError(RetriableException e) {
         }
     }
 
-    final Map<TopicPartition, Long> topicOffsetHighWaterMark;
-    final Map<TopicPartition, OffsetAndMetadata> completedTopicOffsets;
     protected final Consumer<K, V> consumer;
+    private final Map<TopicPartition, Long> topicOffsetHighWaterMark;
+    private final Map<TopicPartition, CompletedOffsets> completedOffsets;
     private final long syncCommitRetryIntervalMs;
     private final int maxAttemptsForEachSyncCommit;
 
     AbstractCommitPolicy(Consumer<K, V> consumer, Duration syncCommitRetryInterval, int maxAttemptsForEachSyncCommit) {
         this.consumer = consumer;
         this.topicOffsetHighWaterMark = new HashMap<>();
-        this.completedTopicOffsets = new HashMap<>();
+        this.completedOffsets = new HashMap<>();
         this.syncCommitRetryIntervalMs = syncCommitRetryInterval.toMillis();
         this.maxAttemptsForEachSyncCommit = maxAttemptsForEachSyncCommit;
     }
 
     @Override
     public void markPendingRecord(ConsumerRecord<K, V> record) {
+        final TopicPartition topicPartition = new TopicPartition(record.topic(), record.partition());
         topicOffsetHighWaterMark.merge(
-                new TopicPartition(record.topic(), record.partition()),
+                topicPartition,
                 record.offset() + 1,
                 Math::max);
+
+        final CompletedOffsets offset = completedOffsets.get(topicPartition);
+        // please note that if offset exists, it could happen for record.offset() >= offset.nextOffsetToCommit()
+        // when there're duplicate records which have lower offset than our next offset to commit consumed from broker
+        if (offset == null) {
+            completedOffsets.put(topicPartition, new CompletedOffsets(record.offset() - 1L));
+        }
     }
 
     @Override
     public void markCompletedRecord(ConsumerRecord<K, V> record) {
-        completedTopicOffsets.merge(
-                new TopicPartition(record.topic(), record.partition()),
-                new OffsetAndMetadata(record.offset() + 1L),
-                maxBy(comparing(OffsetAndMetadata::offset)));
+        final CompletedOffsets offset = completedOffsets.get(new TopicPartition(record.topic(), record.partition()));
+        // offset could be null, when the partition of the record was revoked before its processing was done
+        if (offset != null) {
+            offset.addCompleteOffset(record.offset());
+        }
     }
 
     @Override
-    public Set<TopicPartition> syncPartialCommit() {
-        commitSync(completedTopicOffsets);
-        final Set<TopicPartition> partitions = checkCompletedPartitions();
-        completedTopicOffsets.clear();
-        for (TopicPartition p : partitions) {
-            topicOffsetHighWaterMark.remove(p);
-        }
-        return partitions;
+    public void revokePartitions(Collection<TopicPartition> partitions) {
+        clearProcessingRecordStatesFor(partitions);
     }
 
-    Set<TopicPartition> getCompletedPartitions(boolean noPendingRecords) {
-        final Set<TopicPartition> partitions;
-        if (noPendingRecords) {
-            assert checkCompletedPartitions().equals(topicOffsetHighWaterMark.keySet())
-                    : "expect: " + checkCompletedPartitions() + " actual: " + topicOffsetHighWaterMark.keySet();
-            partitions = new HashSet<>(topicOffsetHighWaterMark.keySet());
-        } else {
-            partitions = checkCompletedPartitions();
+    @Override
+    public Set<TopicPartition> partialCommitSync() {
+        final Map<TopicPartition, OffsetAndMetadata> offsetsToCommit = completedTopicOffsetsToCommit();
+        if (offsetsToCommit.isEmpty()) {
+            return Collections.emptySet();
         }
-        return partitions;
+        commitSyncWithRetry(offsetsToCommit);
+        updatePartialCommittedOffsets(offsetsToCommit);
+
+        return clearProcessingRecordStatesForCompletedPartitions(offsetsToCommit);
     }
 
-    void clearCachedCompletedPartitionsRecords(Set<TopicPartition> completedPartitions, boolean noPendingRecords) {
-        completedTopicOffsets.clear();
-        if (noPendingRecords) {
-            topicOffsetHighWaterMark.clear();
-        } else {
-            for (TopicPartition p : completedPartitions) {
-                topicOffsetHighWaterMark.remove(p);
-            }
-        }
+    Set<TopicPartition> fullCommitSync() {
+        commitSyncWithRetry();
+
+        final Set<TopicPartition> completePartitions = partitionsForAllRecordsStates();
+        clearAllProcessingRecordStates();
+        return completePartitions;
     }
 
     @VisibleForTesting
     Map<TopicPartition, Long> topicOffsetHighWaterMark() {
         return topicOffsetHighWaterMark;
     }
 
-    @VisibleForTesting
-    Map<TopicPartition, OffsetAndMetadata> completedTopicOffsets() {
-        return completedTopicOffsets;
+    Map<TopicPartition, OffsetAndMetadata> completedTopicOffsetsToCommit() {
+        if (noCompletedOffsets()) {
+            return Collections.emptyMap();
+        }
+
+        final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();
+        for (Map.Entry<TopicPartition, CompletedOffsets> entry : completedOffsets.entrySet()) {
+            final CompletedOffsets offset = entry.getValue();
+            if (offset.hasOffsetToCommit()) {
+                offsets.put(entry.getKey(), offset.getOffsetToCommit());
+            }
+        }
+
+        return offsets;
+    }
+
+    boolean noTopicOffsetsToCommit() {
+        if (noCompletedOffsets()) {
+            return true;
+        }
+
+        for (Map.Entry<TopicPartition, CompletedOffsets> entry : completedOffsets.entrySet()) {
+            final CompletedOffsets offset = entry.getValue();
+            if (offset.hasOffsetToCommit()) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    void updatePartialCommittedOffsets(Map<TopicPartition, OffsetAndMetadata> offsets) {
+        for (Map.Entry<TopicPartition, OffsetAndMetadata> entry : offsets.entrySet()) {
+            final CompletedOffsets offset = completedOffsets.get(entry.getKey());
+            offset.updateCommittedOffset(entry.getValue().offset());
+        }
+    }
+
+    boolean noCompletedOffsets() {
+        return completedOffsets.isEmpty();
     }
 
-    void commitSync() {
+    void commitSyncWithRetry() {
         final RetryContext context = context();
         do {
             try {
@@ -135,7 +167,7 @@ void commitSync() {
         } while (true);
     }
 
-    void commitSync(Map<TopicPartition, OffsetAndMetadata> offsets) {
+    void commitSyncWithRetry(Map<TopicPartition, OffsetAndMetadata> offsets) {
         final RetryContext context = context();
         do {
             try {
@@ -147,8 +179,34 @@ void commitSync(Map<TopicPartition, OffsetAndMetadata> offsets) {
         } while (true);
     }
 
-    private Set<TopicPartition> checkCompletedPartitions() {
-        return completedTopicOffsets
+    Set<TopicPartition> partitionsForAllRecordsStates() {
+        return new HashSet<>(topicOffsetHighWaterMark.keySet());
+    }
+
+    void clearAllProcessingRecordStates() {
+        topicOffsetHighWaterMark.clear();
+        completedOffsets.clear();
+    }
+
+    Set<TopicPartition> clearProcessingRecordStatesForCompletedPartitions(Map<TopicPartition, OffsetAndMetadata> committedOffsets) {
+        final Set<TopicPartition> partitions = partitionsToSafeResume(committedOffsets);
+        clearProcessingRecordStatesFor(partitions);
+        return partitions;
+    }
+
+    void clearProcessingRecordStatesFor(Collection<TopicPartition> partitions) {
+        for (TopicPartition p : partitions) {
+            topicOffsetHighWaterMark.remove(p);
+            completedOffsets.remove(p);
+        }
+    }
+
+    Set<TopicPartition> partitionsToSafeResume() {
+        return partitionsToSafeResume(completedTopicOffsetsToCommit());
+    }
+
+    Set<TopicPartition> partitionsToSafeResume(Map<TopicPartition, OffsetAndMetadata> completedOffsets) {
+        return completedOffsets
                 .entrySet()
                 .stream()
                 .filter(entry -> topicOffsetMeetHighWaterMark(entry.getKey(), entry.getValue()))

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

@@ -7,6 +7,7 @@
 import java.time.Duration;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 
 abstract class AbstractRecommitAwareCommitPolicy<K, V> extends AbstractCommitPolicy<K, V> {
     private final Duration recommitInterval;
@@ -21,34 +22,46 @@ abstract class AbstractRecommitAwareCommitPolicy<K, V> extends AbstractCommitPol
         updateNextRecommitTime(System.nanoTime());
     }
 
-    Map<TopicPartition, OffsetAndMetadata> offsetsForRecommit() {
-        assert needRecommit() : "current nanos: " + System.nanoTime() + " nextRecommitNanos:" + nextRecommitNanos;
-
-        final Map<TopicPartition, OffsetAndMetadata> ret = new HashMap<>(completedTopicOffsets);
-        for (TopicPartition partition : consumer.assignment()) {
-            final OffsetAndMetadata offset = consumer.committed(partition);
-            if (offset != null) {
-                ret.putIfAbsent(partition, offset);
-            }
+    @Override
+    public final Set<TopicPartition> tryCommit(boolean noPendingRecords) {
+        if (needRecommit()) {
+            commitSyncWithRetry(offsetsForRecommit());
+            updateNextRecommitTime();
         }
-
-        return ret;
+        return tryCommit0(noPendingRecords);
     }
 
-    boolean needRecommit() {
-        return System.nanoTime() >= nextRecommitNanos;
-    }
+    abstract Set<TopicPartition> tryCommit0(boolean noPendingRecords);
 
     void updateNextRecommitTime() {
         updateNextRecommitTime(System.nanoTime());
     }
 
+    @VisibleForTesting
+    void updateNextRecommitTime(long currentNanos) {
+        nextRecommitNanos = currentNanos + recommitInterval.toNanos();
+    }
+
     @VisibleForTesting
     long nextRecommitNanos() {
         return nextRecommitNanos;
     }
 
-    private void updateNextRecommitTime(long currentNanos) {
-        nextRecommitNanos = currentNanos + recommitInterval.toNanos();
+    private boolean needRecommit() {
+        return System.nanoTime() >= nextRecommitNanos;
+    }
+
+    private Map<TopicPartition, OffsetAndMetadata> offsetsForRecommit() {
+        assert needRecommit() : "current nanos: " + System.nanoTime() + " nextRecommitNanos:" + nextRecommitNanos;
+
+        final Map<TopicPartition, OffsetAndMetadata> ret = new HashMap<>();
+        for (TopicPartition partition : consumer.assignment()) {
+            final OffsetAndMetadata offset = consumer.committed(partition);
+            if (offset != null) {
+                ret.putIfAbsent(partition, offset);
+            }
+        }
+
+        return ret;
     }
 }

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

@@ -8,11 +8,11 @@
 import org.slf4j.LoggerFactory;
 
 import java.time.Duration;
-import java.util.Collections;
-import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 
+import static java.util.Collections.emptySet;
+
 final class AsyncCommitPolicy<K, V> extends AbstractRecommitAwareCommitPolicy<K, V> {
     private static final Logger logger = LoggerFactory.getLogger(AsyncCommitPolicy.class);
 
@@ -32,21 +32,21 @@ final class AsyncCommitPolicy<K, V> extends AbstractRecommitAwareCommitPolicy<K,
     }
 
     @Override
-    public Set<TopicPartition> tryCommit(boolean noPendingRecords) {
-        if (!noPendingRecords || completedTopicOffsets.isEmpty()) {
-            if (needRecommit()) {
-                commit(offsetsForRecommit());
-            }
-            return Collections.emptySet();
+    Set<TopicPartition> tryCommit0(boolean noPendingRecords) {
+        // with forceSync mark it means a previous async commit was failed, so
+        // we do a sync commit no matter if there's any pending records or completed offsets
+        if (!forceSync && (!noPendingRecords || noTopicOffsetsToCommit())) {
+            return emptySet();
         }
 
+        final Set<TopicPartition> partitions = partitionsForAllRecordsStates();
         commit();
 
-        final Set<TopicPartition> partitions = new HashSet<>(completedTopicOffsets.keySet());
-        // it's OK to clear these collections here and we will not left any complete offset without commit even
-        // when this async commit failed because if the async commit failed we will do a sync commit after all
-        completedTopicOffsets.clear();
-        topicOffsetHighWaterMark.clear();
+        // for our commit policy, no matter syncCommit or asyncCommit we are using, 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();
+
         return partitions;
     }
 
@@ -60,44 +60,15 @@ boolean forceSync() {
         return forceSync;
     }
 
-    @VisibleForTesting
-    void setForceSync(boolean forceSync) {
-        this.forceSync = forceSync;
-    }
-
     private void commit() {
-        commit(Collections.emptyMap());
-    }
-
-    private void commit(Map<TopicPartition, OffsetAndMetadata> offsets) {
         if (forceSync || pendingAsyncCommitCounter >= maxPendingAsyncCommits) {
-            syncCommit(offsets);
+            commitSyncWithRetry();
             pendingAsyncCommitCounter = 0;
             forceSync = false;
+            clearAllProcessingRecordStates();
         } else {
-            asyncCommit(offsets);
-        }
-
-        // 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();
-    }
-
-    private void asyncCommit(Map<TopicPartition, OffsetAndMetadata> offsets) {
-        ++pendingAsyncCommitCounter;
-        if (offsets.isEmpty()) {
+            ++pendingAsyncCommitCounter;
             consumer.commitAsync(callback);
-        } else {
-            consumer.commitAsync(offsets, callback);
-        }
-    }
-
-    private void syncCommit(Map<TopicPartition, OffsetAndMetadata> offsets) {
-        if (offsets.isEmpty()) {
-            commitSync();
-        } else {
-            commitSync(offsets);
         }
     }
 
@@ -109,6 +80,8 @@ public void onComplete(Map<TopicPartition, OffsetAndMetadata> offsets, Exception
             if (exception != null) {
                 logger.warn("Failed to commit offsets: " + offsets + " asynchronously", exception);
                 forceSync = true;
+            } else {
+                clearProcessingRecordStatesForCompletedPartitions(offsets);
             }
         }
     }