START: Partition state tracking

step: remove complex dirty/clean management

step: compiles

step: move offset search into PartitionState

step: move offset search into PartitionState

step: replace search with state lookup

Author: astubbs

parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/AbstractOffsetCommitter.java

@@ -3,6 +3,8 @@
 /*-
  * Copyright (C) 2020-2022 Confluent, Inc.
  */
+
+import io.confluent.parallelconsumer.state.PartitionState;
 import io.confluent.parallelconsumer.state.WorkManager;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -11,6 +13,7 @@
 import org.apache.kafka.common.TopicPartition;
 
 import java.util.Map;
+import java.util.stream.Collectors;
 
 @Slf4j
 @RequiredArgsConstructor
@@ -28,15 +31,18 @@ public void retrieveOffsetsAndCommit() {
         // todo shouldn't be removed until commit succeeds (there's no harm in committing the same offset twice)
         preAcquireWork();
         try {
-            Map<TopicPartition, OffsetAndMetadata> offsetsToCommit = wm.findCompletedEligibleOffsetsAndRemove();
+            Map<TopicPartition, PartitionState.OffsetPair> offsetsToCommit = wm.findCompletedEligibleOffsetsAndRemove();
             if (offsetsToCommit.isEmpty()) {
                 log.debug("No offsets ready");
             } else {
                 log.debug("Will commit offsets for {} partition(s): {}", offsetsToCommit.size(), offsetsToCommit);
                 ConsumerGroupMetadata groupMetadata = consumerMgr.groupMetadata();
 
                 log.debug("Begin commit");
-                commitOffsets(offsetsToCommit, groupMetadata);
+                Map<TopicPartition, OffsetAndMetadata> offsetsToCommitOFFSEEETS = offsetsToCommit.entrySet().parallelStream()
+                        .collect(Collectors.toMap(Map.Entry::getKey,
+                                o -> o.getValue().getSync()));
+                commitOffsets(offsetsToCommitOFFSEEETS, groupMetadata);
 
                 log.debug("On commit success");
                 onOffsetCommitSuccess(offsetsToCommit);
@@ -54,7 +60,7 @@ protected void preAcquireWork() {
         // default noop
     }
 
-    private void onOffsetCommitSuccess(final Map<TopicPartition, OffsetAndMetadata> committed) {
+    private void onOffsetCommitSuccess(final Map<TopicPartition, PartitionState.OffsetPair> committed) {
         wm.onOffsetCommitSuccess(committed);
     }
 

parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/internal/ProducerManager.java

@@ -21,7 +21,6 @@
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.time.Duration;
-import java.time.temporal.TemporalUnit;
 import java.util.ConcurrentModificationException;
 import java.util.Map;
 import java.util.concurrent.Future;

parallel-consumer-core/src/main/java/io/confluent/parallelconsumer/state/PartitionMonitor.java

@@ -11,6 +11,7 @@
 import io.confluent.parallelconsumer.internal.InternalRuntimeError;
 import io.confluent.parallelconsumer.offsets.EncodingNotSupportedException;
 import io.confluent.parallelconsumer.offsets.OffsetMapCodecManager;
+import io.confluent.parallelconsumer.state.PartitionState.OffsetPair;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.Setter;
@@ -49,7 +50,7 @@ public class PartitionMonitor<K, V> implements ConsumerRebalanceListener {
      */
     @Getter
     @Setter
-    private double USED_PAYLOAD_THRESHOLD_MULTIPLIER = 0.75;
+    private static double USED_PAYLOAD_THRESHOLD_MULTIPLIER = 0.75;
 
     private final Consumer<K, V> consumer;
 
@@ -172,7 +173,7 @@ public void onPartitionsLost(Collection<TopicPartition> partitions) {
      * When commits are made to broker, we can throw away all the individually tracked offsets before the committed
      * offset.
      */
-    public void onOffsetCommitSuccess(Map<TopicPartition, OffsetAndMetadata> committed) {
+    public void onOffsetCommitSuccess(Map<TopicPartition, OffsetPair> committed) {
         // partitionOffsetHighWaterMarks this will get overwritten in due course
         committed.forEach((tp, meta) -> {
             var partition = getPartitionState(tp);
@@ -408,6 +409,11 @@ public void onSuccess(WorkContainer<K, V> wc) {
         partitionState.onSuccess(wc);
     }
 
+    public void onFailure(WorkContainer<K, V> wc) {
+        PartitionState<K, V> partitionState = getPartitionState(wc.getTopicPartition());
+        partitionState.onFailure(wc);
+    }
+
     /**
      * Takes a record as work and puts it into internal queues, unless it's been previously recorded as completed as per
      * loaded records.
@@ -445,16 +451,47 @@ boolean maybeRegisterNewRecordAsWork(final ConsumerRecord<K, V> rec) {
         }
     }
 
-    public Map<TopicPartition, OffsetAndMetadata> findCompletedEligibleOffsetsAndRemove() {
+    public Map<TopicPartition, OffsetPair> findCompletedEligibleOffsetsAndRemove() {
         return findCompletedEligibleOffsetsAndRemove(true);
     }
 
+    private Map<TopicPartition, OffsetPair> findCompletedEligibleOffsetsAndRemove(boolean remove) {
+
+        var newone = findCompletedEligibleOffsetsAndRemoveNew();
+
+        var oldone = findCompletedEligibleOffsetsAndRemoveOld(remove);
+
+        return newone;
+    }
+
+    private Map<TopicPartition, OffsetPair> findCompletedEligibleOffsetsAndRemoveNew() {
+        Map<TopicPartition, OffsetPair> offsetsToSend = new HashMap<>();
+
+        for (var e : getAssignedPartitions().entrySet()) {
+            Optional<OffsetPair> completedEligibleOffsetsAndRemoveNew = e.getValue().getCompletedEligibleOffsetsAndRemoveNew();
+            completedEligibleOffsetsAndRemoveNew.ifPresent(offsetAndMetadata ->
+                    offsetsToSend.put(e.getKey(), offsetAndMetadata)
+            );
+        }
+
+//        Map<TopicPartition, OffsetAndMetadata> collect = getAssignedPartitions().entrySet().stream()
+//                .map(x -> x.getValue().getCompletedEligibleOffsetsAndRemoveNew(remove))
+//                .filter(Optional::isPresent)
+//                .map(Optional::get)
+//                .collect(Collectors.toMap(Map.Entry::getKey,
+//                        o -> o.getValue().getCompletedEligibleOffsetsAndRemoveNew()));
+
+        log.debug("Scan finished, {} were in flight, offset(s) ({}) to be committed", offsetsToSend.size(), offsetsToSend);
+        return offsetsToSend;
+    }
+
+
     /**
      * Finds eligible offset positions to commit in each assigned partition
      */
     // todo remove completely as state of offsets should be tracked live, no need to scan for them - see #201
     // https://github.com/confluentinc/parallel-consumer/issues/201 Refactor: Live tracking of offsets as they change, so we don't need to scan for them
-    Map<TopicPartition, OffsetAndMetadata> findCompletedEligibleOffsetsAndRemove(boolean remove) {
+    Map<TopicPartition, OffsetAndMetadata> findCompletedEligibleOffsetsAndRemoveOld(boolean remove) {
 
         //
         Map<TopicPartition, OffsetAndMetadata> offsetsToSend = new HashMap<>();
@@ -470,7 +507,7 @@ Map<TopicPartition, OffsetAndMetadata> findCompletedEligibleOffsetsAndRemove(boo
             log.trace("Starting scan of partition: {}", topicPartitionKey);
 
             count += partitionQueue.size();
-            var workToRemove = new LinkedList<WorkContainer<K, V>>();
+            var workToRemove = new LinkedList<WorkContainer<?, ?>>();
             var incompleteOffsets = new LinkedHashSet<Long>();
             long lowWaterMark = -1;
             var highestSucceeded = partitionState.getOffsetHighestSucceeded();
@@ -539,6 +576,14 @@ private Map<TopicPartition, PartitionState<K, V>> getAssignedPartitions() {
                 .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
     }
 
+    boolean isDirty() {
+        return this.partitionStates.values().parallelStream().anyMatch(PartitionState::isDirty);
+    }
+
+    public boolean isClean() {
+        return !isDirty();
+    }
+
     public boolean couldBeTakenAsWork(WorkContainer<?, ?> workContainer) {
         if (checkIfWorkIsStale(workContainer)) {
             log.debug("Work is in queue with stale epoch or no longer assigned. Skipping. Shard it came from will/was removed during partition revocation. WC: {}", workContainer);