This is an automated email from the ASF dual-hosted git repository. Caideyipi pushed a commit to branch fix/subscription-poll-timeout-budget in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit b14c27d0b3a81b9dd127b45c76cbb94bf9feae14 Author: Caideyipi <[email protected]> AuthorDate: Wed Oct 28 12:05:40 2026 +0800 Fix timeout budget propagation in consensus subscription polling --- .../broker/ConsensusSubscriptionBroker.java | 28 ++- .../consensus/ConsensusPrefetchingQueue.java | 55 ++++-- ...onsensusSubscriptionBrokerPayloadLimitTest.java | 23 ++- ...ConsensusSubscriptionBrokerPollTimeoutTest.java | 198 +++++++++++++++++++++ .../ConsensusPrefetchingQueuePollTimeoutTest.java | 180 +++++++++++++++++++ 5 files changed, 464 insertions(+), 20 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java index 22c79ac7303..95e5a04991c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java @@ -25,6 +25,7 @@ import org.apache.iotdb.commons.subscription.meta.consumer.SubscriptionProgressS import org.apache.iotdb.consensus.iot.IoTConsensusServerImpl; import org.apache.iotdb.consensus.iot.SubscriptionWalRetentionPolicy; import org.apache.iotdb.db.i18n.DataNodePipeMessages; +import org.apache.iotdb.db.subscription.agent.SubscriptionAgent; import org.apache.iotdb.db.subscription.broker.consensus.ConsensusLogToTabletConverter; import org.apache.iotdb.db.subscription.broker.consensus.ConsensusPrefetchingQueue; import org.apache.iotdb.db.subscription.broker.consensus.ConsensusRegionRuntimeState; @@ -34,6 +35,7 @@ import org.apache.iotdb.rpc.subscription.payload.poll.RegionProgress; import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; import org.apache.iotdb.rpc.subscription.payload.poll.TopicProgress; import org.apache.iotdb.rpc.subscription.payload.request.SubscriptionSeekReq; +import org.apache.iotdb.session.subscription.util.PollTimer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -110,11 +112,26 @@ public class ConsensusSubscriptionBroker implements ISubscriptionBroker { return poll(consumerId, topicNames, maxBytes, Collections.emptyMap()); } + @Override public List<SubscriptionEvent> poll( final String consumerId, final Set<String> topicNames, final long maxBytes, final Map<String, TopicProgress> progressByTopic) { + return poll( + consumerId, + topicNames, + maxBytes, + progressByTopic, + new PollTimer(System.currentTimeMillis(), SubscriptionAgent.receiver().remainingMs())); + } + + List<SubscriptionEvent> poll( + final String consumerId, + final Set<String> topicNames, + final long maxBytes, + final Map<String, TopicProgress> progressByTopic, + final PollTimer timer) { LOGGER.debug( DataNodePipeMessages .PIPE_LOG_CONSENSUSSUBSCRIPTIONBROKER_POLL_CALLED_CONSUMERID_TOPICNAMES_5F1F5175, @@ -130,6 +147,10 @@ public class ConsensusSubscriptionBroker implements ISubscriptionBroker { boolean responseFull = false; for (final String topicName : topicNames) { + timer.update(); + if (timer.isExpired()) { + break; + } final List<ConsensusPrefetchingQueue> queues = topicNameToConsensusPrefetchingQueues.get(topicName); if (Objects.isNull(queues) || queues.isEmpty()) { @@ -150,6 +171,10 @@ public class ConsensusSubscriptionBroker implements ISubscriptionBroker { final int eventsBeforeTopicPoll = eventsToPoll.size(); for (int queueIndex = 0; queueIndex < pollQueues.size(); queueIndex++) { + timer.update(); + if (timer.isExpired()) { + break; + } final boolean isFallbackQueue = queueIndex >= assignedQueueCount; // Ownership remains the preferred polling path. Only borrow one ready event from another // region when every assigned region is currently empty, preventing a consumer from timing @@ -170,7 +195,8 @@ public class ConsensusSubscriptionBroker implements ISubscriptionBroker { ? topicProgress.getRegionProgress().get(regionIdStr) : null; - final SubscriptionEvent event = consensusQueue.poll(consumerId, regionProgress); + // All topics and regions share the receiver's remaining budget, including queue lock waits. + final SubscriptionEvent event = consensusQueue.poll(consumerId, regionProgress, timer); if (Objects.isNull(event)) { continue; } 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 df756dc2d22..045194f38ea 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 @@ -68,6 +68,7 @@ import org.apache.iotdb.rpc.subscription.payload.poll.TabletsPayload; import org.apache.iotdb.rpc.subscription.payload.poll.WatermarkPayload; import org.apache.iotdb.rpc.subscription.payload.poll.WriterId; import org.apache.iotdb.rpc.subscription.payload.poll.WriterProgress; +import org.apache.iotdb.session.subscription.util.PollTimer; import org.apache.tsfile.utils.Pair; import org.apache.tsfile.write.record.Tablet; @@ -976,12 +977,36 @@ public class ConsensusPrefetchingQueue { } public SubscriptionEvent poll(final String consumerId, final RegionProgress regionProgress) { - if (!retryPendingRuntimeActivation()) { + return poll( + consumerId, + regionProgress, + new PollTimer(System.currentTimeMillis(), SubscriptionAgent.receiver().remainingMs())); + } + + public SubscriptionEvent poll( + final String consumerId, final RegionProgress regionProgress, final PollTimer timer) { + timer.update(); + if (timer.isExpired() || !retryPendingRuntimeActivation()) { return null; } - acquireReadLock(); try { - if (isClosed || closeRequested || !isActive) { + timer.update(); + if (timer.isExpired() + || !lock.readLock().tryLock(timer.remainingMs(), TimeUnit.MILLISECONDS)) { + return null; + } + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + LOGGER.warn( + DataNodePipeMessages + .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_INTERRUPTED_WHILE_POLLING_B7CFF5FD, + this, + e); + return null; + } + try { + timer.update(); + if (timer.isExpired() || isClosed || closeRequested || !isActive) { return null; } final String walReplayError = pendingWalReplayError.getAndSet(null); @@ -994,7 +1019,7 @@ public class ConsensusPrefetchingQueue { if (pendingSeekRequest != null) { return null; } - final SubscriptionEvent event = pollInternal(consumerId); + final SubscriptionEvent event = pollInternal(consumerId, timer); if (Objects.nonNull(event) && !hasReachedPrefetchLimit()) { requestPrefetch(); } else if (Objects.isNull(event) && shouldRecoverPrefetchBindingAfterEmptyPoll()) { @@ -1519,7 +1544,7 @@ public class ConsensusPrefetchingQueue { return buildWriterFrontiers(realtimeEntriesByWriter, this::peekRealtimeEntry); } - private SubscriptionEvent pollInternal(final String consumerId) { + private SubscriptionEvent pollInternal(final String consumerId, final PollTimer timer) { final long size = prefetchingQueue.size(); if (size == 0) { LOGGER.debug( @@ -1545,12 +1570,20 @@ public class ConsensusPrefetchingQueue { SubscriptionEvent event; try { - while (count++ < size - && Objects.nonNull( - event = - prefetchingQueue.poll( - SubscriptionConfig.getInstance().getSubscriptionPollMaxBlockingTimeMs(), - TimeUnit.MILLISECONDS))) { + while (count++ < size) { + timer.update(); + if (timer.isExpired()) { + break; + } + event = + prefetchingQueue.poll( + Math.min( + SubscriptionConfig.getInstance().getSubscriptionPollMaxBlockingTimeMs(), + timer.remainingMs()), + TimeUnit.MILLISECONDS); + if (Objects.isNull(event)) { + break; + } // Metadata events (currently WATERMARK) are fire-and-forget: // skip inFlightEvents tracking so they are not recycled and re-delivered indefinitely. if (event.getCurrentResponse().getResponseType() diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPayloadLimitTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPayloadLimitTest.java index c2cb73d8fd2..56eeb083481 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPayloadLimitTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPayloadLimitTest.java @@ -23,6 +23,7 @@ import org.apache.iotdb.commons.consensus.DataRegionId; import org.apache.iotdb.db.subscription.broker.consensus.ConsensusPrefetchingQueue; import org.apache.iotdb.db.subscription.event.SubscriptionEvent; import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; +import org.apache.iotdb.session.subscription.util.PollTimer; import org.junit.Test; @@ -34,6 +35,9 @@ import java.util.Map; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertSame; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -58,8 +62,8 @@ public class ConsensusSubscriptionBrokerPayloadLimitTest { when(firstQueue.getConsensusGroupId()).thenReturn(new DataRegionId(1)); when(secondQueue.getConsensusGroupId()).thenReturn(new DataRegionId(2)); - when(firstQueue.poll(CONSUMER_ID, null)).thenReturn(firstEvent); - when(secondQueue.poll(CONSUMER_ID, null)).thenReturn(secondEvent); + when(firstQueue.poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class))).thenReturn(firstEvent); + when(secondQueue.poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class))).thenReturn(secondEvent); when(firstEvent.getCurrentResponseSize()).thenReturn(40); when(secondEvent.getCurrentResponseSize()).thenReturn(30); when(firstEvent.getCommitContext()).thenReturn(firstCommitContext); @@ -85,7 +89,8 @@ public class ConsensusSubscriptionBrokerPayloadLimitTest { when(fallbackQueue.getConsensusGroupId()).thenReturn(new DataRegionId(1)); when(assignedQueue.getConsensusGroupId()).thenReturn(new DataRegionId(2)); when(fallbackQueue.getPrefetchedEventCount()).thenReturn(1); - when(fallbackQueue.poll(CONSUMER_ID, null)).thenReturn(fallbackEvent); + when(fallbackQueue.poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class))) + .thenReturn(fallbackEvent); when(fallbackEvent.getCurrentResponseSize()).thenReturn(30); bindQueues(broker, Arrays.asList(fallbackQueue, assignedQueue)); @@ -97,8 +102,8 @@ public class ConsensusSubscriptionBrokerPayloadLimitTest { assertEquals(1, events.size()); assertSame(fallbackEvent, events.get(0)); - verify(assignedQueue).poll(CONSUMER_ID, null); - verify(fallbackQueue).poll(CONSUMER_ID, null); + verify(assignedQueue).poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class)); + verify(fallbackQueue).poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class)); } @Test @@ -112,8 +117,10 @@ public class ConsensusSubscriptionBrokerPayloadLimitTest { when(fallbackQueue.getConsensusGroupId()).thenReturn(new DataRegionId(1)); when(assignedQueue.getConsensusGroupId()).thenReturn(new DataRegionId(2)); when(fallbackQueue.getPrefetchedEventCount()).thenReturn(1); - when(fallbackQueue.poll(CONSUMER_ID, null)).thenReturn(fallbackEvent); - when(assignedQueue.poll(CONSUMER_ID, null)).thenReturn(assignedEvent); + when(fallbackQueue.poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class))) + .thenReturn(fallbackEvent); + when(assignedQueue.poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class))) + .thenReturn(assignedEvent); when(assignedEvent.getCurrentResponseSize()).thenReturn(30); bindQueues(broker, Arrays.asList(fallbackQueue, assignedQueue)); @@ -123,7 +130,7 @@ public class ConsensusSubscriptionBrokerPayloadLimitTest { assertEquals(1, events.size()); assertSame(assignedEvent, events.get(0)); - verify(fallbackQueue, never()).poll(CONSUMER_ID, null); + verify(fallbackQueue, never()).poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class)); } @Test diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPollTimeoutTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPollTimeoutTest.java new file mode 100644 index 00000000000..717fddb1fbf --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPollTimeoutTest.java @@ -0,0 +1,198 @@ +/* + * 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.consensus.DataRegionId; +import org.apache.iotdb.db.subscription.agent.SubscriptionAgent; +import org.apache.iotdb.db.subscription.agent.SubscriptionReceiverAgent; +import org.apache.iotdb.db.subscription.broker.consensus.ConsensusPrefetchingQueue; +import org.apache.iotdb.db.subscription.event.SubscriptionEvent; +import org.apache.iotdb.db.subscription.receiver.SubscriptionReceiver; +import org.apache.iotdb.rpc.subscription.payload.request.PipeSubscribeRequestVersion; +import org.apache.iotdb.session.subscription.util.PollTimer; + +import org.junit.Test; + +import java.lang.reflect.Field; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class ConsensusSubscriptionBrokerPollTimeoutTest { + + private static final String CONSUMER_ID = "consumer"; + + @Test + public void testSharesRemainingBudgetAcrossRegionsAndTopics() throws Exception { + final ConsensusSubscriptionBroker broker = new ConsensusSubscriptionBroker("group"); + final ConsensusPrefetchingQueue first = newQueue(1); + final ConsensusPrefetchingQueue second = newQueue(2); + final ConsensusPrefetchingQueue third = newQueue(3); + bindQueues(broker, "first", first, second); + bindQueues(broker, "second", third); + final PollTimer timer = newManualTimer(1000L); + when(first.poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class))) + .thenAnswer( + invocation -> { + assertSame(timer, invocation.getArgument(2)); + timer.update(timer.currentTimeMs() + 400L); + return null; + }); + when(second.poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class))) + .thenAnswer( + invocation -> { + assertSame(timer, invocation.getArgument(2)); + assertEquals(600L, timer.remainingMs()); + timer.update(timer.currentTimeMs() + 400L); + return null; + }); + when(third.poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class))) + .thenAnswer( + invocation -> { + assertSame(timer, invocation.getArgument(2)); + assertEquals(200L, timer.remainingMs()); + return null; + }); + + assertTrue( + broker + .poll( + CONSUMER_ID, + new LinkedHashSet<>(Arrays.asList("first", "second")), + 100L, + Collections.emptyMap(), + timer) + .isEmpty()); + verify(third).poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class)); + } + + @Test + public void testExpirationStopsLaterRegionsAndTopicsButKeepsCollectedEvents() throws Exception { + final ConsensusSubscriptionBroker broker = new ConsensusSubscriptionBroker("group"); + final ConsensusPrefetchingQueue first = newQueue(1); + final ConsensusPrefetchingQueue second = newQueue(2); + final ConsensusPrefetchingQueue third = newQueue(3); + bindQueues(broker, "first", first, second); + bindQueues(broker, "second", third); + final PollTimer timer = newManualTimer(1000L); + final SubscriptionEvent event = mock(SubscriptionEvent.class); + when(event.getCurrentResponseSize()).thenReturn(10); + when(first.poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class))) + .thenAnswer( + invocation -> { + timer.update(timer.currentTimeMs() + 1000L); + return event; + }); + + assertEquals( + Collections.singletonList(event), + broker.poll( + CONSUMER_ID, + new LinkedHashSet<>(Arrays.asList("first", "second")), + 100L, + Collections.emptyMap(), + timer)); + verify(second, never()).poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class)); + verify(third, never()).poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class)); + } + + @Test + @SuppressWarnings("unchecked") + public void testUsesReceiverRemainingBudget() throws Exception { + final Field field = SubscriptionReceiverAgent.class.getDeclaredField("receiverThreadLocal"); + field.setAccessible(true); + final ThreadLocal<SubscriptionReceiver> receivers = + (ThreadLocal<SubscriptionReceiver>) field.get(SubscriptionAgent.receiver()); + final SubscriptionReceiver previousReceiver = receivers.get(); + final SubscriptionReceiver receiver = mock(SubscriptionReceiver.class); + when(receiver.getVersion()).thenReturn(PipeSubscribeRequestVersion.VERSION_1); + when(receiver.remainingMs()).thenReturn(2000L); + receivers.set(receiver); + try { + final ConsensusSubscriptionBroker broker = new ConsensusSubscriptionBroker("group"); + final ConsensusPrefetchingQueue queue = newQueue(1); + bindQueues(broker, "topic", queue); + when(queue.poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class))) + .thenAnswer( + invocation -> { + final PollTimer timer = invocation.getArgument(2); + assertEquals(2000L, timer.timeoutMs()); + assertTrue(timer.remainingMs() <= 2000L); + return null; + }); + + broker.poll(CONSUMER_ID, Collections.singleton("topic"), 100L); + verify(queue).poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class)); + + when(receiver.remainingMs()).thenReturn(0L); + final ConsensusPrefetchingQueue expiredQueue = newQueue(2); + bindQueues(broker, "expired", expiredQueue); + assertTrue(broker.poll(CONSUMER_ID, Collections.singleton("expired"), 100L).isEmpty()); + verify(expiredQueue, never()).poll(eq(CONSUMER_ID), isNull(), any(PollTimer.class)); + } finally { + if (previousReceiver == null) { + receivers.remove(); + } else { + receivers.set(previousReceiver); + } + } + } + + private static PollTimer newManualTimer(final long timeoutMs) { + return new PollTimer(System.currentTimeMillis(), timeoutMs) { + @Override + public void update() { + // Tests advance the shared clock explicitly without wall-clock timing assertions. + } + }; + } + + private static ConsensusPrefetchingQueue newQueue(final int regionId) { + final ConsensusPrefetchingQueue queue = mock(ConsensusPrefetchingQueue.class); + when(queue.getConsensusGroupId()).thenReturn(new DataRegionId(regionId)); + return queue; + } + + @SuppressWarnings("unchecked") + private static void bindQueues( + final ConsensusSubscriptionBroker broker, + final String topic, + final ConsensusPrefetchingQueue... queues) + throws Exception { + final Field field = + ConsensusSubscriptionBroker.class.getDeclaredField("topicNameToConsensusPrefetchingQueues"); + field.setAccessible(true); + ((Map<String, List<ConsensusPrefetchingQueue>>) field.get(broker)) + .put(topic, Arrays.asList(queues)); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueuePollTimeoutTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueuePollTimeoutTest.java new file mode 100644 index 00000000000..15d9234bb6e --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueuePollTimeoutTest.java @@ -0,0 +1,180 @@ +/* + * 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.consensus.DataRegionId; +import org.apache.iotdb.commons.subscription.config.SubscriptionConfig; +import org.apache.iotdb.consensus.iot.IoTConsensusServerImpl; +import org.apache.iotdb.consensus.iot.SubscriptionWalRetentionPolicy; +import org.apache.iotdb.consensus.iot.log.ConsensusReqReader; +import org.apache.iotdb.db.subscription.event.SubscriptionEvent; +import org.apache.iotdb.rpc.subscription.config.TopicConstant; +import org.apache.iotdb.rpc.subscription.payload.poll.RegionProgress; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponse; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponseType; +import org.apache.iotdb.session.subscription.util.PollTimer; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.lang.reflect.Field; +import java.util.Collections; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.PriorityBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class ConsensusPrefetchingQueuePollTimeoutTest { + + private ConsensusPrefetchingQueue queue; + private PriorityBlockingQueue<SubscriptionEvent> events; + + @Before + @SuppressWarnings("unchecked") + public void setUp() throws Exception { + final IoTConsensusServerImpl server = mock(IoTConsensusServerImpl.class); + when(server.getConsensusReqReader()).thenReturn(mock(ConsensusReqReader.class)); + queue = + new ConsensusPrefetchingQueue( + "group", + "topic", + TopicConstant.ORDER_MODE_MULTI_WRITER_VALUE, + new DataRegionId(1), + server, + mock(SubscriptionWalRetentionPolicy.class), + mock(ConsensusLogToTabletConverter.class), + mock(ConsensusSubscriptionCommitManager.class), + new RegionProgress(Collections.emptyMap()), + 1L, + 1L, + true); + final Field initialized = + ConsensusPrefetchingQueue.class.getDeclaredField("prefetchInitialized"); + initialized.setAccessible(true); + initialized.setBoolean(queue, true); + events = mock(PriorityBlockingQueue.class); + final Field eventQueue = ConsensusPrefetchingQueue.class.getDeclaredField("prefetchingQueue"); + eventQueue.setAccessible(true); + eventQueue.set(queue, events); + } + + @After + public void tearDown() { + queue.close(); + } + + @Test + public void testQueueWaitIsLimitedByRemainingBudget() throws Exception { + when(events.size()).thenReturn(1); + final PollTimer timer = newManualTimer(20L); + when(events.poll(anyLong(), eq(TimeUnit.MILLISECONDS))) + .thenAnswer( + invocation -> { + assertEquals( + Math.min( + 20L, SubscriptionConfig.getInstance().getSubscriptionPollMaxBlockingTimeMs()), + (long) invocation.getArgument(0)); + return null; + }); + + assertNull(queue.poll("consumer", null, timer)); + verify(events).poll(anyLong(), eq(TimeUnit.MILLISECONDS)); + } + + @Test + public void testRetryUsesDecreasingBudgetAndStopsBeforeRemovingAnotherEvent() throws Exception { + final SubscriptionEvent committed = mock(SubscriptionEvent.class); + final SubscriptionPollResponse response = mock(SubscriptionPollResponse.class); + when(response.getResponseType()).thenReturn(SubscriptionPollResponseType.TABLETS.getType()); + when(committed.getCurrentResponse()).thenReturn(response); + when(committed.getCommitContext()).thenReturn(mock(SubscriptionCommitContext.class)); + when(committed.isCommitted()).thenReturn(true); + when(events.size()).thenReturn(3); + final PollTimer timer = newManualTimer(20L); + final AtomicInteger attempts = new AtomicInteger(); + when(events.poll(anyLong(), eq(TimeUnit.MILLISECONDS))) + .thenAnswer( + invocation -> { + final int attempt = attempts.getAndIncrement(); + assertEquals( + Math.min( + attempt == 0 ? 20L : 10L, + SubscriptionConfig.getInstance().getSubscriptionPollMaxBlockingTimeMs()), + (long) invocation.getArgument(0)); + timer.update(timer.currentTimeMs() + 10L); + return committed; + }); + + assertNull(queue.poll("consumer", null, timer)); + assertEquals(2, attempts.get()); + } + + @Test + public void testExpiredBudgetDoesNotRemoveQueuedEvents() throws Exception { + when(events.size()).thenReturn(1); + assertNull(queue.poll("consumer", null, newManualTimer(0L))); + verify(events, never()).poll(anyLong(), eq(TimeUnit.MILLISECONDS)); + } + + @Test + public void testReadLockWaitIsLimitedByRemainingBudget() throws Exception { + final Field field = ConsensusPrefetchingQueue.class.getDeclaredField("lock"); + field.setAccessible(true); + final ReentrantReadWriteLock lock = (ReentrantReadWriteLock) field.get(queue); + final ExecutorService executor = Executors.newSingleThreadExecutor(); + lock.writeLock().lock(); + try { + final Future<SubscriptionEvent> result = + executor.submit( + () -> queue.poll("consumer", null, new PollTimer(System.currentTimeMillis(), 50L))); + // The writer stays locked until the poll returns; an unbounded read lock would time out here. + assertNull(result.get(5L, TimeUnit.SECONDS)); + verify(events, never()).poll(anyLong(), eq(TimeUnit.MILLISECONDS)); + } finally { + lock.writeLock().unlock(); + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5L, TimeUnit.SECONDS)); + } + } + + private static PollTimer newManualTimer(final long timeoutMs) { + return new PollTimer(System.currentTimeMillis(), timeoutMs) { + @Override + public void update() { + // Tests advance the shared clock explicitly without wall-clock timing assertions. + } + }; + } +}
