add multiple caches for accelerating available container count calculation

CHANGELOG.adoc

@@ -13,6 +13,13 @@ NOTE:: Dependency version bumps are not listed here.
 ifndef::github_name[]
 toc::[]
 endif::[]
+
+== 0.5.2.9
+
+=== Improvements
+* improvement: add multiple caches for accelerating available container count calculation （#667）
+
+
 == 0.5.2.8
 
 === Fixes

README.adoc

@@ -1535,6 +1535,9 @@ toc::[]
 endif::[]
 == 0.5.2.8
 
+=== Improvements
+* improvement: add multiple caches for accelerating available container count calculation （#667）
+
 === Fixes
 
 * fix: Fix equality and hash code for ShardKey with array key (#638), resolves (#579)

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

@@ -1,7 +1,7 @@
 package io.confluent.parallelconsumer.internal;
 
 /*-
- * Copyright (C) 2020-2023 Confluent, Inc.
+ * Copyright (C) 2020-2024 Confluent, Inc.
  */
 
 import io.confluent.csid.utils.SupplierUtils;
@@ -1407,20 +1407,8 @@ public void requestCommitAsap() {
         notifySomethingToDo();
     }
 
-    private boolean isCommandedToCommit() {
-        synchronized (commitCommand) {
-            return this.commitCommand.get();
-        }
-    }
 
-    private void clearCommitCommand() {
-        synchronized (commitCommand) {
-            if (commitCommand.get()) {
-                log.debug("Command to commit asap received, clearing");
-                this.commitCommand.set(false);
-            }
-        }
-    }
+
 
     private boolean isTransactionCommittingInProgress() {
         return options.isUsingTransactionCommitMode() &&
@@ -1447,4 +1435,20 @@ public void resumeIfPaused() {
             log.debug("Skipping transition of parallel consumer to state running. Current state is {}.", this.state);
         }
     }
+
+    private boolean isCommandedToCommit() {
+        synchronized (commitCommand) {
+            return this.commitCommand.get();
+        }
+    }
+
+    private void clearCommitCommand() {
+        synchronized (commitCommand) {
+            if (commitCommand.get()) {
+                log.debug("Command to commit asap received, clearing");
+                this.commitCommand.set(false);
+            }
+        }
+    }
+
 }

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

@@ -1,7 +1,7 @@
 package io.confluent.parallelconsumer.state;
 
 /*-
- * Copyright (C) 2020-2023 Confluent, Inc.
+ * Copyright (C) 2020-2024 Confluent, Inc.
  */
 
 import io.confluent.parallelconsumer.ParallelConsumerOptions;
@@ -14,6 +14,7 @@
 import java.time.Duration;
 import java.util.*;
 import java.util.concurrent.ConcurrentSkipListMap;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 
 import static io.confluent.csid.utils.BackportUtils.toSeconds;
