From f753f308f034d0ebc5510c0bc97c9c084a9b6154 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Thu, 13 Aug 2026 18:13:28 +0800 Subject: [PATCH] Fix subscription event loss on poll payload overflow --- .../agent/SubscriptionBrokerAgent.java | 27 ++++- .../broker/ConsensusSubscriptionBroker.java | 23 +++- .../broker/ISubscriptionBroker.java | 3 + .../broker/SubscriptionBroker.java | 21 +++- .../broker/SubscriptionPrefetchingQueue.java | 33 ++++++ .../consensus/ConsensusPrefetchingQueue.java | 38 +++++++ .../subscription/event/SubscriptionEvent.java | 5 + .../receiver/SubscriptionReceiverV1.java | 14 ++- ...susSubscriptionBrokerPayloadLimitTest.java | 93 ++++++++++++++++ ...bscriptionBrokerAgentPayloadLimitTest.java | 104 ++++++++++++++++++ .../ConsensusPrefetchingQueueTest.java | 62 +++++++++++ 11 files changed, 415 insertions(+), 8 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPayloadLimitTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionBrokerAgentPayloadLimitTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionBrokerAgent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionBrokerAgent.java index da649e75b63f..83e816e9cf35 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionBrokerAgent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionBrokerAgent.java @@ -125,12 +125,23 @@ public List poll( } final List events = broker.poll(consumerId, topicNames, remainingBytes, progressByTopic); - allEvents.addAll(events); for (final SubscriptionEvent event : events) { try { - remainingBytes -= event.getCurrentResponseSize(); + final long currentSize = event.getCurrentResponseSize(); + // Each broker preserves the existing handling for its first oversized event. If another + // broker already used part of this response, put the event back so it can be retried with + // the full budget on the next poll. + if (!allEvents.isEmpty() + && currentSize > remainingBytes + && broker.requeue(consumerId, event.getCommitContext())) { + remainingBytes = 0; + break; + } + allEvents.add(event); + remainingBytes -= currentSize; } catch (final IOException ignored) { // best effort + allEvents.add(event); } } } @@ -204,6 +215,18 @@ public List commit( return allSuccessful; } + public boolean requeue( + final ConsumerConfig consumerConfig, final SubscriptionCommitContext commitContext) { + final String consumerGroupId = consumerConfig.getConsumerGroupId(); + final String consumerId = consumerConfig.getConsumerId(); + for (final ISubscriptionBroker broker : getBrokers(consumerGroupId)) { + if (broker.acceptsCommitContext(commitContext) && broker.requeue(consumerId, commitContext)) { + return true; + } + } + return false; + } + public int refreshInFlightEventLeases( final ConsumerConfig consumerConfig, final List processorBufferedCommitContexts) { 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 fb739aef3490..e1696d59b173 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 @@ -124,6 +124,7 @@ public List poll( final List eventsToPoll = new ArrayList<>(); final List eventsToNack = new ArrayList<>(); long totalSize = 0; + boolean responseFull = false; for (final String topicName : topicNames) { final List queues = @@ -169,6 +170,15 @@ public List poll( continue; } + // Preserve the existing handling for a single oversized event. Once this response already + // contains data, defer an event that does not fit instead of returning an oversized batch. + if (totalSize > 0 + && currentSize > maxBytes - totalSize + && consensusQueue.requeue(consumerId, event.getCommitContext())) { + responseFull = true; + break; + } + eventsToPoll.add(event); totalSize += currentSize; @@ -176,7 +186,7 @@ public List poll( break; } } - if (totalSize >= maxBytes) { + if (responseFull || totalSize >= maxBytes) { break; } } @@ -280,6 +290,17 @@ public List commit( return successfulCommitContexts; } + @Override + public boolean requeue(final String consumerId, final SubscriptionCommitContext commitContext) { + final List queues = + topicNameToConsensusPrefetchingQueues.get(commitContext.getTopicName()); + if (Objects.isNull(queues) || queues.isEmpty()) { + return false; + } + final ConsensusPrefetchingQueue queue = getQueueForCommitContext(queues, commitContext); + return Objects.nonNull(queue) && queue.requeue(consumerId, commitContext); + } + @Override public int refreshInFlightEventLeases( final String consumerId, final List commitContexts) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ISubscriptionBroker.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ISubscriptionBroker.java index 7d0ec6deada0..547ffc6b7633 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ISubscriptionBroker.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ISubscriptionBroker.java @@ -48,6 +48,9 @@ List pollTablets( List commit( String consumerId, List commitContexts, boolean nack); + /** Returns an in-flight event to its prefetching queue without incrementing its nack count. */ + boolean requeue(String consumerId, SubscriptionCommitContext commitContext); + default List selectAcceptedCommitContexts( final List commitContexts) { if (Objects.isNull(commitContexts) || commitContexts.isEmpty()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionBroker.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionBroker.java index 1115c465c15f..0353e3d74b1c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionBroker.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/SubscriptionBroker.java @@ -166,6 +166,14 @@ public List poll( continue; } + // Preserve the existing handling for a single oversized event. Once this response already + // contains data, defer an event that does not fit instead of returning an oversized batch. + if (totalSize > 0 + && currentSize > maxBytes - totalSize + && prefetchingQueue.requeue(consumerId, event.getCommitContext())) { + break; + } + // Add the event to the poll list eventsToPoll.add(event); @@ -175,8 +183,8 @@ public List poll( // Update the total size totalSize += currentSize; - // If adding this event exceeds the maxBytes (pessimistic estimation), break the loop - if (totalSize + currentSize > maxBytes) { + // If the response has reached maxBytes, stop polling more events. + if (totalSize >= maxBytes) { break; } } @@ -375,6 +383,15 @@ public List commit( return successfulCommitContexts; } + @Override + public boolean requeue(final String consumerId, final SubscriptionCommitContext commitContext) { + final SubscriptionPrefetchingQueue prefetchingQueue = + topicNameToPrefetchingQueue.get(commitContext.getTopicName()); + return Objects.nonNull(prefetchingQueue) + && !prefetchingQueue.isClosed() + && prefetchingQueue.requeue(consumerId, commitContext); + } + @Override public int refreshInFlightEventLeases( final String consumerId, final List commitContexts) { 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 0bfccd3a78b4..801fc03d3b8f 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 @@ -817,6 +817,39 @@ private boolean refreshInFlightEventLeaseInternal( return refreshed.get(); } + /** + * Returns an event to the prefetching queue without modifying its response or nack count. + * + *

This is used when the server polled the event but cannot fit it in the current response. + */ + public boolean requeue(final String consumerId, final SubscriptionCommitContext commitContext) { + acquireReadLock(); + try { + if (isClosed()) { + return false; + } + final AtomicBoolean requeued = new AtomicBoolean(false); + inFlightEvents.compute( + new Pair<>(consumerId, commitContext), + (key, ev) -> { + if (Objects.isNull(ev)) { + return null; + } + if (ev.isCommitted()) { + ev.cleanUp(false); + return null; + } + ev.resetLastPolledTimestamp(); + prefetchEvent(ev); + requeued.set(true); + return null; + }); + return requeued.get(); + } finally { + releaseReadLock(); + } + } + /** * @return {@code true} if ack successfully */ 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 2d2060e15f75..4dd71c3b8f29 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 @@ -2579,6 +2579,44 @@ private boolean refreshInFlightEventLeaseInternal( return refreshed.get(); } + /** + * Returns an event to the prefetching queue without modifying its response or nack count. + * + *

This is used when the server polled the event but cannot fit it in the current response. + */ + public boolean requeue(final String consumerId, final SubscriptionCommitContext commitContext) { + acquireReadLock(); + try { + if (isClosed || closeRequested || pendingSeekRequest != null || !isActive) { + return false; + } + if (Objects.isNull(commitContext) + || !commitContext.hasWriterProgress() + || isCommitContextOutdated(commitContext)) { + return false; + } + final AtomicBoolean requeued = new AtomicBoolean(false); + inFlightEvents.compute( + new InFlightEventKey(consumerId, commitContext), + (key, ev) -> { + if (Objects.isNull(ev)) { + return null; + } + if (ev.isCommitted()) { + cleanUpEvent(ev, false); + return null; + } + ev.resetLastPolledTimestamp(); + prefetchingQueue.add(ev); + requeued.set(true); + return null; + }); + return requeued.get(); + } finally { + releaseReadLock(); + } + } + private boolean canAcceptCommitContext( final SubscriptionCommitContext commitContext, final String action, final boolean silent) { if (isClosed || closeRequested || pendingSeekRequest != null) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java index 3c99a17f49fc..2e58a68411da 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java @@ -275,6 +275,11 @@ public void nack() { } } + /** Makes this event pollable again without treating local response-size control as a nack. */ + public void resetLastPolledTimestamp() { + lastPolledTimestamp.set(INVALID_TIMESTAMP); + } + /** Returns the current nack count for this event. */ public long getNackCount() { return nackCount.get(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java index dcc92bc22170..b0eec505a5c4 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java @@ -798,10 +798,18 @@ private TPipeSubscribeResp handlePipeSubscribePollInternal(final PipeSubscribePo req.getRequest(), e); } - // nack + // A response-size overflow caused by events already added to this response is + // local batching backpressure, not a consumer rejection. Requeue it without + // increasing the poison-message nack counter. if (!isOutdated) { - SubscriptionAgent.broker() - .commit(consumerConfig, Collections.singletonList(commitContext), true); + final boolean requeued = + e instanceof SubscriptionPayloadExceedException + && totalSize.get() > 0 + && SubscriptionAgent.broker().requeue(consumerConfig, commitContext); + if (!requeued) { + SubscriptionAgent.broker() + .commit(consumerConfig, Collections.singletonList(commitContext), true); + } } return null; } 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 new file mode 100644 index 000000000000..1d2e71375355 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBrokerPayloadLimitTest.java @@ -0,0 +1,93 @@ +/* + * 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.broker.consensus.ConsensusPrefetchingQueue; +import org.apache.iotdb.db.subscription.event.SubscriptionEvent; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; + +import org.junit.Test; + +import java.lang.reflect.Field; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class ConsensusSubscriptionBrokerPayloadLimitTest { + + private static final String CONSUMER_GROUP_ID = "consumerGroup"; + private static final String CONSUMER_ID = "consumer"; + private static final String TOPIC_NAME = "topic"; + + @Test + public void testPollRequeuesEventThatWouldExceedPayloadLimit() throws Exception { + final ConsensusSubscriptionBroker broker = new ConsensusSubscriptionBroker(CONSUMER_GROUP_ID); + final ConsensusPrefetchingQueue firstQueue = mock(ConsensusPrefetchingQueue.class); + final ConsensusPrefetchingQueue secondQueue = mock(ConsensusPrefetchingQueue.class); + final SubscriptionEvent firstEvent = mock(SubscriptionEvent.class); + final SubscriptionEvent secondEvent = mock(SubscriptionEvent.class); + final SubscriptionCommitContext firstCommitContext = newCommitContext(1, 1); + final SubscriptionCommitContext secondCommitContext = newCommitContext(2, 2); + + 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(firstEvent.getCurrentResponseSize()).thenReturn(40); + when(secondEvent.getCurrentResponseSize()).thenReturn(30); + when(firstEvent.getCommitContext()).thenReturn(firstCommitContext); + when(secondEvent.getCommitContext()).thenReturn(secondCommitContext); + when(secondQueue.requeue(CONSUMER_ID, secondCommitContext)).thenReturn(true); + bindQueues(broker, Arrays.asList(firstQueue, secondQueue)); + + final List events = + broker.poll(CONSUMER_ID, Collections.singleton(TOPIC_NAME), 60L); + + assertEquals(1, events.size()); + assertSame(firstEvent, events.get(0)); + verify(secondQueue).requeue(CONSUMER_ID, secondCommitContext); + } + + private static SubscriptionCommitContext newCommitContext( + final int regionId, final int commitId) { + return new SubscriptionCommitContext( + 1, 1, TOPIC_NAME, CONSUMER_GROUP_ID, commitId, "DataRegion[" + regionId + "]", 0L); + } + + @SuppressWarnings("unchecked") + private static void bindQueues( + final ConsensusSubscriptionBroker broker, final List queues) + throws Exception { + final Field field = + ConsensusSubscriptionBroker.class.getDeclaredField("topicNameToConsensusPrefetchingQueues"); + field.setAccessible(true); + final Map> queuesByTopic = + (Map>) field.get(broker); + queuesByTopic.put(TOPIC_NAME, queues); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionBrokerAgentPayloadLimitTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionBrokerAgentPayloadLimitTest.java new file mode 100644 index 000000000000..782f2338e9d2 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/SubscriptionBrokerAgentPayloadLimitTest.java @@ -0,0 +1,104 @@ +/* + * 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.db.subscription.agent.SubscriptionBrokerAgent; +import org.apache.iotdb.db.subscription.event.SubscriptionEvent; +import org.apache.iotdb.rpc.subscription.config.ConsumerConfig; +import org.apache.iotdb.rpc.subscription.config.ConsumerConstant; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponseType; +import org.apache.iotdb.rpc.subscription.payload.poll.TerminationPayload; + +import org.junit.Test; + +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertSame; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class SubscriptionBrokerAgentPayloadLimitTest { + + private static final String CONSUMER_GROUP_ID = "consumerGroup"; + private static final String CONSUMER_ID = "consumer"; + private static final String TOPIC_NAME = "topic"; + + @Test + public void testPollRequeuesFirstEventFromNextBrokerWhenRemainingBudgetIsInsufficient() + throws Exception { + final SubscriptionBrokerAgent agent = new SubscriptionBrokerAgent(); + final ISubscriptionBroker firstBroker = mock(ISubscriptionBroker.class); + final ISubscriptionBroker secondBroker = mock(ISubscriptionBroker.class); + final SubscriptionEvent firstEvent = newEvent(1); + final SubscriptionEvent secondEvent = newEvent(2); + final long firstEventSize = firstEvent.getCurrentResponseSize(); + final long secondEventSize = secondEvent.getCurrentResponseSize(); + final long maxBytes = firstEventSize + secondEventSize - 1L; + final Set topicNames = Collections.singleton(TOPIC_NAME); + + when(firstBroker.poll(CONSUMER_ID, topicNames, maxBytes, Collections.emptyMap())) + .thenReturn(Collections.singletonList(firstEvent)); + when(secondBroker.poll( + CONSUMER_ID, topicNames, maxBytes - firstEventSize, Collections.emptyMap())) + .thenReturn(Collections.singletonList(secondEvent)); + when(secondBroker.requeue(CONSUMER_ID, secondEvent.getCommitContext())).thenReturn(true); + bindBrokers(agent, firstBroker, secondBroker); + + final List events = agent.poll(createConsumerConfig(), topicNames, maxBytes); + + assertEquals(1, events.size()); + assertSame(firstEvent, events.get(0)); + verify(secondBroker).requeue(CONSUMER_ID, secondEvent.getCommitContext()); + assertEquals(0L, secondEvent.getNackCount()); + } + + private static SubscriptionEvent newEvent(final int commitId) { + return new SubscriptionEvent( + SubscriptionPollResponseType.TERMINATION.getType(), + new TerminationPayload(), + new SubscriptionCommitContext(1, 1, TOPIC_NAME, CONSUMER_GROUP_ID, commitId)); + } + + private static ConsumerConfig createConsumerConfig() { + final Map attributes = new HashMap<>(); + attributes.put(ConsumerConstant.CONSUMER_ID_KEY, CONSUMER_ID); + attributes.put(ConsumerConstant.CONSUMER_GROUP_ID_KEY, CONSUMER_GROUP_ID); + return new ConsumerConfig(attributes); + } + + @SuppressWarnings("unchecked") + private static void bindBrokers( + final SubscriptionBrokerAgent agent, final ISubscriptionBroker... brokers) throws Exception { + final Field field = SubscriptionBrokerAgent.class.getDeclaredField("consumerGroupIdToBrokers"); + field.setAccessible(true); + final Map> brokersByConsumerGroup = + (Map>) field.get(agent); + brokersByConsumerGroup.put(CONSUMER_GROUP_ID, new ArrayList<>(List.of(brokers))); + } +} 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 b1fd02cc4c51..5b3fd00b928a 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 @@ -1705,6 +1705,68 @@ public void testLateAckDoesNotStealRecycledEventFromNewConsumer() throws Excepti } } + @Test + public void testRequeueDoesNotIncrementNackCount() throws Exception { + final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir(); + final File systemDir = temporaryFolder.newFolder("system-requeue-without-nack"); + ConsensusPrefetchingQueue queue = null; + try { + final DataRegionId regionId = new DataRegionId(1); + final FakeConsensusReqReader reader = new FakeConsensusReqReader(); + final IoTConsensusServerImpl serverImpl = mock(IoTConsensusServerImpl.class); + when(serverImpl.getConsensusReqReader()).thenReturn(reader); + when(serverImpl.getWriterSafeFrontierTracker()).thenReturn(new WriterSafeFrontierTracker()); + + final ConsensusLogToTabletConverter converter = mock(ConsensusLogToTabletConverter.class); + when(converter.convert(any())) + .thenReturn(Collections.singletonList(createTablet()), Collections.emptyList()); + when(converter.getDatabaseName()).thenReturn("db"); + + queue = + new ConsensusPrefetchingQueue( + "consumerGroup", + "topic", + TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE, + regionId, + serverImpl, + new SubscriptionWalRetentionPolicy( + "topic", + SubscriptionWalRetentionPolicy.UNBOUNDED, + SubscriptionWalRetentionPolicy.UNBOUNDED), + converter, + newCommitManager(systemDir), + new RegionProgress(Collections.emptyMap()), + 1L, + 1L, + true); + + reader.currentSearchIndex = 2L; + assertTrue(pendingEntries(queue).offer(createRequest(1L))); + assertTrue(pendingEntries(queue).offer(createRequest(2L))); + assertNull(queue.poll("consumer")); + queue.drivePrefetchOnce(); + + final SubscriptionEvent event = queue.poll("consumer"); + assertNotNull(event); + assertEquals(1L, queue.getSubscriptionUncommittedEventCount()); + + assertTrue(queue.requeue("consumer", event.getCommitContext())); + assertEquals(0L, event.getNackCount()); + assertEquals(0L, queue.getSubscriptionUncommittedEventCount()); + assertEquals(1, queue.getPrefetchedEventCount()); + + final SubscriptionEvent redeliveredEvent = queue.poll("consumer"); + assertSame(event, redeliveredEvent); + assertEquals(0L, redeliveredEvent.getNackCount()); + assertTrue(queue.ack("consumer", redeliveredEvent.getCommitContext())); + } finally { + if (queue != null) { + queue.close(); + } + IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir); + } + } + @Test public void testDeactivationReleasesMaterializedTabletMemory() throws Exception { final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();