From 74dbb9eae3de589571bb27758f8fed975c4e4077 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 29 Jun 2026 15:27:20 +0800 Subject: [PATCH 1/3] Limit subscription in-flight event accumulation --- .../broker/SubscriptionPrefetchingQueue.java | 29 +++++++++ .../SubscriptionPrefetchingQueueStates.java | 32 ++++++++-- .../consensus/ConsensusPrefetchingQueue.java | 64 +++++++++++++++---- 3 files changed, 106 insertions(+), 19 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java index 8fc7227fd15c4..047532259f71b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java @@ -223,6 +223,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 { @@ -254,6 +258,22 @@ private SubscriptionEvent pollInternal(final String consumerId) { return null; } + private boolean shouldThrottlePollByInFlightEvents() { + if (!states.shouldThrottlePoll()) { + return false; + } + + remapInFlightEventsSnapshot(committedCleaner, pollableNacker); + if (!states.shouldThrottlePoll()) { + return false; + } + + LOGGER.debug( + "Subscription: SubscriptionPrefetchingQueue {} throttles poll because too many events are in flight.", + this); + return true; + } + public SubscriptionEvent pollV2(final String consumerId, final PollTimer timer) { acquireReadLock(); try { @@ -268,6 +288,10 @@ private SubscriptionEvent pollInternalV2(final String consumerId, final PollTime // do-while ensures at least one poll do { + if (shouldThrottlePollByInFlightEvents()) { + return null; + } + SubscriptionEvent event; try { if (prefetchingQueue.isEmpty()) { @@ -976,6 +1000,10 @@ public long getSubscriptionUncommittedEventCount() { return inFlightEvents.size(); } + public long getSubscriptionRetainedEventCount() { + return prefetchingQueue.size() + inFlightEvents.size(); + } + public long getCurrentCommitId() { return commitIdGenerator.get(); } @@ -1023,6 +1051,7 @@ public Map 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)); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStates.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStates.java index 70eacd2eaaae9..d5715e9fa41cc 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStates.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStates.java @@ -106,12 +106,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; } @@ -132,24 +132,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() + return prefetchingQueue.getSubscriptionRetainedEventCount() * SubscriptionAgent.broker().getPrefetchingQueueCount() > PREFETCH_EVENT_GLOBAL_COUNT_THRESHOLD; } + private boolean hasTooManyInFlightLocalEvent() { + return prefetchingQueue.getSubscriptionUncommittedEventCount() + >= PREFETCH_EVENT_LOCAL_COUNT_THRESHOLD; + } + + private boolean hasTooManyInFlightGlobalEvent() { + return prefetchingQueue.getSubscriptionUncommittedEventCount() + * SubscriptionAgent.broker().getPrefetchingQueueCount() + >= PREFETCH_EVENT_GLOBAL_COUNT_THRESHOLD; + } + private boolean isMissingRateTooHigh() { return missingRate() > PREFETCH_MISSING_RATE_THRESHOLD; } @@ -169,6 +185,8 @@ public String toString() { .add("pollRate", pollRate()) .add("missingRate", missingRate()) .add("disorderCause", disorderCauseCounter.getCount()) + .add("retainedEventCount", prefetchingQueue.getSubscriptionRetainedEventCount()) + .add("inFlightEventCount", prefetchingQueue.getSubscriptionUncommittedEventCount()) .toString(); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java index fcd5cc4c3b654..eca5b4f282126 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java @@ -586,8 +586,11 @@ public SubscriptionEvent poll(final String consumerId, final RegionProgress regi if (pendingSeekRequest != null) { return null; } + if (shouldThrottlePollByInFlightEvents()) { + return null; + } final SubscriptionEvent event = pollInternal(consumerId); - if (Objects.nonNull(event) && prefetchingQueue.size() < MAX_PREFETCHING_QUEUE_SIZE) { + if (Objects.nonNull(event) && hasAvailableEventSlot()) { requestPrefetch(); } else if (Objects.isNull(event) && shouldRecoverPrefetchBindingAfterEmptyPoll()) { requestPrefetch(); @@ -598,6 +601,25 @@ public SubscriptionEvent poll(final String consumerId, final RegionProgress regi } } + private boolean shouldThrottlePollByInFlightEvents() { + if (inFlightEvents.size() < MAX_PREFETCHING_QUEUE_SIZE) { + return false; + } + + recycleInFlightEvents(); + if (inFlightEvents.size() < MAX_PREFETCHING_QUEUE_SIZE) { + return false; + } + + LOGGER.debug( + "ConsensusPrefetchingQueue {}: throttles poll because too many events are in flight", this); + return true; + } + + private boolean hasAvailableEventSlot() { + return getSubscriptionRetainedEventCount() < MAX_PREFETCHING_QUEUE_SIZE; + } + private synchronized void initPrefetch(final RegionProgress regionProgress) { if (prefetchInitialized) { return; // double-check under synchronization @@ -1197,7 +1219,7 @@ public PrefetchRoundResult drivePrefetchOnce() { applyPendingSubscriptionWalReset(observedSeekGeneration); recycleInFlightEvents(); - if (!isActive || prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) { + if (!isActive || !hasAvailableEventSlot()) { return computeIdleRoundResult(); } @@ -1244,7 +1266,7 @@ public PrefetchRoundResult drivePrefetchOnce() { 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) { @@ -1331,8 +1353,11 @@ private PrefetchRoundResult computeIdleRoundResult() { if (isClosed || !prefetchInitialized || !isActive) { return PrefetchRoundResult.dormant(); } - if (prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) { - return PrefetchRoundResult.dormant(); + if (!hasAvailableEventSlot()) { + return inFlightEvents.isEmpty() + ? PrefetchRoundResult.dormant() + : PrefetchRoundResult.rescheduleAfter( + SubscriptionConfig.getInstance().getSubscriptionRecycleUncommittedEventIntervalMs()); } if (hasImmediatePrefetchableWork()) { return PrefetchRoundResult.rescheduleNow(); @@ -1598,7 +1623,7 @@ private boolean pumpFromSubscriptionWAL( int entriesRead = 0; while (entriesRead < maxWalEntries && subscriptionWALIterator.hasNext() - && prefetchingQueue.size() < MAX_PREFETCHING_QUEUE_SIZE) { + && hasAvailableEventSlot()) { try { final IndexedConsensusRequest walEntry = subscriptionWALIterator.next(); entriesRead++; @@ -1879,7 +1904,7 @@ private boolean drainBufferedRealtimeWriters( drainRealtimeWriters(batchState, maxTablets, maxBatchBytes); final int bufferedAfter = getRealtimeBufferedEntryCount(); - if (bufferedAfter == 0 || prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) { + if (bufferedAfter == 0 || !hasAvailableEventSlot()) { return true; } @@ -2066,13 +2091,19 @@ private boolean canAcceptCommitContext( } public boolean ack(final String consumerId, final SubscriptionCommitContext commitContext) { + final boolean acked; acquireReadLock(); try { - return canAcceptCommitContext(commitContext, "ack", false) - && ackInternal(consumerId, commitContext); + acked = + canAcceptCommitContext(commitContext, "ack", false) + && ackInternal(consumerId, commitContext); } finally { releaseReadLock(); } + if (acked && hasAvailableEventSlot()) { + requestPrefetch(); + } + return acked; } private boolean ackInternal( @@ -2148,6 +2179,7 @@ public boolean nack(final String consumerId, final SubscriptionCommitContext com * in multi-region iteration where only one queue owns the event. */ public boolean ackSilent(final String consumerId, final SubscriptionCommitContext commitContext) { + final boolean ackedResult; acquireReadLock(); try { if (!canAcceptCommitContext(commitContext, "ack", true)) { @@ -2190,10 +2222,14 @@ public boolean ackSilent(final String consumerId, final SubscriptionCommitContex ev.cleanUp(false); return null; }); - return acked.get(); + ackedResult = acked.get(); } finally { releaseReadLock(); } + if (ackedResult && hasAvailableEventSlot()) { + requestPrefetch(); + } + return ackedResult; } private WriterId extractCommitWriterId(final SubscriptionCommitContext commitContext) { @@ -2754,8 +2790,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) { @@ -3111,6 +3147,10 @@ public long getSubscriptionUncommittedEventCount() { return inFlightEvents.size(); } + public long getSubscriptionRetainedEventCount() { + return prefetchingQueue.size() + inFlightEvents.size(); + } + /** Exposes the current seek generation for runtime tests and metrics. */ public long getCurrentSeekGeneration() { return seekGeneration.get(); From 91787df8166d429815f263862f98b2e87c1debfc Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 29 Jun 2026 15:59:46 +0800 Subject: [PATCH 2/3] Add subscription prefetching queue state tests --- .../SubscriptionPrefetchingQueueStates.java | 14 ++- ...ubscriptionPrefetchingQueueStatesTest.java | 110 ++++++++++++++++++ 2 files changed, 122 insertions(+), 2 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStatesTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStates.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStates.java index d5715e9fa41cc..29beedbae7228 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStates.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStates.java @@ -30,6 +30,8 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.function.IntSupplier; + import static com.google.common.base.MoreObjects.toStringHelper; /** @@ -54,6 +56,7 @@ public class SubscriptionPrefetchingQueueStates { SubscriptionConfig.getInstance().getSubscriptionPrefetchEventGlobalCountThreshold(); private final SubscriptionPrefetchingQueue prefetchingQueue; + private final IntSupplier prefetchingQueueCountSupplier; private volatile long lastPollRequestTimestamp; private final Meter pollRequestMeter; @@ -61,7 +64,14 @@ public class SubscriptionPrefetchingQueueStates { private final Counter disorderCauseCounter; // TODO: use meter 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()); @@ -151,7 +161,7 @@ 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.getSubscriptionRetainedEventCount() - * SubscriptionAgent.broker().getPrefetchingQueueCount() + * prefetchingQueueCountSupplier.getAsInt() > PREFETCH_EVENT_GLOBAL_COUNT_THRESHOLD; } @@ -162,7 +172,7 @@ private boolean hasTooManyInFlightLocalEvent() { private boolean hasTooManyInFlightGlobalEvent() { return prefetchingQueue.getSubscriptionUncommittedEventCount() - * SubscriptionAgent.broker().getPrefetchingQueueCount() + * prefetchingQueueCountSupplier.getAsInt() >= PREFETCH_EVENT_GLOBAL_COUNT_THRESHOLD; } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStatesTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStatesTest.java new file mode 100644 index 0000000000000..d8878f17d9046 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStatesTest.java @@ -0,0 +1,110 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.subscription.broker; + +import org.apache.iotdb.commons.subscription.config.SubscriptionConfig; + +import org.junit.Test; + +import java.lang.reflect.Method; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class SubscriptionPrefetchingQueueStatesTest { + + private static final int LOCAL_COUNT_THRESHOLD = + SubscriptionConfig.getInstance().getSubscriptionPrefetchEventLocalCountThreshold(); + private static final int GLOBAL_COUNT_THRESHOLD = + SubscriptionConfig.getInstance().getSubscriptionPrefetchEventGlobalCountThreshold(); + + @Test + public void testShouldThrottlePollByLocalInFlightEventCount() { + assertFalse(newStates(queueWithCounts(0, LOCAL_COUNT_THRESHOLD - 1L), 1).shouldThrottlePoll()); + assertTrue(newStates(queueWithCounts(0, LOCAL_COUNT_THRESHOLD), 1).shouldThrottlePoll()); + } + + @Test + public void testShouldThrottlePollByGlobalInFlightEventCount() { + final long inFlightEventCount = LOCAL_COUNT_THRESHOLD - 1L; + final int belowGlobalQueueCount = + Math.max(1, (int) ((GLOBAL_COUNT_THRESHOLD - 1L) / inFlightEventCount)); + final int reachingGlobalQueueCount = + (int) ((GLOBAL_COUNT_THRESHOLD + inFlightEventCount - 1L) / inFlightEventCount); + + assertFalse( + newStates(queueWithCounts(0, inFlightEventCount), belowGlobalQueueCount) + .shouldThrottlePoll()); + assertTrue( + newStates(queueWithCounts(0, inFlightEventCount), reachingGlobalQueueCount) + .shouldThrottlePoll()); + } + + @Test + public void testRetainedLocalEventLimitCountsAllRetainedEvents() throws Exception { + assertFalse( + invokeStatePredicate( + newStates(queueWithCounts(LOCAL_COUNT_THRESHOLD, 0), 1), + "hasTooManyRetainedLocalEvent")); + assertTrue( + invokeStatePredicate( + newStates(queueWithCounts(LOCAL_COUNT_THRESHOLD + 1L, 0), 1), + "hasTooManyRetainedLocalEvent")); + } + + @Test + public void testRetainedGlobalEventLimitUsesInjectedQueueCount() throws Exception { + final long retainedEventCount = Math.max(1L, GLOBAL_COUNT_THRESHOLD / 2L); + final int belowOrAtGlobalQueueCount = + Math.max(1, (int) (GLOBAL_COUNT_THRESHOLD / retainedEventCount)); + final int aboveGlobalQueueCount = belowOrAtGlobalQueueCount + 1; + + assertFalse( + invokeStatePredicate( + newStates(queueWithCounts(retainedEventCount, 0), belowOrAtGlobalQueueCount), + "hasTooManyRetainedGlobalEvent")); + assertTrue( + invokeStatePredicate( + newStates(queueWithCounts(retainedEventCount, 0), aboveGlobalQueueCount), + "hasTooManyRetainedGlobalEvent")); + } + + private static SubscriptionPrefetchingQueueStates newStates( + final SubscriptionPrefetchingQueue queue, final int prefetchingQueueCount) { + return new SubscriptionPrefetchingQueueStates(queue, () -> prefetchingQueueCount); + } + + private static SubscriptionPrefetchingQueue queueWithCounts( + final long retainedEventCount, final long inFlightEventCount) { + final SubscriptionPrefetchingQueue queue = mock(SubscriptionPrefetchingQueue.class); + when(queue.getSubscriptionRetainedEventCount()).thenReturn(retainedEventCount); + when(queue.getSubscriptionUncommittedEventCount()).thenReturn(inFlightEventCount); + return queue; + } + + private static boolean invokeStatePredicate( + final SubscriptionPrefetchingQueueStates states, final String methodName) throws Exception { + final Method method = SubscriptionPrefetchingQueueStates.class.getDeclaredMethod(methodName); + method.setAccessible(true); + return (boolean) method.invoke(states); + } +} From 9589f4289ef7050604ddf83d6468e85300b1bbc9 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 26 Oct 2026 12:55:46 +0800 Subject: [PATCH 3/3] Fix subscription retained-event backpressure gaps Serialize poll limit checks with event claims, count retained and lingering events throughout consensus prefetch, and honor topic recycle settings when capacity is full. Preserve bounded realtime admission and add regression coverage for the reproduced failures. --- .../broker/SubscriptionPrefetchingQueue.java | 30 +- .../consensus/ConsensusPrefetchingQueue.java | 54 ++-- ...onPrefetchingQueueRecycleIntervalTest.java | 92 ++++++ ...ensusPrefetchingQueueBackpressureTest.java | 292 ++++++++++++++++++ .../ConsensusPrefetchingQueueTest.java | 12 +- 5 files changed, 439 insertions(+), 41 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueBackpressureTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java index 544a2a5685f49..fe902fb19c880 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueue.java +++ b/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(); } @@ -281,14 +286,7 @@ private boolean shouldThrottlePollByInFlightEvents() { } remapInFlightEventsSnapshot(committedCleaner, pollableNacker); - if (!states.shouldThrottlePoll()) { - return false; - } - - LOGGER.debug( - "Subscription: SubscriptionPrefetchingQueue {} throttles poll because too many events are in flight.", - this); - return true; + return states.shouldThrottlePoll(); } public SubscriptionEvent pollV2(final String consumerId, final PollTimer timer) { @@ -306,13 +304,15 @@ public SubscriptionEvent pollV2(final String consumerId, final PollTimer timer) return null; } - if (shouldThrottlePollByInFlightEvents()) { - return null; - } + synchronized (pollLock) { + if (shouldThrottlePollByInFlightEvents()) { + return null; + } - final SubscriptionEvent event = pollInternalV2(consumerId); - if (Objects.nonNull(event)) { - return event; + final SubscriptionEvent event = pollInternalV2(consumerId); + if (Objects.nonNull(event)) { + return event; + } } } finally { releaseReadLock(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java index d828951ee8959..5f74d33e50ab1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java +++ b/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,10 +996,13 @@ public SubscriptionEvent poll(final String consumerId, final RegionProgress regi if (pendingSeekRequest != null) { return null; } - if (shouldThrottlePollByInFlightEvents()) { - return null; + final SubscriptionEvent event; + synchronized (pollLock) { + if (shouldThrottlePollByInFlightEvents()) { + return null; + } + event = pollInternal(consumerId); } - final SubscriptionEvent event = pollInternal(consumerId); if (Objects.nonNull(event) && hasAvailableEventSlot() && !hasReachedPrefetchLimit()) { requestPrefetch(); } else if (Objects.isNull(event) && shouldRecoverPrefetchBindingAfterEmptyPoll()) { @@ -1035,19 +1040,18 @@ private boolean shouldThrottlePollByInFlightEvents() { } recycleInFlightEvents(); - if (inFlightEvents.size() < MAX_PREFETCHING_QUEUE_SIZE) { - return false; - } - - LOGGER.debug( - "ConsensusPrefetchingQueue {}: throttles poll because too many events are in flight", this); - return true; + 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); @@ -1841,10 +1845,6 @@ private PrefetchRoundResult drivePrefetchOnceUnderReadLock() { return computeIdleRoundResult(walScanBudget); } if (!hasAvailableEventSlot()) { - blockRealtimeAdmission(SubscriptionQueueRejectionReason.QUEUE_CAPACITY); - return computeIdleRoundResult(walScanBudget); - } - if (prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) { return computeIdleRoundResult(walScanBudget); } @@ -2052,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(); } @@ -2173,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 @@ -2594,11 +2596,10 @@ private PrefetchRoundResult computeIdleRoundResult( return PrefetchRoundResult.dormant(); } if (!hasAvailableEventSlot()) { - blockRealtimeAdmission(SubscriptionQueueRejectionReason.QUEUE_CAPACITY); - return inFlightEvents.isEmpty() + final long recycleIntervalMs = getRecycleUncommittedEventIntervalMs(); + return inFlightEvents.isEmpty() || recycleIntervalMs <= 0L ? PrefetchRoundResult.dormant() - : PrefetchRoundResult.rescheduleAfter( - SubscriptionConfig.getInstance().getSubscriptionRecycleUncommittedEventIntervalMs()); + : PrefetchRoundResult.rescheduleAfter(recycleIntervalMs); } if (hasReachedPrefetchLimit()) { return PrefetchRoundResult.dormant(); @@ -3890,6 +3891,9 @@ private boolean drainWriterEntries( final boolean trackLingerTime, final long expectedSeekGeneration) { while (true) { + if (!hasAvailableEventSlot(batchState)) { + return true; + } final PriorityQueue frontiers = frontierSupplier.get(); if (frontiers.isEmpty()) { return true; diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueRecycleIntervalTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueRecycleIntervalTest.java index 3aa26ed54b106..8625e6b05de0d 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueRecycleIntervalTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueRecycleIntervalTest.java @@ -33,14 +33,23 @@ import java.lang.reflect.Field; import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.when; public class SubscriptionPrefetchingQueueRecycleIntervalTest { @@ -55,6 +64,89 @@ public void testRecycleIntervalForV2Queue() throws Exception { verifyRecycleInterval(true); } + @Test + public void testConcurrentPollsRespectInFlightLimitForV1Queue() throws Exception { + verifyConcurrentPollLimit(false); + } + + @Test + public void testConcurrentPollsRespectInFlightLimitForV2Queue() throws Exception { + verifyConcurrentPollLimit(true); + } + + private static void verifyConcurrentPollLimit(final boolean useV2) throws Exception { + final int limit = + CommonDescriptor.getInstance() + .getConfig() + .getSubscriptionPrefetchEventLocalCountThreshold(); + final SubscriptionPrefetchingQueue queue = + new SubscriptionPrefetchingTabletQueue( + "concurrent_group", + "concurrent_poll_topic", + mock(SubscriptionBlockingPendingQueue.class), + new AtomicLong()); + final ExecutorService executor = Executors.newFixedThreadPool(2); + final CountDownLatch claiming = new CountDownLatch(1); + final CountDownLatch releaseClaim = new CountDownLatch(1); + try { + for (int index = 1; index <= limit + 1; index++) { + final SubscriptionEvent event = + new SubscriptionEvent( + SubscriptionPollResponseType.TERMINATION.getType(), + new TerminationPayload(), + new SubscriptionCommitContext( + 1, 1, "concurrent_poll_topic", "concurrent_group", index)); + if (index == limit) { + final SubscriptionEvent blockedEvent = spy(event); + doAnswer( + invocation -> { + claiming.countDown(); + assertTrue(releaseClaim.await(5, TimeUnit.SECONDS)); + return invocation.callRealMethod(); + }) + .when(blockedEvent) + .recordLastPolledTimestamp(); + queue.prefetchEvent(blockedEvent); + } else { + queue.prefetchEvent(event); + } + } + for (int index = 1; index < limit; index++) { + assertNotNull(poll(queue, useV2)); + } + final Future first = executor.submit(() -> poll(queue, useV2)); + assertTrue(claiming.await(5, TimeUnit.SECONDS)); + final AtomicReference secondThread = new AtomicReference<>(); + final Future second = + executor.submit( + () -> { + secondThread.set(Thread.currentThread()); + return poll(queue, useV2); + }); + final long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (!second.isDone() + && (secondThread.get() == null || secondThread.get().getState() != Thread.State.BLOCKED) + && System.nanoTime() < deadline) { + Thread.yield(); + } + assertTrue( + second.isDone() + || (secondThread.get() != null + && secondThread.get().getState() == Thread.State.BLOCKED)); + releaseClaim.countDown(); + assertNotNull(first.get(5, TimeUnit.SECONDS)); + assertNull(second.get(5, TimeUnit.SECONDS)); + assertEquals(limit, queue.getSubscriptionUncommittedEventCount()); + assertEquals(limit + 1L, queue.getSubscriptionRetainedEventCount()); + assertEquals(1, queue.getPrefetchedEventCount()); + } finally { + releaseClaim.countDown(); + executor.shutdownNow(); + executor.awaitTermination(5, TimeUnit.SECONDS); + queue.cleanUp(); + } + } + private static void verifyRecycleInterval(final boolean useV2) throws Exception { final String topicName = "ordinary_recycle_topic_" + useV2; final int originalInterval = diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueBackpressureTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueBackpressureTest.java new file mode 100644 index 0000000000000..9881d4c4e4f48 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueBackpressureTest.java @@ -0,0 +1,292 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.subscription.broker.consensus; + +import org.apache.iotdb.commons.conf.CommonConfig; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.consensus.DataRegionId; +import org.apache.iotdb.commons.subscription.meta.topic.TopicMeta; +import org.apache.iotdb.consensus.common.request.IndexedConsensusRequest; +import org.apache.iotdb.consensus.iot.IoTConsensusServerImpl; +import org.apache.iotdb.consensus.iot.SubscriptionWalRetentionPolicy; +import org.apache.iotdb.consensus.iot.WriterSafeFrontierTracker; +import org.apache.iotdb.consensus.iot.log.ConsensusReqReader; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.plan.statement.StatementTestUtils; +import org.apache.iotdb.db.subscription.agent.SubscriptionAgent; +import org.apache.iotdb.db.subscription.event.SubscriptionEvent; +import org.apache.iotdb.db.subscription.resource.SubscriptionMemoryManager; +import org.apache.iotdb.rpc.subscription.config.TopicConstant; +import org.apache.iotdb.rpc.subscription.payload.poll.RegionProgress; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +@RunWith(Parameterized.class) +public class ConsensusPrefetchingQueueBackpressureTest { + + @Parameterized.Parameters(name = "{0}, WAL={1}") + public static Collection parameters() { + return Arrays.asList( + new Object[] {TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE, false}, + new Object[] {TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE, true}, + new Object[] {TopicConstant.ORDER_MODE_MULTI_WRITER_VALUE, false}); + } + + @Parameterized.Parameter public String orderMode; + + @Parameterized.Parameter(1) + public boolean fromWal; + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + private final CommonConfig config = CommonDescriptor.getInstance().getConfig(); + private final AtomicInteger conversions = new AtomicInteger(); + private final AtomicInteger walPosition = new AtomicInteger(); + private final List requests = new ArrayList<>(); + private final List polledEvents = new ArrayList<>(); + private String originalSystemDir; + private int originalTablets; + private int originalWalEntries; + private int originalDelay; + private long originalWatermarkInterval; + private int capacity; + private ConsensusPrefetchingQueue queue; + + @Before + public void setUp() throws Exception { + originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir(); + originalTablets = config.getSubscriptionConsensusBatchMaxTabletCount(); + originalWalEntries = config.getSubscriptionConsensusBatchMaxWalEntries(); + originalDelay = config.getSubscriptionConsensusBatchMaxDelayInMs(); + originalWatermarkInterval = config.getSubscriptionConsensusWatermarkIntervalMs(); + capacity = config.getSubscriptionConsensusPrefetchingQueueCapacity(); + config.setSubscriptionConsensusBatchMaxTabletCount(1); + config.setSubscriptionConsensusBatchMaxWalEntries(capacity * 2); + config.setSubscriptionConsensusBatchMaxDelayInMs(0); + config.setSubscriptionConsensusWatermarkIntervalMs(0L); + IoTDBDescriptor.getInstance() + .getConfig() + .setSystemDir(temporaryFolder.newFolder("system").getAbsolutePath()); + + for (int index = 1; index <= capacity * 2; index++) { + requests.add( + new IndexedConsensusRequest( + index, Collections.singletonList(StatementTestUtils.genInsertRowNode(index))) + .setNodeId(7) + .setPhysicalTime(1000L + index)); + } + final ConsensusReqReader reader = mock(ConsensusReqReader.class); + when(reader.getCurrentSearchIndex()).thenReturn((long) requests.size()); + final WriterSafeFrontierTracker tracker = new WriterSafeFrontierTracker(); + tracker.recordAppliedProgress(1000L + requests.size(), 7, requests.size()); + tracker.observePendingWriterSafeTimeBarrier(Long.MAX_VALUE, 7, requests.size()); + tracker.observePendingWriterSafeTimeBarrier(Long.MAX_VALUE, 8, 0); + final IoTConsensusServerImpl server = mock(IoTConsensusServerImpl.class); + when(server.getConsensusReqReader()).thenReturn(reader); + when(server.getWriterSafeFrontierTracker()).thenReturn(tracker); + final ConsensusLogToTabletConverter converter = mock(ConsensusLogToTabletConverter.class); + when(converter.getDatabaseName()).thenReturn("db"); + when(converter.convert(any())) + .thenAnswer( + ignored -> { + conversions.incrementAndGet(); + final Tablet tablet = + new Tablet( + "root.test.device", + Collections.singletonList(new MeasurementSchema("s", TSDataType.INT64)), + 1); + tablet.addTimestamp(0, 1L); + tablet.addValue("s", 0, 1L); + tablet.setRowSize(1); + return Collections.singletonList(tablet); + }); + final ProgressWALIterator walIterator = mock(ProgressWALIterator.class); + when(walIterator.hasNext(any(ProgressWALIterator.ScanBudget.class))) + .thenAnswer(ignored -> walPosition.get() < requests.size()); + when(walIterator.next()).thenAnswer(ignored -> requests.get(walPosition.getAndIncrement())); + queue = + new ConsensusPrefetchingQueue( + "backpressure_group", + "backpressure_topic", + orderMode, + new DataRegionId(1), + server, + new SubscriptionWalRetentionPolicy( + "backpressure_topic", + SubscriptionWalRetentionPolicy.UNBOUNDED, + SubscriptionWalRetentionPolicy.UNBOUNDED), + converter, + new ConsensusSubscriptionCommitManager( + (group, topic, id) -> + ConsensusSubscriptionCommitManager.ConfigNodeProgressQueryResult.absent()), + new RegionProgress(Collections.emptyMap()), + 1L, + 1L, + true) { + @Override + protected ProgressWALIterator createSubscriptionWALIterator(final long startSearchIndex) { + return fromWal ? walIterator : null; + } + }; + queue.setSubscriptionMemoryManager(new SubscriptionMemoryManager(64L * 1024 * 1024)); + if (TopicConstant.ORDER_MODE_MULTI_WRITER_VALUE.equals(orderMode)) { + queue.setActiveWriterNodeIds(Set.of(7, 8)); + } + assertNull(queue.poll("consumer")); + if (!fromWal) { + for (final IndexedConsensusRequest request : requests) { + assertTrue(pending().offer(request)); + } + } + } + + @After + public void tearDown() { + if (queue != null) { + queue.close(); + assertEquals(0L, queue.getRetainedRequestBytes()); + assertEquals(0L, queue.getRetainedTabletBytes()); + } + SubscriptionAgent.topic().handleDropTopic("backpressure_topic", false); + config.setSubscriptionConsensusBatchMaxTabletCount(originalTablets); + config.setSubscriptionConsensusBatchMaxWalEntries(originalWalEntries); + config.setSubscriptionConsensusBatchMaxDelayInMs(originalDelay); + config.setSubscriptionConsensusWatermarkIntervalMs(originalWatermarkInterval); + IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir); + } + + @Test + public void testAckResumesWithinRetainedEventCapacity() { + verifyAckResumesWithinRetainedEventCapacity(false); + } + + @Test + public void testAckSilentResumesWithinRetainedEventCapacity() { + verifyAckResumesWithinRetainedEventCapacity(true); + } + + private void verifyAckResumesWithinRetainedEventCapacity(final boolean silent) { + fillRetainedCapacity(); + assertTrue( + silent + ? queue.ackSilent("consumer", polledEvents.get(0).getCommitContext()) + : queue.ack("consumer", polledEvents.get(0).getCommitContext())); + queue.drivePrefetchOnce(); + assertEquals(capacity, queue.getSubscriptionRetainedEventCount()); + assertEquals(2, queue.getPrefetchedEventCount()); + assertEquals(capacity - 2L, queue.getSubscriptionUncommittedEventCount()); + assertEquals(capacity + 1, conversions.get()); + queue.drivePrefetchOnce(); + assertEquals(capacity + 1, conversions.get()); + + for (int index = 1; index < polledEvents.size(); index++) { + assertTrue(queue.ack("consumer", polledEvents.get(index).getCommitContext())); + } + int delivered = polledEvents.size(); + for (int round = 0; round < requests.size() && delivered < requests.size(); round++) { + queue.drivePrefetchOnce(); + assertTrue(queue.getSubscriptionRetainedEventCount() <= capacity); + SubscriptionEvent event; + while ((event = queue.poll("consumer")) != null) { + delivered++; + assertEquals(delivered, event.getCommitContext().getWriterProgress().getLocalSeq()); + assertTrue(queue.ack("consumer", event.getCommitContext())); + } + } + assertEquals(requests.size(), delivered); + assertEquals(requests.size(), conversions.get()); + assertEquals(0L, queue.getSubscriptionRetainedEventCount()); + } + + @Test + public void testFullQueueUsesTopicRecycleInterval() { + fillRetainedCapacity(); + for (final long intervalMs : new long[] {1200000L, 0L, -1L, Long.MIN_VALUE, 300000L}) { + assertNull( + SubscriptionAgent.topic() + .handleSingleTopicMetaChanges( + new TopicMeta( + "backpressure_topic", + 1L, + Collections.singletonMap( + TopicConstant.RECYCLE_UNCOMMITTED_EVENT_INTERVAL_MS_KEY, + Long.toString(intervalMs))))); + final PrefetchRoundResult result = queue.drivePrefetchOnce(); + assertEquals( + intervalMs > 0L + ? PrefetchRoundResult.Type.RESCHEDULE_LATER + : PrefetchRoundResult.Type.DORMANT, + result.getType()); + if (intervalMs > 0L) { + assertEquals(intervalMs, result.getDelayMs()); + } + assertEquals(capacity, queue.getSubscriptionRetainedEventCount()); + assertEquals(capacity, conversions.get()); + } + } + + private void fillRetainedCapacity() { + queue.drivePrefetchOnce(); + assertEquals(capacity, queue.getPrefetchedEventCount()); + for (int index = 1; index < capacity; index++) { + final SubscriptionEvent event = queue.poll("consumer"); + assertNotNull(event); + polledEvents.add(event); + } + queue.drivePrefetchOnce(); + assertEquals(capacity, queue.getSubscriptionRetainedEventCount()); + assertEquals(capacity, conversions.get()); + assertEquals(1, queue.getPrefetchedEventCount()); + } + + @SuppressWarnings("unchecked") + private BlockingQueue pending() throws Exception { + final Field field = ConsensusPrefetchingQueue.class.getDeclaredField("pendingEntries"); + field.setAccessible(true); + return (BlockingQueue) field.get(queue); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java index 8d7edd979c50f..18cfa20714ab6 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueTest.java @@ -3302,7 +3302,17 @@ public void testPrefetchQueueCapacityDoesNotDisableRealtimeAdmission() throws Ex assertEquals("false", queue.coreReportMessage().get("realtimeAdmissionBlocked")); assertTrue(pendingEntries(queue).offer(createRequest(requestCount + 1L))); - assertNotNull(queue.poll("consumer")); + final SubscriptionEvent event = queue.poll("consumer"); + assertNotNull(event); + queue.drivePrefetchOnce(); + + assertEquals(prefetchingQueueCapacity - 1, queue.getPrefetchedEventCount()); + assertEquals(prefetchingQueueCapacity, queue.getSubscriptionRetainedEventCount()); + assertEquals(prefetchingQueueCapacity, conversionCount.get()); + assertEquals(2, pendingEntries(queue).size()); + assertEquals("false", queue.coreReportMessage().get("realtimeAdmissionBlocked")); + + assertTrue(queue.ack("consumer", event.getCommitContext())); queue.drivePrefetchOnce(); assertEquals(prefetchingQueueCapacity, queue.getPrefetchedEventCount());