From 010635c9b6dfde313a917fee058042088d313652 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 24 Jul 2026 13:28:46 +0800 Subject: [PATCH 1/2] [Subscription] Bound consensus prefetch memory usage --- .../iotdb/db/conf/DataNodeMemoryConfig.java | 124 ++- .../memory/QueryEngineMemoryMetrics.java | 26 +- .../consensus/ConsensusPrefetchingQueue.java | 737 ++++++++++++++---- .../SubscriptionDataNodeResourceManager.java | 6 + .../resource/SubscriptionMemoryManager.java | 89 +++ .../db/conf/DataNodeMemoryConfigTest.java | 28 + ...susPrefetchingQueueDataNodeMemoryTest.java | 346 ++++++++ .../ConsensusPrefetchingQueueTest.java | 688 ++++++++++++++++ ...usPrefetchingQueueWalBackpressureTest.java | 336 ++++++++ .../SubscriptionMemoryManagerTest.java | 64 ++ .../conf/iotdb-system.properties.template | 8 +- 11 files changed, 2250 insertions(+), 202 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/resource/SubscriptionMemoryManager.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueDataNodeMemoryTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueWalBackpressureTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/resource/SubscriptionMemoryManagerTest.java 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 73c27a9b805e9..abd7a16765e75 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 @@ 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"; @@ -131,6 +138,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; @@ -464,6 +474,42 @@ private void initCompactionMemoryManager(long compactionMemorySize) { "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(); + } + for (int i = 0; i < proportions.length; i++) { + resolvedProportions[i] = Integer.parseInt(proportions[i].trim()); + if (resolvedProportions[i] < 0) { + throw new IllegalArgumentException(); + } + } + 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(); + } + return resolvedProportions; + } + @SuppressWarnings("squid:S3518") private void initQueryEngineMemoryAllocate( MemoryManager queryEngineMemoryManager, TrimProperties properties) { @@ -494,50 +540,41 @@ private void initQueryEngineMemoryAllocate( 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()) { long sum = @@ -567,6 +604,9 @@ private void initQueryEngineMemoryAllocate( queryEngineMemoryManager.getOrCreateMemoryManager("DataExchange", dataExchangeMemorySize); timeIndexMemoryManager = queryEngineMemoryManager.getOrCreateMemoryManager("TimeIndex", timeIndexMemorySize); + subscriptionMemoryManager = + queryEngineMemoryManager.getOrCreateMemoryManager( + "Subscription", subscriptionMemorySize, true); // must be called after dataExchangeMemoryManager being inited. setQueryThreadCount( @@ -730,6 +770,10 @@ public MemoryManager getTimeIndexMemoryManager() { 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 7138f0daad06c..d336c86bc91f3 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 void bindTo(AbstractMetricService metricService) { 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 void unbindFrom(AbstractMetricService metricService) { 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/consensus/ConsensusPrefetchingQueue.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueue.java index ff04adf344a33..a591ad40db333 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 @@ 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.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.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; @@ -116,6 +120,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. @@ -129,6 +135,26 @@ public class ConsensusPrefetchingQueue { private final Map inFlightEvents; + /** Tablet memory retained by queued and in-flight events, keyed by event identity. */ + private final Map retainedBytesByEvent = new ConcurrentHashMap<>(); + + /** Materialized data events indexed across both queued and in-flight lifecycle stages. */ + private final Map 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(); @@ -143,6 +169,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 activeWriterNodeIds = Collections.emptySet(); @@ -200,6 +231,9 @@ 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; private final Object runtimeActivationLock = new Object(); @@ -271,6 +305,19 @@ protected enum ReplayLocateStatus { 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; @@ -331,15 +378,24 @@ private static final class WakeableIndexedConsensusQueue extends LinkedBlockingDeque { 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(); } @@ -347,9 +403,8 @@ public boolean offer(final IndexedConsensusRequest request) { } @Override - public void put(final IndexedConsensusRequest request) throws InterruptedException { - super.put(request); - wakeupHook.run(); + public synchronized void clear() { + super.clear(); } } @@ -430,6 +485,7 @@ public ConsensusPrefetchingQueue( this.retentionPolicy = retentionPolicy; this.converter = converter; this.commitManager = commitManager; + this.subscriptionMemoryManager = SubscriptionDataNodeResourceManager.memory(); this.fallbackCommittedRegionProgress = fallbackCommittedRegionProgress; this.fallbackTailSearchIndex = tailStartSearchIndex; this.runtimeVersion = initialRuntimeVersion; @@ -445,7 +501,8 @@ public 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( @@ -464,6 +521,16 @@ public 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() { @@ -492,6 +559,24 @@ private void requestPrefetch() { } } + 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; @@ -613,7 +698,13 @@ public SubscriptionEvent poll(final String consumerId, final RegionProgress regi } } - private synchronized boolean initPrefetch(final RegionProgress regionProgress) { + private boolean initPrefetch(final RegionProgress regionProgress) { + synchronized (prefetchInitializationLock) { + return initPrefetchUnderInitializationLock(regionProgress); + } + } + + private boolean initPrefetchUnderInitializationLock(final RegionProgress regionProgress) { if (prefetchInitialized) { return true; // double-check under synchronization } @@ -664,7 +755,7 @@ private synchronized boolean initPrefetch(final RegionProgress regionProgress) { } this.prefetchInitialized = true; this.observedSeekGeneration = seekGeneration.get(); - this.lingerBatch.reset(); + discardBatch(this.lingerBatch); resetBatchWriterProgress(); LOGGER.info( @@ -1054,9 +1145,13 @@ private boolean shouldUseActiveWriterBarriers() { } 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) { @@ -1130,6 +1225,7 @@ private SubscriptionEvent pollInternal(final String consumerId) { .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_POLL_COMMITTED_EVENT_BROKEN_INVARIANT_E478FA3C, this, event); + cleanUpEvent(event, false); continue; } @@ -1140,6 +1236,7 @@ private SubscriptionEvent pollInternal(final String consumerId) { this, event); event.nack(); + prefetchingQueue.add(event); continue; } @@ -1223,6 +1320,7 @@ public SubscriptionEvent pollTablets( // ======================== 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; @@ -1249,6 +1347,7 @@ public PrefetchRoundResult drivePrefetchOnce() { recycleInFlightEvents(); if (!isActive || prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) { + blockRealtimeAdmission(); return computeIdleRoundResult(); } @@ -1258,6 +1357,31 @@ public PrefetchRoundResult drivePrefetchOnce() { 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) { + blockRealtimeAdmission(); + return computeIdleRoundResult(); + } + if (!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 batch = drainPendingEntries(maxWalEntries); if (!batch.isEmpty()) { LOGGER.debug( @@ -1270,22 +1394,38 @@ public PrefetchRoundResult drivePrefetchOnce() { 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( @@ -1323,6 +1463,7 @@ public PrefetchRoundResult drivePrefetchOnce() { fatal.getMessage(), fatal); if (fatal instanceof VirtualMachineError) { + closeCleanupPending = true; markClosed(); return PrefetchRoundResult.dormant(); } @@ -1363,7 +1504,10 @@ private void logPeriodicStatsIfNecessary() { private void resetRoundStateForSeek(final long newSeekGeneration) { restorePendingSubscriptionWalCursor(newSeekGeneration); - lingerBatch.reset(); + discardRealtimeEntries(); + discardBatch(lingerBatch); + memoryBlockedEntryBytes = -1L; + realtimeAdmissionBlocked.set(false); resetBatchWriterProgress(); observedSeekGeneration = newSeekGeneration; } @@ -1371,7 +1515,7 @@ private void resetRoundStateForSeek(final long newSeekGeneration) { private List drainPendingEntries(final int maxWalEntries) { final List 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; @@ -1382,6 +1526,7 @@ private PrefetchRoundResult computeIdleRoundResult() { return PrefetchRoundResult.dormant(); } if (prefetchingQueue.size() >= MAX_PREFETCHING_QUEUE_SIZE) { + blockRealtimeAdmission(); return PrefetchRoundResult.dormant(); } if (hasImmediatePrefetchableWork()) { @@ -1461,7 +1606,7 @@ private void advanceLocalCursorIfPresent(final IndexedConsensusRequest request) } } - private boolean appendRealtimeRequest( + private MaterializationResult appendRealtimeRequest( final IndexedConsensusRequest request, final DeliveryBatchState batchState, final long expectedSeekGeneration, @@ -1470,27 +1615,29 @@ private boolean appendRealtimeRequest( 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 batch, final DeliveryBatchState lingerBatch, final long expectedSeekGeneration, @@ -1514,14 +1661,16 @@ private boolean accumulateFromPending( 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; } } @@ -1541,13 +1690,18 @@ private boolean accumulateFromPending( 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( @@ -1560,7 +1714,7 @@ private boolean accumulateFromPending( lingerBatch.tablets.size(), nextExpectedSearchIndex.get()); - return true; + return MaterializationResult.SUCCESS; } /** @@ -1572,10 +1726,9 @@ private boolean accumulateFromPending( * 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, @@ -1585,18 +1738,20 @@ private boolean fillGapFromWAL( 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(); @@ -1619,41 +1774,53 @@ private boolean fillGapFromWAL( } 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(); @@ -1679,9 +1846,14 @@ private boolean pumpFromSubscriptionWAL( 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); @@ -1703,7 +1875,7 @@ private boolean pumpFromSubscriptionWAL( entriesRead, nextExpectedSearchIndex.get()); } - return true; + return MaterializationResult.SUCCESS; } private void ensureSubscriptionWalReadable() { @@ -1797,6 +1969,7 @@ private PreparedEntry prepareEntry(final IndexedConsensusRequest indexedRequest) final InsertNode insertNode = ConsensusLogToTabletConverter.deserializeToInsertNode(indexedRequest); if (Objects.isNull(insertNode)) { + memoryBlockedEntryBytes = -1L; return null; } @@ -1818,24 +1991,24 @@ private PreparedEntry prepareEntry(final IndexedConsensusRequest indexedRequest) maxObservedTimestamp = maxTs; } final List 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 tablets, final long startSearchIndex, final long endSearchIndex) { - createAndEnqueueEvent( - tablets, startSearchIndex, endSearchIndex, endSearchIndex, seekGeneration.get()); - } - private boolean createAndEnqueueEvent( final List tablets, final long startSearchIndex, final long endSearchIndex, final long commitLocalSeq, + final long retainedBytes, final long expectedSeekGeneration) { if (seekGeneration.get() != expectedSeekGeneration) { LOGGER.debug( @@ -1874,6 +2047,8 @@ private boolean createAndEnqueueEvent( 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( @@ -1947,6 +2122,173 @@ private long estimateTabletsBytes(final List tablets) { 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 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, @@ -1954,7 +2296,8 @@ private boolean appendPreparedEntryViaRealtimeWriter( final int maxTablets, final long maxBatchBytes) { bufferRealtimeEntry(preparedEntry); - return drainRealtimeWriters(batchState, expectedSeekGeneration, maxTablets, maxBatchBytes); + return drainBufferedRealtimeWriters( + batchState, expectedSeekGeneration, maxTablets, maxBatchBytes); } private int getRealtimeBufferedEntryCount() { @@ -2069,7 +2412,7 @@ private boolean drainWriterEntries( continue; } - final long entryEstimatedBytes = estimateTabletsBytes(writerHead.getTablets()); + final long entryEstimatedBytes = writerHead.getEstimatedBytes(); if (!canAppendWriterEntry( batchState, writerHead, entryEstimatedBytes, maxEntries, maxTablets, maxBatchBytes)) { return true; @@ -2095,6 +2438,7 @@ private boolean commitEmptyWriterEntry( entry.getSearchIndex(), entry.getSearchIndex(), entry.getLocalSeq(), + 0L, expectedSeekGeneration); resetBatchWriterProgress(); return committed; @@ -2108,6 +2452,7 @@ private boolean flushBatch( batchState.startSearchIndex, batchState.endSearchIndex, batchState.lastLocalSeq, + batchState.estimatedBytes, expectedSeekGeneration)) { return false; } @@ -2144,7 +2489,7 @@ private boolean refreshInFlightEventLeaseInternal( return null; } if (ev.isCommitted()) { - ev.cleanUp(false); + cleanUpEvent(ev, false); return null; } ev.recordLastPolledTimestamp(); @@ -2177,6 +2522,9 @@ private boolean canAcceptCommitContext( } return false; } + if (isCommitContextOutdated(commitContext)) { + return false; + } if (!isActive) { if (silent) { LOGGER.debug( @@ -2201,49 +2549,51 @@ private boolean canAcceptCommitContext( } 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; } @@ -2255,22 +2605,25 @@ private boolean ackInternal( 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) { @@ -2288,52 +2641,24 @@ public boolean nack(final String consumerId, final SubscriptionCommitContext com * 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) { @@ -2375,7 +2700,7 @@ public boolean nackSilent( ev); ev.ack(); ev.recordCommittedTimestamp(); - ev.cleanUp(false); + cleanUpEvent(ev, false); return null; } prefetchingQueue.add(ev); @@ -2413,7 +2738,7 @@ private boolean nackInternal( ev); ev.ack(); ev.recordCommittedTimestamp(); - ev.cleanUp(false); + cleanUpEvent(ev, false); return null; } prefetchingQueue.add(ev); @@ -2435,7 +2760,7 @@ private void recycleInFlightEvents() { return null; } if (ev.isCommitted()) { - ev.cleanUp(false); + cleanUpEvent(ev, false); return null; } if (ev.pollable()) { @@ -2449,7 +2774,7 @@ private void recycleInFlightEvents() { ev); ev.ack(); ev.recordCommittedTimestamp(); - ev.cleanUp(false); + cleanUpEvent(ev, false); return null; } prefetchingQueue.add(ev); @@ -2468,20 +2793,26 @@ private void recycleInFlightEvents() { // ======================== 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; @@ -2747,17 +3078,21 @@ private void failPendingSeekBeforeScheduling(final PendingSeekRequest request) { } 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(); @@ -2771,7 +3106,8 @@ private void applySeekResetUnderWriteLock(final PendingSeekRequest request) { // first replayable searchIndex here. nextExpectedSearchIndex.set(request.targetSearchIndex); requestSubscriptionWalReset(request.targetSearchIndex, seekGeneration.get()); - lingerBatch.reset(); + discardBatch(lingerBatch); + reconcileRetainedTabletMemoryAfterCleanup(); resetBatchWriterProgress(); observedSeekGeneration = seekGeneration.get(); pendingWalGapRetryRequested = false; @@ -2999,10 +3335,11 @@ public void close() { acquireWriteLock(); try { - if (isClosed || closeRequested) { + if (closeRequested || (isClosed && !closeCleanupPending)) { return; } closeRequested = true; + closeCleanupPending = true; seekRequestToFail = pendingSeekRequest; pendingSeekRequest = null; } finally { @@ -3060,6 +3397,7 @@ public void close() { // Persist progress before closing commitManager.persistAll(); } + closeCleanupPending = false; } finally { closeRequested = false; } @@ -3090,7 +3428,7 @@ private void flushLingeringBatchOnCloseUnderWriteLock() { DataNodePipeMessages .PIPE_LOG_CONSENSUSPREFETCHINGQUEUE_FAILED_TO_FLUSH_LINGERING_BATCH_F97D8AA7, this); - lingerBatch.reset(); + discardBatch(lingerBatch); resetBatchWriterProgress(); } } @@ -3166,12 +3504,71 @@ public long getEpochChangeCount() { * 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(); } @@ -3369,6 +3766,14 @@ public long getSubscriptionUncommittedEventCount() { 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(); @@ -3433,6 +3838,12 @@ public Map coreReportMessage() { 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())); @@ -3490,6 +3901,8 @@ public int hashCode() { private interface WriterBufferedEntry { List getTablets(); + long getEstimatedBytes(); + long getPhysicalTime(); int getWriterNodeId(); @@ -3569,6 +3982,7 @@ private static final class WriterState { private static final class PreparedEntry implements WriterBufferedEntry { private final List tablets; + private final long estimatedBytes; private final long physicalTime; private final int writerNodeId; private final long localSeq; @@ -3576,11 +3990,13 @@ private static final class PreparedEntry implements WriterBufferedEntry { private PreparedEntry( final List 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; @@ -3592,6 +4008,11 @@ public List getTablets() { 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 347299df10f1a..bf1a7de30f7a9 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 @@ 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 0000000000000..bf71fad6c4b72 --- /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. + * + *

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 7d227abcd0f2b..ae79ab7072ae9 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 @@ -24,11 +24,39 @@ 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.assertTrue; 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(expected = IllegalArgumentException.class) + public void testRejectInvalidSubscriptionQueryMemoryProportions() { + DataNodeMemoryConfig.resolveQueryMemoryProportions("1:100:200:50:200:200:200", true); + } + + @Test(expected = IllegalArgumentException.class) + public void testRejectNegativeSubscriptionQueryMemoryProportion() { + DataNodeMemoryConfig.resolveQueryMemoryProportions("1:100:200:50:200:200:200:50:-1", true); + } + @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 0000000000000..60f30834cad0a --- /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 pendingEntries( + final ConsensusPrefetchingQueue queue) throws Exception { + final Field field = ConsensusPrefetchingQueue.class.getDeclaredField("pendingEntries"); + field.setAccessible(true); + return (BlockingQueue) field.get(queue); + } + + private static Tablet createTablet() { + final List columnNames = Arrays.asList("device", "temperature"); + final List dataTypes = Arrays.asList(TSDataType.STRING, TSDataType.DOUBLE); + final List 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 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 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 4902554dc9043..1c8643f2a7024 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 @@ -36,6 +36,7 @@ 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; @@ -53,12 +54,18 @@ import java.io.IOException; 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; @@ -75,6 +82,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 offered = new AtomicReference<>(); + final AtomicReference 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 queue = + (BlockingQueue) + 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(); @@ -749,6 +857,550 @@ public void testActivationInstallsRuntimeStateBeforeRefreshingAuthoritativeProgr } } + @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 emptyRequest = + new IndexedConsensusRequest( + 2L, Collections.singletonList(StatementTestUtils.genInsertRowNode(2))) + .setPhysicalTime(1001L) + .setNodeId(7); + reader.currentSearchIndex = 2L; + pendingEntries(queue).offer(dataRequest); + pendingEntries(queue).offer(emptyRequest); + + 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; + 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 < 20; 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 < 20; 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 pendingEntries( final ConsensusPrefetchingQueue queue) throws Exception { @@ -769,6 +1421,42 @@ private static Tablet createTablet() { 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 0000000000000..b3783beab43dd --- /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 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 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 columnNames = Arrays.asList("device", "temperature"); + final List dataTypes = Arrays.asList(TSDataType.STRING, TSDataType.DOUBLE); + final List 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 pendingEntries( + final ConsensusPrefetchingQueue queue) throws Exception { + final Field field = ConsensusPrefetchingQueue.class.getDeclaredField("pendingEntries"); + field.setAccessible(true); + return (BlockingQueue) field.get(queue); + } + + private static ConsensusSubscriptionCommitManager newCommitManager(final File systemDir) + throws Exception { + IoTDBDescriptor.getInstance().getConfig().setSystemDir(systemDir.getAbsolutePath()); + final Constructor 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 0000000000000..1635a6e0ecf83 --- /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 02dcb1b3b5de7..c4a06b68b1497 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 @@ -1156,10 +1156,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 From d41d8bc013d339bf3b50be0b3e37b4bdf56d6a80 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 24 Jul 2026 19:15:41 +0800 Subject: [PATCH 2/2] [Subscription] Accept fenced consensus commits as no-ops --- .../broker/ConsensusSubscriptionBroker.java | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/ConsensusSubscriptionBroker.java index 32598129e1866..fb739aef34900 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 List commit( 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) {