Limit subscription in-flight event accumulation

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java

@@ -109,6 +109,9 @@ public abstract class SubscriptionPrefetchingQueue {
 
   private final SubscriptionPrefetchingQueueStates states;
 
+  // Keep the in-flight limit check and event claim atomic across concurrent polls.
+  private final Object pollLock = new Object();
+
   private long lastStateReportTimestamp = System.currentTimeMillis();
 
   private volatile boolean isCompleted = false;
@@ -230,7 +233,9 @@ protected void executeReceiverSubtask(
   public SubscriptionEvent poll(final String consumerId) {
     acquireReadLock();
     try {
-      return isClosed() ? null : pollInternal(consumerId);
+      synchronized (pollLock) {
+        return isClosed() ? null : pollInternal(consumerId);
+      }
     } finally {
       releaseReadLock();
     }
@@ -239,6 +244,10 @@ public SubscriptionEvent poll(final String consumerId) {
   private SubscriptionEvent pollInternal(final String consumerId) {
     states.markPollRequest();
 
+    if (shouldThrottlePollByInFlightEvents()) {
+      return null;
+    }
+
     if (prefetchingQueue.isEmpty()) {
       states.markMissingPrefetch();
       try {
@@ -271,6 +280,15 @@ private SubscriptionEvent pollInternal(final String consumerId) {
     return null;
   }
 
+  private boolean shouldThrottlePollByInFlightEvents() {
+    if (!states.shouldThrottlePoll()) {
+      return false;
+    }
+
+    remapInFlightEventsSnapshot(committedCleaner, pollableNacker);
+    return states.shouldThrottlePoll();
+  }
+
   public SubscriptionEvent pollV2(final String consumerId, final PollTimer timer) {
     states.markPollRequest();
 
@@ -286,9 +304,15 @@ public SubscriptionEvent pollV2(final String consumerId, final PollTimer timer)
           return null;
         }
 
-        final SubscriptionEvent event = pollInternalV2(consumerId);
-        if (Objects.nonNull(event)) {
-          return event;
+        synchronized (pollLock) {
+          if (shouldThrottlePollByInFlightEvents()) {
+            return null;
+          }
+
+          final SubscriptionEvent event = pollInternalV2(consumerId);
+          if (Objects.nonNull(event)) {
+            return event;
+          }
         }
       } finally {
         releaseReadLock();
@@ -1119,6 +1143,10 @@ public long getSubscriptionUncommittedEventCount() {
     return inFlightEvents.size();
   }
 
+  public long getSubscriptionRetainedEventCount() {
+    return prefetchingQueue.size() + inFlightEvents.size();
+  }
+
   public long getCurrentCommitId() {
     return commitIdGenerator.get();
   }
@@ -1166,6 +1194,7 @@ public Map<String, String> coreReportMessage() {
     result.put("size of inputPendingQueue", String.valueOf(inputPendingQueue.size()));
     result.put("size of prefetchingQueue", String.valueOf(prefetchingQueue.size()));
     result.put("size of inFlightEvents", String.valueOf(inFlightEvents.size()));
+    result.put("size of retainedEvents", String.valueOf(getSubscriptionRetainedEventCount()));
     result.put("commitIdGenerator", commitIdGenerator.toString());
     result.put("states", states.toString());
     result.put("isCompleted", String.valueOf(isCompleted));

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStates.java

@@ -30,6 +30,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.IntSupplier;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 
@@ -55,14 +56,22 @@ public class SubscriptionPrefetchingQueueStates {
       SubscriptionConfig.getInstance().getSubscriptionPrefetchEventGlobalCountThreshold();
 
   private final SubscriptionPrefetchingQueue prefetchingQueue;
+  private final IntSupplier prefetchingQueueCountSupplier;
 
   private volatile long lastPollRequestTimestamp;
   private final Meter pollRequestMeter;
   private final Meter missingPrefechMeter;
   private final AtomicLong disorderCauseCounter;
 
   public SubscriptionPrefetchingQueueStates(final SubscriptionPrefetchingQueue prefetchingQueue) {
+    this(prefetchingQueue, () -> SubscriptionAgent.broker().getPrefetchingQueueCount());
+  }
+
+  SubscriptionPrefetchingQueueStates(
+      final SubscriptionPrefetchingQueue prefetchingQueue,
+      final IntSupplier prefetchingQueueCountSupplier) {
     this.prefetchingQueue = prefetchingQueue;
+    this.prefetchingQueueCountSupplier = prefetchingQueueCountSupplier;
 
     this.lastPollRequestTimestamp = -1;
     this.pollRequestMeter = new Meter(new IoTDBMovingAverage(), Clock.defaultClock());
@@ -109,12 +118,12 @@ public boolean shouldPrefetch() {
     }
 
     // 1.3. local event count
-    if (hasTooManyPrefetchedLocalEvent()) {
+    if (hasTooManyRetainedLocalEvent()) {
       return false;
     }
 
     // 1.4. global event count
-    if (hasTooManyPrefetchedGlobalEvent()) {
+    if (hasTooManyRetainedGlobalEvent()) {
       return false;
     }
 
@@ -142,24 +151,40 @@ public boolean shouldPrefetch() {
     return (System.currentTimeMillis() - lastPollRequestTimestamp) * pollRate() > 1000;
   }
 
+  public boolean shouldThrottlePoll() {
+    return hasTooManyInFlightLocalEvent() || hasTooManyInFlightGlobalEvent();
+  }
+
   private boolean isMemoryEnough() {
     return PipeDataNodeResourceManager.memory().getTotalNonFloatingMemorySizeInBytes()
             * PREFETCH_MEMORY_THRESHOLD
         > PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes();
   }
 
-  private boolean hasTooManyPrefetchedLocalEvent() {
-    return prefetchingQueue.getPrefetchedEventCount() > PREFETCH_EVENT_LOCAL_COUNT_THRESHOLD;
+  private boolean hasTooManyRetainedLocalEvent() {
+    return prefetchingQueue.getSubscriptionRetainedEventCount()
+        > PREFETCH_EVENT_LOCAL_COUNT_THRESHOLD;
   }
 
-  private boolean hasTooManyPrefetchedGlobalEvent() {
-    // The number of prefetched events in the current prefetching queue > floor(t / number of
+  private boolean hasTooManyRetainedGlobalEvent() {
+    // The number of retained events in the current prefetching queue > floor(t / number of
     // prefetching queues), where t is an adjustable parameter.
-    return prefetchingQueue.getPrefetchedEventCount()
-            * SubscriptionAgent.broker().getPrefetchingQueueCount()
+    return prefetchingQueue.getSubscriptionRetainedEventCount()
+            * prefetchingQueueCountSupplier.getAsInt()
         > PREFETCH_EVENT_GLOBAL_COUNT_THRESHOLD;
   }
 
+  private boolean hasTooManyInFlightLocalEvent() {
+    return prefetchingQueue.getSubscriptionUncommittedEventCount()
+        >= PREFETCH_EVENT_LOCAL_COUNT_THRESHOLD;
+  }
+
+  private boolean hasTooManyInFlightGlobalEvent() {
+    return prefetchingQueue.getSubscriptionUncommittedEventCount()
+            * prefetchingQueueCountSupplier.getAsInt()
+        >= PREFETCH_EVENT_GLOBAL_COUNT_THRESHOLD;
+  }
+
   private boolean isMissingRateTooHigh() {
     return missingRate() > PREFETCH_MISSING_RATE_THRESHOLD;
   }
@@ -183,6 +208,8 @@ public String toString() {
         .add("pollRate", pollRate())
         .add("missingRate", missingRate())
         .add("disorderCause", disorderCauseCounter.get())
+        .add("retainedEventCount", prefetchingQueue.getSubscriptionRetainedEventCount())
+        .add("inFlightEventCount", prefetchingQueue.getSubscriptionUncommittedEventCount())
         .toString();
   }
 }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java

@@ -145,6 +145,9 @@ public class ConsensusPrefetchingQueue {
 
   private final WakeableIndexedConsensusQueue pendingEntries;
 
+  // Keep the in-flight limit check and event claim atomic across concurrent polls.
+  private final Object pollLock = new Object();
+
   private static final int PENDING_QUEUE_CAPACITY = 4096;
 
   private static final long FIRST_CONSENSUS_SEARCH_INDEX = 1L;
@@ -825,7 +828,7 @@ private boolean hasReachedPrefetchLimit() {
     // include in-flight events, so polling without ACK cannot bypass the memory target. Allow an
     // empty delivery queue to make progress even when a single entry exceeds the target; the
     // shared memory manager still enforces the hard allocation limit.
-    return prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE
+    return !hasAvailableEventSlot()
         || (!prefetchingQueue.isEmpty()
             && retainedTabletBytes.get()
                 >= Math.max(1L, subscriptionMemoryHandle.getMemoryQuotaInBytes() / 2L));
@@ -834,8 +837,7 @@ private boolean hasReachedPrefetchLimit() {
   private boolean hasReachedPrefetchLimit(final DeliveryBatchState batchState) {
     // The current batch becomes one more event when flushed. Stop conversion before that event
     // fills the delivery queue so the unprocessed pending suffix stays queued.
-    return hasReachedPrefetchLimit()
-        || (!batchState.isEmpty() && prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE - 1);
+    return hasReachedPrefetchLimit() || !hasAvailableEventSlot(batchState);
   }
 
   private SubscriptionQueueRejectionReason getRealtimeAdmissionRejectionReason() {
@@ -994,8 +996,14 @@ public SubscriptionEvent poll(final String consumerId, final RegionProgress regi
       if (pendingSeekRequest != null) {
         return null;
       }
-      final SubscriptionEvent event = pollInternal(consumerId);
-      if (Objects.nonNull(event) && !hasReachedPrefetchLimit()) {
+      final SubscriptionEvent event;
+      synchronized (pollLock) {
+        if (shouldThrottlePollByInFlightEvents()) {
+          return null;
+        }
+        event = pollInternal(consumerId);
+      }
+      if (Objects.nonNull(event) && hasAvailableEventSlot() && !hasReachedPrefetchLimit()) {
         requestPrefetch();
       } else if (Objects.isNull(event) && shouldRecoverPrefetchBindingAfterEmptyPoll()) {
         requestPrefetch();
@@ -1026,6 +1034,24 @@ public SubscriptionEvent poll(final String consumerId, final RegionProgress regi
     }
   }
 
+  private boolean shouldThrottlePollByInFlightEvents() {
+    if (inFlightEvents.size() < MAX_PREFETCHING_QUEUE_SIZE) {
+      return false;
+    }
+
+    recycleInFlightEvents();
+    return inFlightEvents.size() >= MAX_PREFETCHING_QUEUE_SIZE;
+  }
+
+  private boolean hasAvailableEventSlot() {
+    return getSubscriptionRetainedEventCount() < MAX_PREFETCHING_QUEUE_SIZE;
+  }
+
+  private boolean hasAvailableEventSlot(final DeliveryBatchState batchState) {
+    return getSubscriptionRetainedEventCount() + (batchState.isEmpty() ? 0L : 1L)
+        < MAX_PREFETCHING_QUEUE_SIZE;
+  }
+
   private boolean initPrefetch(final RegionProgress regionProgress) {
     synchronized (prefetchInitializationLock) {
       return initPrefetchUnderInitializationLock(regionProgress);
@@ -1818,7 +1844,7 @@ private PrefetchRoundResult drivePrefetchOnceUnderReadLock() {
         blockRealtimeAdmission(SubscriptionQueueRejectionReason.INACTIVE_OR_CLOSED);
         return computeIdleRoundResult(walScanBudget);
       }
-      if (prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) {
+      if (!hasAvailableEventSlot()) {
         return computeIdleRoundResult(walScanBudget);
       }
 
@@ -1937,14 +1963,16 @@ private PrefetchRoundResult drivePrefetchOnceUnderReadLock() {
 
       // Publish reserved Tablets when read-ahead pauses. Leaving this batch lingering at the
       // delivery limit would keep rescheduling pending work that cannot yet be materialized.
-      if (!lingerBatch.isEmpty() && hasReachedPrefetchLimit(lingerBatch)) {
+      if (!lingerBatch.isEmpty()
+          && hasReachedPrefetchLimit(lingerBatch)
+          && hasAvailableEventSlot()) {
         if (!flushBatch(lingerBatch, observedSeekGeneration)) {
           resetRoundStateForSeek(seekGeneration.get());
           return PrefetchRoundResult.rescheduleNow();
         }
       }
 
-      if (!lingerBatch.isEmpty() && lingerBatch.firstTabletTimeMs > 0L) {
+      if (!lingerBatch.isEmpty() && lingerBatch.firstTabletTimeMs > 0L && hasAvailableEventSlot()) {
         final long lingerElapsedMs = System.currentTimeMillis() - lingerBatch.firstTabletTimeMs;
         if (lingerElapsedMs >= batchMaxDelayMs) {
           if (seekGeneration.get() != observedSeekGeneration) {
@@ -2024,7 +2052,7 @@ private PrefetchRoundResult driveOrderedWriterReplay(
     // Bound entries processed across all writer readers in this round, including skipped entries.
     final int[] remainingEntryBudget = {Math.max(1, maxWalEntries)};
     final int roundRobinStart = orderedWriterRoundRobinStart++;
-    while (prefetchingQueue.size() < MAX_PREFETCHING_QUEUE_SIZE) {
+    while (hasAvailableEventSlot(lingerBatch)) {
       if (seekGeneration.get() != expectedSeekGeneration || closeRequested) {
         return PrefetchRoundResult.rescheduleNow();
       }
@@ -2145,14 +2173,16 @@ private PrefetchRoundResult driveOrderedWriterReplay(
       }
     }
     if (!lingerBatch.isEmpty()
-        && System.currentTimeMillis() - lingerBatch.firstTabletTimeMs >= batchMaxDelayMs) {
+        && (!hasAvailableEventSlot(lingerBatch)
+            || System.currentTimeMillis() - lingerBatch.firstTabletTimeMs >= batchMaxDelayMs)
+        && hasAvailableEventSlot()) {
       if (!flushBatch(lingerBatch, expectedSeekGeneration)) {
         return PrefetchRoundResult.rescheduleNow();
       }
     }
     maybeInjectWatermark();
-    if (prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) {
-      return PrefetchRoundResult.dormant();
+    if (!hasAvailableEventSlot()) {
+      return computeIdleRoundResult(walScanBudget);
     }
     if (walScanBudget.hasYielded()
         || remainingEntryBudget[0] <= 0
@@ -2565,6 +2595,12 @@ private PrefetchRoundResult computeIdleRoundResult(
     if (isClosed || !prefetchInitialized || !isActive) {
       return PrefetchRoundResult.dormant();
     }
+    if (!hasAvailableEventSlot()) {
+      final long recycleIntervalMs = getRecycleUncommittedEventIntervalMs();
+      return inFlightEvents.isEmpty() || recycleIntervalMs <= 0L
+          ? PrefetchRoundResult.dormant()
+          : PrefetchRoundResult.rescheduleAfter(recycleIntervalMs);
+    }
     if (hasReachedPrefetchLimit()) {
       return PrefetchRoundResult.dormant();
     }
@@ -2961,6 +2997,7 @@ private MaterializationResult pumpFromSubscriptionWAL(
     int entriesRead = 0;
     while (entriesRead < maxWalEntries
         && !hasReachedPrefetchLimit(batchState)
+        && hasAvailableEventSlot()
         && subscriptionWALIterator.hasNext(walScanBudget)) {
       reportUnreadableWalReplayIfNecessary();
       try {
@@ -3779,7 +3816,7 @@ private boolean drainBufferedRealtimeWriters(
       }
 
       final int bufferedAfter = getRealtimeBufferedEntryCount();
-      if (bufferedAfter == 0 || prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) {
+      if (bufferedAfter == 0 || !hasAvailableEventSlot()) {
         return true;
       }
 
@@ -3854,6 +3891,9 @@ private <T extends WriterBufferedEntry> boolean drainWriterEntries(
       final boolean trackLingerTime,
       final long expectedSeekGeneration) {
     while (true) {
+      if (!hasAvailableEventSlot(batchState)) {
+        return true;
+      }
       final PriorityQueue<WriterFrontier> frontiers = frontierSupplier.get();
       if (frontiers.isEmpty()) {
         return true;
@@ -4082,7 +4122,9 @@ public boolean ack(final String consumerId, final SubscriptionCommitContext comm
                 && ackMissingInFlightEvent(commitContext, false));
     if (acked) {
       serverImpl.checkAndUpdateSafeDeletedSearchIndex();
-      requestPrefetch();
+      if (hasAvailableEventSlot()) {
+        requestPrefetch();
+      }
     }
     return acked;
   }
@@ -4175,7 +4217,9 @@ public boolean ackSilent(final String consumerId, final SubscriptionCommitContex
                 && ackMissingInFlightEvent(commitContext, true));
     if (ackedResult) {
       serverImpl.checkAndUpdateSafeDeletedSearchIndex();
-      requestPrefetch();
+      if (hasAvailableEventSlot()) {
+        requestPrefetch();
+      }
     }
     return ackedResult;
   }
@@ -5108,8 +5152,8 @@ private void maybeInjectWatermark() {
     }
     final long intervalMs =
         SubscriptionConfig.getInstance().getSubscriptionConsensusWatermarkIntervalMs();
-    if (intervalMs <= 0) {
-      return; // Watermark disabled
+    if (intervalMs <= 0 || !hasAvailableEventSlot()) {
+      return; // Watermark disabled or the queue is full
     }
     final long now = System.currentTimeMillis();
     if (now - lastWatermarkEmitTimeMs >= intervalMs) {
@@ -5858,6 +5902,10 @@ public long getSubscriptionUncommittedEventCount() {
     return inFlightEvents.size();
   }
 
+  public long getSubscriptionRetainedEventCount() {
+    return prefetchingQueue.size() + inFlightEvents.size();
+  }
+
   public long getRetainedTabletBytes() {
     return retainedTabletBytes.get();
   }