omg - hashsets vs queues, wow

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

@@ -889,7 +889,7 @@ private void transitionToClosing() {
      */
     private void processWorkCompleteMailBox() {
         log.trace("Processing mailbox (might block waiting for results)...");
-        Set<ActionItem> results = new HashSet<>();
+        Queue<ActionItem> results = new ArrayDeque<>();
 
         final Duration timeToBlockFor = getTimeToBlockFor();
 
@@ -927,7 +927,7 @@ private void processWorkCompleteMailBox() {
         for (var action : results) {
             WorkContainer<K, V> work = action.getWorkContainer();
             if (work == null) {
-                EpochAndRecords consumerRecords = action.getConsumerRecords();
+                EpochAndRecords<?, ?> consumerRecords = action.getConsumerRecords();
                 wm.registerWork(consumerRecords);
             } else {
                 MDC.put("offset", work.toString());
@@ -1138,6 +1138,7 @@ protected void addToMailbox(WorkContainer<K, V> wc) {
 
     public void registerWork(EpochAndRecords polledRecords) {
         log.debug("Adding {} to mailbox...", polledRecords);
+
         workMailBox.add(new ActionItem(null, polledRecords));
     }
 

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

@@ -186,9 +186,15 @@ private void updateHighestSucceededOffsetSoFar(WorkContainer<K, V> work) {
     }
 
     public void addWorkContainer(WorkContainer<K, V> wc) {
+        long offsetHighestSeen = getOffsetHighestSeen();
+        if (wc.offset() != offsetHighestSeen + 1) {
+            log.error("");
+        }
+
         maybeRaiseHighestSeenOffset(wc.offset());
         commitQueue.put(wc.offset(), wc);
         incompleteOffsets.add(wc.offset());
+
     }
 
     /**

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

@@ -35,6 +35,7 @@ public class ProcessingShard<K, V> {
      * Uses a ConcurrentSkipListMap instead of a TreeMap as under high pressure there appears to be some concurrency
      * errors (missing WorkContainers).
      */
+    @Getter
     private final NavigableMap<Long, WorkContainer<K, V>> entries = new ConcurrentSkipListMap<>();
 
     @Getter(PRIVATE)
@@ -74,9 +75,19 @@ public Optional<WorkContainer<K, V>> getWorkForOffset(long offset) {
     }
 
     public long getCountOfWorkAwaitingSelection() {
-        return entries.values().parallelStream()
+        return entries.values().stream()
                 // todo missing pm.isBlocked(topicPartition) ?
-                .filter(kvWorkContainer -> kvWorkContainer.isAvailableToTakeAsWork())
+                .filter(WorkContainer::isAvailableToTakeAsWork)
+                .count();
+    }
+
+    public long getCountOfWorkTracked() {
+        return entries.size();
+    }
+
+    public long getCountWorkInFlight() {
+        return entries.values().stream()
+                .filter(WorkContainer::isInFlight)
                 .count();
     }
 

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

@@ -58,6 +58,7 @@ public class ShardManager<K, V> {
      * @see WorkManager#getWorkIfAvailable()
      */
     // todo performance: disable/remove if using partition order
+    @Getter
     private final Map<Object, ProcessingShard<K, V>> processingShards = new ConcurrentHashMap<>();
 
     private final NavigableSet<WorkContainer<?, ?>> retryQueue = new TreeSet<>(Comparator.comparing(wc -> wc.getDelayUntilRetryDue()));

parallel-consumer-core/src/test/java/io/confluent/parallelconsumer/AbstractParallelEoSStreamProcessorTestBase.java

@@ -79,7 +79,7 @@ public abstract class AbstractParallelEoSStreamProcessorTestBase {
 
     protected AbstractParallelEoSStreamProcessor<String, String> parentParallelConsumer;
 
-    public static int defaultTimeoutSeconds = 30;
+    public static int defaultTimeoutSeconds = 3000;
 
     public static Duration defaultTimeout = ofSeconds(defaultTimeoutSeconds);
     protected static long defaultTimeoutMs = defaultTimeout.toMillis();
@@ -287,7 +287,7 @@ private void blockingLoopLatchTrigger(int waitForCount) {
         loopLatchV = new CountDownLatch(waitForCount);
         try {
             boolean timeout = !loopLatchV.await(defaultTimeoutSeconds, SECONDS);
-            if (timeout)
+            if (timeout || parentParallelConsumer.isClosedOrFailed())
                 throw new TimeoutException(msg("Timeout of {}, waiting for {} counts, on latch with {} left", defaultTimeout, waitForCount, loopLatchV.getCount()));
         } catch (InterruptedException e) {
             log.error("Interrupted while waiting for loop latch - timeout was {}", defaultTimeout);

parallel-consumer-core/src/test/java/io/confluent/parallelconsumer/offsets/OffsetEncodingBackPressureTest.java

@@ -9,10 +9,9 @@
 import io.confluent.parallelconsumer.FakeRuntimeError;
 import io.confluent.parallelconsumer.ParallelEoSStreamProcessorTestBase;
 import io.confluent.parallelconsumer.offsets.OffsetMapCodecManager.HighestOffsetAndIncompletes;
-import io.confluent.parallelconsumer.state.PartitionMonitor;
-import io.confluent.parallelconsumer.state.WorkContainer;
-import io.confluent.parallelconsumer.state.WorkManager;
+import io.confluent.parallelconsumer.state.*;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
 import org.apache.kafka.common.TopicPartition;
 import org.junit.jupiter.api.AfterAll;
@@ -23,12 +22,11 @@
 import pl.tlinkowski.unij.api.UniLists;
 
 import java.time.Duration;
-import java.util.List;
-import java.util.Map;
-import java.util.Optional;
-import java.util.Set;
+import java.util.*;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
 
 import static io.confluent.csid.utils.JavaUtils.getLast;
 import static io.confluent.csid.utils.JavaUtils.getOnlyOne;
@@ -83,9 +81,12 @@ void backPressureShouldPreventTooManyMessagesBeingQueuedForProcessing() throws O
         OffsetMapCodecManager.DefaultMaxMetadataSize = 40; // reduce available to make testing easier
         OffsetMapCodecManager.forcedCodec = Optional.of(OffsetEncoding.BitSetV2); // force one that takes a predictable large amount of space
 
-        ktu.send(consumerSpy, ktu.generateRecords(numberOfRecords));
+        List<ConsumerRecord<String, String>> records = ktu.generateRecords(numberOfRecords);
+        ktu.send(consumerSpy, records);
+
+        AtomicInteger userFuncFinishedCount = new AtomicInteger();
+        AtomicInteger userFuncStartCount = new AtomicInteger();
 
-        AtomicInteger userFuncFinishedCount = new AtomicInteger(0);
         CountDownLatch msgLock = new CountDownLatch(1);
         CountDownLatch msgLockTwo = new CountDownLatch(1);
         CountDownLatch msgLockThree = new CountDownLatch(1);
@@ -94,32 +95,40 @@ void backPressureShouldPreventTooManyMessagesBeingQueuedForProcessing() throws O
         List<Long> blockedOffsets = UniLists.of(0L, 2L);
         final int numberOfBlockedMessages = blockedOffsets.size();
 
-        parallelConsumer.poll((rec) -> {
+        WorkManager<String, String> wm = parallelConsumer.getWm();
+        final PartitionState<String, String> partitionState = wm.getPm().getPartitionState(topicPartition);
+
+        ConcurrentLinkedQueue<Long> seen = new ConcurrentLinkedQueue<>();
 
+        parallelConsumer.poll((rec) -> {
+            seen.add(rec.offset());
+            userFuncStartCount.incrementAndGet();
             // block the partition to create bigger and bigger offset encoding blocks
             // don't let offset 0 finish
             if (rec.offset() == offsetToBlock) {
                 int attemptNumber = attempts.incrementAndGet();
                 if (attemptNumber == 1) {
                     log.debug("Force first message to 'never' complete, causing a large offset encoding (lots of messages completing above the low water mark. Waiting for msgLock countdown.");
-                    int sleepFor = 120;
-                    awaitLatch(msgLock, sleepFor);
+                    int timeout = 120;
+                    awaitLatch(msgLock, timeout);
                     log.debug("Very slow message awoken, throwing exception");
                     throw new FakeRuntimeError("Fake error");
                 } else {
                     log.debug("Second attempt, waiting for msgLockTwo countdown");
                     awaitLatch(msgLockTwo, 60);
                     log.debug("Second attempt, unlocked, succeeding");
                 }
-            } else if (rec.offset() == 2l) {
+            } else if (rec.offset() == 2L) {
                 awaitLatch(msgLockThree);
                 log.debug("// msg 2L unblocked");
             } else {
                 sleepQuietly(1);
             }
-            userFuncFinishedCount.getAndIncrement();
+            userFuncFinishedCount.incrementAndGet();
         });
 
+        ShardManager<String, String> sm = wm.getSm();
+
         try {
 
             // wait for all pre-produced messages to be processed and produced
@@ -129,7 +138,17 @@ void backPressureShouldPreventTooManyMessagesBeingQueuedForProcessing() throws O
                     //, () -> parallelConsumer.getFailureCause()) // requires https://github.com/awaitility/awaitility/issues/178#issuecomment-734769761
                     .pollInterval(1, SECONDS)
                     .untilAsserted(() -> {
+                        ProcessingShard<String, String> stringStringProcessingShard = sm.getProcessingShards().get(topicPartition);
+                        if (stringStringProcessingShard != null) {
+                            long countOfWorkAwaitingSelection = stringStringProcessingShard.getCountOfWorkAwaitingSelection();
+                            NavigableMap<Long, WorkContainer<String, String>> entries = stringStringProcessingShard.getEntries();
+                            boolean b = sm.workIsWaitingToBeProcessed();
+                            long countWorkInFlight = stringStringProcessingShard.getCountWorkInFlight();
+                            long countOfWorkTracked = stringStringProcessingShard.getCountOfWorkTracked();
+                            long numberOfWorkQueuedInShardsAwaitingSelection = sm.getNumberOfWorkQueuedInShardsAwaitingSelection();
+                        }
                         assertThat(userFuncFinishedCount.get()).isEqualTo(numberOfRecords - numberOfBlockedMessages);
+//                        Truth.assertThat(numberOfWorkQueuedInShardsAwaitingSelection).isEqualTo(-4);
                     });
 
             // # assert commit ok - nothing blocked
@@ -159,7 +178,6 @@ void backPressureShouldPreventTooManyMessagesBeingQueuedForProcessing() throws O
                 assertThat(highestSeenOffset).as("offset 99 is encoded as having been seen").isEqualTo(expectedHighestSeenOffset);
             }
 
-            WorkManager<String, String> wm = parallelConsumer.getWm();
 
             // partition not blocked
             {
@@ -202,25 +220,30 @@ void backPressureShouldPreventTooManyMessagesBeingQueuedForProcessing() throws O
                                     .deserialiseIncompleteOffsetMapFromBase64(0L, meta);
                             Truth.assertWithMessage("The only incomplete record now is offset zero, which we are blocked on")
                                     .that(incompletes.getIncompleteOffsets()).containsExactlyElementsIn(blockedOffsets);
-                            Truth8.assertThat(incompletes.getHighestSeenOffset()).hasValue(numberOfRecords + extraRecordsToBlockWithThresholdBlocks - 1);
+                            int expectedHighestSeen = numberOfRecords + extraRecordsToBlockWithThresholdBlocks - 1;
+                            Truth8.assertThat(incompletes.getHighestSeenOffset()).hasValue(expectedHighestSeen);
                         }
                 );
             }
 
             // recreates the situation where the payload size is too large and must be dropped
             log.debug("// test max payload exceeded, payload dropped");
-            int processedBeforePartitionBlock = userFuncFinishedCount.get();
-            int extraMessages = numberOfRecords + extraRecordsToBlockWithThresholdBlocks / 2;
-            log.debug("// messages already sent {}, sending {} more", processedBeforePartitionBlock, extraMessages);
+
+//            log.debug("// messages already sent {}, sending {} more", processedBeforePartitionBlock, extraMessages);
             {
+                long numberOfWorkQueuedInShardsAwaitingSelection = sm.getNumberOfWorkQueuedInShardsAwaitingSelection();
                 log.debug("// force system to allow more records (i.e. the actual system attempts to never allow the payload to grow this big)");
-                PartitionMonitor.setUSED_PAYLOAD_THRESHOLD_MULTIPLIER(2);
+                PartitionMonitor.setUSED_PAYLOAD_THRESHOLD_MULTIPLIER(30);
                 parallelConsumer.requestCommitAsap();
                 awaitForOneLoopCycle();
 
+
                 //
                 log.debug("// unlock to make state dirty to get a commit");
+
                 msgLockThree.countDown();
+                int processedBeforePartitionBlock = userFuncFinishedCount.get();
+                int extraMessages = numberOfRecords + extraRecordsToBlockWithThresholdBlocks / 2;
                 log.debug("// send {} more messages", extraMessages);
                 ktu.send(consumerSpy, ktu.generateRecords(extraMessages));
 
@@ -229,7 +252,18 @@ void backPressureShouldPreventTooManyMessagesBeingQueuedForProcessing() throws O
 
                 log.debug("// wait for the new message to be processed");
                 await().atMost(defaultTimeout).untilAsserted(() ->
-                        assertThat(userFuncFinishedCount.get()).isEqualTo(processedBeforePartitionBlock + extraMessages + 1)
+                        {
+                            long numberOfWorkQueuedInShardsAwaitingSelection1 = sm.getNumberOfWorkQueuedInShardsAwaitingSelection();
+                            ShardManager<String, String> sm1 = sm;
+                            List<Long> seen1 = seen.stream().sorted().collect(Collectors.toList());
+                            long offsetHighestSucceeded = partitionState.getOffsetHighestSucceeded();
+                            long offsetHighestSeen = partitionState.getOffsetHighestSeen();
+                            long offsetHighestSequentialSucceeded = partitionState.getOffsetHighestSequentialSucceeded();
+                            int i = userFuncStartCount.get();
+                            int i1 = userFuncFinishedCount.get();
+                            int expectedUserFunctionFinishedCount = processedBeforePartitionBlock + extraMessages + 1;
+                            assertThat(userFuncFinishedCount.get()).isEqualTo(expectedUserFunctionFinishedCount);
+                        }
                 );
 
                 log.debug("// assert payload missing from commit now");

parallel-consumer-core/src/test/java/io/confluent/parallelconsumer/offsets/OffsetEncodingBackPressureUnitTest.java

@@ -4,6 +4,7 @@
  * Copyright (C) 2020-2022 Confluent, Inc.
  */
 
+import com.google.common.truth.Truth;
 import io.confluent.parallelconsumer.ParallelEoSStreamProcessorTestBase;
 import io.confluent.parallelconsumer.state.PartitionMonitor;
 import io.confluent.parallelconsumer.state.PartitionState;
@@ -207,6 +208,7 @@ private void sendRecordsToWM(int numberOfRecords, WorkManager<String, String> wm
         log.debug("~Sending {} more records", numberOfRecords);
         List<ConsumerRecord<String, String>> records = ktu.generateRecords(numberOfRecords);
         wm.registerWork(new ConsumerRecords<>(UniMaps.of(topicPartition, records)));
+        Truth.assertThat(wm.getTotalWorkAwaitingIngestion()).isEqualTo(numberOfRecords);
     }