This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 3154bffd068 [Subscription] Prevent unsubscribe from blocking on
prefetch cleanup (#18742)
3154bffd068 is described below
commit 3154bffd068ee1e3d0206218add6b63ca6ecf003
Author: Caideyipi <[email protected]>
AuthorDate: Tue Sep 29 15:46:05 2026 +0800
[Subscription] Prevent unsubscribe from blocking on prefetch cleanup
(#18742)
---
.../agent/SubscriptionConsumerAgent.java | 147 +++++++++++++++------
.../broker/ConsensusSubscriptionBroker.java | 8 +-
.../consensus/ConsensusPrefetchingQueue.java | 44 +++++-
.../agent/SubscriptionConsumerAgentTest.java | 114 ++++++++++++++++
.../consensus/ConsensusPrefetchingQueueTest.java | 90 +++++++++++++
5 files changed, 352 insertions(+), 51 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionConsumerAgent.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionConsumerAgent.java
index a268f515a6a..79f0c59ec7e 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionConsumerAgent.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/agent/SubscriptionConsumerAgent.java
@@ -34,6 +34,8 @@ import java.util.LinkedHashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.locks.ReentrantLock;
public class SubscriptionConsumerAgent {
@@ -41,6 +43,9 @@ public class SubscriptionConsumerAgent {
private final ConsumerGroupMetaKeeper consumerGroupMetaKeeper;
+ private final ConcurrentHashMap<String, ConsumerGroupOperationLock>
consumerGroupOperationLocks =
+ new ConcurrentHashMap<>();
+
public SubscriptionConsumerAgent() {
this.consumerGroupMetaKeeper = new ConsumerGroupMetaKeeper();
}
@@ -67,16 +72,19 @@ public class SubscriptionConsumerAgent {
public TPushConsumerGroupMetaRespExceptionMessage
handleSingleConsumerGroupMetaChanges(
final ConsumerGroupMeta consumerGroupMetaFromCoordinator) {
- acquireWriteLock();
+ final String consumerGroupId =
consumerGroupMetaFromCoordinator.getConsumerGroupId();
try {
- if (consumerGroupMetaFromCoordinator.isEmpty()) {
-
handleDropConsumerGroupInternal(consumerGroupMetaFromCoordinator.getConsumerGroupId());
- } else {
-
handleSingleConsumerGroupMetaChangesInternal(consumerGroupMetaFromCoordinator);
- }
+ executeWithConsumerGroupLock(
+ consumerGroupId,
+ () -> {
+ if (consumerGroupMetaFromCoordinator.isEmpty()) {
+ handleDropConsumerGroupInternal(consumerGroupId);
+ } else {
+
handleSingleConsumerGroupMetaChangesInternal(consumerGroupMetaFromCoordinator);
+ }
+ });
return null;
} catch (final Exception e) {
- final String consumerGroupId =
consumerGroupMetaFromCoordinator.getConsumerGroupId();
LOGGER.warn(
DataNodePipeMessages
.PIPE_LOG_EXCEPTION_OCCURRED_WHEN_HANDLING_SINGLE_CONSUMER_GROUP_META_10E7688C,
@@ -88,16 +96,13 @@ public class SubscriptionConsumerAgent {
consumerGroupId, e);
return new TPushConsumerGroupMetaRespExceptionMessage(
consumerGroupId, exceptionMessage, System.currentTimeMillis());
- } finally {
- releaseWriteLock();
}
}
private void handleSingleConsumerGroupMetaChangesInternal(
final ConsumerGroupMeta metaFromCoordinator) {
final String consumerGroupId = metaFromCoordinator.getConsumerGroupId();
- final ConsumerGroupMeta metaInAgent =
- consumerGroupMetaKeeper.getConsumerGroupMeta(consumerGroupId);
+ final ConsumerGroupMeta metaInAgent =
getConsumerGroupMeta(consumerGroupId);
// if consumer group meta does not exist on local agent
if (Objects.isNull(metaInAgent)) {
@@ -106,7 +111,7 @@ public class SubscriptionConsumerAgent {
consumerGroupId,
metaFromCoordinator.getSubscribedTopicNames(),
metaFromCoordinator.visibleUnder(true));
- consumerGroupMetaKeeper.addConsumerGroupMeta(consumerGroupId,
metaFromCoordinator);
+ replaceConsumerGroupMeta(consumerGroupId, metaFromCoordinator);
return;
}
@@ -133,8 +138,7 @@ public class SubscriptionConsumerAgent {
consumerGroupId,
metaFromCoordinator.getSubscribedTopicNames(),
metaFromCoordinator.visibleUnder(true));
- consumerGroupMetaKeeper.removeConsumerGroupMeta(consumerGroupId);
- consumerGroupMetaKeeper.addConsumerGroupMeta(consumerGroupId,
metaFromCoordinator);
+ replaceConsumerGroupMeta(consumerGroupId, metaFromCoordinator);
// no need to create broker manually
return;
}
@@ -182,8 +186,7 @@ public class SubscriptionConsumerAgent {
() -> {
// TODO: Currently we fully replace the entire ConsumerGroupMeta
without carefully
// checking the changes in its fields.
- consumerGroupMetaKeeper.removeConsumerGroupMeta(consumerGroupId);
- consumerGroupMetaKeeper.addConsumerGroupMeta(consumerGroupId,
metaFromCoordinator);
+ replaceConsumerGroupMeta(consumerGroupId, metaFromCoordinator);
});
}
@@ -198,38 +201,35 @@ public class SubscriptionConsumerAgent {
public TPushConsumerGroupMetaRespExceptionMessage
handleConsumerGroupMetaChanges(
final List<ConsumerGroupMeta> consumerGroupMetasFromCoordinator) {
- acquireWriteLock();
- try {
- for (final ConsumerGroupMeta consumerGroupMetaFromCoordinator :
- consumerGroupMetasFromCoordinator) {
- try {
-
handleSingleConsumerGroupMetaChangesInternal(consumerGroupMetaFromCoordinator);
- } catch (final Exception e) {
- final String consumerGroupId =
consumerGroupMetaFromCoordinator.getConsumerGroupId();
- LOGGER.warn(
- DataNodePipeMessages
-
.PIPE_LOG_EXCEPTION_OCCURRED_WHEN_HANDLING_SINGLE_CONSUMER_GROUP_META_10E7688C,
- consumerGroupId,
- e);
- final String exceptionMessage =
- String.format(
- "Subscription: Failed to handle single consumer group meta
changes for consumer group %s, because %s",
- consumerGroupId, e);
- return new TPushConsumerGroupMetaRespExceptionMessage(
- consumerGroupId, exceptionMessage, System.currentTimeMillis());
- }
+ for (final ConsumerGroupMeta consumerGroupMetaFromCoordinator :
+ consumerGroupMetasFromCoordinator) {
+ final String consumerGroupId =
consumerGroupMetaFromCoordinator.getConsumerGroupId();
+ try {
+ executeWithConsumerGroupLock(
+ consumerGroupId,
+ () ->
handleSingleConsumerGroupMetaChangesInternal(consumerGroupMetaFromCoordinator));
+ } catch (final Exception e) {
+ LOGGER.warn(
+ DataNodePipeMessages
+
.PIPE_LOG_EXCEPTION_OCCURRED_WHEN_HANDLING_SINGLE_CONSUMER_GROUP_META_10E7688C,
+ consumerGroupId,
+ e);
+ final String exceptionMessage =
+ String.format(
+ "Subscription: Failed to handle single consumer group meta
changes for consumer group %s, because %s",
+ consumerGroupId, e);
+ return new TPushConsumerGroupMetaRespExceptionMessage(
+ consumerGroupId, exceptionMessage, System.currentTimeMillis());
}
- return null;
- } finally {
- releaseWriteLock();
}
+ return null;
}
public TPushConsumerGroupMetaRespExceptionMessage handleDropConsumerGroup(
final String consumerGroupId) {
- acquireWriteLock();
try {
- handleDropConsumerGroupInternal(consumerGroupId);
+ executeWithConsumerGroupLock(
+ consumerGroupId, () ->
handleDropConsumerGroupInternal(consumerGroupId));
return null;
} catch (final Exception e) {
LOGGER.warn(DataNodeMiscMessages.EXCEPTION_DROPPING_CONSUMER_GROUP,
consumerGroupId, e);
@@ -238,8 +238,6 @@ public class SubscriptionConsumerAgent {
"Subscription: Failed to drop consumer group %s, because %s",
consumerGroupId, e);
return new TPushConsumerGroupMetaRespExceptionMessage(
consumerGroupId, exceptionMessage, System.currentTimeMillis());
- } finally {
- releaseWriteLock();
}
}
@@ -258,7 +256,68 @@ public class SubscriptionConsumerAgent {
consumerGroupId);
}
- consumerGroupMetaKeeper.removeConsumerGroupMeta(consumerGroupId);
+ removeConsumerGroupMeta(consumerGroupId);
+ }
+
+ private ConsumerGroupMeta getConsumerGroupMeta(final String consumerGroupId)
{
+ acquireReadLock();
+ try {
+ return consumerGroupMetaKeeper.getConsumerGroupMeta(consumerGroupId);
+ } finally {
+ releaseReadLock();
+ }
+ }
+
+ private void replaceConsumerGroupMeta(
+ final String consumerGroupId, final ConsumerGroupMeta consumerGroupMeta)
{
+ acquireWriteLock();
+ try {
+ consumerGroupMetaKeeper.addConsumerGroupMeta(consumerGroupId,
consumerGroupMeta);
+ } finally {
+ releaseWriteLock();
+ }
+ }
+
+ private void removeConsumerGroupMeta(final String consumerGroupId) {
+ acquireWriteLock();
+ try {
+ consumerGroupMetaKeeper.removeConsumerGroupMeta(consumerGroupId);
+ } finally {
+ releaseWriteLock();
+ }
+ }
+
+ void executeWithConsumerGroupLock(final String consumerGroupId, final
Runnable operation) {
+ final ConsumerGroupOperationLock operationLock =
+ consumerGroupOperationLocks.compute(
+ consumerGroupId,
+ (ignored, existingLock) -> {
+ if (Objects.isNull(existingLock)) {
+ return new ConsumerGroupOperationLock();
+ }
+ existingLock.referenceCount++;
+ return existingLock;
+ });
+ operationLock.lock.lock();
+ try {
+ operation.run();
+ } finally {
+ operationLock.lock.unlock();
+ consumerGroupOperationLocks.computeIfPresent(
+ consumerGroupId,
+ (ignored, existingLock) -> {
+ if (existingLock != operationLock) {
+ return existingLock;
+ }
+ return --existingLock.referenceCount == 0 ? null : existingLock;
+ });
+ }
+ }
+
+ private static final class ConsumerGroupOperationLock {
+
+ private final ReentrantLock lock = new ReentrantLock(true);
+ private int referenceCount = 1;
}
public boolean isConsumerExisted(final String consumerGroupId, final String
consumerId) {
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 4df735a7b76..9765cdc4ea0 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
@@ -710,7 +710,7 @@ public class ConsensusSubscriptionBroker implements
ISubscriptionBroker {
}
public void unbindConsensusPrefetchingQueue(final String topicName) {
- closeAndRemoveConsensusPrefetchingQueues(topicName, true);
+ closeAndRemoveConsensusPrefetchingQueues(topicName, true, true);
}
@Override
@@ -816,11 +816,11 @@ public class ConsensusSubscriptionBroker implements
ISubscriptionBroker {
topicName,
brokerId);
}
- closeAndRemoveConsensusPrefetchingQueues(topicName, false);
+ closeAndRemoveConsensusPrefetchingQueues(topicName, false, false);
}
private void closeAndRemoveConsensusPrefetchingQueues(
- final String topicName, final boolean warnIfMissing) {
+ final String topicName, final boolean warnIfMissing, final boolean
removeProgressAfterClose) {
final List<ConsensusPrefetchingQueue> queuesToClose;
synchronized (queueLifecycleLock) {
final List<ConsensusPrefetchingQueue> queues =
@@ -841,7 +841,7 @@ public class ConsensusSubscriptionBroker implements
ISubscriptionBroker {
}
for (final ConsensusPrefetchingQueue q : queuesToClose) {
- q.close();
+ q.close(removeProgressAfterClose);
}
LOGGER.info(
DataNodePipeMessages
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 7f9ba6afd24..79cb048afac 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
@@ -19,6 +19,8 @@
package org.apache.iotdb.db.subscription.broker.consensus;
+import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
+import org.apache.iotdb.commons.concurrent.ThreadName;
import org.apache.iotdb.commons.consensus.ConsensusGroupId;
import org.apache.iotdb.commons.subscription.config.SubscriptionConfig;
import
org.apache.iotdb.commons.subscription.meta.consumer.SubscriptionProgressSnapshot;
@@ -84,6 +86,7 @@ import java.util.PriorityQueue;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.PriorityBlockingQueue;
import java.util.concurrent.TimeUnit;
@@ -103,6 +106,10 @@ public class ConsensusPrefetchingQueue {
private static final Logger LOGGER =
LoggerFactory.getLogger(ConsensusPrefetchingQueue.class);
+ private static final ExecutorService CLOSE_CLEANUP_EXECUTOR =
+ IoTDBThreadPoolFactory.newCachedThreadPoolWithDaemon(
+ ThreadName.SUBSCRIPTION_CONSENSUS_PREFETCH_EXECUTOR_POOL.getName() +
"-Close-Cleanup");
+
private final String consumerGroupId;
private final String topicName;
private final ConsensusGroupId consensusGroupId;
@@ -3726,6 +3733,10 @@ public class ConsensusPrefetchingQueue {
}
public void close() {
+ close(false);
+ }
+
+ public void close(final boolean removeProgressAfterClose) {
final PendingSeekRequest seekRequestToFail;
final Pair<ConsensusSubscriptionPrefetchExecutor,
ConsensusPrefetchSubtask> prefetchBinding;
@@ -3760,9 +3771,32 @@ public class ConsensusPrefetchingQueue {
if (Objects.nonNull(prefetchBinding.right)) {
prefetchBinding.right.cancelPendingExecution();
- prefetchBinding.right.awaitIdle();
+ if (prefetchBinding.right.isScheduledOrRunning()) {
+ CLOSE_CLEANUP_EXECUTOR.execute(
+ () -> finishCloseInBackground(prefetchBinding,
removeProgressAfterClose));
+ return;
+ }
+ }
+
+ finishClose(prefetchBinding, removeProgressAfterClose);
+ }
+
+ private void finishCloseInBackground(
+ final Pair<ConsensusSubscriptionPrefetchExecutor,
ConsensusPrefetchSubtask> prefetchBinding,
+ final boolean removeProgressAfterClose) {
+ try {
+ finishClose(prefetchBinding, removeProgressAfterClose);
+ } catch (final Exception e) {
+ LOGGER.warn("Failed to finish closing consensus prefetching queue {}",
this, e);
}
+ }
+ private void finishClose(
+ final Pair<ConsensusSubscriptionPrefetchExecutor,
ConsensusPrefetchSubtask> prefetchBinding,
+ final boolean removeProgressAfterClose) {
+ if (Objects.nonNull(prefetchBinding.right)) {
+ prefetchBinding.right.awaitIdle();
+ }
try {
acquireWriteLock();
try {
@@ -3790,8 +3824,12 @@ public class ConsensusPrefetchingQueue {
try {
cleanUp();
} finally {
- // Persist progress before closing
- commitManager.persistAll();
+ if (removeProgressAfterClose) {
+ commitManager.removeAllStatesForTopic(consumerGroupId, topicName);
+ } else {
+ // Persist progress before closing.
+ commitManager.persistAll();
+ }
}
closeCleanupPending = false;
} finally {
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/agent/SubscriptionConsumerAgentTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/agent/SubscriptionConsumerAgentTest.java
index 8ad28911e08..0d564ed0c72 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/agent/SubscriptionConsumerAgentTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/agent/SubscriptionConsumerAgentTest.java
@@ -25,12 +25,21 @@ import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
+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 static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
public class SubscriptionConsumerAgentTest {
+ private static final long TEST_TIMEOUT_SECONDS = 5L;
+
@Test
public void testTopicDiffSetsUpBeforeTeardownAndMetaPublication() {
final List<String> operations = new ArrayList<>();
@@ -62,4 +71,109 @@ public class SubscriptionConsumerAgentTest {
assertEquals(Collections.singletonList("setup"), operations);
}
+
+ @Test
+ public void testSameConsumerGroupMetaChangesAreSerialized() throws Exception
{
+ final SubscriptionConsumerAgent agent = new SubscriptionConsumerAgent();
+ final ExecutorService executor = Executors.newFixedThreadPool(2);
+ final CountDownLatch firstEntered = new CountDownLatch(1);
+ final CountDownLatch releaseFirst = new CountDownLatch(1);
+ final CountDownLatch secondEntered = new CountDownLatch(1);
+ try {
+ final Future<?> first =
+ executor.submit(
+ () ->
+ agent.executeWithConsumerGroupLock(
+ "group",
+ () -> {
+ firstEntered.countDown();
+ await(releaseFirst);
+ }));
+ assertTrue(firstEntered.await(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS));
+
+ final Future<?> second =
+ executor.submit(
+ () -> agent.executeWithConsumerGroupLock("group",
secondEntered::countDown));
+ assertFalse(secondEntered.await(200L, TimeUnit.MILLISECONDS));
+
+ releaseFirst.countDown();
+ first.get(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+ second.get(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+ assertEquals(0L, secondEntered.getCount());
+ } finally {
+ releaseFirst.countDown();
+ executor.shutdownNow();
+ executor.awaitTermination(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+ }
+ }
+
+ @Test
+ public void testDifferentConsumerGroupMetaChangesRunIndependently() throws
Exception {
+ final SubscriptionConsumerAgent agent = new SubscriptionConsumerAgent();
+ final ExecutorService executor = Executors.newFixedThreadPool(2);
+ final CountDownLatch firstEntered = new CountDownLatch(1);
+ final CountDownLatch releaseFirst = new CountDownLatch(1);
+ final CountDownLatch secondEntered = new CountDownLatch(1);
+ try {
+ final Future<?> first =
+ executor.submit(
+ () ->
+ agent.executeWithConsumerGroupLock(
+ "group-1",
+ () -> {
+ firstEntered.countDown();
+ await(releaseFirst);
+ }));
+ assertTrue(firstEntered.await(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS));
+
+ final Future<?> second =
+ executor.submit(
+ () -> agent.executeWithConsumerGroupLock("group-2",
secondEntered::countDown));
+ assertTrue(secondEntered.await(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS));
+
+ releaseFirst.countDown();
+ first.get(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+ second.get(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+ } finally {
+ releaseFirst.countDown();
+ executor.shutdownNow();
+ executor.awaitTermination(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+ }
+ }
+
+ @Test
+ public void testConsumerMetaReadsRemainAvailableDuringGroupOperation()
throws Exception {
+ final SubscriptionConsumerAgent agent = new SubscriptionConsumerAgent();
+ final ExecutorService executor = Executors.newSingleThreadExecutor();
+ final CountDownLatch operationEntered = new CountDownLatch(1);
+ final CountDownLatch releaseOperation = new CountDownLatch(1);
+ try {
+ final Future<?> operation =
+ executor.submit(
+ () ->
+ agent.executeWithConsumerGroupLock(
+ "group",
+ () -> {
+ operationEntered.countDown();
+ await(releaseOperation);
+ }));
+ assertTrue(operationEntered.await(TEST_TIMEOUT_SECONDS,
TimeUnit.SECONDS));
+ assertFalse(agent.isConsumerExisted("group", "consumer"));
+
+ releaseOperation.countDown();
+ operation.get(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+ } finally {
+ releaseOperation.countDown();
+ executor.shutdownNow();
+ executor.awaitTermination(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+ }
+ }
+
+ private static void await(final CountDownLatch latch) {
+ try {
+ latch.await();
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
}
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 6dabc8e2a19..69ece3840ef 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
@@ -43,6 +43,8 @@ import
org.apache.iotdb.db.subscription.agent.SubscriptionAgent;
import org.apache.iotdb.db.subscription.agent.SubscriptionConsumerAgent;
import org.apache.iotdb.db.subscription.event.SubscriptionEvent;
import org.apache.iotdb.db.subscription.resource.SubscriptionMemoryManager;
+import
org.apache.iotdb.db.subscription.task.execution.ConsensusSubscriptionPrefetchExecutor;
+import org.apache.iotdb.db.subscription.task.subtask.ConsensusPrefetchSubtask;
import org.apache.iotdb.rpc.subscription.config.TopicConstant;
import org.apache.iotdb.rpc.subscription.payload.poll.ErrorPayload;
import org.apache.iotdb.rpc.subscription.payload.poll.RegionProgress;
@@ -92,14 +94,88 @@ import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
public class ConsensusPrefetchingQueueTest {
+ private static final long TEST_TIMEOUT_SECONDS = 5L;
+
@Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+ @Test
+ public void testCloseDetachesBlockedPrefetchCleanup() throws Exception {
+ final String originalSystemDir =
IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final File systemDir = temporaryFolder.newFolder("async-prefetch-close");
+ final CountDownLatch cleanupStarted = new CountDownLatch(1);
+ final CountDownLatch releaseCleanup = new CountDownLatch(1);
+ ConsensusPrefetchingQueue queue = null;
+ try {
+ 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.getDatabaseName()).thenReturn("db");
+ when(converter.isTableModel()).thenReturn(true);
+
+ queue =
+ new ConsensusPrefetchingQueue(
+ "consumerGroup",
+ "topic",
+ TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE,
+ new DataRegionId(1),
+ serverImpl,
+ new SubscriptionWalRetentionPolicy(
+ "topic",
+ SubscriptionWalRetentionPolicy.UNBOUNDED,
+ SubscriptionWalRetentionPolicy.UNBOUNDED),
+ converter,
+ newCommitManager(systemDir),
+ new RegionProgress(Collections.emptyMap()),
+ 1L,
+ 1L,
+ true);
+
+ final ConsensusPrefetchSubtask subtask =
mock(ConsensusPrefetchSubtask.class);
+ final ConsensusSubscriptionPrefetchExecutor prefetchExecutor =
+ mock(ConsensusSubscriptionPrefetchExecutor.class);
+ when(subtask.isScheduledOrRunning()).thenReturn(true);
+ when(prefetchExecutor.isShutdown()).thenReturn(true);
+ doAnswer(
+ invocation -> {
+ cleanupStarted.countDown();
+ releaseCleanup.await(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS);
+ return null;
+ })
+ .when(subtask)
+ .awaitIdle();
+ setPrefetchBinding(queue, prefetchExecutor, subtask);
+
+ final long closeStartNanos = System.nanoTime();
+ queue.close();
+ final long closeElapsedMillis =
+ TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - closeStartNanos);
+
+ assertTrue(closeElapsedMillis < 1000L);
+ assertTrue(cleanupStarted.await(TEST_TIMEOUT_SECONDS, TimeUnit.SECONDS));
+ assertFalse(queue.isClosed());
+
+ releaseCleanup.countDown();
+ verify(subtask,
timeout(TimeUnit.SECONDS.toMillis(TEST_TIMEOUT_SECONDS))).close();
+ assertTrue(queue.isClosed());
+ } finally {
+ releaseCleanup.countDown();
+ if (queue != null && !queue.isClosed()) {
+ queue.close();
+ }
+
IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
@Test
public void testInitializationAndActivationUseIndependentMonitors() throws
Exception {
assertFalse(
@@ -2305,6 +2381,20 @@ public class ConsensusPrefetchingQueueTest {
return (ReentrantReadWriteLock) field.get(queue);
}
+ private static void setPrefetchBinding(
+ final ConsensusPrefetchingQueue queue,
+ final ConsensusSubscriptionPrefetchExecutor executor,
+ final ConsensusPrefetchSubtask subtask)
+ throws Exception {
+ final Field executorField =
+ ConsensusPrefetchingQueue.class.getDeclaredField("prefetchExecutor");
+ executorField.setAccessible(true);
+ executorField.set(queue, executor);
+ final Field subtaskField =
ConsensusPrefetchingQueue.class.getDeclaredField("prefetchSubtask");
+ subtaskField.setAccessible(true);
+ subtaskField.set(queue, subtask);
+ }
+
private static ProgressWALIterator subscriptionWalIterator(final
ConsensusPrefetchingQueue queue)
throws Exception {
final Field field =
ConsensusPrefetchingQueue.class.getDeclaredField("subscriptionWALIterator");