@@ -44,6 +45,7 @@ public class ProcessingShard<K, V> {
     @Getter
     private final NavigableMap<Long, WorkContainer<K, V>> entries = new ConcurrentSkipListMap<>();
 
+
     @Getter(PRIVATE)
     private final ShardKey key;
 
@@ -53,9 +55,12 @@ public class ProcessingShard<K, V> {
 
     private final RateLimiter slowWarningRateLimit = new RateLimiter(5);
 
+    private AtomicLong availableWorkContainerCnt = new AtomicLong(0);
+
+    private AtomicLong expiredRetryContainerCnt = new AtomicLong(0);
+
     public boolean workIsWaitingToBeProcessed() {
-        return entries.values().parallelStream()
-                .anyMatch(kvWorkContainer -> kvWorkContainer.isAvailableToTakeAsWork());
+        return availableWorkContainerCnt.get() > 0L;
     }
 
     public void addWorkContainer(WorkContainer<K, V> wc) {
@@ -64,23 +69,36 @@ public void addWorkContainer(WorkContainer<K, V> wc) {
             log.debug("Entry for {} already exists in shard queue, dropping record", wc);
         } else {
             entries.put(key, wc);
+            availableWorkContainerCnt.incrementAndGet();
         }
     }
 
     public void onSuccess(WorkContainer<?, ?> wc) {
         // remove work from shard's queue
         entries.remove(wc.offset());
+
+        // decrease for retry only
+        if (wc.getLastFailedAt().isPresent()) {
+            expiredRetryContainerCnt.decrementAndGet();
+        }
     }
 
+    public void onFailure() {
+        // increase available cnt first to let retry expired calculated later
+        availableWorkContainerCnt.incrementAndGet();
+    }
+
+
     public boolean isEmpty() {
         return entries.isEmpty();
     }
 
     public long getCountOfWorkAwaitingSelection() {
-        return entries.values().stream()
-                // todo missing pm.isBlocked(topicPartition) ?
-                .filter(WorkContainer::isAvailableToTakeAsWork)
-                .count();
+        return availableWorkContainerCnt.get();
+    }
+
+    public long getExpiredRetryContainerCnt() {
+        return expiredRetryContainerCnt.get();
     }
 
     public long getCountOfWorkTracked() {
@@ -94,10 +112,16 @@ public long getCountWorkInFlight() {
     }
 
     public WorkContainer<K, V> remove(long offset) {
+        // from onPartitionsRemoved callback, need to deduce the available worker count for the revoked partition
+        WorkContainer<K, V> toRemovedWorker = entries.get(offset);
+        if (toRemovedWorker != null && toRemovedWorker.isAvailableToTakeAsWork()) {
+            dcrAvailableWorkContainerCntByDelta(1);
+        }
         return entries.remove(offset);
     }
 
 
+
     // remove staled WorkContainer otherwise when the partition is reassigned, the staled messages will:
     // 1. block the new work containers to be picked and processed
     // 2. will cause the consumer to paused consuming new messages indefinitely
@@ -109,7 +133,6 @@ public boolean removeStaleWorkContainersFromShard() {
                 });
     }
 
-
     ArrayList<WorkContainer<K, V>> getWorkIfAvailable(int workToGetDelta) {
         log.trace("Looking for work on shardQueueEntry: {}", getKey());
 
@@ -123,6 +146,11 @@ ArrayList<WorkContainer<K, V>> getWorkIfAvailable(int workToGetDelta) {
             if (pm.couldBeTakenAsWork(workContainer)) {
                 if (workContainer.isAvailableToTakeAsWork()) {
                     log.trace("Taking {} as work", workContainer);
+
+                    // only increase the ExpiredRetryContainerCnt when this is retry due since already added in the new container creation
+                    if (workContainer.isDelayExistsExpired()) {
+                        expiredRetryContainerCnt.incrementAndGet();
+                    }
                     workContainer.onQueueingForExecution();
                     workTaken.add(workContainer);
                 } else {
@@ -154,6 +182,8 @@ ArrayList<WorkContainer<K, V>> getWorkIfAvailable(int workToGetDelta) {
 
         logSlowWork(slowWork);
 
+        dcrAvailableWorkContainerCntByDelta(workTaken.size());
+
         return workTaken;
     }
 
@@ -200,4 +230,12 @@ private boolean isOrderRestricted() {
     private boolean isWorkContainerStale(WorkContainer<K, V> workContainer) {
         return pm.getPartitionState(workContainer).checkIfWorkIsStale(workContainer);
     }
+
+    private void dcrAvailableWorkContainerCntByDelta(int ByNum) {
+        availableWorkContainerCnt.getAndAdd(-1 * ByNum);
+        // in case of possible race condition
+        if (availableWorkContainerCnt.get() < 0L) {
+            availableWorkContainerCnt.set(0L);
+        }
+    }
 }

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

@@ -1,7 +1,7 @@
 package io.confluent.parallelconsumer.state;
 
 /*-
- * Copyright (C) 2020-2023 Confluent, Inc.
+ * Copyright (C) 2020-2024 Confluent, Inc.
  */
 
 import io.confluent.csid.utils.LoopingResumingIterator;
@@ -22,7 +22,11 @@
 import java.time.Duration;
 import java.time.Instant;
 import java.util.*;
+import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.concurrent.PriorityBlockingQueue;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 
 import static io.confluent.parallelconsumer.ParallelConsumerOptions.ProcessingOrder.KEY;
@@ -92,6 +96,9 @@ public class ShardManager<K, V> {
     @Getter(AccessLevel.PACKAGE) // visible for testing
     private final NavigableSet<WorkContainer<?, ?>> retryQueue = new TreeSet<>(retryQueueWorkContainerComparator);
 
+    @Getter
+    private final AtomicLong retryItemCnt = new AtomicLong(0);
+
     /**
      * Iteration resume point, to ensure fairness (prevent shard starvation) when we can't process messages from every
      * shard.
@@ -132,9 +139,13 @@ ShardKey computeShardKey(ConsumerRecord<?, ?> wc) {
      * @return Work ready in the processing shards, awaiting selection as work to do
      */
     public long getNumberOfWorkQueuedInShardsAwaitingSelection() {
+        // all available container count - (still pending for running retry containers count)
+        // => all_available_count - (retryCnt - all_expired_retry_cnt) => all_available_count + all_expired_retry_cnt - retryCnt
+
         return processingShards.values().stream()
-                .mapToLong(ProcessingShard::getCountOfWorkAwaitingSelection)
-                .sum();
+                .mapToLong(processingShard ->
+                        processingShard.getCountOfWorkAwaitingSelection() + processingShard.getExpiredRetryContainerCnt())
+                .sum() - retryItemCnt.get();
     }
 
     public boolean workIsWaitingToBeProcessed() {
@@ -170,7 +181,9 @@ private void removeWorkFromShardFor(ConsumerRecord<K, V> consumerRecord) {
             WorkContainer<K, V> removedWC = shard.remove(consumerRecord.offset());
 
             // remove if in retry queue
-            this.retryQueue.remove(removedWC);
+            if (this.retryQueue.remove(removedWC)) {
+                retryItemCnt.decrementAndGet();
+            }
 
             // remove the shard if empty
             removeShardIfEmpty(shardKey);
@@ -204,7 +217,9 @@ void removeShardIfEmpty(ShardKey key) {
 
     public void onSuccess(WorkContainer<?, ?> wc) {
         // remove from the retry queue if it's contained
-        this.retryQueue.remove(wc);
+        if (this.retryQueue.remove(wc)) {
+            retryItemCnt.decrementAndGet();
+        }
 
         // remove from processing queues
         var key = computeShardKey(wc);
@@ -224,7 +239,16 @@ public void onSuccess(WorkContainer<?, ?> wc) {
      */
     public void onFailure(WorkContainer<?, ?> wc) {
         log.debug("Work FAILED");
-        this.retryQueue.add(wc);
+        if (this.retryQueue.add(wc)) {
+            retryItemCnt.incrementAndGet();
+        }
+        var key = computeShardKey(wc);
+        var shardOptional = getShard(key);
+
+        if (shardOptional.isPresent()) {
+            shardOptional.get().onFailure();
+        }
+
     }
 
     /**

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

@@ -1,7 +1,7 @@
 package io.confluent.parallelconsumer.state;
 
 /*-
- * Copyright (C) 2020-2023 Confluent, Inc.
+ * Copyright (C) 2020-2024 Confluent, Inc.
  */
 
 import io.confluent.parallelconsumer.PollContextInternal;
@@ -109,6 +109,18 @@ public boolean isDelayPassed() {
         return negative;
     }
 
+    /**
+     *
+     * @return whether this is retry and should be executed
+     */
+    public boolean isDelayExistsExpired() {
+        if (!lastFailedAt.isPresent()) {
+            return false;
+        }
+        Duration delay = getDelayUntilRetryDue();
+        return delay.isNegative() || delay.isZero();
+    }
+
     /**
      * @return time until it should be retried
      */

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

@@ -1,7 +1,7 @@
 package io.confluent.parallelconsumer.state;
 
 /*-
- * Copyright (C) 2020-2022 Confluent, Inc.
+ * Copyright (C) 2020-2024 Confluent, Inc.
  */
 
 import io.confluent.parallelconsumer.internal.PCModuleTestEnv;

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

@@ -1,7 +1,7 @@
 package io.confluent.parallelconsumer.state;
 
 /*-
- * Copyright (C) 2020-2022 Confluent, Inc.
+ * Copyright (C) 2020-2024 Confluent, Inc.
  */
 
 import com.google.common.truth.Truth;
@@ -131,6 +131,7 @@ private ConsumerRecord<String, String> makeRec(String value, String key, int par
         return stringStringConsumerRecord;
     }
 
+
     @ParameterizedTest
     @EnumSource
     void basic(ParallelConsumerOptions.ProcessingOrder order) {
@@ -738,4 +739,41 @@ void starvation() {
 
     }
 
+    /**
+     * Tests available worker cnt
+     */
+    @Test
+    void testAvailableWorkerCnt() {
+        ParallelConsumerOptions<?, ?> build = ParallelConsumerOptions.builder().ordering(PARTITION).build();
+        setupWorkManager(build);
+        // sanity
+        assertThat(wm.getOptions().getOrdering()).isEqualTo(PARTITION);
+
+        registerSomeWork();
+
+        int total = 3;
+        int maxWorkToGet = 2;
+
+        var works = wm.getWorkIfAvailable(maxWorkToGet);
+
+        assertThat(wm.getNumberOfWorkQueuedInShardsAwaitingSelection()).isEqualTo(total - works.size());
+
+        // fail the work
+        var wc = works.get(0);
+        fail(wc);
+
+
+        // advance clock to make delay pass
+        advanceClockByDelay();
+        try {
+            Thread.sleep(1000);
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+        // work should now be ready to take
+        works = wm.getWorkIfAvailable(maxWorkToGet);
+        assertThat(wm.getNumberOfWorkQueuedInShardsAwaitingSelection()).isEqualTo(total - works.size());
+
+    }
+
 }