You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@iotdb.apache.org by ju...@apache.org on 2026/08/25 03:09:25 UTC
(iotdb) 01/01: [Subscription] Bound consensus prefetch memory usage (#18298)
This is an automated email from the ASF dual-hosted git repository.
Caideyipi pushed a commit to branch rc-to-20111
in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 4baf19f21a2579583e83f82da568454206a07240
Author: Caideyipi <87...@users.noreply.github.com>
AuthorDate: Tue Jul 28 18:35:32 2026 +0800
[Subscription] Bound consensus prefetch memory usage (#18298)
* [Subscription] Bound consensus prefetch memory usage
* [Subscription] Accept fenced consensus commits as no-ops
* [Subscription] Address review comments
(cherry picked from commit 44b2cf5676aa21f9aac8a6c7d28773f03659db72)
---
.../apache/iotdb/db/i18n/DataNodeMiscMessages.java | 9 +
.../apache/iotdb/db/i18n/DataNodeMiscMessages.java | 9 +
.../apache/iotdb/db/conf/DataNodeMemoryConfig.java | 137 ++--
.../metrics/memory/QueryEngineMemoryMetrics.java | 26 +-
.../broker/ConsensusSubscriptionBroker.java | 21 +-
.../consensus/ConsensusPrefetchingQueue.java | 737 ++++++++++++++++-----
.../SubscriptionDataNodeResourceManager.java | 6 +
.../resource/SubscriptionMemoryManager.java | 89 +++
.../iotdb/db/conf/DataNodeMemoryConfigTest.java | 69 ++
...onsensusPrefetchingQueueDataNodeMemoryTest.java | 346 ++++++++++
.../consensus/ConsensusPrefetchingQueueTest.java | 697 ++++++++++++++++++-
...nsensusPrefetchingQueueWalBackpressureTest.java | 336 ++++++++++
.../resource/SubscriptionMemoryManagerTest.java | 64 ++
.../conf/iotdb-system.properties.template | 8 +-
14 files changed, 2343 insertions(+), 211 deletions(-)
diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
index b5011edc8bd..a74d9cc7676 100644
--- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
+++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
@@ -1392,6 +1392,15 @@ public final class DataNodeMiscMessages {
public static final String MISC_EXCEPTION_EACH_SUBSECTION_OF_CONFIGURATION_ITEM_UDF_READER_TRANSFORMER_97CA8962 =
"Each subsection of configuration item udf_reader_transformer_collector_memory_proportion "
+ "should be an integer, which is %s";
+ public static final String
+ EXCEPTION_QUERY_MEMORY_PROPORTIONS_MUST_CONTAIN_8_OR_9_COLON_SEPARATED_VALUES_BUT_FOUND_ARG_03A03941 =
+ "Query memory proportions must contain 8 or 9 colon-separated values, but found %d";
+ public static final String
+ EXCEPTION_QUERY_MEMORY_PROPORTION_AT_POSITION_ARG_MUST_BE_NON_NEGATIVE_BUT_FOUND_ARG_DC69BC75 =
+ "Query memory proportion at position %d must be non-negative, but found %d";
+ public static final String
+ EXCEPTION_THE_SUM_OF_QUERY_MEMORY_PROPORTIONS_MUST_BE_POSITIVE_BUT_WAS_ARG_407092B6 =
+ "The sum of query memory proportions must be positive, but was %d";
public static final String MISC_EXCEPTION_EACH_SUBSECTION_OF_CONFIGURATION_ITEM_CHUNKMETA_CHUNK_TIMESERIESMETA_77A43CE2 =
"Each subsection of configuration item chunkmeta_chunk_timeseriesmeta_free_memory_proportion "
+ "should be an integer, which is %s";
diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
index 8983fd61dc4..31ab68978df 100644
--- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
+++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
@@ -1362,6 +1362,15 @@ public final class DataNodeMiscMessages {
"无法重新加载配置文件 %s，原因：%s";
public static final String MISC_EXCEPTION_EACH_SUBSECTION_OF_CONFIGURATION_ITEM_UDF_READER_TRANSFORMER_97CA8962 =
"配置项 udf_reader_transformer_collector_memory_proportion 的每个子项都应为整数，当前为 %s";
+ public static final String
+ EXCEPTION_QUERY_MEMORY_PROPORTIONS_MUST_CONTAIN_8_OR_9_COLON_SEPARATED_VALUES_BUT_FOUND_ARG_03A03941 =
+ "查询内存比例必须包含 8 个或 9 个以冒号分隔的值，但实际有 %d 个";
+ public static final String
+ EXCEPTION_QUERY_MEMORY_PROPORTION_AT_POSITION_ARG_MUST_BE_NON_NEGATIVE_BUT_FOUND_ARG_DC69BC75 =
+ "查询内存比例中第 %d 个值必须为非负数，但实际为 %d";
+ public static final String
+ EXCEPTION_THE_SUM_OF_QUERY_MEMORY_PROPORTIONS_MUST_BE_POSITIVE_BUT_WAS_ARG_407092B6 =
+ "查询内存比例之和必须为正数，但实际为 %d";
public static final String MISC_EXCEPTION_EACH_SUBSECTION_OF_CONFIGURATION_ITEM_CHUNKMETA_CHUNK_TIMESERIESMETA_77A43CE2 =
"配置项 chunkmeta_chunk_timeseriesmeta_free_memory_proportion 的每个子项都应为整数，当前为 %s";
public static final String MISC_EXCEPTION_ILLEGAL_DEFAULTDATABASELEVEL_D_SHOULD_1_03088B38 =
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/DataNodeMemoryConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/DataNodeMemoryConfig.java
index 9476e84e15e..350f5d78019 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/DataNodeMemoryConfig.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/DataNodeMemoryConfig.java
@@ -33,6 +33,13 @@ import org.slf4j.LoggerFactory;
public class DataNodeMemoryConfig {
private static final Logger LOGGER = LoggerFactory.getLogger(DataNodeMemoryConfig.class);
+ private static final int LEGACY_QUERY_MEMORY_COMPONENT_COUNT = 8;
+ private static final int QUERY_MEMORY_COMPONENT_COUNT = 9;
+ private static final int SUBSCRIPTION_MEMORY_INDEX = 8;
+ private static final int[] DEFAULT_QUERY_MEMORY_PROPORTIONS = {
+ 1, 100, 200, 50, 200, 200, 200, 50, 250
+ };
+
public static final String SCHEMA_CACHE = "SchemaCache";
public static final String SCHEMA_REGION = "SchemaRegion";
public static final String PARTITION_CACHE = "PartitionCache";
@@ -128,6 +135,9 @@ public class DataNodeMemoryConfig {
/** Memory manager proportion for timeIndex */
private MemoryManager timeIndexMemoryManager;
+ /** Memory manager for subscription */
+ private MemoryManager subscriptionMemoryManager;
+
/** Memory manager for the mtree */
private MemoryManager schemaEngineMemoryManager;
@@ -429,6 +439,55 @@ public class DataNodeMemoryConfig {
"Compaction", compactionMemorySize - fixedMemoryCost);
}
+ static int[] resolveQueryMemoryProportions(
+ final String configuredProportions, final boolean subscriptionEnabled) {
+ final int[] resolvedProportions = DEFAULT_QUERY_MEMORY_PROPORTIONS.clone();
+ if (configuredProportions != null) {
+ final String[] proportions = configuredProportions.split(":");
+ if (proportions.length != LEGACY_QUERY_MEMORY_COMPONENT_COUNT
+ && proportions.length != QUERY_MEMORY_COMPONENT_COUNT) {
+ throw new IllegalArgumentException(
+ String.format(
+ DataNodeMiscMessages
+ .EXCEPTION_QUERY_MEMORY_PROPORTIONS_MUST_CONTAIN_8_OR_9_COLON_SEPARATED_VALUES_BUT_FOUND_ARG_03A03941,
+ proportions.length));
+ }
+ for (int i = 0; i < proportions.length; i++) {
+ resolvedProportions[i] = Integer.parseInt(proportions[i].trim());
+ if (resolvedProportions[i] < 0) {
+ throw new IllegalArgumentException(
+ String.format(
+ DataNodeMiscMessages
+ .EXCEPTION_QUERY_MEMORY_PROPORTION_AT_POSITION_ARG_MUST_BE_NON_NEGATIVE_BUT_FOUND_ARG_DC69BC75,
+ i + 1,
+ resolvedProportions[i]));
+ }
+ }
+ if (proportions.length == LEGACY_QUERY_MEMORY_COMPONENT_COUNT) {
+ int legacyProportionSum = 0;
+ for (int i = 0; i < LEGACY_QUERY_MEMORY_COMPONENT_COUNT; i++) {
+ legacyProportionSum += resolvedProportions[i];
+ }
+ resolvedProportions[SUBSCRIPTION_MEMORY_INDEX] = legacyProportionSum / 4;
+ }
+ }
+ if (!subscriptionEnabled) {
+ resolvedProportions[SUBSCRIPTION_MEMORY_INDEX] = 0;
+ }
+ int proportionSum = 0;
+ for (final int proportion : resolvedProportions) {
+ proportionSum += proportion;
+ }
+ if (proportionSum <= 0) {
+ throw new IllegalArgumentException(
+ String.format(
+ DataNodeMiscMessages
+ .EXCEPTION_THE_SUM_OF_QUERY_MEMORY_PROPORTIONS_MUST_BE_POSITIVE_BUT_WAS_ARG_407092B6,
+ proportionSum));
+ }
+ return resolvedProportions;
+ }
+
@SuppressWarnings("squid:S3518")
private void initQueryEngineMemoryAllocate(
MemoryManager queryEngineMemoryManager, TrimProperties properties) {
@@ -459,49 +518,40 @@ public class DataNodeMemoryConfig {
LOGGER.error(String.format(DataNodeMiscMessages.FAIL_RELOAD_CONFIGURATION_FMT, e));
}
+ long maxMemoryAvailable = queryEngineMemoryManager.getTotalMemorySizeInBytes();
String queryMemoryAllocateProportion =
properties.getProperty("chunk_timeseriesmeta_free_memory_proportion");
- long maxMemoryAvailable = queryEngineMemoryManager.getTotalMemorySizeInBytes();
-
- long bloomFilterCacheMemorySize = maxMemoryAvailable / 1001;
- long chunkCacheMemorySize = maxMemoryAvailable * 100 / 1001;
- long timeSeriesMetaDataCacheMemorySize = maxMemoryAvailable * 200 / 1001;
- long coordinatorMemorySize = maxMemoryAvailable * 50 / 1001;
- long operatorsMemorySize = maxMemoryAvailable * 200 / 1001;
- long dataExchangeMemorySize = maxMemoryAvailable * 200 / 1001;
- long timeIndexMemorySize = maxMemoryAvailable * 200 / 1001;
- if (queryMemoryAllocateProportion != null) {
- String[] proportions = queryMemoryAllocateProportion.split(":");
- int proportionSum = 0;
- for (String proportion : proportions) {
- proportionSum += Integer.parseInt(proportion.trim());
- }
- if (proportionSum != 0) {
- try {
- bloomFilterCacheMemorySize =
- maxMemoryAvailable * Integer.parseInt(proportions[0].trim()) / proportionSum;
- chunkCacheMemorySize =
- maxMemoryAvailable * Integer.parseInt(proportions[1].trim()) / proportionSum;
- timeSeriesMetaDataCacheMemorySize =
- maxMemoryAvailable * Integer.parseInt(proportions[2].trim()) / proportionSum;
- coordinatorMemorySize =
- maxMemoryAvailable * Integer.parseInt(proportions[3].trim()) / proportionSum;
- operatorsMemorySize =
- maxMemoryAvailable * Integer.parseInt(proportions[4].trim()) / proportionSum;
- dataExchangeMemorySize =
- maxMemoryAvailable * Integer.parseInt(proportions[5].trim()) / proportionSum;
- timeIndexMemorySize =
- maxMemoryAvailable * Integer.parseInt(proportions[6].trim()) / proportionSum;
- } catch (Exception e) {
- throw new IllegalArgumentException(
- String.format(
- DataNodeMiscMessages
- .MISC_EXCEPTION_EACH_SUBSECTION_OF_CONFIGURATION_ITEM_CHUNKMETA_CHUNK_TIMESERIESMETA_77A43CE2,
- queryMemoryAllocateProportion),
- e);
- }
- }
+ boolean subscriptionEnabled =
+ Boolean.parseBoolean(
+ properties.getProperty("subscription_enabled", Boolean.TRUE.toString()));
+ final int[] queryMemoryProportions;
+ try {
+ queryMemoryProportions =
+ resolveQueryMemoryProportions(queryMemoryAllocateProportion, subscriptionEnabled);
+ } catch (Exception e) {
+ throw new IllegalArgumentException(
+ String.format(
+ DataNodeMiscMessages
+ .MISC_EXCEPTION_EACH_SUBSECTION_OF_CONFIGURATION_ITEM_CHUNKMETA_CHUNK_TIMESERIESMETA_77A43CE2,
+ queryMemoryAllocateProportion),
+ e);
}
+ int proportionSum = 0;
+ for (int proportion : queryMemoryProportions) {
+ proportionSum += proportion;
+ }
+
+ long bloomFilterCacheMemorySize =
+ maxMemoryAvailable * queryMemoryProportions[0] / proportionSum;
+ long chunkCacheMemorySize = maxMemoryAvailable * queryMemoryProportions[1] / proportionSum;
+ long timeSeriesMetaDataCacheMemorySize =
+ maxMemoryAvailable * queryMemoryProportions[2] / proportionSum;
+ long coordinatorMemorySize = maxMemoryAvailable * queryMemoryProportions[3] / proportionSum;
+ long operatorsMemorySize = maxMemoryAvailable * queryMemoryProportions[4] / proportionSum;
+ long dataExchangeMemorySize = maxMemoryAvailable * queryMemoryProportions[5] / proportionSum;
+ long timeIndexMemorySize = maxMemoryAvailable * queryMemoryProportions[6] / proportionSum;
+ long subscriptionMemorySize =
+ maxMemoryAvailable * queryMemoryProportions[SUBSCRIPTION_MEMORY_INDEX] / proportionSum;
// metadata cache is disabled, we need to move all their allocated memory to other parts
if (!isMetaDataCacheEnable()) {
@@ -532,6 +582,9 @@ public class DataNodeMemoryConfig {
queryEngineMemoryManager.getOrCreateMemoryManager("DataExchange", dataExchangeMemorySize);
timeIndexMemoryManager =
queryEngineMemoryManager.getOrCreateMemoryManager("TimeIndex", timeIndexMemorySize);
+ subscriptionMemoryManager =
+ queryEngineMemoryManager.getOrCreateMemoryManager(
+ "Subscription", subscriptionMemorySize, true);
// must be called after dataExchangeMemoryManager being inited.
setQueryThreadCount(
@@ -695,6 +748,10 @@ public class DataNodeMemoryConfig {
return timeIndexMemoryManager;
}
+ public MemoryManager getSubscriptionMemoryManager() {
+ return subscriptionMemoryManager;
+ }
+
public MemoryManager getSchemaEngineMemoryManager() {
return schemaEngineMemoryManager;
}
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/memory/QueryEngineMemoryMetrics.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/memory/QueryEngineMemoryMetrics.java
index 7138f0daad0..d336c86bc91 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/memory/QueryEngineMemoryMetrics.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/service/metrics/memory/QueryEngineMemoryMetrics.java
@@ -43,6 +43,7 @@ public class QueryEngineMemoryMetrics implements IMetricSet {
private static final String QUERY_ENGINE_DATA_EXCHANGE = "QueryEngine-DataExchange";
private static final String QUERY_ENGINE_TIME_INDEX = "QueryEngine-TimeIndex";
private static final String QUERY_ENGINE_COORDINATOR = "QueryEngine-Coordinator";
+ private static final String QUERY_ENGINE_SUBSCRIPTION = "QueryEngine-Subscription";
@Override
public void bindTo(AbstractMetricService metricService) {
@@ -211,6 +212,28 @@ public class QueryEngineMemoryMetrics implements IMetricSet {
GlobalMemoryMetrics.ON_HEAP,
Tag.LEVEL.toString(),
GlobalMemoryMetrics.LEVELS[2]);
+ metricService.createAutoGauge(
+ Metric.MEMORY_THRESHOLD_SIZE.toString(),
+ MetricLevel.IMPORTANT,
+ memoryConfig.getSubscriptionMemoryManager(),
+ MemoryManager::getTotalMemorySizeInBytes,
+ Tag.NAME.toString(),
+ QUERY_ENGINE_SUBSCRIPTION,
+ Tag.TYPE.toString(),
+ GlobalMemoryMetrics.ON_HEAP,
+ Tag.LEVEL.toString(),
+ GlobalMemoryMetrics.LEVELS[2]);
+ metricService.createAutoGauge(
+ Metric.MEMORY_ACTUAL_SIZE.toString(),
+ MetricLevel.IMPORTANT,
+ memoryConfig.getSubscriptionMemoryManager(),
+ MemoryManager::getUsedMemorySizeInBytes,
+ Tag.NAME.toString(),
+ QUERY_ENGINE_SUBSCRIPTION,
+ Tag.TYPE.toString(),
+ GlobalMemoryMetrics.ON_HEAP,
+ Tag.LEVEL.toString(),
+ GlobalMemoryMetrics.LEVELS[2]);
}
@Override
@@ -231,7 +254,8 @@ public class QueryEngineMemoryMetrics implements IMetricSet {
QUERY_ENGINE_OPERATORS,
QUERY_ENGINE_DATA_EXCHANGE,
QUERY_ENGINE_TIME_INDEX,
- QUERY_ENGINE_COORDINATOR)
+ QUERY_ENGINE_COORDINATOR,
+ QUERY_ENGINE_SUBSCRIPTION)
.forEach(
name -> {
metricService.remove(
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 e089a8426d3..98dda3709f3 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
@@ -249,15 +249,22 @@ public class ConsensusSubscriptionBroker implements ISubscriptionBroker {
getQueueForCommitContext(queues, commitContext);
boolean handled = false;
if (Objects.nonNull(assignedQueue)) {
- final boolean success;
- if (!nack) {
- success = assignedQueue.ackSilent(consumerId, commitContext);
- } else {
- success = assignedQueue.nackSilent(consumerId, commitContext);
- }
- if (success) {
+ if (assignedQueue.isCommitContextOutdated(commitContext)) {
+ // A seek or reboot has already fenced this context. Accept the delayed commit as a no-op
+ // so the client does not retry it, while keeping the current commit frontier unchanged.
successfulCommitContexts.add(commitContext);
handled = true;
+ } else {
+ final boolean success;
+ if (!nack) {
+ success = assignedQueue.ackSilent(consumerId, commitContext);
+ } else {
+ success = assignedQueue.nackSilent(consumerId, commitContext);
+ }
+ if (success) {
+ successfulCommitContexts.add(commitContext);
+ handled = true;
+ }
}
}
if (!handled) {
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 cc03f1941bc..e00a3c5b21d 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
@@ -21,6 +21,7 @@ package org.apache.iotdb.db.subscription.broker.consensus;
import org.apache.iotdb.commons.consensus.ConsensusGroupId;
import org.apache.iotdb.commons.subscription.config.SubscriptionConfig;
+import org.apache.iotdb.commons.utils.TestOnly;
import org.apache.iotdb.consensus.common.request.IndexedConsensusRequest;
import org.apache.iotdb.consensus.iot.IoTConsensusServerImpl;
import org.apache.iotdb.consensus.iot.SubscriptionWalRetentionPolicy;
@@ -44,6 +45,8 @@ import org.apache.iotdb.db.subscription.agent.SubscriptionAgent;
import org.apache.iotdb.db.subscription.columnfilter.ColumnFilterMatcher;
import org.apache.iotdb.db.subscription.event.SubscriptionEvent;
import org.apache.iotdb.db.subscription.metric.ConsensusSubscriptionPrefetchingQueueMetrics;
+import org.apache.iotdb.db.subscription.resource.SubscriptionDataNodeResourceManager;
+import org.apache.iotdb.db.subscription.resource.SubscriptionMemoryManager;
import org.apache.iotdb.db.subscription.task.execution.ConsensusSubscriptionPrefetchExecutor;
import org.apache.iotdb.db.subscription.task.execution.ConsensusSubscriptionPrefetchExecutorManager;
import org.apache.iotdb.db.subscription.task.subtask.ConsensusPrefetchSubtask;
@@ -86,6 +89,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.function.BiConsumer;
+import java.util.function.BooleanSupplier;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.function.Supplier;
@@ -114,6 +118,8 @@ public class ConsensusPrefetchingQueue {
private final ConsensusSubscriptionCommitManager commitManager;
+ private SubscriptionMemoryManager subscriptionMemoryManager;
+
/**
* Incremented on each seek to distinguish batches/events created before and after the seek. The
* value is copied into commit contexts so stale events cannot be committed after a later seek.
@@ -127,6 +133,26 @@ public class ConsensusPrefetchingQueue {
private final Map<InFlightEventKey, SubscriptionEvent> inFlightEvents;
+ /** Tablet memory retained by queued and in-flight events, keyed by event identity. */
+ private final Map<SubscriptionEvent, Long> retainedBytesByEvent = new ConcurrentHashMap<>();
+
+ /** Materialized data events indexed across both queued and in-flight lifecycle stages. */
+ private final Map<SubscriptionCommitContext, SubscriptionEvent> outstandingEventsByCommitContext =
+ new ConcurrentHashMap<>();
+
+ /** Tablet memory retained by this queue across all materialized lifecycle stages. */
+ private final AtomicLong retainedTabletBytes = new AtomicLong(0L);
+
+ /**
+ * Size of the first request that could not be materialized because of the shared memory limit.
+ * This prevents repeatedly converting the same request while some memory is free but still less
+ * than the request requires.
+ */
+ private volatile long memoryBlockedEntryBytes = -1L;
+
+ /** Rejects the real-time fast path while this queue cannot safely materialize more data. */
+ private final AtomicBoolean realtimeAdmissionBlocked = new AtomicBoolean(false);
+
private static final int MAX_PREFETCHING_QUEUE_SIZE =
SubscriptionConfig.getInstance().getSubscriptionConsensusPrefetchingQueueCapacity();
@@ -141,6 +167,11 @@ public class ConsensusPrefetchingQueue {
private volatile boolean closeRequested = false;
+ /**
+ * Allows a later close call to finish cleanup after a fatal prefetch failure marked the queue.
+ */
+ private volatile boolean closeCleanupPending = false;
+
private volatile boolean isActive = true;
private volatile Set<Integer> activeWriterNodeIds = Collections.emptySet();
@@ -198,8 +229,14 @@ public class ConsensusPrefetchingQueue {
*/
private volatile boolean prefetchInitialized = false;
+ /** Serializes lazy initialization without participating in queue lifecycle lock ordering. */
+ private final Object prefetchInitializationLock = new Object();
+
private volatile PendingSeekRequest pendingSeekRequest;
+ /** Serializes lifecycle transitions independently from lazy prefetch initialization. */
+ private final Object runtimeActivationLock = new Object();
+
private final DeliveryBatchState lingerBatch = new DeliveryBatchState();
private volatile long observedSeekGeneration;
@@ -264,6 +301,19 @@ public class ConsensusPrefetchingQueue {
LOCATE_MISS
}
+ private enum MaterializationResult {
+ SUCCESS,
+ MEMORY_BLOCKED,
+ WAL_GAP,
+ STALE
+ }
+
+ private enum InFlightAckResult {
+ ACKED,
+ REJECTED,
+ MISSING
+ }
+
protected static final class ReplayLocateDecision {
private final ReplayLocateStatus status;
private final long startSearchIndex;
@@ -324,15 +374,24 @@ public class ConsensusPrefetchingQueue {
extends LinkedBlockingDeque<IndexedConsensusRequest> {
private final Runnable wakeupHook;
+ private final BooleanSupplier admissionSupplier;
- private WakeableIndexedConsensusQueue(final int capacity, final Runnable wakeupHook) {
+ private WakeableIndexedConsensusQueue(
+ final int capacity, final Runnable wakeupHook, final BooleanSupplier admissionSupplier) {
super(capacity);
this.wakeupHook = wakeupHook;
+ this.admissionSupplier = admissionSupplier;
}
@Override
public boolean offer(final IndexedConsensusRequest request) {
- final boolean offered = super.offer(request);
+ final boolean offered;
+ synchronized (this) {
+ if (!admissionSupplier.getAsBoolean()) {
+ return false;
+ }
+ offered = super.offer(request);
+ }
if (offered) {
wakeupHook.run();
}
@@ -340,9 +399,8 @@ public class ConsensusPrefetchingQueue {
}
@Override
- public void put(final IndexedConsensusRequest request) throws InterruptedException {
- super.put(request);
- wakeupHook.run();
+ public synchronized void clear() {
+ super.clear();
}
}
@@ -423,6 +481,7 @@ public class ConsensusPrefetchingQueue {
this.retentionPolicy = retentionPolicy;
this.converter = converter;
this.commitManager = commitManager;
+ this.subscriptionMemoryManager = SubscriptionDataNodeResourceManager.memory();
this.fallbackCommittedRegionProgress = fallbackCommittedRegionProgress;
this.fallbackTailSearchIndex = tailStartSearchIndex;
this.runtimeVersion = initialRuntimeVersion;
@@ -438,7 +497,8 @@ public class ConsensusPrefetchingQueue {
// Register pending queue early so we don't miss real-time writes
this.pendingEntries =
- new WakeableIndexedConsensusQueue(PENDING_QUEUE_CAPACITY, this::requestPrefetch);
+ new WakeableIndexedConsensusQueue(
+ PENDING_QUEUE_CAPACITY, this::requestPrefetch, this::canAcceptRealtimeEntry);
serverImpl.registerSubscriptionQueue(pendingEntries, retentionPolicy);
LOGGER.info(
@@ -457,6 +517,16 @@ public class ConsensusPrefetchingQueue {
ConsensusSubscriptionPrefetchingQueueMetrics.getInstance().register(this);
}
+ @TestOnly
+ void setSubscriptionMemoryManager(final SubscriptionMemoryManager subscriptionMemoryManager) {
+ if (retainedTabletBytes.get() != 0L) {
+ throw new IllegalStateException();
+ }
+ this.subscriptionMemoryManager = Objects.requireNonNull(subscriptionMemoryManager);
+ memoryBlockedEntryBytes = -1L;
+ realtimeAdmissionBlocked.set(false);
+ }
+
// ======================== Lock Operations ========================
private void acquireReadLock() {
@@ -485,6 +555,24 @@ public class ConsensusPrefetchingQueue {
}
}
+ private boolean canAcceptRealtimeEntry() {
+ return isActive
+ && !closeRequested
+ && !isClosed
+ && !realtimeAdmissionBlocked.get()
+ && prefetchingQueue.size() < MAX_PREFETCHING_QUEUE_SIZE
+ && subscriptionMemoryManager.getFreeMemorySizeInBytes() > 0L;
+ }
+
+ private void blockRealtimeAdmission() {
+ realtimeAdmissionBlocked.set(true);
+ pendingEntries.clear();
+ }
+
+ private void unblockRealtimeAdmission() {
+ realtimeAdmissionBlocked.set(false);
+ }
+
private ConsensusPrefetchSubtask ensurePrefetchSubtaskBound() {
if (closeRequested || isClosed) {
return null;
@@ -603,7 +691,13 @@ public class ConsensusPrefetchingQueue {
}
}
- private synchronized void initPrefetch(final RegionProgress regionProgress) {
+ private void initPrefetch(final RegionProgress regionProgress) {
+ synchronized (prefetchInitializationLock) {
+ initPrefetchUnderInitializationLock(regionProgress);
+ }
+ }
+
+ private void initPrefetchUnderInitializationLock(final RegionProgress regionProgress) {
if (prefetchInitialized) {
return; // double-check under synchronization
}
@@ -641,7 +735,7 @@ public class ConsensusPrefetchingQueue {
}
this.prefetchInitialized = true;
this.observedSeekGeneration = seekGeneration.get();
- this.lingerBatch.reset();
+ discardBatch(this.lingerBatch);
resetBatchWriterProgress();
LOGGER.info(
@@ -1021,9 +1115,13 @@ public class ConsensusPrefetchingQueue {
}
private void bufferRealtimeEntry(final PreparedEntry entry) {
- realtimeEntriesByWriter
- .computeIfAbsent(entry.writerNodeId, ignored -> new TreeMap<>())
- .put(entry.localSeq, entry);
+ final PreparedEntry replaced =
+ realtimeEntriesByWriter
+ .computeIfAbsent(entry.writerNodeId, ignored -> new TreeMap<>())
+ .put(entry.localSeq, entry);
+ if (Objects.nonNull(replaced)) {
+ releaseTabletMemory(replaced.estimatedBytes);
+ }
}
private PreparedEntry peekRealtimeEntry(final int writerNodeId) {
@@ -1097,6 +1195,7 @@ public class ConsensusPrefetchingQueue {
.PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_POLL_COMMITTED_EVENT_BROKEN_INVARIANT_E478FA3C,
this,
event);
+ cleanUpEvent(event, false);
continue;
}
@@ -1107,6 +1206,7 @@ public class ConsensusPrefetchingQueue {
this,
event);
event.nack();
+ prefetchingQueue.add(event);
continue;
}
@@ -1190,6 +1290,7 @@ public class ConsensusPrefetchingQueue {
// ======================== Prefetch Round Drive ========================
private static final long WAL_GAP_RETRY_SLEEP_MS = 10L;
+ private static final long MEMORY_RETRY_SLEEP_MS = 100L;
private static final long WAL_GAP_WAIT_LOG_INTERVAL_MS = 5_000L;
private static final long PREFETCH_STATS_LOG_INTERVAL_MS = 5_000L;
@@ -1216,6 +1317,7 @@ public class ConsensusPrefetchingQueue {
recycleInFlightEvents();
if (!isActive || prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) {
+ blockRealtimeAdmission();
return computeIdleRoundResult();
}
@@ -1225,6 +1327,28 @@ public class ConsensusPrefetchingQueue {
final int maxTablets = config.getSubscriptionConsensusBatchMaxTabletCount();
final long maxBatchBytes = config.getSubscriptionConsensusBatchMaxSizeInBytes();
+ // Always consume already materialized entries before converting more WAL requests into
+ // Tablets. Otherwise a full delivery batch can leave an unbounded hidden writer backlog.
+ if (!drainBufferedRealtimeWriters(
+ lingerBatch, observedSeekGeneration, maxTablets, maxBatchBytes)) {
+ resetRoundStateForSeek(seekGeneration.get());
+ return PrefetchRoundResult.rescheduleNow();
+ }
+ if (prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE
+ || !realtimeEntriesByWriter.isEmpty()) {
+ blockRealtimeAdmission();
+ return computeIdleRoundResult();
+ }
+ if (shouldWaitForSubscriptionMemory()) {
+ blockRealtimeAdmission();
+ if (!lingerBatch.isEmpty() && !flushBatch(lingerBatch, observedSeekGeneration)) {
+ resetRoundStateForSeek(seekGeneration.get());
+ return PrefetchRoundResult.rescheduleNow();
+ }
+ return PrefetchRoundResult.rescheduleAfter(MEMORY_RETRY_SLEEP_MS);
+ }
+ unblockRealtimeAdmission();
+
final List<IndexedConsensusRequest> batch = drainPendingEntries(maxWalEntries);
if (!batch.isEmpty()) {
LOGGER.debug(
@@ -1237,22 +1361,38 @@ public class ConsensusPrefetchingQueue {
nextExpectedSearchIndex.get(),
prefetchingQueue.size());
- final boolean batchAccepted =
+ final MaterializationResult batchResult =
accumulateFromPending(
batch, lingerBatch, observedSeekGeneration, maxTablets, maxBatchBytes);
- if (!batchAccepted) {
- if (pendingWalGapRetryRequested) {
- // Once a drained batch hits an unresolved WAL gap, the affected suffix falls back to
- // the WAL path on later rounds instead of being requeued into the bounded pending path.
+ if (batchResult != MaterializationResult.SUCCESS) {
+ if (batchResult == MaterializationResult.WAL_GAP) {
return PrefetchRoundResult.rescheduleAfter(WAL_GAP_RETRY_SLEEP_MS);
}
+ if (batchResult == MaterializationResult.MEMORY_BLOCKED) {
+ // Publish already reserved Tablets immediately so ACK can release their memory. The
+ // blocked request and the drained suffix are recovered from WAL on a later round.
+ blockRealtimeAdmission();
+ if (!lingerBatch.isEmpty() && !flushBatch(lingerBatch, observedSeekGeneration)) {
+ resetRoundStateForSeek(seekGeneration.get());
+ return PrefetchRoundResult.rescheduleNow();
+ }
+ return PrefetchRoundResult.rescheduleAfter(MEMORY_RETRY_SLEEP_MS);
+ }
resetRoundStateForSeek(seekGeneration.get());
return PrefetchRoundResult.rescheduleNow();
}
}
if (batch.isEmpty() && lingerBatch.isEmpty()) {
- tryCatchUpFromWAL(observedSeekGeneration);
+ final MaterializationResult walResult = tryCatchUpFromWAL(observedSeekGeneration);
+ if (walResult == MaterializationResult.MEMORY_BLOCKED) {
+ blockRealtimeAdmission();
+ return PrefetchRoundResult.rescheduleAfter(MEMORY_RETRY_SLEEP_MS);
+ }
+ if (walResult == MaterializationResult.STALE) {
+ resetRoundStateForSeek(seekGeneration.get());
+ return PrefetchRoundResult.rescheduleNow();
+ }
}
if (!drainBufferedRealtimeWriters(
@@ -1290,6 +1430,7 @@ public class ConsensusPrefetchingQueue {
fatal.getMessage(),
fatal);
if (fatal instanceof VirtualMachineError) {
+ closeCleanupPending = true;
markClosed();
return PrefetchRoundResult.dormant();
}
@@ -1330,7 +1471,10 @@ public class ConsensusPrefetchingQueue {
private void resetRoundStateForSeek(final long newSeekGeneration) {
restorePendingSubscriptionWalCursor(newSeekGeneration);
- lingerBatch.reset();
+ discardRealtimeEntries();
+ discardBatch(lingerBatch);
+ memoryBlockedEntryBytes = -1L;
+ realtimeAdmissionBlocked.set(false);
resetBatchWriterProgress();
observedSeekGeneration = newSeekGeneration;
}
@@ -1338,7 +1482,7 @@ public class ConsensusPrefetchingQueue {
private List<IndexedConsensusRequest> drainPendingEntries(final int maxWalEntries) {
final List<IndexedConsensusRequest> batch = new ArrayList<>();
IndexedConsensusRequest next;
- while (batch.size() < maxWalEntries && (next = pendingEntries.poll()) != null) {
+ while (batch.size() < maxWalEntries && Objects.nonNull(next = pendingEntries.poll())) {
batch.add(next);
}
return batch;
@@ -1349,6 +1493,7 @@ public class ConsensusPrefetchingQueue {
return PrefetchRoundResult.dormant();
}
if (prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) {
+ blockRealtimeAdmission();
return PrefetchRoundResult.dormant();
}
if (hasImmediatePrefetchableWork()) {
@@ -1428,7 +1573,7 @@ public class ConsensusPrefetchingQueue {
}
}
- private boolean appendRealtimeRequest(
+ private MaterializationResult appendRealtimeRequest(
final IndexedConsensusRequest request,
final DeliveryBatchState batchState,
final long expectedSeekGeneration,
@@ -1437,27 +1582,29 @@ public class ConsensusPrefetchingQueue {
final boolean fromPending) {
final PreparedEntry preparedEntry = prepareEntry(request);
if (Objects.isNull(preparedEntry)) {
- return true;
+ return memoryBlockedEntryBytes > 0L
+ ? MaterializationResult.MEMORY_BLOCKED
+ : MaterializationResult.SUCCESS;
}
if (!appendPreparedEntryViaRealtimeWriter(
batchState, preparedEntry, expectedSeekGeneration, maxTablets, maxBatchBytes)) {
- return false;
+ return MaterializationResult.STALE;
}
if (fromPending) {
markAcceptedFromPending();
} else {
markAcceptedFromWal();
}
- return true;
+ return MaterializationResult.SUCCESS;
}
/**
* Accumulates tablets from pending entries into the linger buffer. When pending replay outruns
* the local WAL reader, this method backfills the local-index gap from WAL before continuing.
*
- * @return false if the batch became stale because seek generation changed while flushing
+ * @return the materialization result for the drained pending batch
*/
- private boolean accumulateFromPending(
+ private MaterializationResult accumulateFromPending(
final List<IndexedConsensusRequest> batch,
final DeliveryBatchState lingerBatch,
final long expectedSeekGeneration,
@@ -1481,14 +1628,16 @@ public class ConsensusPrefetchingQueue {
expected,
searchIndex,
searchIndex - expected);
- if (!fillGapFromWAL(
- expected,
- searchIndex,
- lingerBatch,
- expectedSeekGeneration,
- maxTablets,
- maxBatchBytes)) {
- return false;
+ final MaterializationResult gapFillResult =
+ fillGapFromWAL(
+ expected,
+ searchIndex,
+ lingerBatch,
+ expectedSeekGeneration,
+ maxTablets,
+ maxBatchBytes);
+ if (gapFillResult != MaterializationResult.SUCCESS) {
+ return gapFillResult;
}
}
@@ -1508,13 +1657,18 @@ public class ConsensusPrefetchingQueue {
continue;
}
- if (!appendRealtimeRequest(
- request, lingerBatch, expectedSeekGeneration, maxTablets, maxBatchBytes, true)) {
- return false;
+ final MaterializationResult appendResult =
+ appendRealtimeRequest(
+ request, lingerBatch, expectedSeekGeneration, maxTablets, maxBatchBytes, true);
+ if (appendResult != MaterializationResult.SUCCESS) {
+ return appendResult;
}
markMaterializedProgress(request);
processedCount++;
advanceLocalCursorIfPresent(request);
+ if (prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) {
+ break;
+ }
}
LOGGER.debug(
@@ -1527,7 +1681,7 @@ public class ConsensusPrefetchingQueue {
lingerBatch.tablets.size(),
nextExpectedSearchIndex.get());
- return true;
+ return MaterializationResult.SUCCESS;
}
/**
@@ -1539,10 +1693,9 @@ public class ConsensusPrefetchingQueue {
* affected suffix fall back to the WAL path on later rounds. This keeps replay contiguous without
* requeueing the drained batch back into the bounded pending queue.
*
- * @return false if gap fill had to stop because the current batch became stale or the queue was
- * interrupted/closed
+ * @return the materialization result for the WAL gap
*/
- private boolean fillGapFromWAL(
+ private MaterializationResult fillGapFromWAL(
final long fromIndex,
final long toIndex,
final DeliveryBatchState batchState,
@@ -1552,18 +1705,20 @@ public class ConsensusPrefetchingQueue {
pendingWalGapRetryRequested = false;
resetSubscriptionWALPosition(fromIndex);
if (seekGeneration.get() != expectedSeekGeneration || isClosed) {
- return false;
+ return MaterializationResult.STALE;
}
- if (!pumpFromSubscriptionWAL(
- batchState, expectedSeekGeneration, Integer.MAX_VALUE, maxTablets, maxBatchBytes)) {
- return false;
+ final MaterializationResult pumpResult =
+ pumpFromSubscriptionWAL(
+ batchState, expectedSeekGeneration, Integer.MAX_VALUE, maxTablets, maxBatchBytes);
+ if (pumpResult != MaterializationResult.SUCCESS) {
+ return pumpResult;
}
final long nextExpected = nextExpectedSearchIndex.get();
if (nextExpected >= toIndex) {
walGapWaitStartTimeMs = 0L;
lastWalGapWaitLogTimeMs = 0L;
- return true;
+ return MaterializationResult.SUCCESS;
}
final long nowMs = System.currentTimeMillis();
@@ -1586,41 +1741,53 @@ public class ConsensusPrefetchingQueue {
}
onWalGapRetryScheduled();
pendingWalGapRetryRequested = true;
- return false;
+ return MaterializationResult.WAL_GAP;
}
/**
* Try catch-up from WAL when the pending queue was empty. This handles cold-start or scenarios
* where the subscription started after data was already written.
*/
- private void tryCatchUpFromWAL(final long expectedSeekGeneration) {
+ private MaterializationResult tryCatchUpFromWAL(final long expectedSeekGeneration) {
final SubscriptionConfig config = SubscriptionConfig.getInstance();
final int maxTablets = config.getSubscriptionConsensusBatchMaxTabletCount();
final long maxBatchBytes = config.getSubscriptionConsensusBatchMaxSizeInBytes();
final int maxWalEntries = config.getSubscriptionConsensusBatchMaxWalEntries();
- final DeliveryBatchState batchState = new DeliveryBatchState();
+ // Use the persistent linger batch so an unexpected runtime failure cannot orphan already
+ // reserved Tablets or advance replay progress past data that has become unreachable.
+ final DeliveryBatchState batchState = lingerBatch;
resetSubscriptionWALPosition(nextExpectedSearchIndex.get());
- final boolean accepted =
+ final MaterializationResult materializationResult =
pumpFromSubscriptionWAL(
batchState, expectedSeekGeneration, maxWalEntries, maxTablets, maxBatchBytes);
- if (!accepted) {
- return;
+ if (materializationResult != MaterializationResult.SUCCESS) {
+ if (materializationResult == MaterializationResult.MEMORY_BLOCKED && !batchState.isEmpty()) {
+ if (!flushBatch(batchState, expectedSeekGeneration)) {
+ discardBatch(batchState);
+ return MaterializationResult.STALE;
+ }
+ } else {
+ discardBatch(batchState);
+ }
+ return materializationResult;
}
- if (!batchState.isEmpty()) {
- flushBatch(batchState, expectedSeekGeneration);
+ if (!batchState.isEmpty() && !flushBatch(batchState, expectedSeekGeneration)) {
+ discardBatch(batchState);
+ return MaterializationResult.STALE;
}
+ return MaterializationResult.SUCCESS;
}
- private boolean pumpFromSubscriptionWAL(
+ private MaterializationResult pumpFromSubscriptionWAL(
final DeliveryBatchState batchState,
final long expectedSeekGeneration,
final int maxWalEntries,
final int maxTablets,
final long maxBatchBytes) {
if (Objects.isNull(subscriptionWALIterator)) {
- return true;
+ return MaterializationResult.SUCCESS;
}
subscriptionWALIterator.refresh();
@@ -1646,9 +1813,14 @@ public class ConsensusPrefetchingQueue {
continue;
}
- if (!appendRealtimeRequest(
- walEntry, batchState, expectedSeekGeneration, maxTablets, maxBatchBytes, false)) {
- return false;
+ final MaterializationResult appendResult =
+ appendRealtimeRequest(
+ walEntry, batchState, expectedSeekGeneration, maxTablets, maxBatchBytes, false);
+ if (appendResult != MaterializationResult.SUCCESS) {
+ if (appendResult == MaterializationResult.MEMORY_BLOCKED) {
+ resetSubscriptionWALPosition(nextExpectedSearchIndex.get());
+ }
+ return appendResult;
}
markMaterializedProgress(walEntry);
advanceLocalCursorIfPresent(walEntry);
@@ -1670,7 +1842,7 @@ public class ConsensusPrefetchingQueue {
entriesRead,
nextExpectedSearchIndex.get());
}
- return true;
+ return MaterializationResult.SUCCESS;
}
private void ensureSubscriptionWalReadable() {
@@ -1764,6 +1936,7 @@ public class ConsensusPrefetchingQueue {
final InsertNode insertNode =
ConsensusLogToTabletConverter.deserializeToInsertNode(indexedRequest);
if (Objects.isNull(insertNode)) {
+ memoryBlockedEntryBytes = -1L;
return null;
}
@@ -1785,24 +1958,24 @@ public class ConsensusPrefetchingQueue {
maxObservedTimestamp = maxTs;
}
final List<Tablet> tablets = converter.convert(insertNode);
- return new PreparedEntry(tablets, physicalTime, writerNodeId, localSeq, searchIndex);
+ final long estimatedBytes = estimateTabletsBytes(tablets);
+ if (!tryReserveTabletMemory(estimatedBytes)) {
+ return null;
+ }
+ return new PreparedEntry(
+ tablets, estimatedBytes, physicalTime, writerNodeId, localSeq, searchIndex);
}
private static long estimateTabletSize(final Tablet tablet) {
return PipeMemoryWeightUtil.calculateTabletSizeInBytes(tablet);
}
- private void createAndEnqueueEvent(
- final List<Tablet> tablets, final long startSearchIndex, final long endSearchIndex) {
- createAndEnqueueEvent(
- tablets, startSearchIndex, endSearchIndex, endSearchIndex, seekGeneration.get());
- }
-
private boolean createAndEnqueueEvent(
final List<Tablet> tablets,
final long startSearchIndex,
final long endSearchIndex,
final long commitLocalSeq,
+ final long retainedBytes,
final long expectedSeekGeneration) {
if (seekGeneration.get() != expectedSeekGeneration) {
LOGGER.debug(
@@ -1841,6 +2014,8 @@ public class ConsensusPrefetchingQueue {
SubscriptionAgent.broker().getColumnFilterMatcher(topicName).isTimeSelected(),
getTimeSelectedByTable(converter.getDatabaseName(), tablets));
+ // Install the ownership record before exposing the event to concurrent poll/ack threads.
+ retainEventMemory(event, retainedBytes);
prefetchingQueue.add(event);
LOGGER.debug(
@@ -1914,6 +2089,173 @@ public class ConsensusPrefetchingQueue {
return estimatedBytes;
}
+ private boolean tryReserveTabletMemory(final long bytes) {
+ if (bytes <= 0L) {
+ memoryBlockedEntryBytes = -1L;
+ return true;
+ }
+ if (!subscriptionMemoryManager.tryAllocate(bytes)) {
+ memoryBlockedEntryBytes = bytes;
+ return false;
+ }
+ memoryBlockedEntryBytes = -1L;
+ retainedTabletBytes.addAndGet(bytes);
+ return true;
+ }
+
+ private boolean shouldWaitForSubscriptionMemory() {
+ final long blockedEntryBytes = memoryBlockedEntryBytes;
+ if (blockedEntryBytes <= 0L) {
+ return subscriptionMemoryManager.getFreeMemorySizeInBytes() <= 0L;
+ }
+
+ final long totalMemorySizeInBytes = subscriptionMemoryManager.getTotalMemorySizeInBytes();
+ final long usedMemorySizeInBytes = subscriptionMemoryManager.getUsedMemorySizeInBytes();
+ final boolean canAllocateWithinLimit =
+ blockedEntryBytes <= totalMemorySizeInBytes
+ && totalMemorySizeInBytes - usedMemorySizeInBytes >= blockedEntryBytes;
+ final boolean canAllocateSingleOversizedEntry =
+ totalMemorySizeInBytes > 0L
+ && usedMemorySizeInBytes == 0L
+ && blockedEntryBytes > totalMemorySizeInBytes;
+ if (canAllocateWithinLimit || canAllocateSingleOversizedEntry) {
+ memoryBlockedEntryBytes = -1L;
+ return false;
+ }
+ return true;
+ }
+
+ private void releaseTabletMemory(final long bytes) {
+ if (bytes <= 0L) {
+ return;
+ }
+ long currentRetainedBytes;
+ long releasedBytes;
+ do {
+ currentRetainedBytes = retainedTabletBytes.get();
+ if (currentRetainedBytes <= 0L) {
+ return;
+ }
+ releasedBytes = Math.min(bytes, currentRetainedBytes);
+ } while (!retainedTabletBytes.compareAndSet(
+ currentRetainedBytes, currentRetainedBytes - releasedBytes));
+ subscriptionMemoryManager.release(releasedBytes);
+ }
+
+ private void reconcileRetainedTabletMemoryAfterCleanup() {
+ // All indexed lifecycle owners have already been cleared. Any remaining bytes were reserved
+ // before a fatal failure interrupted the transfer into a tracked container.
+ releaseTabletMemory(retainedTabletBytes.get());
+ }
+
+ private void retainEventMemory(final SubscriptionEvent event, final long bytes) {
+ outstandingEventsByCommitContext.put(event.getCommitContext(), event);
+ if (bytes > 0L) {
+ retainedBytesByEvent.put(event, bytes);
+ }
+ }
+
+ private void cleanUpEvent(final SubscriptionEvent event, final boolean force) {
+ synchronized (event) {
+ final boolean indexed =
+ outstandingEventsByCommitContext.remove(event.getCommitContext(), event);
+ final Long retainedBytes = retainedBytesByEvent.remove(event);
+ // A materialized data event can be observed by a regular ACK and a late ACK concurrently.
+ // Only the thread that still owns either lifecycle index may clean and release it.
+ if (!indexed
+ && Objects.isNull(retainedBytes)
+ && event.getCommitContext().hasWriterProgress()) {
+ return;
+ }
+ try {
+ event.cleanUp(force);
+ } finally {
+ if (Objects.nonNull(retainedBytes)) {
+ releaseTabletMemory(retainedBytes);
+ }
+ }
+ }
+ }
+
+ private boolean ackMissingInFlightEvent(
+ final SubscriptionCommitContext commitContext, final boolean silent) {
+ acquireWriteLock();
+ try {
+ if (!canAcceptCommitContext(commitContext, "ack", silent)) {
+ return false;
+ }
+
+ final SubscriptionEvent event = outstandingEventsByCommitContext.get(commitContext);
+ // A recycled event may already have been claimed by another consumer. Only an event that is
+ // still queued can be committed by the previous consumer's late ACK.
+ if (Objects.nonNull(event) && !prefetchingQueue.remove(event)) {
+ return false;
+ }
+
+ final WriterId commitWriterId = extractCommitWriterId(commitContext);
+ final WriterProgress commitWriterProgress = extractCommitWriterProgress(commitContext);
+ boolean committed = false;
+ try {
+ committed =
+ commitManager.commitWithoutOutstanding(
+ consumerGroupId, topicName, consensusGroupId, commitWriterId, commitWriterProgress);
+ } finally {
+ if (!committed
+ && Objects.nonNull(event)
+ && outstandingEventsByCommitContext.get(commitContext) == event
+ && !event.isCommitted()) {
+ prefetchingQueue.add(event);
+ }
+ }
+
+ if (!committed) {
+ if (!silent) {
+ LOGGER.warn(
+ DataNodePipeMessages
+ .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_COMMIT_CONTEXT_DOES_NOT_EXIST_99B8A8F3,
+ this,
+ commitContext);
+ }
+ return false;
+ }
+
+ if (Objects.nonNull(event)) {
+ event.ack();
+ event.recordCommittedTimestamp();
+ cleanUpEvent(event, false);
+ }
+ return true;
+ } finally {
+ releaseWriteLock();
+ }
+ }
+
+ private void discardBatch(final DeliveryBatchState batchState) {
+ releaseTabletMemory(batchState.estimatedBytes);
+ batchState.reset();
+ }
+
+ private void discardRealtimeEntries() {
+ long retainedBytes = 0L;
+ for (final NavigableMap<Long, PreparedEntry> writerEntries : realtimeEntriesByWriter.values()) {
+ for (final PreparedEntry entry : writerEntries.values()) {
+ retainedBytes += entry.estimatedBytes;
+ }
+ }
+ realtimeEntriesByWriter.clear();
+ releaseTabletMemory(retainedBytes);
+ }
+
+ private void discardRetainedEventMemory() {
+ long retainedBytes = 0L;
+ for (final Long eventRetainedBytes : retainedBytesByEvent.values()) {
+ retainedBytes += eventRetainedBytes;
+ }
+ retainedBytesByEvent.clear();
+ outstandingEventsByCommitContext.clear();
+ releaseTabletMemory(retainedBytes);
+ }
+
private boolean appendPreparedEntryViaRealtimeWriter(
final DeliveryBatchState batchState,
final PreparedEntry preparedEntry,
@@ -1921,7 +2263,8 @@ public class ConsensusPrefetchingQueue {
final int maxTablets,
final long maxBatchBytes) {
bufferRealtimeEntry(preparedEntry);
- return drainRealtimeWriters(batchState, expectedSeekGeneration, maxTablets, maxBatchBytes);
+ return drainBufferedRealtimeWriters(
+ batchState, expectedSeekGeneration, maxTablets, maxBatchBytes);
}
private int getRealtimeBufferedEntryCount() {
@@ -2036,7 +2379,7 @@ public class ConsensusPrefetchingQueue {
continue;
}
- final long entryEstimatedBytes = estimateTabletsBytes(writerHead.getTablets());
+ final long entryEstimatedBytes = writerHead.getEstimatedBytes();
if (!canAppendWriterEntry(
batchState, writerHead, entryEstimatedBytes, maxEntries, maxTablets, maxBatchBytes)) {
return true;
@@ -2062,6 +2405,7 @@ public class ConsensusPrefetchingQueue {
entry.getSearchIndex(),
entry.getSearchIndex(),
entry.getLocalSeq(),
+ 0L,
expectedSeekGeneration);
resetBatchWriterProgress();
return committed;
@@ -2075,6 +2419,7 @@ public class ConsensusPrefetchingQueue {
batchState.startSearchIndex,
batchState.endSearchIndex,
batchState.lastLocalSeq,
+ batchState.estimatedBytes,
expectedSeekGeneration)) {
return false;
}
@@ -2111,7 +2456,7 @@ public class ConsensusPrefetchingQueue {
return null;
}
if (ev.isCommitted()) {
- ev.cleanUp(false);
+ cleanUpEvent(ev, false);
return null;
}
ev.recordLastPolledTimestamp();
@@ -2144,6 +2489,9 @@ public class ConsensusPrefetchingQueue {
}
return false;
}
+ if (isCommitContextOutdated(commitContext)) {
+ return false;
+ }
if (!isActive) {
if (silent) {
LOGGER.debug(
@@ -2168,49 +2516,51 @@ public class ConsensusPrefetchingQueue {
}
public boolean ack(final String consumerId, final SubscriptionCommitContext commitContext) {
+ InFlightAckResult result = InFlightAckResult.REJECTED;
acquireReadLock();
try {
- return canAcceptCommitContext(commitContext, "ack", false)
- && ackInternal(consumerId, commitContext);
+ if (canAcceptCommitContext(commitContext, "ack", false)) {
+ result = ackInFlightEvent(consumerId, commitContext, false);
+ }
} finally {
releaseReadLock();
}
+
+ final boolean acked =
+ result == InFlightAckResult.ACKED
+ || (result == InFlightAckResult.MISSING
+ && ackMissingInFlightEvent(commitContext, false));
+ if (acked) {
+ requestPrefetch();
+ }
+ return acked;
}
- private boolean ackInternal(
- final String consumerId, final SubscriptionCommitContext commitContext) {
+ private InFlightAckResult ackInFlightEvent(
+ final String consumerId,
+ final SubscriptionCommitContext commitContext,
+ final boolean silent) {
final WriterId commitWriterId = extractCommitWriterId(commitContext);
final WriterProgress commitWriterProgress = extractCommitWriterProgress(commitContext);
+ final AtomicBoolean found = new AtomicBoolean(false);
final AtomicBoolean acked = new AtomicBoolean(false);
inFlightEvents.compute(
new InFlightEventKey(consumerId, commitContext),
(key, ev) -> {
if (Objects.isNull(ev)) {
- final boolean directCommitted =
- commitManager.commitWithoutOutstanding(
- consumerGroupId,
- topicName,
- consensusGroupId,
- commitWriterId,
- commitWriterProgress);
- acked.set(directCommitted);
- if (!acked.get()) {
+ return null;
+ }
+ found.set(true);
+
+ if (ev.isCommitted()) {
+ if (!silent) {
LOGGER.warn(
DataNodePipeMessages
- .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_COMMIT_CONTEXT_DOES_NOT_EXIST_99B8A8F3,
+ .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_EVENT_ALREADY_COMMITTED_AC34E829,
this,
commitContext);
}
- return null;
- }
-
- if (ev.isCommitted()) {
- LOGGER.warn(
- DataNodePipeMessages
- .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_EVENT_ALREADY_COMMITTED_AC34E829,
- this,
- commitContext);
- ev.cleanUp(false);
+ cleanUpEvent(ev, false);
return null;
}
@@ -2222,22 +2572,25 @@ public class ConsensusPrefetchingQueue {
commitWriterId,
commitWriterProgress);
if (!committed) {
- LOGGER.warn(
- DataNodePipeMessages
- .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_FAILED_TO_ADVANCE_COMMIT_FRONTIER_56E606C0,
- this,
- commitContext);
+ if (!silent) {
+ LOGGER.warn(
+ DataNodePipeMessages
+ .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_FAILED_TO_ADVANCE_COMMIT_FRONTIER_56E606C0,
+ this,
+ commitContext);
+ }
return ev;
}
ev.ack();
ev.recordCommittedTimestamp();
acked.set(true);
- ev.cleanUp(false);
+ cleanUpEvent(ev, false);
return null;
});
-
- return acked.get();
+ return !found.get()
+ ? InFlightAckResult.MISSING
+ : acked.get() ? InFlightAckResult.ACKED : InFlightAckResult.REJECTED;
}
public boolean nack(final String consumerId, final SubscriptionCommitContext commitContext) {
@@ -2255,52 +2608,24 @@ public class ConsensusPrefetchingQueue {
* in multi-region iteration where only one queue owns the event.
*/
public boolean ackSilent(final String consumerId, final SubscriptionCommitContext commitContext) {
+ InFlightAckResult result = InFlightAckResult.REJECTED;
acquireReadLock();
try {
- if (!canAcceptCommitContext(commitContext, "ack", true)) {
- return false;
+ if (canAcceptCommitContext(commitContext, "ack", true)) {
+ result = ackInFlightEvent(consumerId, commitContext, true);
}
- final WriterId commitWriterId = extractCommitWriterId(commitContext);
- final WriterProgress commitWriterProgress = extractCommitWriterProgress(commitContext);
- final AtomicBoolean acked = new AtomicBoolean(false);
- inFlightEvents.compute(
- new InFlightEventKey(consumerId, commitContext),
- (key, ev) -> {
- if (Objects.isNull(ev)) {
- final boolean directCommitted =
- commitManager.commitWithoutOutstanding(
- consumerGroupId,
- topicName,
- consensusGroupId,
- commitWriterId,
- commitWriterProgress);
- acked.set(directCommitted);
- return null;
- }
- if (ev.isCommitted()) {
- ev.cleanUp(false);
- return null;
- }
- final boolean committed =
- commitManager.commit(
- consumerGroupId,
- topicName,
- consensusGroupId,
- commitWriterId,
- commitWriterProgress);
- if (!committed) {
- return ev;
- }
- ev.ack();
- ev.recordCommittedTimestamp();
- acked.set(true);
- ev.cleanUp(false);
- return null;
- });
- return acked.get();
} finally {
releaseReadLock();
}
+
+ final boolean ackedResult =
+ result == InFlightAckResult.ACKED
+ || (result == InFlightAckResult.MISSING
+ && ackMissingInFlightEvent(commitContext, true));
+ if (ackedResult) {
+ requestPrefetch();
+ }
+ return ackedResult;
}
private WriterId extractCommitWriterId(final SubscriptionCommitContext commitContext) {
@@ -2342,7 +2667,7 @@ public class ConsensusPrefetchingQueue {
ev);
ev.ack();
ev.recordCommittedTimestamp();
- ev.cleanUp(false);
+ cleanUpEvent(ev, false);
return null;
}
prefetchingQueue.add(ev);
@@ -2380,7 +2705,7 @@ public class ConsensusPrefetchingQueue {
ev);
ev.ack();
ev.recordCommittedTimestamp();
- ev.cleanUp(false);
+ cleanUpEvent(ev, false);
return null;
}
prefetchingQueue.add(ev);
@@ -2402,7 +2727,7 @@ public class ConsensusPrefetchingQueue {
return null;
}
if (ev.isCommitted()) {
- ev.cleanUp(false);
+ cleanUpEvent(ev, false);
return null;
}
if (ev.pollable()) {
@@ -2416,7 +2741,7 @@ public class ConsensusPrefetchingQueue {
ev);
ev.ack();
ev.recordCommittedTimestamp();
- ev.cleanUp(false);
+ cleanUpEvent(ev, false);
return null;
}
prefetchingQueue.add(ev);
@@ -2435,20 +2760,26 @@ public class ConsensusPrefetchingQueue {
// ======================== Cleanup ========================
public void cleanUp() {
+ blockRealtimeAdmission();
acquireWriteLock();
try {
- prefetchingQueue.forEach(event -> event.cleanUp(true));
+ prefetchingQueue.forEach(event -> cleanUpEvent(event, true));
prefetchingQueue.clear();
- inFlightEvents.values().forEach(event -> event.cleanUp(true));
+ inFlightEvents.values().forEach(event -> cleanUpEvent(event, true));
inFlightEvents.clear();
+ discardRetainedEventMemory();
- realtimeEntriesByWriter.clear();
+ discardRealtimeEntries();
writerStates.clear();
clearRecoveryWriterProgress();
materializedProgressByWriter.clear();
- pendingEntries.clear();
- lingerBatch.reset();
+ // A prefetch round that already held the read lock may have reopened admission while this
+ // cleanup was waiting for the write lock. Fence and clear once more under the write lock.
+ blockRealtimeAdmission();
+ discardBatch(lingerBatch);
+ reconcileRetainedTabletMemoryAfterCleanup();
+ memoryBlockedEntryBytes = -1L;
resetBatchWriterProgress();
pendingWalGapRetryRequested = false;
walGapWaitStartTimeMs = 0L;
@@ -2714,17 +3045,21 @@ public class ConsensusPrefetchingQueue {
}
private void applySeekResetUnderWriteLock(final PendingSeekRequest request) {
+ blockRealtimeAdmission();
+
// 1. Clean up all queued and in-flight events
- prefetchingQueue.forEach(event -> event.cleanUp(true));
+ prefetchingQueue.forEach(event -> cleanUpEvent(event, true));
prefetchingQueue.clear();
- inFlightEvents.values().forEach(event -> event.cleanUp(true));
+ inFlightEvents.values().forEach(event -> cleanUpEvent(event, true));
inFlightEvents.clear();
+ discardRetainedEventMemory();
// 2. Discard stale pending entries from in-memory queue
- pendingEntries.clear();
+ memoryBlockedEntryBytes = -1L;
+ unblockRealtimeAdmission();
// 3. Reset per-writer release state and source-level dedup frontiers.
- realtimeEntriesByWriter.clear();
+ discardRealtimeEntries();
writerStates.clear();
clearRecoveryWriterProgress();
materializedProgressByWriter.clear();
@@ -2738,7 +3073,8 @@ public class ConsensusPrefetchingQueue {
// first replayable searchIndex here.
nextExpectedSearchIndex.set(request.targetSearchIndex);
requestSubscriptionWalReset(request.targetSearchIndex, seekGeneration.get());
- lingerBatch.reset();
+ discardBatch(lingerBatch);
+ reconcileRetainedTabletMemoryAfterCleanup();
resetBatchWriterProgress();
observedSeekGeneration = seekGeneration.get();
pendingWalGapRetryRequested = false;
@@ -2938,10 +3274,11 @@ public class ConsensusPrefetchingQueue {
acquireWriteLock();
try {
- if (isClosed || closeRequested) {
+ if (closeRequested || (isClosed && !closeCleanupPending)) {
return;
}
closeRequested = true;
+ closeCleanupPending = true;
seekRequestToFail = pendingSeekRequest;
pendingSeekRequest = null;
} finally {
@@ -2999,6 +3336,7 @@ public class ConsensusPrefetchingQueue {
// Persist progress before closing
commitManager.persistAll();
}
+ closeCleanupPending = false;
} finally {
closeRequested = false;
}
@@ -3029,7 +3367,7 @@ public class ConsensusPrefetchingQueue {
DataNodePipeMessages
.PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_FAILED_TO_FLUSH_LINGERING_BATCH_F97D8AA7,
this);
- lingerBatch.reset();
+ discardBatch(lingerBatch);
resetBatchWriterProgress();
}
}
@@ -3105,12 +3443,71 @@ public class ConsensusPrefetchingQueue {
* active. Inactive queues skip prefetching and return null on poll.
*/
public void setActive(final boolean active) {
- this.isActive = active;
+ synchronized (runtimeActivationLock) {
+ setActiveUnderRuntimeLock(active);
+ }
+ }
+
+ private void setActiveUnderRuntimeLock(final boolean active) {
+ PendingSeekRequest seekRequestToFail = null;
+ if (this.isActive != active) {
+ if (active) {
+ this.isActive = true;
+ unblockRealtimeAdmission();
+ } else {
+ // Fence admission before waiting for the write lock so no new realtime reference can race
+ // with the lifecycle cleanup below.
+ this.isActive = false;
+ realtimeAdmissionBlocked.set(true);
+ acquireWriteLock();
+ try {
+ seekRequestToFail = pendingSeekRequest;
+ pendingSeekRequest = null;
+ seekGeneration.incrementAndGet();
+
+ prefetchingQueue.forEach(event -> cleanUpEvent(event, true));
+ prefetchingQueue.clear();
+ inFlightEvents.values().forEach(event -> cleanUpEvent(event, true));
+ inFlightEvents.clear();
+ discardRetainedEventMemory();
+
+ pendingEntries.clear();
+ discardRealtimeEntries();
+ discardBatch(lingerBatch);
+ reconcileRetainedTabletMemoryAfterCleanup();
+ writerStates.clear();
+ clearRecoveryWriterProgress();
+ materializedProgressByWriter.clear();
+ resetBatchWriterProgress();
+
+ memoryBlockedEntryBytes = -1L;
+ prefetchInitialized = false;
+ observedSeekGeneration = seekGeneration.get();
+ pendingWalGapRetryRequested = false;
+ walGapWaitStartTimeMs = 0L;
+ lastWalGapWaitLogTimeMs = 0L;
+ pendingSubscriptionWalResetSearchIndex = Long.MIN_VALUE;
+ pendingSubscriptionWalResetGeneration = Long.MIN_VALUE;
+ closeSubscriptionWALIterator();
+ } finally {
+ releaseWriteLock();
+ }
+ }
+ }
LOGGER.info(
DataNodePipeMessages.PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_ISACTIVE_SET_TO_REGION_EC0AD7BA,
this,
active,
consensusGroupId);
+ if (Objects.nonNull(seekRequestToFail)) {
+ seekRequestToFail.fail(
+ new IllegalStateException(
+ String.format(
+ DataNodePipeMessages
+ .PIPE_EXCEPTION_CONSENSUSPREFETCHINGQUEUE_S_RUNTIME_STOPPED_BEFORE_SEEK_7BCB4F4B,
+ this,
+ seekRequestToFail.seekReason)));
+ }
if (active) {
requestPrefetch();
}
@@ -3259,6 +3656,14 @@ public class ConsensusPrefetchingQueue {
return inFlightEvents.size();
}
+ public long getRetainedTabletBytes() {
+ return retainedTabletBytes.get();
+ }
+
+ public long getSubscriptionMemoryLimitInBytes() {
+ return subscriptionMemoryManager.getTotalMemorySizeInBytes();
+ }
+
/** Exposes the current seek generation for runtime tests and metrics. */
public long getCurrentSeekGeneration() {
return seekGeneration.get();
@@ -3323,6 +3728,12 @@ public class ConsensusPrefetchingQueue {
result.put("prefetchingQueueSize", String.valueOf(prefetchingQueue.size()));
result.put("inFlightEventsSize", String.valueOf(inFlightEvents.size()));
result.put("pendingEntriesSize", String.valueOf(pendingEntries.size()));
+ result.put("retainedTabletBytes", String.valueOf(retainedTabletBytes.get()));
+ result.put("memoryBlockedEntryBytes", String.valueOf(memoryBlockedEntryBytes));
+ result.put("realtimeAdmissionBlocked", String.valueOf(realtimeAdmissionBlocked.get()));
+ result.put(
+ "subscriptionMemoryLimitInBytes",
+ String.valueOf(subscriptionMemoryManager.getTotalMemorySizeInBytes()));
result.put("pendingPathAcceptedEntries", String.valueOf(getPendingPathAcceptedEntries()));
result.put("walPathAcceptedEntries", String.valueOf(getWalPathAcceptedEntries()));
result.put("seekGeneration", String.valueOf(seekGeneration.get()));
@@ -3380,6 +3791,8 @@ public class ConsensusPrefetchingQueue {
private interface WriterBufferedEntry {
List<Tablet> getTablets();
+ long getEstimatedBytes();
+
long getPhysicalTime();
int getWriterNodeId();
@@ -3459,6 +3872,7 @@ public class ConsensusPrefetchingQueue {
private static final class PreparedEntry implements WriterBufferedEntry {
private final List<Tablet> tablets;
+ private final long estimatedBytes;
private final long physicalTime;
private final int writerNodeId;
private final long localSeq;
@@ -3466,11 +3880,13 @@ public class ConsensusPrefetchingQueue {
private PreparedEntry(
final List<Tablet> tablets,
+ final long estimatedBytes,
final long physicalTime,
final int writerNodeId,
final long localSeq,
final long searchIndex) {
this.tablets = tablets;
+ this.estimatedBytes = estimatedBytes;
this.physicalTime = physicalTime;
this.writerNodeId = writerNodeId;
this.localSeq = localSeq;
@@ -3482,6 +3898,11 @@ public class ConsensusPrefetchingQueue {
return tablets;
}
+ @Override
+ public long getEstimatedBytes() {
+ return estimatedBytes;
+ }
+
@Override
public long getPhysicalTime() {
return physicalTime;
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/resource/SubscriptionDataNodeResourceManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/resource/SubscriptionDataNodeResourceManager.java
index 347299df10f..bf1a7de30f7 100644
--- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/resource/SubscriptionDataNodeResourceManager.java
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/resource/SubscriptionDataNodeResourceManager.java
@@ -24,15 +24,21 @@ import org.apache.iotdb.db.subscription.resource.log.SubscriptionLogManager;
public class SubscriptionDataNodeResourceManager {
private final SubscriptionLogManager subscriptionLogManager;
+ private final SubscriptionMemoryManager subscriptionMemoryManager;
public static SubscriptionLogManager log() {
return SubscriptionDataNodeResourceManagerHolder.INSTANCE.subscriptionLogManager;
}
+ public static SubscriptionMemoryManager memory() {
+ return SubscriptionDataNodeResourceManagerHolder.INSTANCE.subscriptionMemoryManager;
+ }
+
///////////////////////////// SINGLETON /////////////////////////////
private SubscriptionDataNodeResourceManager() {
subscriptionLogManager = new SubscriptionLogManager();
+ subscriptionMemoryManager = new SubscriptionMemoryManager();
}
private static class SubscriptionDataNodeResourceManagerHolder {
diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/resource/SubscriptionMemoryManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/resource/SubscriptionMemoryManager.java
new file mode 100644
index 00000000000..bf71fad6c4b
--- /dev/null
+++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/resource/SubscriptionMemoryManager.java
@@ -0,0 +1,89 @@
+/*
+ * 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.resource;
+
+import org.apache.iotdb.commons.memory.AtomicLongMemoryBlock;
+import org.apache.iotdb.commons.memory.IMemoryBlock;
+import org.apache.iotdb.commons.memory.MemoryBlockType;
+import org.apache.iotdb.commons.utils.TestOnly;
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+
+public class SubscriptionMemoryManager {
+
+ private static final String MEMORY_BLOCK_NAME = "Subscription";
+
+ private final IMemoryBlock memoryBlock;
+
+ SubscriptionMemoryManager() {
+ memoryBlock =
+ IoTDBDescriptor.getInstance()
+ .getMemoryConfig()
+ .getSubscriptionMemoryManager()
+ .exactAllocate(MEMORY_BLOCK_NAME, MemoryBlockType.DYNAMIC);
+ }
+
+ @TestOnly
+ public SubscriptionMemoryManager(final long totalMemorySizeInBytes) {
+ memoryBlock =
+ new AtomicLongMemoryBlock(
+ MEMORY_BLOCK_NAME, null, totalMemorySizeInBytes, MemoryBlockType.DYNAMIC);
+ }
+
+ /**
+ * Reserves memory for materialized subscription data.
+ *
+ * <p>A single entry larger than the whole budget is allowed only while the budget is otherwise
+ * empty. This avoids permanently blocking progress while keeping the overrun bounded by one
+ * consensus entry.
+ */
+ public synchronized boolean tryAllocate(final long sizeInBytes) {
+ if (sizeInBytes <= 0L) {
+ return true;
+ }
+ if (memoryBlock.allocate(sizeInBytes)) {
+ return true;
+ }
+ if (memoryBlock.getTotalMemorySizeInBytes() > 0L
+ && memoryBlock.getUsedMemoryInBytes() == 0L
+ && sizeInBytes > memoryBlock.getTotalMemorySizeInBytes()) {
+ memoryBlock.forceAllocateWithoutLimitation(sizeInBytes);
+ return true;
+ }
+ return false;
+ }
+
+ public synchronized void release(final long sizeInBytes) {
+ if (sizeInBytes > 0L) {
+ memoryBlock.release(sizeInBytes);
+ }
+ }
+
+ public long getTotalMemorySizeInBytes() {
+ return memoryBlock.getTotalMemorySizeInBytes();
+ }
+
+ public long getUsedMemorySizeInBytes() {
+ return memoryBlock.getUsedMemoryInBytes();
+ }
+
+ public long getFreeMemorySizeInBytes() {
+ return memoryBlock.getFreeMemoryInBytes();
+ }
+}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/DataNodeMemoryConfigTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/DataNodeMemoryConfigTest.java
index 0d3447bbefb..54f4b849ce7 100644
--- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/DataNodeMemoryConfigTest.java
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/DataNodeMemoryConfigTest.java
@@ -20,12 +20,81 @@
package org.apache.iotdb.db.conf;
import org.apache.iotdb.commons.memory.MemoryConfig;
+import org.apache.iotdb.db.i18n.DataNodeMiscMessages;
import org.junit.Assert;
import org.junit.Test;
+import java.util.Arrays;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertThrows;
+
public class DataNodeMemoryConfigTest {
+ @Test
+ public void testResolveSubscriptionQueryMemoryProportions() {
+ final int[] defaultProportions = DataNodeMemoryConfig.resolveQueryMemoryProportions(null, true);
+ assertArrayEquals(new int[] {1, 100, 200, 50, 200, 200, 200, 50, 250}, defaultProportions);
+ assertEquals(
+ 0.2, (double) defaultProportions[8] / Arrays.stream(defaultProportions).sum(), 0.001);
+ assertArrayEquals(
+ new int[] {1, 100, 200, 50, 200, 200, 200, 50, 250},
+ DataNodeMemoryConfig.resolveQueryMemoryProportions("1:100:200:50:200:200:200:50", true));
+ assertArrayEquals(
+ new int[] {1, 100, 200, 50, 200, 200, 200, 50, 0},
+ DataNodeMemoryConfig.resolveQueryMemoryProportions(
+ "1:100:200:50:200:200:200:50:1000", false));
+ }
+
+ @Test
+ public void testRejectInvalidSubscriptionQueryMemoryProportions() {
+ final IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ DataNodeMemoryConfig.resolveQueryMemoryProportions(
+ "1:100:200:50:200:200:200", true));
+ assertEquals(
+ String.format(
+ DataNodeMiscMessages
+ .EXCEPTION_QUERY_MEMORY_PROPORTIONS_MUST_CONTAIN_8_OR_9_COLON_SEPARATED_VALUES_BUT_FOUND_ARG_03A03941,
+ 7),
+ exception.getMessage());
+ }
+
+ @Test
+ public void testRejectNegativeSubscriptionQueryMemoryProportion() {
+ final IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ DataNodeMemoryConfig.resolveQueryMemoryProportions(
+ "1:100:200:50:200:200:200:50:-1", true));
+ assertEquals(
+ String.format(
+ DataNodeMiscMessages
+ .EXCEPTION_QUERY_MEMORY_PROPORTION_AT_POSITION_ARG_MUST_BE_NON_NEGATIVE_BUT_FOUND_ARG_DC69BC75,
+ 9,
+ -1),
+ exception.getMessage());
+ }
+
+ @Test
+ public void testRejectNonPositiveSubscriptionQueryMemoryProportionSum() {
+ final IllegalArgumentException exception =
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> DataNodeMemoryConfig.resolveQueryMemoryProportions("0:0:0:0:0:0:0:0:0", true));
+ assertEquals(
+ String.format(
+ DataNodeMiscMessages
+ .EXCEPTION_THE_SUM_OF_QUERY_MEMORY_PROPORTIONS_MUST_BE_POSITIVE_BUT_WAS_ARG_407092B6,
+ 0),
+ exception.getMessage());
+ }
+
@Test
public void testRpcMemoryControlIsActivatedOnlyExplicitly() {
DataNodeMemoryConfig memoryConfig = IoTDBDescriptor.getInstance().getMemoryConfig();
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueDataNodeMemoryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueDataNodeMemoryTest.java
new file mode 100644
index 00000000000..60f30834cad
--- /dev/null
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueDataNodeMemoryTest.java
@@ -0,0 +1,346 @@
+/*
+ * 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.CommonDescriptor;
+import org.apache.iotdb.commons.consensus.DataRegionId;
+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.pipe.resource.memory.PipeMemoryWeightUtil;
+import org.apache.iotdb.db.queryengine.plan.statement.StatementTestUtils;
+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.ColumnCategory;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.utils.Pair;
+import org.apache.tsfile.write.record.Tablet;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.File;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Field;
+import java.util.Arrays;
+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.assertFalse;
+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;
+
+public class ConsensusPrefetchingQueueDataNodeMemoryTest {
+
+ @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ @Test
+ public void testQueuesShareDataNodeMemoryBudget() throws Exception {
+ final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final int originalBatchMaxDelay =
+ CommonDescriptor.getInstance().getConfig().getSubscriptionConsensusBatchMaxDelayInMs();
+ final File systemDir = temporaryFolder.newFolder("shared-datanode-memory");
+ ConsensusPrefetchingQueue queueA = null;
+ ConsensusPrefetchingQueue queueB = null;
+ try {
+ CommonDescriptor.getInstance().getConfig().setSubscriptionConsensusBatchMaxDelayInMs(0);
+ final Tablet tablet = createTablet();
+ final long oneTabletBytes = PipeMemoryWeightUtil.calculateTabletSizeInBytes(tablet);
+ final SubscriptionMemoryManager memoryManager = new SubscriptionMemoryManager(oneTabletBytes);
+ final ConsensusSubscriptionCommitManager commitManager = newCommitManager(systemDir);
+
+ final FakeConsensusReqReader readerA = new FakeConsensusReqReader();
+ final AtomicInteger conversionCountA = new AtomicInteger();
+ queueA =
+ newQueue(
+ "consumerGroupA",
+ new DataRegionId(1),
+ readerA,
+ newConverter(conversionCountA),
+ commitManager,
+ TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE);
+ queueA.setSubscriptionMemoryManager(memoryManager);
+
+ final FakeConsensusReqReader readerB = new FakeConsensusReqReader();
+ final AtomicInteger conversionCountB = new AtomicInteger();
+ queueB =
+ newQueue(
+ "consumerGroupB",
+ new DataRegionId(2),
+ readerB,
+ newConverter(conversionCountB),
+ commitManager,
+ TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE);
+ queueB.setSubscriptionMemoryManager(memoryManager);
+
+ assertNull(queueA.poll("consumerA"));
+ assertNull(queueB.poll("consumerB"));
+
+ readerA.currentSearchIndex = 1L;
+ assertTrue(pendingEntries(queueA).offer(createRequest(1L)));
+ queueA.drivePrefetchOnce();
+
+ assertEquals(1, conversionCountA.get());
+ assertEquals(0, conversionCountB.get());
+ assertEquals(oneTabletBytes, queueA.getRetainedTabletBytes());
+ assertEquals(0L, queueB.getRetainedTabletBytes());
+ assertEquals(oneTabletBytes, memoryManager.getUsedMemorySizeInBytes());
+ assertTrue(
+ queueA.getRetainedTabletBytes() + queueB.getRetainedTabletBytes()
+ <= memoryManager.getTotalMemorySizeInBytes());
+
+ readerB.currentSearchIndex = 1L;
+ assertFalse(pendingEntries(queueB).offer(createRequest(1L)));
+ queueB.drivePrefetchOnce();
+ assertEquals(0, conversionCountB.get());
+ assertEquals("true", queueB.coreReportMessage().get("realtimeAdmissionBlocked"));
+ assertEquals(oneTabletBytes, memoryManager.getUsedMemorySizeInBytes());
+
+ final SubscriptionEvent eventA = queueA.poll("consumerA");
+ assertNotNull(eventA);
+ assertTrue(queueA.ack("consumerA", eventA.getCommitContext()));
+ assertEquals(0L, queueA.getRetainedTabletBytes());
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+
+ queueB.drivePrefetchOnce();
+ assertEquals("false", queueB.coreReportMessage().get("realtimeAdmissionBlocked"));
+ assertTrue(pendingEntries(queueB).offer(createRequest(1L)));
+ queueB.drivePrefetchOnce();
+
+ assertEquals(1, conversionCountB.get());
+ assertEquals(oneTabletBytes, queueB.getRetainedTabletBytes());
+ assertEquals(oneTabletBytes, memoryManager.getUsedMemorySizeInBytes());
+ assertTrue(
+ queueA.getRetainedTabletBytes() + queueB.getRetainedTabletBytes()
+ <= memoryManager.getTotalMemorySizeInBytes());
+
+ final SubscriptionEvent eventB = queueB.poll("consumerB");
+ assertNotNull(eventB);
+ assertTrue(queueB.ack("consumerB", eventB.getCommitContext()));
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ } finally {
+ if (queueA != null) {
+ queueA.close();
+ }
+ if (queueB != null) {
+ queueB.close();
+ }
+ CommonDescriptor.getInstance()
+ .getConfig()
+ .setSubscriptionConsensusBatchMaxDelayInMs(originalBatchMaxDelay);
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
+ @Test
+ public void testRealtimeBacklogIsDrainedBeforeMorePendingEntriesAreConverted() throws Exception {
+ final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final int originalBatchMaxDelay =
+ CommonDescriptor.getInstance().getConfig().getSubscriptionConsensusBatchMaxDelayInMs();
+ final File systemDir = temporaryFolder.newFolder("realtime-backlog-first");
+ ConsensusPrefetchingQueue queue = null;
+ try {
+ CommonDescriptor.getInstance().getConfig().setSubscriptionConsensusBatchMaxDelayInMs(60_000);
+ final FakeConsensusReqReader reader = new FakeConsensusReqReader();
+ final AtomicInteger conversionCount = new AtomicInteger();
+ queue =
+ newQueue(
+ "consumerGroup",
+ new DataRegionId(1),
+ reader,
+ newConverter(conversionCount),
+ newCommitManager(systemDir),
+ TopicConstant.ORDER_MODE_MULTI_WRITER_VALUE);
+ queue.setSubscriptionMemoryManager(
+ new SubscriptionMemoryManager(
+ PipeMemoryWeightUtil.calculateTabletSizeInBytes(createTablet()) * 4L));
+ queue.setActiveWriterNodeIds(Set.of(7, 8));
+
+ assertNull(queue.poll("consumer"));
+ reader.currentSearchIndex = 2L;
+ assertTrue(pendingEntries(queue).offer(createRequest(1L)));
+ queue.drivePrefetchOnce();
+
+ assertEquals(1, conversionCount.get());
+ assertEquals(2L, queue.getCurrentReadSearchIndex());
+ assertEquals("1", queue.coreReportMessage().get("bufferedRealtimeEntryCount"));
+ assertEquals(0, queue.getPrefetchedEventCount());
+
+ assertTrue(pendingEntries(queue).offer(createRequest(2L)));
+ queue.drivePrefetchOnce();
+
+ assertEquals(1, conversionCount.get());
+ assertEquals(2L, queue.getCurrentReadSearchIndex());
+ assertEquals("1", queue.coreReportMessage().get("bufferedRealtimeEntryCount"));
+ assertEquals("0", queue.coreReportMessage().get("pendingEntriesSize"));
+ assertEquals("true", queue.coreReportMessage().get("realtimeAdmissionBlocked"));
+
+ queue.setActiveWriterNodeIds(Collections.singleton(7));
+ queue.drivePrefetchOnce();
+
+ assertEquals(1, conversionCount.get());
+ assertEquals("0", queue.coreReportMessage().get("bufferedRealtimeEntryCount"));
+ assertEquals("false", queue.coreReportMessage().get("realtimeAdmissionBlocked"));
+
+ assertTrue(pendingEntries(queue).offer(createRequest(2L)));
+ queue.drivePrefetchOnce();
+
+ assertEquals(2, conversionCount.get());
+ assertEquals(3L, queue.getCurrentReadSearchIndex());
+ assertEquals("0", queue.coreReportMessage().get("bufferedRealtimeEntryCount"));
+ assertEquals("0", queue.coreReportMessage().get("pendingEntriesSize"));
+ } finally {
+ if (queue != null) {
+ queue.close();
+ }
+ CommonDescriptor.getInstance()
+ .getConfig()
+ .setSubscriptionConsensusBatchMaxDelayInMs(originalBatchMaxDelay);
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
+ private static ConsensusPrefetchingQueue newQueue(
+ final String consumerGroupId,
+ final DataRegionId regionId,
+ final FakeConsensusReqReader reader,
+ final ConsensusLogToTabletConverter converter,
+ final ConsensusSubscriptionCommitManager commitManager,
+ final String orderMode) {
+ final IoTConsensusServerImpl serverImpl = mock(IoTConsensusServerImpl.class);
+ when(serverImpl.getConsensusReqReader()).thenReturn(reader);
+ when(serverImpl.getWriterSafeFrontierTracker()).thenReturn(new WriterSafeFrontierTracker());
+ return new ConsensusPrefetchingQueue(
+ consumerGroupId,
+ "topic",
+ orderMode,
+ regionId,
+ serverImpl,
+ new SubscriptionWalRetentionPolicy(
+ "topic",
+ SubscriptionWalRetentionPolicy.UNBOUNDED,
+ SubscriptionWalRetentionPolicy.UNBOUNDED),
+ converter,
+ commitManager,
+ new RegionProgress(Collections.emptyMap()),
+ 1L,
+ 1L,
+ true);
+ }
+
+ private static ConsensusLogToTabletConverter newConverter(final AtomicInteger conversionCount) {
+ final ConsensusLogToTabletConverter converter = mock(ConsensusLogToTabletConverter.class);
+ when(converter.convert(any()))
+ .thenAnswer(
+ ignored -> {
+ conversionCount.incrementAndGet();
+ return Collections.singletonList(createTablet());
+ });
+ when(converter.getDatabaseName()).thenReturn("db");
+ return converter;
+ }
+
+ @SuppressWarnings("unchecked")
+ private static BlockingQueue<IndexedConsensusRequest> pendingEntries(
+ final ConsensusPrefetchingQueue queue) throws Exception {
+ final Field field = ConsensusPrefetchingQueue.class.getDeclaredField("pendingEntries");
+ field.setAccessible(true);
+ return (BlockingQueue<IndexedConsensusRequest>) field.get(queue);
+ }
+
+ private static Tablet createTablet() {
+ final List<String> columnNames = Arrays.asList("device", "temperature");
+ final List<TSDataType> dataTypes = Arrays.asList(TSDataType.STRING, TSDataType.DOUBLE);
+ final List<ColumnCategory> categories = Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD);
+ final Tablet tablet = new Tablet("sensors", columnNames, dataTypes, categories, 1);
+ tablet.addTimestamp(0, 1L);
+ tablet.addValue(0, 0, "d1");
+ tablet.addValue(0, 1, 36.5);
+ tablet.setRowSize(1);
+ return tablet;
+ }
+
+ private static IndexedConsensusRequest createRequest(final long searchIndex) {
+ return new IndexedConsensusRequest(
+ searchIndex,
+ Collections.singletonList(
+ StatementTestUtils.genInsertRowNode(Math.toIntExact(searchIndex))))
+ .setPhysicalTime(1000L + searchIndex)
+ .setNodeId(7);
+ }
+
+ private static ConsensusSubscriptionCommitManager newCommitManager(final File systemDir)
+ throws Exception {
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(systemDir.getAbsolutePath());
+ final Constructor<ConsensusSubscriptionCommitManager> constructor =
+ ConsensusSubscriptionCommitManager.class.getDeclaredConstructor();
+ constructor.setAccessible(true);
+ return constructor.newInstance();
+ }
+
+ private static final class FakeConsensusReqReader implements ConsensusReqReader {
+
+ private long currentSearchIndex;
+
+ @Override
+ public void setSafelyDeletedSearchIndex(final long safelyDeletedSearchIndex) {
+ // no-op
+ }
+
+ @Override
+ public ReqIterator getReqIterator(final long startIndex) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public long getCurrentSearchIndex() {
+ return currentSearchIndex;
+ }
+
+ @Override
+ public long getCurrentWALFileVersion() {
+ return 0;
+ }
+
+ @Override
+ public long getTotalSize() {
+ return 0;
+ }
+
+ @Override
+ public Pair<Long, Long> getDeletionBoundToFreeAtLeast(final long bytesToFree) {
+ return new Pair<>(DEFAULT_SAFELY_DELETED_SEARCH_INDEX, 0L);
+ }
+ }
+}
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 1ad473cb564..8460fe319f8 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
@@ -29,6 +29,7 @@ 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.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.iotdb.rpc.subscription.payload.poll.WriterId;
@@ -45,10 +46,18 @@ import org.junit.rules.TemporaryFolder;
import java.io.File;
import java.lang.reflect.Constructor;
import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.lang.reflect.Modifier;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BooleanSupplier;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
@@ -63,6 +72,107 @@ public class ConsensusPrefetchingQueueTest {
@Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+ @Test
+ public void testInitializationAndActivationUseIndependentMonitors() throws Exception {
+ assertFalse(
+ Modifier.isSynchronized(
+ ConsensusPrefetchingQueue.class
+ .getDeclaredMethod("initPrefetch", RegionProgress.class)
+ .getModifiers()));
+ assertFalse(
+ Modifier.isSynchronized(
+ ConsensusPrefetchingQueue.class
+ .getDeclaredMethod("setActive", boolean.class)
+ .getModifiers()));
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ public void testAdmissionClearCannotLeaveEntryEnqueuedAfterFence() throws Exception {
+ final Class<?> queueClass =
+ Class.forName(ConsensusPrefetchingQueue.class.getName() + "$WakeableIndexedConsensusQueue");
+ final Constructor<?> constructor =
+ queueClass.getDeclaredConstructor(int.class, Runnable.class, BooleanSupplier.class);
+ constructor.setAccessible(true);
+
+ final AtomicBoolean admissionEnabled = new AtomicBoolean(true);
+ final CountDownLatch admissionChecked = new CountDownLatch(1);
+ final CountDownLatch finishAdmissionCheck = new CountDownLatch(1);
+ final CountDownLatch clearInvoked = new CountDownLatch(1);
+ final CountDownLatch clearCompleted = new CountDownLatch(1);
+ final AtomicReference<Boolean> offered = new AtomicReference<>();
+ final AtomicReference<Throwable> asyncFailure = new AtomicReference<>();
+ final BooleanSupplier admissionSupplier =
+ () -> {
+ final boolean admitted = admissionEnabled.get();
+ admissionChecked.countDown();
+ try {
+ if (!finishAdmissionCheck.await(5, TimeUnit.SECONDS)) {
+ throw new AssertionError();
+ }
+ } catch (final InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError(e);
+ }
+ return admitted;
+ };
+ final BlockingQueue<IndexedConsensusRequest> queue =
+ (BlockingQueue<IndexedConsensusRequest>)
+ constructor.newInstance(8, (Runnable) () -> {}, admissionSupplier);
+
+ final Thread offerThread =
+ new Thread(
+ () -> {
+ try {
+ offered.set(queue.offer(createRequest(1L)));
+ } catch (final Throwable t) {
+ asyncFailure.compareAndSet(null, t);
+ }
+ });
+ final Thread clearThread =
+ new Thread(
+ () -> {
+ try {
+ clearInvoked.countDown();
+ queue.clear();
+ } catch (final Throwable t) {
+ asyncFailure.compareAndSet(null, t);
+ } finally {
+ clearCompleted.countDown();
+ }
+ });
+ offerThread.setDaemon(true);
+ clearThread.setDaemon(true);
+
+ offerThread.start();
+ try {
+ assertTrue(admissionChecked.await(5, TimeUnit.SECONDS));
+ admissionEnabled.set(false);
+ clearThread.start();
+ assertTrue(clearInvoked.await(5, TimeUnit.SECONDS));
+
+ final long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+ while (clearCompleted.getCount() > 0
+ && clearThread.getState() != Thread.State.BLOCKED
+ && System.nanoTime() < deadlineNanos) {
+ Thread.yield();
+ }
+ assertTrue(clearCompleted.getCount() == 0 || clearThread.getState() == Thread.State.BLOCKED);
+ } finally {
+ finishAdmissionCheck.countDown();
+ offerThread.join(TimeUnit.SECONDS.toMillis(5));
+ clearThread.join(TimeUnit.SECONDS.toMillis(5));
+ }
+
+ assertFalse(offerThread.isAlive());
+ assertFalse(clearThread.isAlive());
+ if (asyncFailure.get() != null) {
+ throw new AssertionError(asyncFailure.get());
+ }
+ assertTrue(Boolean.TRUE.equals(offered.get()));
+ assertTrue(queue.isEmpty());
+ }
+
@Test
public void testLagIncludesLingeringBatchUntilCommitted() throws Exception {
final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
@@ -234,7 +344,7 @@ public class ConsensusPrefetchingQueueTest {
1L, Collections.singletonList(StatementTestUtils.genInsertRowNode(1)))
.setPhysicalTime(1000L)
.setNodeId(7);
- final IndexedConsensusRequest emptyRequest =
+ final IndexedConsensusRequest requestWithEmptyConversionResult =
new IndexedConsensusRequest(
2L, Collections.singletonList(StatementTestUtils.genInsertRowNode(2)))
.setPhysicalTime(1001L)
@@ -243,7 +353,7 @@ public class ConsensusPrefetchingQueueTest {
assertNull(queue.poll("consumer"));
pendingEntries(queue).offer(dataRequest);
- pendingEntries(queue).offer(emptyRequest);
+ pendingEntries(queue).offer(requestWithEmptyConversionResult);
queue.drivePrefetchOnce();
final WriterId writerId = new WriterId(regionId.toString(), 7);
@@ -272,6 +382,553 @@ public class ConsensusPrefetchingQueueTest {
}
}
+ @Test
+ public void testTabletMemoryReleasedAfterAck() throws Exception {
+ final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final File systemDir = temporaryFolder.newFolder("system-memory-release");
+ 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);
+
+ final IndexedConsensusRequest dataRequest =
+ new IndexedConsensusRequest(
+ 1L, Collections.singletonList(StatementTestUtils.genInsertRowNode(1)))
+ .setPhysicalTime(1000L)
+ .setNodeId(7);
+ final IndexedConsensusRequest requestWithEmptyConversionResult =
+ new IndexedConsensusRequest(
+ 2L, Collections.singletonList(StatementTestUtils.genInsertRowNode(2)))
+ .setPhysicalTime(1001L)
+ .setNodeId(7);
+ reader.currentSearchIndex = 2L;
+ pendingEntries(queue).offer(dataRequest);
+ pendingEntries(queue).offer(requestWithEmptyConversionResult);
+
+ assertNull(queue.poll("consumer"));
+ queue.drivePrefetchOnce();
+
+ final long retainedBytes = queue.getRetainedTabletBytes();
+ assertTrue(retainedBytes > 0L);
+ final SubscriptionEvent event = queue.poll("consumer");
+ assertNotNull(event);
+ assertEquals(retainedBytes, queue.getRetainedTabletBytes());
+
+ assertTrue(queue.ack("consumer", event.getCommitContext()));
+ assertEquals(0L, queue.getRetainedTabletBytes());
+ } finally {
+ if (queue != null) {
+ queue.close();
+ }
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
+ @Test
+ public void testCleanupReconcilesUnindexedTabletReservation() throws Exception {
+ final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final File systemDir = temporaryFolder.newFolder("system-orphan-memory-release");
+ 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());
+
+ queue =
+ new ConsensusPrefetchingQueue(
+ "consumerGroup",
+ "topic",
+ TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE,
+ regionId,
+ serverImpl,
+ new SubscriptionWalRetentionPolicy(
+ "topic",
+ SubscriptionWalRetentionPolicy.UNBOUNDED,
+ SubscriptionWalRetentionPolicy.UNBOUNDED),
+ mock(ConsensusLogToTabletConverter.class),
+ newCommitManager(systemDir),
+ new RegionProgress(Collections.emptyMap()),
+ 1L,
+ 1L,
+ true);
+ final SubscriptionMemoryManager memoryManager = new SubscriptionMemoryManager(1024L);
+ queue.setSubscriptionMemoryManager(memoryManager);
+
+ final Method reserveTabletMemory =
+ ConsensusPrefetchingQueue.class.getDeclaredMethod("tryReserveTabletMemory", long.class);
+ reserveTabletMemory.setAccessible(true);
+ assertTrue((Boolean) reserveTabletMemory.invoke(queue, 256L));
+ assertEquals(256L, queue.getRetainedTabletBytes());
+ assertEquals(256L, memoryManager.getUsedMemorySizeInBytes());
+
+ queue.cleanUp();
+
+ assertEquals(0L, queue.getRetainedTabletBytes());
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ } finally {
+ if (queue != null) {
+ queue.close();
+ }
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
+ @Test
+ public void testMemoryBackpressureStopsFurtherMaterializationUntilAck() throws Exception {
+ final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final File systemDir = temporaryFolder.newFolder("system-memory-backpressure");
+ 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 AtomicInteger conversionCount = new AtomicInteger();
+ final ConsensusLogToTabletConverter converter = mock(ConsensusLogToTabletConverter.class);
+ when(converter.convert(any()))
+ .thenAnswer(
+ ignored -> {
+ conversionCount.incrementAndGet();
+ return Collections.singletonList(createTablet());
+ });
+ 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);
+ final long oneTabletBytes = createTablet().ramBytesUsed();
+ final SubscriptionMemoryManager memoryManager =
+ new SubscriptionMemoryManager(oneTabletBytes + Math.max(1L, oneTabletBytes / 2L));
+ queue.setSubscriptionMemoryManager(memoryManager);
+
+ final IndexedConsensusRequest firstRequest = createRequest(1L);
+ final IndexedConsensusRequest secondRequest = createRequest(2L);
+ final IndexedConsensusRequest thirdRequest = createRequest(3L);
+ assertTrue(pendingEntries(queue).offer(firstRequest));
+ assertTrue(pendingEntries(queue).offer(secondRequest));
+ assertTrue(pendingEntries(queue).offer(thirdRequest));
+ reader.currentSearchIndex = 3L;
+
+ assertNull(queue.poll("consumer"));
+ queue.drivePrefetchOnce();
+
+ assertEquals(2, conversionCount.get());
+ assertEquals(2L, queue.getCurrentReadSearchIndex());
+ assertEquals(oneTabletBytes, queue.getRetainedTabletBytes());
+ assertTrue(memoryManager.getFreeMemorySizeInBytes() > 0L);
+ assertEquals(1, queue.getPrefetchedEventCount());
+ assertTrue(pendingEntries(queue).isEmpty());
+ assertEquals("true", queue.coreReportMessage().get("realtimeAdmissionBlocked"));
+ assertFalse(pendingEntries(queue).offer(createRequest(4L)));
+
+ queue.drivePrefetchOnce();
+ assertEquals(2, conversionCount.get());
+ assertEquals(oneTabletBytes, queue.getRetainedTabletBytes());
+
+ final SubscriptionEvent event = queue.poll("consumer");
+ assertNotNull(event);
+ assertTrue(queue.ack("consumer", event.getCommitContext()));
+ assertEquals(0L, queue.getRetainedTabletBytes());
+
+ queue.drivePrefetchOnce();
+ assertEquals("false", queue.coreReportMessage().get("realtimeAdmissionBlocked"));
+ assertTrue(pendingEntries(queue).offer(secondRequest));
+ assertTrue(pendingEntries(queue).offer(thirdRequest));
+ queue.drivePrefetchOnce();
+ assertEquals(4, conversionCount.get());
+ assertEquals(3L, queue.getCurrentReadSearchIndex());
+ assertEquals(oneTabletBytes, queue.getRetainedTabletBytes());
+ assertEquals(1, queue.getPrefetchedEventCount());
+
+ final SubscriptionEvent secondEvent = queue.poll("consumer");
+ assertNotNull(secondEvent);
+ assertTrue(queue.ack("consumer", secondEvent.getCommitContext()));
+ queue.drivePrefetchOnce();
+ assertTrue(pendingEntries(queue).offer(thirdRequest));
+ queue.drivePrefetchOnce();
+
+ assertEquals(5, conversionCount.get());
+ assertEquals(4L, queue.getCurrentReadSearchIndex());
+ assertEquals("0", queue.coreReportMessage().get("pendingEntriesSize"));
+ assertEquals("0", queue.coreReportMessage().get("bufferedRealtimeEntryCount"));
+ assertEquals(oneTabletBytes, queue.getRetainedTabletBytes());
+
+ queue.close();
+ queue = null;
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ } finally {
+ if (queue != null) {
+ queue.close();
+ }
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
+ @Test
+ public void testWideTablePausedConsumerKeepsMaterializedMemoryBounded() throws Exception {
+ final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final File systemDir = temporaryFolder.newFolder("system-wide-table-memory-bound");
+ 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 int columnCount = 128;
+ final int rowCount = 64;
+ final long oneTabletBytes = createWideTablet(columnCount, rowCount).ramBytesUsed();
+ final AtomicInteger conversionCount = new AtomicInteger();
+ final ConsensusLogToTabletConverter converter = mock(ConsensusLogToTabletConverter.class);
+ when(converter.convert(any()))
+ .thenAnswer(
+ ignored -> {
+ conversionCount.incrementAndGet();
+ return Collections.singletonList(createWideTablet(columnCount, rowCount));
+ });
+ 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);
+ final long memoryLimit = oneTabletBytes * 2L + oneTabletBytes / 2L;
+ final SubscriptionMemoryManager memoryManager = new SubscriptionMemoryManager(memoryLimit);
+ queue.setSubscriptionMemoryManager(memoryManager);
+
+ final int writeCount = 256;
+ // One scheduling attempt per submitted write proves repeated prefetch cannot bypass
+ // backpressure.
+ final int backpressureVerificationRounds = writeCount;
+ for (long searchIndex = 1L; searchIndex <= writeCount; searchIndex++) {
+ assertTrue(pendingEntries(queue).offer(createRequest(searchIndex)));
+ }
+ reader.currentSearchIndex = writeCount;
+
+ assertNull(queue.poll("pausedConsumer"));
+ queue.drivePrefetchOnce();
+
+ assertEquals(3, conversionCount.get());
+ assertEquals(3L, queue.getCurrentReadSearchIndex());
+ assertEquals(oneTabletBytes * 2L, queue.getRetainedTabletBytes());
+ assertTrue(queue.getRetainedTabletBytes() <= queue.getSubscriptionMemoryLimitInBytes());
+ assertEquals(1, queue.getPrefetchedEventCount());
+ assertEquals("0", queue.coreReportMessage().get("pendingEntriesSize"));
+ assertEquals("0", queue.coreReportMessage().get("bufferedRealtimeEntryCount"));
+ assertEquals("true", queue.coreReportMessage().get("realtimeAdmissionBlocked"));
+
+ for (int round = 0; round < backpressureVerificationRounds; round++) {
+ queue.drivePrefetchOnce();
+ }
+ assertEquals(3, conversionCount.get());
+ assertEquals(oneTabletBytes * 2L, queue.getRetainedTabletBytes());
+ assertFalse(pendingEntries(queue).offer(createRequest(writeCount + 1L)));
+
+ final SubscriptionEvent pausedEvent = queue.poll("pausedConsumer");
+ assertNotNull(pausedEvent);
+ assertEquals(0, queue.getPrefetchedEventCount());
+ assertEquals(1L, queue.getSubscriptionUncommittedEventCount());
+ for (int round = 0; round < backpressureVerificationRounds; round++) {
+ queue.drivePrefetchOnce();
+ }
+ assertEquals(3, conversionCount.get());
+ assertEquals(oneTabletBytes * 2L, queue.getRetainedTabletBytes());
+
+ queue.close();
+ queue = null;
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ } finally {
+ if (queue != null) {
+ queue.close();
+ }
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
+ @Test
+ public void testLateAckAfterRecycleReleasesQueuedTabletMemory() throws Exception {
+ final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final int originalRecycleInterval =
+ CommonDescriptor.getInstance()
+ .getConfig()
+ .getSubscriptionRecycleUncommittedEventIntervalMs();
+ final File systemDir = temporaryFolder.newFolder("system-late-ack-memory-release");
+ ConsensusPrefetchingQueue queue = null;
+ try {
+ CommonDescriptor.getInstance()
+ .getConfig()
+ .setSubscriptionRecycleUncommittedEventIntervalMs(-1);
+ 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);
+ final SubscriptionMemoryManager memoryManager =
+ new SubscriptionMemoryManager(createTablet().ramBytesUsed() * 2L);
+ queue.setSubscriptionMemoryManager(memoryManager);
+
+ 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);
+ assertTrue(queue.getRetainedTabletBytes() > 0L);
+ assertEquals(1L, queue.getSubscriptionUncommittedEventCount());
+
+ queue.drivePrefetchOnce();
+ assertEquals(0L, queue.getSubscriptionUncommittedEventCount());
+ assertEquals(1, queue.getPrefetchedEventCount());
+
+ assertTrue(queue.ack("consumer", event.getCommitContext()));
+ assertEquals(0, queue.getPrefetchedEventCount());
+ assertEquals(0L, queue.getRetainedTabletBytes());
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ } finally {
+ if (queue != null) {
+ queue.close();
+ }
+ CommonDescriptor.getInstance()
+ .getConfig()
+ .setSubscriptionRecycleUncommittedEventIntervalMs(originalRecycleInterval);
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
+ @Test
+ public void testLateAckDoesNotStealRecycledEventFromNewConsumer() throws Exception {
+ final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final int originalRecycleInterval =
+ CommonDescriptor.getInstance()
+ .getConfig()
+ .getSubscriptionRecycleUncommittedEventIntervalMs();
+ final File systemDir = temporaryFolder.newFolder("system-late-ack-new-owner");
+ ConsensusPrefetchingQueue queue = null;
+ try {
+ CommonDescriptor.getInstance()
+ .getConfig()
+ .setSubscriptionRecycleUncommittedEventIntervalMs(-1);
+ 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);
+ final SubscriptionMemoryManager memoryManager =
+ new SubscriptionMemoryManager(createTablet().ramBytesUsed() * 2L);
+ queue.setSubscriptionMemoryManager(memoryManager);
+
+ reader.currentSearchIndex = 2L;
+ assertTrue(pendingEntries(queue).offer(createRequest(1L)));
+ assertTrue(pendingEntries(queue).offer(createRequest(2L)));
+
+ assertNull(queue.poll("oldConsumer"));
+ queue.drivePrefetchOnce();
+ final SubscriptionEvent event = queue.poll("oldConsumer");
+ assertNotNull(event);
+ final long retainedBytes = queue.getRetainedTabletBytes();
+ assertTrue(retainedBytes > 0L);
+
+ queue.drivePrefetchOnce();
+ assertEquals(0L, queue.getSubscriptionUncommittedEventCount());
+ assertEquals(1, queue.getPrefetchedEventCount());
+
+ final SubscriptionEvent redeliveredEvent = queue.poll("newConsumer");
+ assertTrue(redeliveredEvent == event);
+ assertFalse(queue.ack("oldConsumer", event.getCommitContext()));
+ assertEquals(1L, queue.getSubscriptionUncommittedEventCount());
+ assertEquals(0, queue.getPrefetchedEventCount());
+ assertEquals(retainedBytes, queue.getRetainedTabletBytes());
+ assertTrue(memoryManager.getUsedMemorySizeInBytes() > 0L);
+
+ assertTrue(queue.ack("newConsumer", event.getCommitContext()));
+ assertEquals(0L, queue.getSubscriptionUncommittedEventCount());
+ assertEquals(0L, queue.getRetainedTabletBytes());
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ } finally {
+ if (queue != null) {
+ queue.close();
+ }
+ CommonDescriptor.getInstance()
+ .getConfig()
+ .setSubscriptionRecycleUncommittedEventIntervalMs(originalRecycleInterval);
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
+ @Test
+ public void testDeactivationReleasesMaterializedTabletMemory() throws Exception {
+ final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final File systemDir = temporaryFolder.newFolder("system-deactivation-memory-release");
+ 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);
+ final SubscriptionMemoryManager memoryManager =
+ new SubscriptionMemoryManager(createTablet().ramBytesUsed() * 2L);
+ queue.setSubscriptionMemoryManager(memoryManager);
+
+ 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);
+ assertTrue(queue.getRetainedTabletBytes() > 0L);
+ assertEquals(1L, queue.getSubscriptionUncommittedEventCount());
+ final long previousSeekGeneration = queue.getCurrentSeekGeneration();
+
+ queue.setActive(false);
+
+ assertFalse(queue.isActive());
+ assertEquals(0L, queue.getInitializedStatus());
+ assertEquals(previousSeekGeneration + 1L, queue.getCurrentSeekGeneration());
+ assertEquals(0, queue.getPrefetchedEventCount());
+ assertEquals(0L, queue.getSubscriptionUncommittedEventCount());
+ assertEquals(0L, queue.getRetainedTabletBytes());
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ assertFalse(pendingEntries(queue).offer(createRequest(3L)));
+ } finally {
+ if (queue != null) {
+ queue.close();
+ }
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
@SuppressWarnings("unchecked")
private static BlockingQueue<IndexedConsensusRequest> pendingEntries(
final ConsensusPrefetchingQueue queue) throws Exception {
@@ -292,6 +949,42 @@ public class ConsensusPrefetchingQueueTest {
return tablet;
}
+ private static Tablet createWideTablet(final int columnCount, final int rowCount) {
+ final String[] columnNames = new String[columnCount];
+ final TSDataType[] dataTypes = new TSDataType[columnCount];
+ final ColumnCategory[] categories = new ColumnCategory[columnCount];
+ Arrays.fill(dataTypes, TSDataType.DOUBLE);
+ Arrays.fill(categories, ColumnCategory.FIELD);
+ for (int columnIndex = 0; columnIndex < columnCount; columnIndex++) {
+ columnNames[columnIndex] = "field_" + columnIndex;
+ }
+
+ final Tablet tablet =
+ new Tablet(
+ "wide_table",
+ Arrays.asList(columnNames),
+ Arrays.asList(dataTypes),
+ Arrays.asList(categories),
+ rowCount);
+ for (int rowIndex = 0; rowIndex < rowCount; rowIndex++) {
+ tablet.addTimestamp(rowIndex, rowIndex);
+ for (int columnIndex = 0; columnIndex < columnCount; columnIndex++) {
+ tablet.addValue(rowIndex, columnIndex, (double) (rowIndex + columnIndex));
+ }
+ }
+ tablet.setRowSize(rowCount);
+ return tablet;
+ }
+
+ private static IndexedConsensusRequest createRequest(final long searchIndex) {
+ return new IndexedConsensusRequest(
+ searchIndex,
+ Collections.singletonList(
+ StatementTestUtils.genInsertRowNode(Math.toIntExact(searchIndex))))
+ .setPhysicalTime(1000L + searchIndex)
+ .setNodeId(7);
+ }
+
private static ConsensusSubscriptionCommitManager newCommitManager(final File systemDir)
throws Exception {
IoTDBDescriptor.getInstance().getConfig().setSystemDir(systemDir.getAbsolutePath());
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueWalBackpressureTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueWalBackpressureTest.java
new file mode 100644
index 00000000000..b3783beab43
--- /dev/null
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueWalBackpressureTest.java
@@ -0,0 +1,336 @@
+/*
+ * 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.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.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertRowNode;
+import org.apache.iotdb.db.queryengine.plan.statement.StatementTestUtils;
+import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALInfoEntry;
+import org.apache.iotdb.db.storageengine.dataregion.wal.io.WALFileVersion;
+import org.apache.iotdb.db.storageengine.dataregion.wal.io.WALMetaData;
+import org.apache.iotdb.db.storageengine.dataregion.wal.io.WALWriter;
+import org.apache.iotdb.db.storageengine.dataregion.wal.node.WALNode;
+import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALByteBufferForTest;
+import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALFileStatus;
+import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALFileUtils;
+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.iotdb.rpc.subscription.payload.poll.WriterId;
+import org.apache.iotdb.rpc.subscription.payload.poll.WriterProgress;
+
+import org.apache.tsfile.enums.ColumnCategory;
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.write.record.Tablet;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+import java.io.File;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Field;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.BlockingQueue;
+
+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.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class ConsensusPrefetchingQueueWalBackpressureTest {
+
+ private static final int WRITER_NODE_ID = 7;
+ private static final int REQUEST_COUNT = 3;
+
+ @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ @Test
+ public void testAckRecoversDrainedSuffixFromWalWithoutReoffer() throws Exception {
+ final String originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir();
+ final CommonConfig config = CommonDescriptor.getInstance().getConfig();
+ final int originalBatchMaxWalEntries = config.getSubscriptionConsensusBatchMaxWalEntries();
+ final int originalBatchMaxTabletCount = config.getSubscriptionConsensusBatchMaxTabletCount();
+ final long originalBatchMaxSize = config.getSubscriptionConsensusBatchMaxSizeInBytes();
+ final int originalBatchMaxDelay = config.getSubscriptionConsensusBatchMaxDelayInMs();
+ final File systemDir = temporaryFolder.newFolder("system-wal-backpressure-recovery");
+ final File walDirectory = temporaryFolder.newFolder("wal-backpressure-recovery");
+ ConsensusPrefetchingQueue queue = null;
+ try {
+ final Tablet sampleTablet = createTablet();
+ final long oneTabletBytes = sampleTablet.ramBytesUsed();
+ config.setSubscriptionConsensusBatchMaxWalEntries(128);
+ config.setSubscriptionConsensusBatchMaxTabletCount(128);
+ config.setSubscriptionConsensusBatchMaxSizeInBytes(oneTabletBytes * 16L);
+ config.setSubscriptionConsensusBatchMaxDelayInMs(60_000);
+
+ writeSealedWal(walDirectory);
+
+ final WALNode walNode = mock(WALNode.class);
+ when(walNode.getLogDirectory()).thenReturn(walDirectory);
+ when(walNode.getCurrentSearchIndex()).thenReturn((long) REQUEST_COUNT);
+ when(walNode.getCurrentWALFileVersion()).thenReturn(1L);
+ when(walNode.getCurrentWALMetaDataSnapshot()).thenReturn(new WALMetaData());
+
+ final IoTConsensusServerImpl serverImpl = mock(IoTConsensusServerImpl.class);
+ when(serverImpl.getConsensusReqReader()).thenReturn(walNode);
+ when(serverImpl.getWriterSafeFrontierTracker()).thenReturn(new WriterSafeFrontierTracker());
+
+ final List<Long> conversionAttempts = new ArrayList<>();
+ final ConsensusLogToTabletConverter converter = mock(ConsensusLogToTabletConverter.class);
+ when(converter.convert(any()))
+ .thenAnswer(
+ invocation -> {
+ final InsertNode insertNode = (InsertNode) invocation.getArgument(0);
+ conversionAttempts.add(insertNode.getSearchIndex());
+ return Collections.singletonList(createTablet());
+ });
+ when(converter.getDatabaseName()).thenReturn("db");
+
+ final DataRegionId regionId = new DataRegionId(1);
+ final ConsensusSubscriptionCommitManager commitManager = newCommitManager(systemDir);
+ queue =
+ new ConsensusPrefetchingQueue(
+ "consumerGroup",
+ "topic",
+ TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE,
+ regionId,
+ serverImpl,
+ new SubscriptionWalRetentionPolicy(
+ "topic",
+ SubscriptionWalRetentionPolicy.UNBOUNDED,
+ SubscriptionWalRetentionPolicy.UNBOUNDED),
+ converter,
+ commitManager,
+ new RegionProgress(Collections.emptyMap()),
+ 1L,
+ 1L,
+ true);
+ final long memoryLimit = oneTabletBytes + Math.max(1L, oneTabletBytes / 2L);
+ final TrackingSubscriptionMemoryManager memoryManager =
+ new TrackingSubscriptionMemoryManager(memoryLimit);
+ queue.setSubscriptionMemoryManager(memoryManager);
+
+ assertNull(queue.poll("consumer"));
+ for (long searchIndex = 1L; searchIndex <= REQUEST_COUNT; searchIndex++) {
+ assertTrue(pendingEntries(queue).offer(createRequest(searchIndex)));
+ }
+
+ queue.drivePrefetchOnce();
+
+ assertEquals(Arrays.asList(1L, 2L), conversionAttempts);
+ assertEquals(2L, queue.getCurrentReadSearchIndex());
+ assertEquals(1, queue.getPrefetchedEventCount());
+ assertTrue(pendingEntries(queue).isEmpty());
+ assertEquals("true", queue.coreReportMessage().get("realtimeAdmissionBlocked"));
+ assertEquals(oneTabletBytes, memoryManager.getUsedMemorySizeInBytes());
+ assertMemoryBounded(queue, memoryManager);
+
+ final List<Long> deliveredLocalSequences = new ArrayList<>();
+ for (long expectedLocalSequence = 1L;
+ expectedLocalSequence <= REQUEST_COUNT;
+ expectedLocalSequence++) {
+ final SubscriptionEvent event = queue.poll("consumer");
+ assertNotNull(event);
+ final long deliveredLocalSequence =
+ event.getCommitContext().getWriterProgress().getLocalSeq();
+ assertEquals(expectedLocalSequence, deliveredLocalSequence);
+ deliveredLocalSequences.add(deliveredLocalSequence);
+
+ assertTrue(queue.ack("consumer", event.getCommitContext()));
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ assertMemoryBounded(queue, memoryManager);
+
+ queue.drivePrefetchOnce();
+ assertMemoryBounded(queue, memoryManager);
+ if (expectedLocalSequence < REQUEST_COUNT) {
+ assertEquals(1, queue.getPrefetchedEventCount());
+ assertEquals(oneTabletBytes, memoryManager.getUsedMemorySizeInBytes());
+ }
+ }
+
+ assertEquals(Arrays.asList(1L, 2L, 3L), deliveredLocalSequences);
+ assertEquals(Arrays.asList(1L, 2L, 2L, 3L, 3L), conversionAttempts);
+ assertEquals(1L, queue.getPendingPathAcceptedEntries());
+ assertEquals(2L, queue.getWalPathAcceptedEntries());
+ assertEquals(4L, queue.getCurrentReadSearchIndex());
+ assertEquals(0, queue.getPrefetchedEventCount());
+ assertEquals(0L, queue.getSubscriptionUncommittedEventCount());
+ assertNull(queue.poll("consumer"));
+ assertEquals(
+ new WriterProgress(1000L + REQUEST_COUNT, REQUEST_COUNT),
+ commitManager
+ .getCommittedRegionProgress("consumerGroup", "topic", regionId)
+ .getWriterPositions()
+ .get(new WriterId(regionId.toString(), WRITER_NODE_ID)));
+ assertEquals(oneTabletBytes, memoryManager.getMaximumUsedMemorySizeInBytes());
+ assertFalse(memoryManager.hasExceededBudget());
+ } finally {
+ if (queue != null) {
+ queue.close();
+ }
+ config.setSubscriptionConsensusBatchMaxWalEntries(originalBatchMaxWalEntries);
+ config.setSubscriptionConsensusBatchMaxTabletCount(originalBatchMaxTabletCount);
+ config.setSubscriptionConsensusBatchMaxSizeInBytes(originalBatchMaxSize);
+ config.setSubscriptionConsensusBatchMaxDelayInMs(originalBatchMaxDelay);
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir);
+ }
+ }
+
+ private static void writeSealedWal(final File walDirectory) throws Exception {
+ final File historicalWal =
+ new File(
+ walDirectory, WALFileUtils.getLogFileName(0L, 0L, WALFileStatus.CONTAINS_SEARCH_INDEX));
+ final File liveWal =
+ new File(
+ walDirectory,
+ WALFileUtils.getLogFileName(1L, REQUEST_COUNT, WALFileStatus.CONTAINS_SEARCH_INDEX));
+
+ try (WALWriter writer = new WALWriter(historicalWal, WALFileVersion.V3)) {
+ for (long searchIndex = 1L; searchIndex <= REQUEST_COUNT; searchIndex++) {
+ writeWalEntry(writer, searchIndex);
+ }
+ }
+ try (WALWriter ignored = new WALWriter(liveWal, WALFileVersion.V3)) {
+ // Keep an empty live successor so the WAL containing the requests is historical and readable.
+ }
+ }
+
+ private static void writeWalEntry(final WALWriter writer, final long searchIndex)
+ throws Exception {
+ final RelationalInsertRowNode insertNode =
+ StatementTestUtils.genInsertRowNode(Math.toIntExact(searchIndex));
+ insertNode.setSearchIndex(searchIndex);
+ insertNode.setLastFragment(true);
+ final WALInfoEntry walEntry = new WALInfoEntry(1L, insertNode);
+ final ByteBuffer buffer = ByteBuffer.allocate(walEntry.serializedSize());
+ walEntry.serialize(new WALByteBufferForTest(buffer));
+
+ final WALMetaData metaData = new WALMetaData();
+ metaData.add(
+ walEntry.serializedSize(),
+ searchIndex,
+ 1L,
+ 1000L + searchIndex,
+ WRITER_NODE_ID,
+ searchIndex);
+ writer.write(buffer, metaData);
+ }
+
+ private static IndexedConsensusRequest createRequest(final long searchIndex) {
+ return new IndexedConsensusRequest(
+ searchIndex,
+ Collections.singletonList(
+ StatementTestUtils.genInsertRowNode(Math.toIntExact(searchIndex))))
+ .setPhysicalTime(1000L + searchIndex)
+ .setNodeId(WRITER_NODE_ID);
+ }
+
+ private static Tablet createTablet() {
+ final List<String> columnNames = Arrays.asList("device", "temperature");
+ final List<TSDataType> dataTypes = Arrays.asList(TSDataType.STRING, TSDataType.DOUBLE);
+ final List<ColumnCategory> categories = Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD);
+ final Tablet tablet = new Tablet("sensors", columnNames, dataTypes, categories, 1);
+ tablet.addTimestamp(0, 1L);
+ tablet.addValue(0, 0, "d1");
+ tablet.addValue(0, 1, 36.5);
+ tablet.setRowSize(1);
+ return tablet;
+ }
+
+ private static void assertMemoryBounded(
+ final ConsensusPrefetchingQueue queue,
+ final TrackingSubscriptionMemoryManager memoryManager) {
+ assertEquals(queue.getRetainedTabletBytes(), memoryManager.getUsedMemorySizeInBytes());
+ assertTrue(
+ memoryManager.getUsedMemorySizeInBytes() <= memoryManager.getTotalMemorySizeInBytes());
+ assertFalse(memoryManager.hasExceededBudget());
+ }
+
+ @SuppressWarnings("unchecked")
+ private static BlockingQueue<IndexedConsensusRequest> pendingEntries(
+ final ConsensusPrefetchingQueue queue) throws Exception {
+ final Field field = ConsensusPrefetchingQueue.class.getDeclaredField("pendingEntries");
+ field.setAccessible(true);
+ return (BlockingQueue<IndexedConsensusRequest>) field.get(queue);
+ }
+
+ private static ConsensusSubscriptionCommitManager newCommitManager(final File systemDir)
+ throws Exception {
+ IoTDBDescriptor.getInstance().getConfig().setSystemDir(systemDir.getAbsolutePath());
+ final Constructor<ConsensusSubscriptionCommitManager> constructor =
+ ConsensusSubscriptionCommitManager.class.getDeclaredConstructor();
+ constructor.setAccessible(true);
+ return constructor.newInstance();
+ }
+
+ private static final class TrackingSubscriptionMemoryManager extends SubscriptionMemoryManager {
+
+ private long maximumUsedMemorySizeInBytes;
+ private boolean exceededBudget;
+
+ private TrackingSubscriptionMemoryManager(final long totalMemorySizeInBytes) {
+ super(totalMemorySizeInBytes);
+ }
+
+ @Override
+ public synchronized boolean tryAllocate(final long sizeInBytes) {
+ final boolean allocated = super.tryAllocate(sizeInBytes);
+ observeUsage();
+ return allocated;
+ }
+
+ @Override
+ public synchronized void release(final long sizeInBytes) {
+ super.release(sizeInBytes);
+ observeUsage();
+ }
+
+ private void observeUsage() {
+ final long usedMemorySizeInBytes = getUsedMemorySizeInBytes();
+ maximumUsedMemorySizeInBytes = Math.max(maximumUsedMemorySizeInBytes, usedMemorySizeInBytes);
+ exceededBudget |= usedMemorySizeInBytes > getTotalMemorySizeInBytes();
+ }
+
+ private long getMaximumUsedMemorySizeInBytes() {
+ return maximumUsedMemorySizeInBytes;
+ }
+
+ private boolean hasExceededBudget() {
+ return exceededBudget;
+ }
+ }
+}
diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/resource/SubscriptionMemoryManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/resource/SubscriptionMemoryManagerTest.java
new file mode 100644
index 00000000000..1635a6e0ecf
--- /dev/null
+++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/resource/SubscriptionMemoryManagerTest.java
@@ -0,0 +1,64 @@
+/*
+ * 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.resource;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+public class SubscriptionMemoryManagerTest {
+
+ @Test
+ public void testAllocateAndReleaseWithinBudget() {
+ final SubscriptionMemoryManager memoryManager = new SubscriptionMemoryManager(10L);
+
+ assertTrue(memoryManager.tryAllocate(6L));
+ assertFalse(memoryManager.tryAllocate(5L));
+ assertEquals(6L, memoryManager.getUsedMemorySizeInBytes());
+ assertEquals(4L, memoryManager.getFreeMemorySizeInBytes());
+
+ memoryManager.release(6L);
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ assertEquals(10L, memoryManager.getFreeMemorySizeInBytes());
+ }
+
+ @Test
+ public void testAllowOnlyOneOversizedEntryWhenBudgetIsEmpty() {
+ final SubscriptionMemoryManager memoryManager = new SubscriptionMemoryManager(10L);
+
+ assertTrue(memoryManager.tryAllocate(11L));
+ assertFalse(memoryManager.tryAllocate(1L));
+ assertEquals(11L, memoryManager.getUsedMemorySizeInBytes());
+
+ memoryManager.release(11L);
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ assertTrue(memoryManager.tryAllocate(10L));
+ }
+
+ @Test
+ public void testZeroBudgetRejectsAllocation() {
+ final SubscriptionMemoryManager memoryManager = new SubscriptionMemoryManager(0L);
+
+ assertFalse(memoryManager.tryAllocate(1L));
+ assertEquals(0L, memoryManager.getUsedMemorySizeInBytes());
+ }
+}
diff --git a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template
index c54878bfe1b..5ccbf54da96 100644
--- a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template
+++ b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template
@@ -1131,10 +1131,12 @@ meta_data_cache_enable=true
# Datatype: boolean
may_cache_nonexist_series=true
-# Read memory Allocation Ratio: BloomFilterCache : ChunkCache : TimeSeriesMetadataCache : Coordinator : Operators : DataExchange : timeIndex in TsFileResourceList : others.
-# The parameter form is a:b:c:d:e:f:g:h, where a, b, c, d, e, f, g and h are integers. for example: 1:1:1:1:1:1:1:1 , 1:100:200:50:200:200:200:50
+# Read memory Allocation Ratio: BloomFilterCache : ChunkCache : TimeSeriesMetadataCache : Coordinator : Operators : DataExchange : timeIndex in TsFileResourceList : others : Subscription.
+# The parameter form is a:b:c:d:e:f:g:h:i, where a, b, c, d, e, f, g, h and i are integers. for example: 1:1:1:1:1:1:1:1:1 , 1:100:200:50:200:200:200:50:250
+# The legacy eight-part form is still accepted and defaults Subscription to 20% of query memory.
+# When subscription is disabled, its memory proportion is 0.
# effectiveMode: restart
-chunk_timeseriesmeta_free_memory_proportion=1:100:200:50:200:200:200:50
+chunk_timeseriesmeta_free_memory_proportion=1:100:200:50:200:200:200:50:250
# Whether to enable LAST cache
# effectiveMode: restart