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 755b75df3dc3..f110c08d9082 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(); } @@ -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(); @@ -1169,6 +1193,10 @@ public long getSubscriptionUncommittedEventCount() { return inFlightEvents.size(); } + public long getSubscriptionRetainedEventCount() { + return prefetchingQueue.size() + inFlightEvents.size(); + } + public long getCurrentCommitId() { return commitIdGenerator.get(); } @@ -1216,6 +1244,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 a9167b7b4755..9eb83482ea86 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,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,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; @@ -62,7 +64,14 @@ public class SubscriptionPrefetchingQueueStates { 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(); } } 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 90d102fbbd87..7da446c093d2 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 @@ -146,6 +146,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; @@ -826,7 +829,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)); @@ -835,8 +838,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() { @@ -1019,8 +1021,14 @@ public SubscriptionEvent poll( if (pendingSeekRequest != null) { return null; } - final SubscriptionEvent event = pollInternal(consumerId, timer); - if (Objects.nonNull(event) && !hasReachedPrefetchLimit()) { + final SubscriptionEvent event; + synchronized (pollLock) { + if (shouldThrottlePollByInFlightEvents()) { + return null; + } + event = pollInternal(consumerId, timer); + } + if (Objects.nonNull(event) && hasAvailableEventSlot() && !hasReachedPrefetchLimit()) { requestPrefetch(); } else if (Objects.isNull(event) && shouldRecoverPrefetchBindingAfterEmptyPoll()) { requestPrefetch(); @@ -1051,6 +1059,24 @@ public SubscriptionEvent poll( } } + 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); @@ -1876,7 +1902,7 @@ private PrefetchRoundResult drivePrefetchOnceUnderReadLock() { blockRealtimeAdmission(SubscriptionQueueRejectionReason.INACTIVE_OR_CLOSED); return computeIdleRoundResult(walScanBudget); } - if (prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) { + if (!hasAvailableEventSlot()) { return computeIdleRoundResult(walScanBudget); } @@ -1995,14 +2021,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) { @@ -2082,7 +2110,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(); } @@ -2203,13 +2231,15 @@ 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) { + if (!hasAvailableEventSlot()) { return computeIdleRoundResult(walScanBudget); } if (walScanBudget.hasYielded() @@ -2623,6 +2653,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()) { // Backpressure is temporary: writes may have stopped, and another queue releasing its // memory quota does not wake this queue. Keep a delayed retry without scanning more WAL @@ -3022,6 +3058,7 @@ private MaterializationResult pumpFromSubscriptionWAL( int entriesRead = 0; while (entriesRead < maxWalEntries && !hasReachedPrefetchLimit(batchState) + && hasAvailableEventSlot() && subscriptionWALIterator.hasNext(walScanBudget)) { reportUnreadableWalReplayIfNecessary(); try { @@ -3840,7 +3877,7 @@ private boolean drainBufferedRealtimeWriters( } final int bufferedAfter = getRealtimeBufferedEntryCount(); - if (bufferedAfter == 0 || prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) { + if (bufferedAfter == 0 || !hasAvailableEventSlot()) { return true; } @@ -3915,6 +3952,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; @@ -4144,7 +4184,9 @@ public boolean ack(final String consumerId, final SubscriptionCommitContext comm && ackMissingInFlightEvent(commitContext, false)); if (acked) { serverImpl.checkAndUpdateSafeDeletedSearchIndex(); - requestPrefetch(); + if (hasAvailableEventSlot()) { + requestPrefetch(); + } } return acked; } @@ -4237,7 +4279,9 @@ public boolean ackSilent(final String consumerId, final SubscriptionCommitContex && ackMissingInFlightEvent(commitContext, true)); if (ackedResult) { serverImpl.checkAndUpdateSafeDeletedSearchIndex(); - requestPrefetch(); + if (hasAvailableEventSlot()) { + requestPrefetch(); + } } return ackedResult; } @@ -5175,8 +5219,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) { @@ -5925,6 +5969,10 @@ public long getSubscriptionUncommittedEventCount() { return inFlightEvents.size(); } + public long getSubscriptionRetainedEventCount() { + return prefetchingQueue.size() + inFlightEvents.size(); + } + public long getRetainedTabletBytes() { return retainedTabletBytes.get(); } 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 3aa26ed54b10..8625e6b05de0 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/SubscriptionPrefetchingQueueStatesTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionPrefetchingQueueStatesTest.java new file mode 100644 index 000000000000..d8878f17d904 --- /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); + } +} 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 000000000000..9881d4c4e4f4 --- /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 84ec91d34bce..149b400160bb 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 @@ -3319,7 +3319,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());