This is an automated email from the ASF dual-hosted git repository. Caideyipi pushed a commit to branch fix/subscription-wal-backed-inflight in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit d1c70f5ba78dc2d5d8b6fbc8b392d76de7f90841 Author: Caideyipi <[email protected]> AuthorDate: Wed Oct 14 12:40:36 2026 +0800 [Subscription] Replay uncommitted delivery payloads from WAL --- .../apache/iotdb/db/i18n/DataNodePipeMessages.java | 2 + .../apache/iotdb/db/i18n/DataNodePipeMessages.java | 2 + .../consensus/ConsensusLogToTabletConverter.java | 28 + .../consensus/ConsensusPrefetchingQueue.java | 393 ++++++++++++- .../db/subscription/event/SubscriptionEvent.java | 70 ++- .../cache/CachedSubscriptionPollResponse.java | 1 + .../response/SubscriptionEventSingleResponse.java | 15 +- .../receiver/SubscriptionReceiverV1.java | 7 +- .../ConsensusLogToTabletConverterTest.java | 25 + .../ConsensusPrefetchingQueueDeliveryTest.java | 627 +++++++++++++++++++++ 10 files changed, 1157 insertions(+), 13 deletions(-) diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java index 836ac88d84f..bec917531be 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java @@ -2704,4 +2704,6 @@ public final class DataNodePipeMessages { "Failed to read OPC UA server operation limits, use defaults: maxNodesPerWrite={}, maxNodesPerNodeManagement={}"; public static final String MESSAGE_ARG_SUBSCRIPTION_ENTRY_REQUIRES_ARG_BYTES_ABOVE_THE_CURRENT_PER_QUEUE_MAXIMUM_ARG_BYTES_DATANODE_BUDGET_ARG_BYTES_REDUCE_THE_WRITE_BATCH_FIELD_SIZE_OR_INCREASE_SUBSCRIPTION_MATERIALIZATION_MEMORY_WAL_PROGRESS_HAS_NOT_ADVANCED_AFCBC7FC = "[%s] Subscription entry requires %d bytes, above the current per-queue maximum %d bytes (DataNode budget %d bytes). Reduce the write batch/field size or increase subscription materialization memory. WAL progress has not advanced."; + public static final String LOG_SUBSCRIPTION_FAILED_TO_REPLAY_DELIVERY_BATCH_ARG_FROM_WAL_RETAINING_UNCOMMITTED_STATE_0654EACD = + "Subscription failed to replay delivery batch {} from WAL; retaining uncommitted state"; } diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java index 5840da12daf..5b300e79b4b 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java @@ -2529,4 +2529,6 @@ public final class DataNodePipeMessages { "读取 OPC UA 服务器操作限制失败，使用默认值：maxNodesPerWrite={}，maxNodesPerNodeManagement={}"; public static final String MESSAGE_ARG_SUBSCRIPTION_ENTRY_REQUIRES_ARG_BYTES_ABOVE_THE_CURRENT_PER_QUEUE_MAXIMUM_ARG_BYTES_DATANODE_BUDGET_ARG_BYTES_REDUCE_THE_WRITE_BATCH_FIELD_SIZE_OR_INCREASE_SUBSCRIPTION_MATERIALIZATION_MEMORY_WAL_PROGRESS_HAS_NOT_ADVANCED_AFCBC7FC = "[%s] 订阅 entry 需要 %d 字节，超过当前队列可用上限 %d 字节（DataNode 预算 %d 字节）。请减小写入 batch/字段大小或增加订阅物化内存。WAL 进度未推进。"; + public static final String LOG_SUBSCRIPTION_FAILED_TO_REPLAY_DELIVERY_BATCH_ARG_FROM_WAL_RETAINING_UNCOMMITTED_STATE_0654EACD = + "Subscription 无法从 WAL 重放交付批次 {}；保留未提交状态"; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusLogToTabletConverter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusLogToTabletConverter.java index ab80b997061..8fd79a6a08b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusLogToTabletConverter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusLogToTabletConverter.java @@ -104,6 +104,18 @@ public class ConsensusLogToTabletConverter { return databaseName; } + /** Freezes the column projection used by a delivered batch for later WAL replay. */ + ConsensusLogToTabletConverter snapshotForReplay() { + return new ConsensusLogToTabletConverter( + treePattern, tablePattern, getColumnFilterMatcher(), databaseName); + } + + boolean hasSameReplayProjection(final ConsensusLogToTabletConverter other) { + return treePattern == other.treePattern + && tablePattern == other.tablePattern + && fallbackColumnFilterMatcher == other.fallbackColumnFilterMatcher; + } + boolean isTableModel() { return Objects.nonNull(tablePattern); } @@ -126,6 +138,16 @@ public class ConsensusLogToTabletConverter { * DataRegionStateMachine.grabPlanNode()}. */ static InsertNode deserializeToInsertNode(final IndexedConsensusRequest indexedRequest) { + return deserializeToInsertNode(indexedRequest, false); + } + + /** A payload can be discarded only after every fragment, including the seal, is readable. */ + static InsertNode deserializeCompleteInsertNode(final IndexedConsensusRequest indexedRequest) { + return deserializeToInsertNode(indexedRequest, true); + } + + private static InsertNode deserializeToInsertNode( + final IndexedConsensusRequest indexedRequest, final boolean requireComplete) { final List<SearchNode> searchNodes = new ArrayList<>(); PlanNode nonSearchNode = null; @@ -151,6 +173,9 @@ public class ConsensusLogToTabletConverter { indexedRequest.getSearchIndex(), e.getMessage(), e); + if (requireComplete) { + return null; + } continue; } @@ -173,6 +198,9 @@ public class ConsensusLogToTabletConverter { } if (!searchNodes.isEmpty()) { + if (requireComplete && !searchNodes.get(searchNodes.size() - 1).isLastFragment()) { + return null; + } final PlanNode merged = searchNodes.get(0).merge(searchNodes); if (merged instanceof InsertNode) { final InsertNode mergedInsert = (InsertNode) merged; 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 63c06f06960..db4aa2a96b9 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 @@ -61,6 +61,7 @@ import org.apache.iotdb.rpc.subscription.config.TopicConstant; import org.apache.iotdb.rpc.subscription.payload.poll.ErrorPayload; import org.apache.iotdb.rpc.subscription.payload.poll.RegionProgress; import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponse; import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponseType; import org.apache.iotdb.rpc.subscription.payload.poll.TabletsPayload; import org.apache.iotdb.rpc.subscription.payload.poll.WatermarkPayload; @@ -72,8 +73,10 @@ import org.apache.tsfile.write.record.Tablet; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.io.ByteArrayOutputStream; import java.io.File; import java.io.IOException; +import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; @@ -88,6 +91,7 @@ import java.util.PriorityQueue; import java.util.Set; import java.util.TreeMap; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentSkipListMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.LinkedBlockingDeque; import java.util.concurrent.PriorityBlockingQueue; @@ -195,6 +199,15 @@ public class ConsensusPrefetchingQueue { private final Map<SubscriptionCommitContext, SubscriptionEvent> outstandingEventsByCommitContext = new ConcurrentHashMap<>(); + /** + * Compact batch boundaries and source positions; no Tablet or serialized response is retained. + */ + private final Map<SubscriptionEvent, DeliveryBatchRecord> deliveryBatches = + new ConcurrentHashMap<>(); + + private final ConcurrentSkipListMap<Long, Integer> deliveryStartIndices = + new ConcurrentSkipListMap<>(); + /** Tablet memory retained by this queue across all materialized lifecycle stages. */ private final AtomicLong retainedTabletBytes = new AtomicLong(0L); @@ -1462,6 +1475,25 @@ public class ConsensusPrefetchingQueue { continue; } + synchronized (event) { + // ACK may have removed a recycled event while it was being claimed from the queue. + if (event.isCommitted() + || (event.getCommitContext().hasWriterProgress() + && outstandingEventsByCommitContext.get(event.getCommitContext()) != event)) { + continue; + } + if (!restorePayloadForDelivery(event)) { + prefetchingQueue.add(event); + continue; + } + final DeliveryBatchRecord batchRecord = deliveryBatches.get(event); + if (Objects.nonNull(batchRecord)) { + batchRecord.delivered = false; + // A fresh response identity fences a late delivery callback from an earlier poll. + event.restorePayload(event.getCurrentResponse().getPayload()); + } + } + // Mark as polled before updating inFlightEvents event.recordLastPolledTimestamp(); timeTracker.recordDelivery(); @@ -1638,6 +1670,7 @@ public class ConsensusPrefetchingQueue { } applyPendingSubscriptionWalReset(observedSeekGeneration); + reclaimDeliveredPayloads(); recycleInFlightEvents(); if (!isActive) { @@ -2450,7 +2483,10 @@ public class ConsensusPrefetchingQueue { if (maxTs > maxObservedTimestamp) { maxObservedTimestamp = maxTs; } - final List<Tablet> tablets = converter.convert(insertNode); + final ConsensusLogToTabletConverter snapshot = converter.snapshotForReplay(); + final ConsensusLogToTabletConverter replayConverter = + Objects.nonNull(snapshot) ? snapshot : converter; + final List<Tablet> tablets = replayConverter.convert(insertNode); final long estimatedBytes = estimateTabletsBytes(tablets); if (debugEnabled) { LOGGER.debug( @@ -2468,7 +2504,13 @@ public class ConsensusPrefetchingQueue { return null; } return new PreparedEntry( - tablets, estimatedBytes, physicalTime, writerNodeId, localSeq, searchIndex); + tablets, + estimatedBytes, + physicalTime, + writerNodeId, + localSeq, + searchIndex, + replayConverter); } private static long estimateTabletSize(final Tablet tablet) { @@ -2482,6 +2524,24 @@ public class ConsensusPrefetchingQueue { final long commitLocalSeq, final long retainedBytes, final long expectedSeekGeneration) { + return createAndEnqueueEvent( + tablets, + startSearchIndex, + endSearchIndex, + commitLocalSeq, + retainedBytes, + expectedSeekGeneration, + null); + } + + private boolean createAndEnqueueEvent( + final List<Tablet> tablets, + final long startSearchIndex, + final long endSearchIndex, + final long commitLocalSeq, + final long retainedBytes, + final long expectedSeekGeneration, + final DeliveryBatchRecord batchRecord) { if (closeRequested || isClosed) { return false; } @@ -2525,6 +2585,11 @@ public class ConsensusPrefetchingQueue { // Install the ownership record before exposing the event to concurrent poll/ack threads. retainEventMemory(event, retainedBytes); + if (Objects.nonNull(batchRecord)) { + deliveryBatches.put(event, batchRecord); + deliveryStartIndices.merge(batchRecord.minimumSearchIndex, 1, Integer::sum); + event.setResponseDeliveryListener(response -> onResponseDelivered(event, response)); + } prefetchingQueue.add(event); LOGGER.debug( @@ -2696,11 +2761,161 @@ public class ConsensusPrefetchingQueue { } } + private void onResponseDelivered( + final SubscriptionEvent event, final SubscriptionPollResponse response) { + acquireReadLock(); + try { + synchronized (event) { + final DeliveryBatchRecord record = deliveryBatches.get(event); + if (isClosed + || closeRequested + || pendingSeekRequest != null + || Objects.isNull(record) + || event.isCommitted() + || event.getCurrentResponse() != response + || outstandingEventsByCommitContext.get(event.getCommitContext()) != event) { + return; + } + record.delivered = true; + reclaimDeliveredPayload(event, record); + } + } finally { + releaseReadLock(); + } + requestPrefetch(); + } + + private void reclaimDeliveredPayloads() { + // Only materialized events need a check. The number of these is bounded by the memory budget. + for (final SubscriptionEvent event : retainedBytesByEvent.keySet()) { + synchronized (event) { + final DeliveryBatchRecord record = deliveryBatches.get(event); + if (Objects.nonNull(record) && record.delivered && !event.isCommitted()) { + reclaimDeliveredPayload(event, record); + } + } + } + } + + private void reclaimDeliveredPayload( + final SubscriptionEvent event, final DeliveryBatchRecord record) { + if (event.isPayloadDetached() || !record.delivered) { + return; + } + // Pending entries can precede durable WAL visibility. Verify every member and fragment before + // releasing the only payload copy. This read does not change replay or committed progress. + if (!record.replayable && Objects.isNull(readDeliveryBatch(event, record, false))) { + return; + } + record.replayable = true; + if (event.detachCurrentPayload(event.getCurrentResponse())) { + final Long retainedBytes = retainedBytesByEvent.remove(event); + if (Objects.nonNull(retainedBytes)) { + releaseTabletMemory(retainedBytes); + } + } + } + + private boolean restorePayloadForDelivery(final SubscriptionEvent event) { + if (!event.isPayloadDetached()) { + return true; + } + final DeliveryBatchRecord record = deliveryBatches.get(event); + if (Objects.isNull(record) || !tryReserveTabletMemory(record.estimatedBytes)) { + return false; + } + boolean restored = false; + try { + final List<Tablet> tablets = readDeliveryBatch(event, record, true); + if (Objects.isNull(tablets) + || tablets.size() != record.tabletCount + || estimateTabletsBytes(tablets) > record.estimatedBytes) { + return false; + } + event.restorePayload( + new TabletsPayload( + Collections.singletonMap(record.databaseName, tablets), -record.tabletCount)); + retainedBytesByEvent.put(event, record.estimatedBytes); + restored = true; + return true; + } finally { + if (!restored) { + releaseTabletMemory(record.estimatedBytes); + } + } + } + + /** Reads exactly the recorded members in delivery order, including non-local writer records. */ + private List<Tablet> readDeliveryBatch( + final SubscriptionEvent event, final DeliveryBatchRecord record, final boolean materialize) { + final List<ReplayEntryPosition> positions = record.decodePositions(); + final Map<WriterProgress, Integer> members = new HashMap<>(); + for (int index = 0; index < positions.size(); index++) { + final ReplayEntryPosition position = positions.get(index); + members.put(new WriterProgress(position.physicalTime, position.localSeq), index); + } + final List<List<Tablet>> tabletsByMember = new ArrayList<>(positions.size()); + for (int index = 0; index < positions.size(); index++) { + tabletsByMember.add(null); + } + final ProgressWALIterator iterator = createSubscriptionWALIterator(record.minimumSearchIndex); + if (Objects.isNull(iterator)) { + return null; + } + try (final ProgressWALIterator ignored = iterator) { + while (!members.isEmpty() && iterator.hasNext()) { + if (closeRequested || isClosed || isCommitContextOutdated(event.getCommitContext())) { + return null; + } + final IndexedConsensusRequest request = iterator.next(); + if (request.getNodeId() != event.getCommitContext().getWriterId().getNodeId()) { + continue; + } + final WriterProgress progress = toWriterProgress(request); + final Integer memberIndex = members.get(progress); + if (Objects.isNull(memberIndex)) { + continue; + } + final ReplayEntryPosition position = positions.get(memberIndex); + if (request.getPhysicalTime() != position.physicalTime + || request.getSearchIndex() != position.searchIndex) { + continue; + } + final InsertNode node = + ConsensusLogToTabletConverter.deserializeCompleteInsertNode(request); + if (Objects.isNull(node)) { + return null; + } + tabletsByMember.set( + memberIndex, materialize ? record.converter.convert(node) : Collections.emptyList()); + members.remove(progress); + } + if (!members.isEmpty() || iterator.hasIncompleteScan()) { + return null; + } + final List<Tablet> tablets = new ArrayList<>(); + tabletsByMember.forEach(tablets::addAll); + return tablets; + } catch (final IOException | RuntimeException e) { + LOGGER.warn( + DataNodePipeMessages + .LOG_SUBSCRIPTION_FAILED_TO_REPLAY_DELIVERY_BATCH_ARG_FROM_WAL_RETAINING_UNCOMMITTED_STATE_0654EACD, + event.getCommitContext(), + e); + return null; + } + } + 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); + final DeliveryBatchRecord batchRecord = deliveryBatches.remove(event); + if (Objects.nonNull(batchRecord)) { + deliveryStartIndices.computeIfPresent( + batchRecord.minimumSearchIndex, (index, count) -> count == 1 ? null : count - 1); + } // 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 @@ -2796,6 +3011,8 @@ public class ConsensusPrefetchingQueue { } retainedBytesByEvent.clear(); outstandingEventsByCommitContext.clear(); + deliveryBatches.clear(); + deliveryStartIndices.clear(); releaseTabletMemory(retainedBytes); } @@ -2862,10 +3079,17 @@ public class ConsensusPrefetchingQueue { // SubscriptionCommitContext and RegionProgress remain per-writer. final boolean writerChanged = !batchState.isEmpty() && batchState.writerNodeId != entry.getWriterNodeId(); + final boolean projectionChanged = + !batchState.isEmpty() + && Objects.nonNull(batchState.replayConverter) + && Objects.nonNull(entry.getReplayConverter()) + && batchState.replayConverter != entry.getReplayConverter() + && !batchState.replayConverter.hasSameReplayProjection(entry.getReplayConverter()); return !(wouldExceedEntryLimit || wouldExceedTabletLimit || wouldExceedByteLimit - || writerChanged); + || writerChanged + || projectionChanged); } private boolean drainRealtimeWriters( @@ -2963,7 +3187,8 @@ public class ConsensusPrefetchingQueue { batchState.endSearchIndex, batchState.lastLocalSeq, batchState.estimatedBytes, - expectedSeekGeneration)) { + expectedSeekGeneration, + batchState.createDeliveryRecord(converter))) { return false; } resetBatchWriterProgress(); @@ -3122,6 +3347,7 @@ public class ConsensusPrefetchingQueue { || (result == InFlightAckResult.MISSING && ackMissingInFlightEvent(commitContext, false)); if (acked) { + serverImpl.checkAndUpdateSafeDeletedSearchIndex(); requestPrefetch(); } return acked; @@ -3209,6 +3435,7 @@ public class ConsensusPrefetchingQueue { || (result == InFlightAckResult.MISSING && ackMissingInFlightEvent(commitContext, true)); if (ackedResult) { + serverImpl.checkAndUpdateSafeDeletedSearchIndex(); requestPrefetch(); } return ackedResult; @@ -3719,7 +3946,20 @@ public class ConsensusPrefetchingQueue { // Committed progress can advance on another consumer or replica while this queue is still // replaying an older local search index. Keep both boundaries so WAL deletion can never pass // the earliest entry this queue has not inspected yet. - return Math.min(getCommittedRetainedMinVersionId(), getReplayRetainedMinVersionId()); + long requiredVersion = + Math.min(getCommittedRetainedMinVersionId(), getReplayRetainedMinVersionId()); + // A replica's committed frontier may advance independently. Local delivery records remain + // replayable until their lifecycle is cleaned, even after another replica reports progress. + if (consensusReqReader instanceof WALNode) { + final Map.Entry<Long, Integer> firstDelivery = deliveryStartIndices.firstEntry(); + if (Objects.nonNull(firstDelivery)) { + final File[] walFiles = ((WALNode) consensusReqReader).getSortedWalFilesSnapshot(); + requiredVersion = + Math.min( + requiredVersion, findReplayRetainedMinVersionId(walFiles, firstDelivery.getKey())); + } + } + return requiredVersion; } private static boolean isProgressAtLeast( @@ -5012,12 +5252,22 @@ public class ConsensusPrefetchingQueue { long getSearchIndex(); + ConsensusLogToTabletConverter getReplayConverter(); + OrderingKey getOrderingKey(); } private static final class DeliveryBatchState { private final List<Tablet> tablets = new ArrayList<>(); + private final ByteArrayOutputStream replayPositions = new ByteArrayOutputStream(); + private long previousSearchIndex; + private long previousPhysicalTime; + private long previousLocalSeq; + private long minimumSearchIndex; + private boolean hasNonLocalEntry; + private boolean replayEligible; + private ConsensusLogToTabletConverter replayConverter; private long startSearchIndex; private long endSearchIndex; private long estimatedBytes; @@ -5048,6 +5298,7 @@ public class ConsensusPrefetchingQueue { firstTabletTimeMs = System.currentTimeMillis(); } writerNodeId = entry.getWriterNodeId(); + replayConverter = entry.getReplayConverter(); } if (entry.getSearchIndex() >= 0) { if (startSearchIndex < 0) { @@ -5056,6 +5307,17 @@ public class ConsensusPrefetchingQueue { endSearchIndex = entry.getSearchIndex(); } tablets.addAll(entry.getTablets()); + DeliveryBatchRecord.writeDelta(replayPositions, entry.getSearchIndex() - previousSearchIndex); + DeliveryBatchRecord.writeDelta( + replayPositions, entry.getPhysicalTime() - previousPhysicalTime); + DeliveryBatchRecord.writeDelta(replayPositions, entry.getLocalSeq() - previousLocalSeq); + previousSearchIndex = entry.getSearchIndex(); + previousPhysicalTime = entry.getPhysicalTime(); + previousLocalSeq = entry.getLocalSeq(); + minimumSearchIndex = Math.min(minimumSearchIndex, entry.getSearchIndex()); + hasNonLocalEntry |= entry.getSearchIndex() < 0; + replayEligible &= + entry.getPhysicalTime() >= 0 && entry.getLocalSeq() >= 0 && entry.getWriterNodeId() >= 0; estimatedBytes += entryEstimatedBytes; physicalTime = entry.getPhysicalTime(); lastLocalSeq = entry.getLocalSeq(); @@ -5063,8 +5325,31 @@ public class ConsensusPrefetchingQueue { entryCount++; } + private DeliveryBatchRecord createDeliveryRecord( + final ConsensusLogToTabletConverter converter) { + if (!replayEligible || entryCount == 0) { + return null; + } + return new DeliveryBatchRecord( + replayPositions.toByteArray(), + entryCount, + hasNonLocalEntry ? Long.MIN_VALUE : minimumSearchIndex, + tablets.size(), + estimatedBytes, + converter.getDatabaseName(), + Objects.nonNull(replayConverter) ? replayConverter : converter); + } + private void reset() { tablets.clear(); + replayPositions.reset(); + previousSearchIndex = 0L; + previousPhysicalTime = 0L; + previousLocalSeq = 0L; + minimumSearchIndex = Long.MAX_VALUE; + hasNonLocalEntry = false; + replayEligible = true; + replayConverter = null; startSearchIndex = -1L; endSearchIndex = -1L; estimatedBytes = 0L; @@ -5076,6 +5361,86 @@ public class ConsensusPrefetchingQueue { } } + private static final class ReplayEntryPosition { + private final long searchIndex; + private final long physicalTime; + private final long localSeq; + + private ReplayEntryPosition( + final long searchIndex, final long physicalTime, final long localSeq) { + this.searchIndex = searchIndex; + this.physicalTime = physicalTime; + this.localSeq = localSeq; + } + } + + private static final class DeliveryBatchRecord { + // Delta + zigzag varints keep contiguous indices/sequences and timestamp deltas compact. + // Positions are decoded only while verifying or materializing one memory-bounded batch. + private final byte[] positions; + private final int entryCount; + private final long minimumSearchIndex; + private final int tabletCount; + private final long estimatedBytes; + private final String databaseName; + private final ConsensusLogToTabletConverter converter; + private boolean delivered; + private boolean replayable; + + private DeliveryBatchRecord( + final byte[] positions, + final int entryCount, + final long minimumSearchIndex, + final int tabletCount, + final long estimatedBytes, + final String databaseName, + final ConsensusLogToTabletConverter converter) { + this.positions = positions; + this.entryCount = entryCount; + this.minimumSearchIndex = minimumSearchIndex; + this.tabletCount = tabletCount; + this.estimatedBytes = estimatedBytes; + this.databaseName = databaseName; + this.converter = converter; + } + + private List<ReplayEntryPosition> decodePositions() { + final ByteBuffer buffer = ByteBuffer.wrap(positions); + final List<ReplayEntryPosition> result = new ArrayList<>(entryCount); + long searchIndex = 0L; + long physicalTime = 0L; + long localSeq = 0L; + for (int index = 0; index < entryCount; index++) { + searchIndex += readDelta(buffer); + physicalTime += readDelta(buffer); + localSeq += readDelta(buffer); + result.add(new ReplayEntryPosition(searchIndex, physicalTime, localSeq)); + } + return result; + } + + private static void writeDelta(final ByteArrayOutputStream stream, final long delta) { + long encoded = (delta << 1) ^ (delta >> 63); + while ((encoded & ~0x7FL) != 0L) { + stream.write((int) (encoded & 0x7F) | 0x80); + encoded >>>= 7; + } + stream.write((int) encoded); + } + + private static long readDelta(final ByteBuffer buffer) { + long encoded = 0L; + int shift = 0; + byte value; + do { + value = buffer.get(); + encoded |= (long) (value & 0x7F) << shift; + shift += 7; + } while (value < 0); + return (encoded >>> 1) ^ -(encoded & 1L); + } + } + private static final class WriterState { private long effectiveSafePt = 0L; private boolean closed = false; @@ -5088,6 +5453,7 @@ public class ConsensusPrefetchingQueue { private final int writerNodeId; private final long localSeq; private final long searchIndex; + private final ConsensusLogToTabletConverter replayConverter; private PreparedEntry( final List<Tablet> tablets, @@ -5096,12 +5462,24 @@ public class ConsensusPrefetchingQueue { final int writerNodeId, final long localSeq, final long searchIndex) { + this(tablets, estimatedBytes, physicalTime, writerNodeId, localSeq, searchIndex, null); + } + + private PreparedEntry( + final List<Tablet> tablets, + final long estimatedBytes, + final long physicalTime, + final int writerNodeId, + final long localSeq, + final long searchIndex, + final ConsensusLogToTabletConverter replayConverter) { this.tablets = tablets; this.estimatedBytes = estimatedBytes; this.physicalTime = physicalTime; this.writerNodeId = writerNodeId; this.localSeq = localSeq; this.searchIndex = searchIndex; + this.replayConverter = replayConverter; } @Override @@ -5134,6 +5512,11 @@ public class ConsensusPrefetchingQueue { return searchIndex; } + @Override + public ConsensusLogToTabletConverter getReplayConverter() { + return replayConverter; + } + @Override public OrderingKey getOrderingKey() { return new OrderingKey(physicalTime, writerNodeId, localSeq); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java index 2e58a68411d..6457a61811d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/SubscriptionEvent.java @@ -35,6 +35,7 @@ import org.apache.iotdb.db.subscription.event.response.SubscriptionEventTsFileRe import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionCommitContext; import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollPayload; import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponse; +import org.apache.iotdb.rpc.subscription.payload.poll.TabletsPayload; import org.apache.thrift.annotation.Nullable; import org.slf4j.Logger; @@ -43,10 +44,12 @@ import org.slf4j.LoggerFactory; import java.io.File; import java.io.IOException; import java.nio.ByteBuffer; +import java.util.Collections; import java.util.Map; import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Consumer; import static com.google.common.base.MoreObjects.toStringHelper; @@ -60,6 +63,9 @@ public class SubscriptionEvent implements Comparable<SubscriptionEvent> { private final SubscriptionEventResponse<CachedSubscriptionPollResponse> response; private final SubscriptionCommitContext commitContext; + private volatile boolean payloadDetached; + private volatile Consumer<SubscriptionPollResponse> responseDeliveryListener; + // lastPolledConsumerId is not used as a criterion for determining pollability private volatile String lastPolledConsumerId = null; private final AtomicLong lastPolledTimestamp = new AtomicLong(INVALID_TIMESTAMP); @@ -210,13 +216,53 @@ public class SubscriptionEvent implements Comparable<SubscriptionEvent> { * {@link ConcurrentHashMap#compute} method of inFlightEvents in {@link * SubscriptionPrefetchingQueue} or {@link SubscriptionPrefetchingQueue#cleanUp}. */ - public void cleanUp(final boolean force) { + public synchronized void cleanUp(final boolean force) { pipeEvents.cleanUp(force); response.cleanUp(); + responseDeliveryListener = null; // TODO: clean more fields } + /** Marks this event's serialized response as safe for a broker to detach, if supported. */ + public void markResponseDelivered(final SubscriptionPollResponse deliveredResponse) { + final Consumer<SubscriptionPollResponse> listener = responseDeliveryListener; + if (listener != null) { + listener.accept(deliveredResponse); + } + } + + public void setResponseDeliveryListener(final Consumer<SubscriptionPollResponse> listener) { + responseDeliveryListener = listener; + } + + public boolean isPayloadDetached() { + return payloadDetached; + } + + /** Detaches the tablet payload while preserving commit and response metadata. */ + public synchronized boolean detachCurrentPayload( + final SubscriptionPollResponse deliveredResponse) { + final CachedSubscriptionPollResponse currentResponse = response.getCurrentResponse(); + if (payloadDetached + || !(response instanceof SubscriptionEventSingleResponse) + || currentResponse != deliveredResponse + || !(currentResponse.getPayload() instanceof TabletsPayload)) { + return false; + } + final int nextOffset = ((TabletsPayload) currentResponse.getPayload()).getNextOffset(); + ((SubscriptionEventSingleResponse) response) + .replaceCurrentPayload(new TabletsPayload(Collections.emptyMap(), nextOffset)); + payloadDetached = true; + return true; + } + + /** Restores a detached payload before redelivery. */ + public synchronized void restorePayload(final SubscriptionPollPayload payload) { + ((SubscriptionEventSingleResponse) response).replaceCurrentPayload(payload); + payloadDetached = false; + } + //////////////////////////// pollable //////////////////////////// public void recordLastPolledTimestamp() { @@ -262,7 +308,7 @@ public class SubscriptionEvent implements Comparable<SubscriptionEvent> { > SubscriptionConfig.getInstance().getSubscriptionRecycleUncommittedEventIntervalMs(); } - public void nack() { + public synchronized void nack() { // nack response response.nack(); @@ -313,14 +359,28 @@ public class SubscriptionEvent implements Comparable<SubscriptionEvent> { response.trySerializeRemainingResponses(); } - public void trySerializeCurrentResponse() { - response.trySerializeCurrentResponse(); + public synchronized void trySerializeCurrentResponse() { + if (!payloadDetached) { + response.trySerializeCurrentResponse(); + } } - public ByteBuffer getCurrentResponseByteBuffer() throws IOException { + public synchronized ByteBuffer getCurrentResponseByteBuffer() throws IOException { + if (payloadDetached) { + throw new IOException(DataNodeMiscMessages.NULL_RESPONSE_WHEN_SERIALIZING); + } return response.getCurrentResponseByteBuffer(); } + /** Serializes the selected response while fencing payload replacement and cleanup. */ + public synchronized ByteBuffer getResponseByteBuffer( + final SubscriptionPollResponse selectedResponse) throws IOException { + if (selectedResponse != response.getCurrentResponse()) { + throw new IOException(DataNodeMiscMessages.NULL_RESPONSE_WHEN_SERIALIZING); + } + return getCurrentResponseByteBuffer(); + } + public void invalidateCurrentResponseByteBuffer() { response.invalidateCurrentResponseByteBuffer(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/cache/CachedSubscriptionPollResponse.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/cache/CachedSubscriptionPollResponse.java index 63f1fb347b1..2c25792b66b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/cache/CachedSubscriptionPollResponse.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/cache/CachedSubscriptionPollResponse.java @@ -84,6 +84,7 @@ public class CachedSubscriptionPollResponse extends SubscriptionPollResponse { public void closeMemoryBlock() { if (Objects.nonNull(memoryBlock)) { memoryBlock.close(); + memoryBlock = null; } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/response/SubscriptionEventSingleResponse.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/response/SubscriptionEventSingleResponse.java index f9cb9db971c..d16044c6f2d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/response/SubscriptionEventSingleResponse.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/event/response/SubscriptionEventSingleResponse.java @@ -41,7 +41,7 @@ import java.util.Objects; public class SubscriptionEventSingleResponse implements SubscriptionEventResponse<CachedSubscriptionPollResponse> { - private final CachedSubscriptionPollResponse response; + private volatile CachedSubscriptionPollResponse response; public SubscriptionEventSingleResponse( final short responseType, @@ -109,6 +109,19 @@ public class SubscriptionEventSingleResponse SubscriptionPollResponseCache.getInstance().invalidate(response); } + public void replaceCurrentPayload(final SubscriptionPollPayload payload) { + final CachedSubscriptionPollResponse previous = response; + SubscriptionPollResponseCache.getInstance().invalidate(previous); + previous.closeMemoryBlock(); + response = + new CachedSubscriptionPollResponse( + previous.getResponseType(), + payload, + previous.getCommitContext(), + previous.isTimeSelected(), + previous.getTimeSelectedByTable()); + } + @Override public void nack() { invalidateCurrentResponseByteBuffer(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java index dd897f3a162..7e9b94b508e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/subscription/receiver/SubscriptionReceiverV1.java @@ -825,10 +825,10 @@ public class SubscriptionReceiverV1 implements SubscriptionReceiver { } try { - final ByteBuffer byteBuffer = event.getCurrentResponseByteBuffer(); + final ByteBuffer byteBuffer = event.getResponseByteBuffer(response); // payload size control - final long size = event.getCurrentResponseSize(); + final long size = byteBuffer.remaining(); if (totalSize.get() + size > maxBytes) { throw new SubscriptionPayloadExceedException( String.format( @@ -868,6 +868,9 @@ public class SubscriptionReceiverV1 implements SubscriptionReceiver { response, req.getRequest()); } + // The buffer now owns the serialized bytes. Notify outside the event monitor + // so queue lifecycle locks always precede the event's payload lock. + event.markResponseDelivered(response); return byteBuffer; } catch (final Exception e) { final boolean isOutdated = diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusLogToTabletConverterTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusLogToTabletConverterTest.java index e7c1f21cdb8..1cb1b05b3a5 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusLogToTabletConverterTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusLogToTabletConverterTest.java @@ -62,6 +62,31 @@ public class ConsensusLogToTabletConverterTest { private static final String DATABASE_NAME = "db"; + @Test + public void testReplayRequiresCompleteAndDecodableWalFragments() { + final List<IConsensusRequest> fragments = new ArrayList<>(); + final RelationalInsertRowNode node = StatementTestUtils.genInsertRowNode(1); + node.setSearchIndex(1); + final WALInfoEntry entry = new WALInfoEntry(1L, node); + final ByteBuffer buffer = ByteBuffer.allocate(entry.serializedSize()); + entry.serialize(new WALByteBufferForTest(buffer)); + buffer.flip(); + fragments.add(new IoTConsensusRequest(buffer)); + final IndexedConsensusRequest request = new IndexedConsensusRequest(1, fragments); + Assert.assertNull(ConsensusLogToTabletConverter.deserializeCompleteInsertNode(request)); + + node.setLastFragment(true); + final WALInfoEntry completeEntry = new WALInfoEntry(1L, node); + final ByteBuffer complete = ByteBuffer.allocate(completeEntry.serializedSize()); + completeEntry.serialize(new WALByteBufferForTest(complete)); + complete.flip(); + request.getRequests().set(0, new IoTConsensusRequest(complete)); + Assert.assertNotNull(ConsensusLogToTabletConverter.deserializeCompleteInsertNode(request)); + + request.getRequests().add(0, new IoTConsensusRequest(ByteBuffer.allocate(1))); + Assert.assertNull(ConsensusLogToTabletConverter.deserializeCompleteInsertNode(request)); + } + @Test public void testRepeatedWalDeserializationPreservesAllFragments() { final List<IConsensusRequest> fragments = new ArrayList<>(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueDeliveryTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueDeliveryTest.java new file mode 100644 index 00000000000..3cf1f6556f1 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/subscription/broker/consensus/ConsensusPrefetchingQueueDeliveryTest.java @@ -0,0 +1,627 @@ +/* + * 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.pipe.resource.memory.PipeMemoryWeightUtil; +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.SubscriptionCommitContext; +import org.apache.iotdb.rpc.subscription.payload.poll.SubscriptionPollResponse; +import org.apache.iotdb.rpc.subscription.payload.poll.TabletsPayload; +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.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.io.File; +import java.io.IOException; +import java.lang.reflect.Constructor; +import java.lang.reflect.Field; +import java.lang.reflect.Method; +import java.nio.ByteBuffer; +import java.nio.file.Files; +import java.util.ArrayList; +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.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; + +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.assertSame; +import static org.junit.Assert.assertThrows; +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 ConsensusPrefetchingQueueDeliveryTest { + + private static final int WRITER = 7; + private static final String CONSUMER = "consumer"; + private static final DataRegionId REGION = new DataRegionId(1); + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + private final CommonConfig config = CommonDescriptor.getInstance().getConfig(); + private final List<ConsensusPrefetchingQueue> queues = new ArrayList<>(); + private String originalSystemDir; + private int originalMaxDelay; + private int originalMaxTablets; + private int originalMaxEntries; + private long originalMaxBytes; + private File walDirectory; + private WALNode walNode; + private ConsensusLogToTabletConverter converter; + private ConsensusSubscriptionCommitManager commitManager; + private SubscriptionMemoryManager memoryManager; + private AtomicLong currentSearchIndex; + private long tabletBytes; + + @Before + public void setUp() throws Exception { + originalSystemDir = IoTDBDescriptor.getInstance().getConfig().getSystemDir(); + originalMaxDelay = config.getSubscriptionConsensusBatchMaxDelayInMs(); + originalMaxTablets = config.getSubscriptionConsensusBatchMaxTabletCount(); + originalMaxEntries = config.getSubscriptionConsensusBatchMaxWalEntries(); + originalMaxBytes = config.getSubscriptionConsensusBatchMaxSizeInBytes(); + config.setSubscriptionConsensusBatchMaxDelayInMs(0); + config.setSubscriptionConsensusBatchMaxTabletCount(1); + config.setSubscriptionConsensusBatchMaxWalEntries(128); + config.setSubscriptionConsensusBatchMaxSizeInBytes(16 * 1024 * 1024); + IoTDBDescriptor.getInstance() + .getConfig() + .setSystemDir(temporaryFolder.newFolder("system").getAbsolutePath()); + walDirectory = temporaryFolder.newFolder("wal"); + currentSearchIndex = new AtomicLong(); + walNode = mock(WALNode.class); + when(walNode.getLogDirectory()).thenReturn(walDirectory); + when(walNode.getCurrentSearchIndex()).thenAnswer(ignored -> currentSearchIndex.get()); + when(walNode.getCurrentWALFileVersion()).thenReturn(1L); + when(walNode.getCurrentWALMetaDataSnapshot()).thenReturn(new WALMetaData()); + when(walNode.getSortedWalFilesSnapshot()) + .thenAnswer( + ignored -> { + final File[] files = WALFileUtils.listAllWALFiles(walDirectory); + WALFileUtils.ascSortByVersionId(files); + return files; + }); + converter = mock(ConsensusLogToTabletConverter.class); + when(converter.getDatabaseName()).thenReturn("db"); + when(converter.isTableModel()).thenReturn(true); + when(converter.convert(any())) + .thenAnswer( + invocation -> + Collections.singletonList( + tablet(((InsertNode) invocation.getArgument(0)).getPhysicalTime()))); + tabletBytes = PipeMemoryWeightUtil.calculateTabletSizeInBytes(tablet(1001)); + memoryManager = new SubscriptionMemoryManager(tabletBytes * 2); + commitManager = + new ConsensusSubscriptionCommitManager( + (group, topic, region) -> + ConsensusSubscriptionCommitManager.ConfigNodeProgressQueryResult.absent()); + } + + @After + public void tearDown() { + try { + queues.forEach(ConsensusPrefetchingQueue::close); + } finally { + config.setSubscriptionConsensusBatchMaxDelayInMs(originalMaxDelay); + config.setSubscriptionConsensusBatchMaxTabletCount(originalMaxTablets); + config.setSubscriptionConsensusBatchMaxWalEntries(originalMaxEntries); + config.setSubscriptionConsensusBatchMaxSizeInBytes(originalMaxBytes); + IoTDBDescriptor.getInstance().getConfig().setSystemDir(originalSystemDir); + } + } + + @Test + public void testTenDeliveriesWithTwoBatchBudgetAndDelayedCommit() throws Exception { + writeWal(10, true); + final ConsensusPrefetchingQueue queue = newQueue(); + assertNull(queue.poll(CONSUMER)); + final List<SubscriptionCommitContext> contexts = new ArrayList<>(); + for (long index = 1; index <= 10; index++) { + queue.drivePrefetchOnce(); + final SubscriptionEvent event = queue.poll(CONSUMER); + assertNotNull(event); + assertEquals(index, event.getCommitContext().getWriterProgress().getLocalSeq()); + assertEquals(Collections.singletonList(1000 + index), timestamps(deliver(event))); + assertTrue(event.isPayloadDetached()); + assertTrue(((TabletsPayload) event.getCurrentResponse().getPayload()).getTablets().isEmpty()); + assertThrows(IOException.class, event::getCurrentResponseByteBuffer); + contexts.add(event.getCommitContext()); + assertTrue(memoryManager.getUsedMemorySizeInBytes() <= tabletBytes * 2); + assertEquals(queue.getRetainedTabletBytes(), memoryManager.getUsedMemorySizeInBytes()); + assertTrue(committedProgress().getWriterPositions().isEmpty()); + assertEquals(index, queue.getSubscriptionUncommittedEventCount()); + assertEquals(0L, requiredRetainedVersion(queue)); + } + + // An out-of-order ACK still waits for the preceding nine batches. + assertTrue(queue.ack(CONSUMER, contexts.get(9))); + assertTrue(committedProgress().getWriterPositions().isEmpty()); + for (int index = 0; index < 9; index++) { + assertTrue(queue.ack(CONSUMER, contexts.get(index))); + } + assertEquals( + new WriterProgress(1010, 10), committedProgress().getWriterPositions().get(writerId())); + assertEquals(0L, queue.getSubscriptionUncommittedEventCount()); + assertEquals(0L, queue.getRetainedTabletBytes()); + assertEquals(1L, requiredRetainedVersion(queue)); + assertEquals(0, deliveryRecordCount(queue)); + } + + @Test + public void testNackTimeoutDisconnectAndLateAckAfterReclamation() throws Exception { + writeWal(1, true); + final ConsensusPrefetchingQueue queue = newQueue(); + assertNull(queue.poll(CONSUMER)); + queue.drivePrefetchOnce(); + final SubscriptionEvent event = queue.poll(CONSUMER); + final SubscriptionCommitContext context = event.getCommitContext(); + final SubscriptionPollResponse first = deliver(event); + assertEquals(0L, queue.getRetainedTabletBytes()); + + assertTrue(queue.nack(CONSUMER, context)); + assertSame(event, queue.poll(CONSUMER)); + assertFalse(event.isPayloadDetached()); + assertEquals(tabletBytes, queue.getRetainedTabletBytes()); + final SubscriptionPollResponse second = deliver(event); + assertEquals(first.getPayload(), second.getPayload()); + assertEquals(first.getCommitContext(), second.getCommitContext()); + + // Force the timeout without sleeping or changing the process-wide recycle interval. + final Field timestamp = SubscriptionEvent.class.getDeclaredField("lastPolledTimestamp"); + timestamp.setAccessible(true); + ((AtomicLong) timestamp.get(event)).set(0L); + queue.drivePrefetchOnce(); + assertSame(event, queue.poll("replacement")); + assertEquals(first.getPayload(), deliver(event).getPayload()); + + assertEquals(1, queue.requeueInFlightEvents("replacement")); + assertSame(event, queue.poll("replacement")); + assertEquals(first.getPayload(), deliver(event).getPayload()); + assertEquals(1, queue.requeueInFlightEvents("replacement")); + // A late ACK may commit a queued, detached event without reading WAL again. + assertTrue(queue.ack("replacement", context)); + assertNull(queue.poll(CONSUMER)); + assertEquals(0, deliveryRecordCount(queue)); + assertEquals(0L, memoryManager.getUsedMemorySizeInBytes()); + } + + @Test + public void testIncompleteOrNotYetVisibleWalKeepsPayloadUntilReplayable() throws Exception { + currentSearchIndex.set(1); + final ConsensusPrefetchingQueue queue = newQueue(); + assertNull(queue.poll(CONSUMER)); + assertTrue(pendingEntries(queue).offer(request(1, 1, WRITER))); + queue.drivePrefetchOnce(); + final SubscriptionEvent event = queue.poll(CONSUMER); + deliver(event); + assertFalse(event.isPayloadDetached()); + assertEquals(tabletBytes, queue.getRetainedTabletBytes()); + + writeWal(1, false); + queue.drivePrefetchOnce(); + assertFalse(event.isPayloadDetached()); + assertEquals(tabletBytes, queue.getRetainedTabletBytes()); + + writeWal(1, true); + queue.drivePrefetchOnce(); + assertTrue(event.isPayloadDetached()); + assertEquals(0L, queue.getRetainedTabletBytes()); + assertTrue(committedProgress().getWriterPositions().isEmpty()); + assertTrue(queue.ack(CONSUMER, event.getCommitContext())); + } + + @Test + public void testUnreadableReplayKeepsRecordAndRetriesWithoutAdvancingCommit() throws Exception { + writeWal(1, true); + final ConsensusPrefetchingQueue queue = newQueue(); + assertNull(queue.poll(CONSUMER)); + queue.drivePrefetchOnce(); + final SubscriptionEvent event = queue.poll(CONSUMER); + final SubscriptionPollResponse first = deliver(event); + assertTrue(queue.nack(CONSUMER, event.getCommitContext())); + final File hidden = new File(temporaryFolder.getRoot(), "hidden.wal"); + Files.move(dataWal().toPath(), hidden.toPath()); + assertNull(queue.poll(CONSUMER)); + assertTrue(event.isPayloadDetached()); + assertEquals(1, deliveryRecordCount(queue)); + assertEquals(0L, queue.getRetainedTabletBytes()); + assertTrue(committedProgress().getWriterPositions().isEmpty()); + Files.move(hidden.toPath(), dataWal().toPath()); + assertSame(event, queue.poll(CONSUMER)); + assertEquals(first.getPayload(), deliver(event).getPayload()); + assertTrue(queue.ack(CONSUMER, event.getCommitContext())); + } + + @Test + public void testSparseBatchReplaysOnlyItsOriginalMembers() throws Exception { + config.setSubscriptionConsensusBatchMaxTabletCount(8); + config.setSubscriptionConsensusBatchMaxDelayInMs(60_000); + memoryManager = new SubscriptionMemoryManager(tabletBytes * 8); + writeWal(3, true); + final ConsensusPrefetchingQueue queue = newQueue(); + assertNull(queue.poll(CONSUMER)); + // Build a batch with members 1 and 3, as writer reordering can do for interleaved WAL records. + // Replaying the inclusive [1, 3] range would silently add entry 2 to the redelivery. + final Field linger = ConsensusPrefetchingQueue.class.getDeclaredField("lingerBatch"); + linger.setAccessible(true); + final Object batch = linger.get(queue); + final Class<?> entryClass = + Class.forName(ConsensusPrefetchingQueue.class.getName() + "$PreparedEntry"); + final Constructor<?> constructor = + entryClass.getDeclaredConstructor( + List.class, long.class, long.class, int.class, long.class, long.class); + constructor.setAccessible(true); + final Method append = + batch + .getClass() + .getDeclaredMethod( + "append", + Class.forName(ConsensusPrefetchingQueue.class.getName() + "$WriterBufferedEntry"), + long.class, + boolean.class); + append.setAccessible(true); + final Method reserve = + ConsensusPrefetchingQueue.class.getDeclaredMethod("tryReserveTabletMemory", long.class); + reserve.setAccessible(true); + for (final long index : new long[] {1, 3}) { + assertTrue((boolean) reserve.invoke(queue, tabletBytes)); + append.invoke( + batch, + constructor.newInstance( + Collections.singletonList(tablet(1000 + index)), + tabletBytes, + 1000 + index, + WRITER, + index, + index), + tabletBytes, + true); + } + final Method flush = + ConsensusPrefetchingQueue.class.getDeclaredMethod( + "flushBatch", batch.getClass(), long.class); + flush.setAccessible(true); + assertTrue((boolean) flush.invoke(queue, batch, queue.getCurrentSeekGeneration())); + final SubscriptionEvent event = queue.poll(CONSUMER); + final SubscriptionPollResponse first = deliver(event); + assertEquals(Arrays.asList(1001L, 1003L), timestamps(first)); + assertTrue(event.isPayloadDetached()); + assertTrue(queue.nack(CONSUMER, event.getCommitContext())); + assertSame(event, queue.poll(CONSUMER)); + assertEquals(first.getPayload(), deliver(event).getPayload()); + } + + @Test + public void testRestartReplaysAllUncommittedBatchesFromCommittedProgress() throws Exception { + writeWal(10, true); + ConsensusPrefetchingQueue queue = newQueue(); + assertNull(queue.poll(CONSUMER)); + for (int index = 1; index <= 10; index++) { + queue.drivePrefetchOnce(); + final SubscriptionEvent event = queue.poll(CONSUMER); + deliver(event); + if (index <= 2) { + assertTrue(queue.ack(CONSUMER, event.getCommitContext())); + } + } + final RegionProgress committed = committedProgress(); + commitManager.persistAll(); + queue.close(); + assertEquals(0, deliveryRecordCount(queue)); + + commitManager = + new ConsensusSubscriptionCommitManager( + (group, topic, region) -> + ConsensusSubscriptionCommitManager.ConfigNodeProgressQueryResult.unavailable()); + queue = newQueue(); + assertNull(queue.poll(CONSUMER)); + final List<Long> received = new ArrayList<>(); + for (int index = 3; index <= 10; index++) { + queue.drivePrefetchOnce(); + final SubscriptionEvent event = queue.poll(CONSUMER); + assertNotNull(event); + received.addAll(timestamps(deliver(event))); + assertTrue(queue.ack(CONSUMER, event.getCommitContext())); + } + assertEquals(Arrays.asList(1003L, 1004L, 1005L, 1006L, 1007L, 1008L, 1009L, 1010L), received); + assertEquals(new WriterProgress(1002, 2), committed.getWriterPositions().get(writerId())); + assertEquals( + new WriterProgress(1010, 10), committedProgress().getWriterPositions().get(writerId())); + } + + @Test + public void testReplicatedEntryWithoutLocalIndexReplaysByWriterPosition() throws Exception { + writeWal(1, true, false); + final ConsensusPrefetchingQueue queue = newQueue(); + assertNull(queue.poll(CONSUMER)); + queue.drivePrefetchOnce(); + final SubscriptionEvent event = queue.poll(CONSUMER); + final SubscriptionPollResponse first = deliver(event); + assertTrue(event.isPayloadDetached()); + assertEquals(0L, queue.getRetainedTabletBytes()); + assertEquals(0L, requiredRetainedVersion(queue)); + assertTrue(queue.nack(CONSUMER, event.getCommitContext())); + assertSame(event, queue.poll(CONSUMER)); + assertEquals(first.getPayload(), deliver(event).getPayload()); + assertTrue(queue.ack(CONSUMER, event.getCommitContext())); + } + + @Test + public void testSeekResetDiscardsDetachedStateAndRejectsOldCallbacksAndAcks() throws Exception { + writeWal(1, true); + final ConsensusPrefetchingQueue queue = newQueue(); + assertNull(queue.poll(CONSUMER)); + queue.drivePrefetchOnce(); + final SubscriptionEvent event = queue.poll(CONSUMER); + final SubscriptionPollResponse oldResponse = event.getCurrentResponse(); + deliver(event); + assertTrue(event.isPayloadDetached()); + + // Install the request a public seek hands to the serial worker. Drive it synchronously so the + // test does not depend on a process-wide executor being enabled. + final long oldGeneration = queue.getCurrentSeekGeneration(); + final Class<?> requestClass = + Class.forName(ConsensusPrefetchingQueue.class.getName() + "$PendingSeekRequest"); + final Constructor<?> constructor = + requestClass.getDeclaredConstructor( + long.class, RegionProgress.class, String.class, boolean.class, long.class, long.class); + constructor.setAccessible(true); + final Field generation = ConsensusPrefetchingQueue.class.getDeclaredField("seekGeneration"); + generation.setAccessible(true); + ((AtomicLong) generation.get(queue)).incrementAndGet(); + final Field pendingSeek = + ConsensusPrefetchingQueue.class.getDeclaredField("pendingSeekRequest"); + pendingSeek.setAccessible(true); + pendingSeek.set( + queue, + constructor.newInstance( + 1L, + new RegionProgress(Collections.emptyMap()), + "beginning", + true, + oldGeneration, + oldGeneration + 1)); + queue.drivePrefetchOnce(); + + assertEquals(0, deliveryRecordCount(queue)); + assertEquals(0L, memoryManager.getUsedMemorySizeInBytes()); + event.markResponseDelivered(oldResponse); + assertFalse(queue.ack(CONSUMER, event.getCommitContext())); + queue.drivePrefetchOnce(); + final SubscriptionEvent replayed = queue.poll(CONSUMER); + assertEquals(Collections.singletonList(1001L), timestamps(deliver(replayed))); + assertTrue(queue.ack(CONSUMER, replayed.getCommitContext())); + } + + @Test + public void testLateDeliveryCallbackCannotDetachAnotherAttemptAndCloseWaitsForSerialization() + throws Exception { + writeWal(1, true); + final ConsensusPrefetchingQueue queue = newQueue(); + assertNull(queue.poll(CONSUMER)); + queue.drivePrefetchOnce(); + final SubscriptionEvent event = queue.poll(CONSUMER); + final SubscriptionPollResponse oldResponse = event.getCurrentResponse(); + event.getResponseByteBuffer(oldResponse); + assertTrue(queue.requeue(CONSUMER, event.getCommitContext())); + assertSame(event, queue.poll(CONSUMER)); + event.markResponseDelivered(oldResponse); + assertFalse(event.isPayloadDetached()); + assertEquals(tabletBytes, queue.getRetainedTabletBytes()); + + final ExecutorService executor = Executors.newSingleThreadExecutor(); + final CountDownLatch started = new CountDownLatch(1); + try { + final Future<?> cleanup; + synchronized (event) { + cleanup = + executor.submit( + () -> { + started.countDown(); + queue.cleanUp(); + }); + assertTrue(started.await(5, TimeUnit.SECONDS)); + assertFalse(cleanup.isDone()); + assertNotNull(event.getCurrentResponseByteBuffer()); + } + cleanup.get(5, TimeUnit.SECONDS); + event.markResponseDelivered(oldResponse); + assertEquals(0, deliveryRecordCount(queue)); + assertEquals(0L, memoryManager.getUsedMemorySizeInBytes()); + } finally { + executor.shutdownNow(); + } + } + + private ConsensusPrefetchingQueue newQueue() { + final IoTConsensusServerImpl server = mock(IoTConsensusServerImpl.class); + when(server.getConsensusReqReader()).thenReturn(walNode); + when(server.getWriterSafeFrontierTracker()).thenReturn(new WriterSafeFrontierTracker()); + final ConsensusPrefetchingQueue queue = + new ConsensusPrefetchingQueue( + "group", + "topic", + TopicConstant.ORDER_MODE_LEADER_ONLY_VALUE, + REGION, + server, + new SubscriptionWalRetentionPolicy( + "topic", + SubscriptionWalRetentionPolicy.UNBOUNDED, + SubscriptionWalRetentionPolicy.UNBOUNDED), + converter, + commitManager, + new RegionProgress(Collections.emptyMap()), + 1L, + 1L, + true); + queue.setSubscriptionMemoryManager(memoryManager); + queues.add(queue); + return queue; + } + + private void writeWal(final int count, final boolean complete) throws Exception { + writeWal(count, complete, true); + } + + private void writeWal(final int count, final boolean complete, final boolean local) + throws Exception { + currentSearchIndex.set(local ? count : 0); + // WALWriter appends to existing files; replacing a visibility fixture needs a fresh segment. + Files.deleteIfExists(dataWal().toPath()); + try (WALWriter writer = new WALWriter(dataWal(), WALFileVersion.V3)) { + for (long index = 1; index <= count; index++) { + final RelationalInsertRowNode node = StatementTestUtils.genInsertRowNode((int) index); + node.setSearchIndex(local ? index : -1L); + node.setLastFragment(complete); + final WALInfoEntry entry = new WALInfoEntry(1L, node); + final ByteBuffer buffer = ByteBuffer.allocate(entry.serializedSize()); + entry.serialize(new WALByteBufferForTest(buffer)); + final WALMetaData metadata = new WALMetaData(); + metadata.add(entry.serializedSize(), local ? index : -1L, 1L, 1000 + index, WRITER, index); + writer.write(buffer, metadata); + } + } + try (WALWriter ignored = + new WALWriter( + new File( + walDirectory, + WALFileUtils.getLogFileName(1L, count, WALFileStatus.CONTAINS_SEARCH_INDEX)), + WALFileVersion.V3)) { + // The data file is sealed; the live successor contains no requests. + } + } + + private File dataWal() { + return new File( + walDirectory, WALFileUtils.getLogFileName(0L, 0L, WALFileStatus.CONTAINS_SEARCH_INDEX)); + } + + private static IndexedConsensusRequest request( + final long index, final long sequence, final int writer) { + return new IndexedConsensusRequest( + index, + sequence, + Collections.singletonList(StatementTestUtils.genInsertRowNode((int) sequence))) + .setPhysicalTime(1000 + sequence) + .setNodeId(writer); + } + + private static Tablet tablet(final long timestamp) { + final Tablet tablet = + new Tablet( + "sensors", + Arrays.asList("device", "temperature"), + Arrays.asList(TSDataType.STRING, TSDataType.DOUBLE), + Arrays.asList(ColumnCategory.TAG, ColumnCategory.FIELD), + 1); + tablet.addTimestamp(0, timestamp); + tablet.addValue(0, 0, "d1"); + tablet.addValue(0, 1, (double) timestamp); + tablet.setRowSize(1); + return tablet; + } + + private static SubscriptionPollResponse deliver(final SubscriptionEvent event) throws Exception { + final SubscriptionPollResponse response = event.getCurrentResponse(); + final ByteBuffer bytes = event.getResponseByteBuffer(response); + event.invalidateCurrentResponseByteBuffer(); + event.markResponseDelivered(response); + return SubscriptionPollResponse.deserialize(bytes.duplicate()); + } + + private static List<Long> timestamps(final SubscriptionPollResponse response) { + final List<Long> result = new ArrayList<>(); + for (final Tablet tablet : ((TabletsPayload) response.getPayload()).getTablets()) { + result.add(tablet.getTimestamps()[0]); + } + return result; + } + + private RegionProgress committedProgress() { + return commitManager.getCommittedRegionProgress("group", "topic", REGION); + } + + private static WriterId writerId() { + return new WriterId(REGION.toString(), WRITER); + } + + private static long requiredRetainedVersion(final ConsensusPrefetchingQueue queue) + throws Exception { + final Method method = + ConsensusPrefetchingQueue.class.getDeclaredMethod("getRequiredRetainedMinVersionId"); + method.setAccessible(true); + return (long) method.invoke(queue); + } + + private static int deliveryRecordCount(final ConsensusPrefetchingQueue queue) throws Exception { + final Field field = ConsensusPrefetchingQueue.class.getDeclaredField("deliveryBatches"); + field.setAccessible(true); + return ((java.util.Map<?, ?>) field.get(queue)).size(); + } + + @SuppressWarnings("unchecked") + private static BlockingQueue<IndexedConsensusRequest> pendingEntries( + final ConsensusPrefetchingQueue queue) throws Exception { + final Field field = ConsensusPrefetchingQueue.class.getDeclaredField("pendingEntries"); + field.setAccessible(true); + return (BlockingQueue<IndexedConsensusRequest>) field.get(queue); + } +}
