This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 576d6a7b3cf [Fix] Harden split LOAD recovery, deletion dispatch and
memory bounds (#18930)
576d6a7b3cf is described below
commit 576d6a7b3cf48c3e912a4e598c54e99ddaad7856
Author: Caideyipi <[email protected]>
AuthorDate: Tue Oct 20 19:19:01 2026 +0800
[Fix] Harden split LOAD recovery, deletion dispatch and memory bounds
(#18930)
* Fix LOAD partition recovery and PREPARE completeness
* Fix LOAD recovery, deletion delivery and buffer eviction
---
.../plan/scheduler/load/PieceDispatcher.java | 63 ++---
.../plan/scheduler/load/TsFileSplitConsumer.java | 15 +-
.../db/storageengine/load/TsFileWriterManager.java | 98 +++++---
.../load/splitter/TsFileSplitter.java | 8 +-
.../write/TsFilePrecalculatedChunkWriter.java | 8 +-
.../plan/scheduler/load/PieceDispatcherTest.java | 146 ++++++++++++
.../scheduler/load/TsFileSplitConsumerTest.java | 233 +++++++++++++++++++
.../db/storageengine/load/TsFileSplitterTest.java | 41 ++++
.../load/TsFileWriterManagerRecoveryTest.java | 254 +++++++++++++++++++++
9 files changed, 805 insertions(+), 61 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java
index 970b337b9ca..a56423ec2f8 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java
@@ -35,10 +35,14 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
+import java.util.Comparator;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.Map;
+import java.util.NavigableSet;
import java.util.Objects;
-import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.TreeSet;
/**
* Buffers chunks and deletions of one source TsFile per target region
partition, managing
@@ -63,9 +67,13 @@ class PieceDispatcher {
private final Map<PartitionKey, ChunkOffsetCalculator>
partition2OffsetCalculator =
new HashMap<>();
- /** Max-heap tracking (partitionKey, bufferedSize) for memory budget
eviction. */
- private final PriorityQueue<HeapEntry> largestPieceRegions =
- new PriorityQueue<>((a, b) -> Long.compare(b.dataSize(), a.dataSize()));
+ /** One ordered entry per non-empty partition, so tracking stays bounded as
chunks stream in. */
+ private final NavigableSet<PieceSize> largestPieceRegions =
+ new TreeSet<>(
+ Comparator.comparingLong(PieceSize::dataSize)
+ .reversed()
+ .thenComparing(entry -> entry.key().regionId())
+ .thenComparingLong(entry -> entry.key().timePartitionStart()));
PieceDispatcher(
final LoadSingleTsFileNode singleTsFileNode,
@@ -116,23 +124,29 @@ class PieceDispatcher {
partition2ReplicaSetAndNode.computeIfAbsent(
key, k -> new Pair<>(replicaSet, createPieceNode(k)));
+ largestPieceRegions.remove(new PieceSize(key,
targetPair.getRight().getDataSize()));
targetPair.getRight().addTsFileData(chunkData);
- trackHeapEntry(key, targetPair.getRight().getDataSize());
+ trackPieceSize(key, targetPair.getRight().getDataSize());
}
- /** Broadcasts deletion data across every active partition buffer. */
- void addDeletionToAll(final DeletionData deletionData) {
+ /** Buffers one copy per region; its receiver applies the deletion to all
staged partitions. */
+ void addDeletionToAllRegions(final DeletionData deletionData) {
if (deletionData == null || partition2ReplicaSetAndNode.isEmpty()) {
return;
}
final long deletionSize = deletionData.getDataSize();
+ final Set<TConsensusGroupId> regions = new HashSet<>();
for (final Map.Entry<PartitionKey, Pair<TRegionReplicaSet,
LoadTsFilePieceNode>> entry :
partition2ReplicaSetAndNode.entrySet()) {
+ if (!regions.add(entry.getKey().regionId())) {
+ continue;
+ }
memoryBuffer.add(deletionSize);
final LoadTsFilePieceNode pieceNode = entry.getValue().getRight();
+ largestPieceRegions.remove(new PieceSize(entry.getKey(),
pieceNode.getDataSize()));
pieceNode.addTsFileData(deletionData);
- trackHeapEntry(entry.getKey(), pieceNode.getDataSize());
+ trackPieceSize(entry.getKey(), pieceNode.getDataSize());
}
}
@@ -143,7 +157,7 @@ class PieceDispatcher {
/** Progressively dispatches largest buffered pieces until memory usage is
back within budget. */
boolean dispatchLargestUntilMemoryEnough() throws LoadFileException {
while (!memoryBuffer.isMemoryEnough()) {
- final PartitionKey key = pollLargestValidPartition();
+ final PartitionKey key = pollLargestPartition();
if (key == null) {
break;
}
@@ -158,6 +172,7 @@ class PieceDispatcher {
final long pieceSize = pieceNode.getDataSize();
if (!dispatchCallback.dispatch(pieceNode, pair.getLeft())) {
+ trackPieceSize(key, pieceSize);
return false;
}
@@ -192,10 +207,12 @@ class PieceDispatcher {
}
// -------------------------------------------------------------------------
- // Helpers & Eviction Heap Management
+ // Helpers & Eviction Tracking
// -------------------------------------------------------------------------
private void replacePieceNode(final PartitionKey key, final
TRegionReplicaSet replicaSet) {
+ largestPieceRegions.remove(
+ new PieceSize(key,
partition2ReplicaSetAndNode.get(key).getRight().getDataSize()));
partition2ReplicaSetAndNode.put(key, new Pair<>(replicaSet,
createPieceNode(key)));
}
@@ -208,25 +225,15 @@ class PieceDispatcher {
pieceIndex);
}
- private void trackHeapEntry(final PartitionKey key, final long size) {
- largestPieceRegions.offer(new HeapEntry(key, size));
+ private void trackPieceSize(final PartitionKey key, final long size) {
+ if (size > 0) {
+ largestPieceRegions.add(new PieceSize(key, size));
+ }
}
- private PartitionKey pollLargestValidPartition() {
- while (!largestPieceRegions.isEmpty()) {
- final HeapEntry entry = largestPieceRegions.poll();
- final Pair<TRegionReplicaSet, LoadTsFilePieceNode> pair =
- partition2ReplicaSetAndNode.get(entry.key());
- if (pair == null) {
- continue;
- }
-
- final long actualSize = pair.getRight().getDataSize();
- if (entry.dataSize() == actualSize && actualSize > 0) {
- return entry.key();
- }
- }
- return null;
+ private PartitionKey pollLargestPartition() {
+ final PieceSize entry = largestPieceRegions.pollFirst();
+ return entry == null ? null : entry.key();
}
void clear() {
@@ -242,5 +249,5 @@ class PieceDispatcher {
private record PartitionKey(TConsensusGroupId regionId, long
timePartitionStart) {}
- private record HeapEntry(PartitionKey key, long dataSize) {}
+ private record PieceSize(PartitionKey key, long dataSize) {}
}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java
index 0481d0ce42d..fcb704431f3 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java
@@ -49,6 +49,7 @@ public class TsFileSplitConsumer implements
TsFileSplitter.TsFileDataConsumer {
private final Consumer<TTimePartitionSlot> progressIndexCallback;
private final List<ChunkData> nonDirectionalChunkData = new ArrayList<>();
+ private boolean isProcessingDeletions;
public TsFileSplitConsumer(
final LoadSingleTsFileNode singleTsFileNode,
@@ -111,10 +112,17 @@ public class TsFileSplitConsumer implements
TsFileSplitter.TsFileDataConsumer {
}
private boolean addOrSendDeletionData(final DeletionData deletionData)
throws LoadFileException {
- // Route pending chunks first to ensure deletions never precede the chunks
they apply to
+ // The receiver broadcasts each deletion to all staged partitions of its
region. Stage all
+ // chunks first, then send each deletion once per region, including
partitions already evicted.
routeChunkData();
- dispatcher.addDeletionToAll(deletionData);
- return true;
+ if (!isProcessingDeletions) {
+ if (!dispatcher.flushAll()) {
+ return false;
+ }
+ isProcessingDeletions = true;
+ }
+ dispatcher.addDeletionToAllRegions(deletionData);
+ return dispatcher.dispatchLargestUntilMemoryEnough();
}
// -------------------------------------------------------------------------
@@ -157,6 +165,7 @@ public class TsFileSplitConsumer implements
TsFileSplitter.TsFileDataConsumer {
/** Drops all unrouted chunks, clears partition buffers, and releases cached
memory accounting. */
void clear() {
nonDirectionalChunkData.clear();
+ isProcessingDeletions = false;
dispatcher.clear();
memoryBuffer.clear();
}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java
index d3d9efb1345..d72fd02e54b 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java
@@ -45,6 +45,7 @@ import
org.apache.iotdb.db.storageengine.load.splitter.ChunkPayloadRef;
import org.apache.iotdb.db.storageengine.load.splitter.DeletionData;
import org.apache.iotdb.db.storageengine.load.splitter.TsFileData;
+import org.apache.tsfile.common.conf.TSFileConfig;
import org.apache.tsfile.common.constant.TsFileConstant;
import org.apache.tsfile.enums.TSDataType;
import org.apache.tsfile.exception.write.PageException;
@@ -54,6 +55,7 @@ import org.apache.tsfile.file.metadata.IChunkMetadata;
import org.apache.tsfile.file.metadata.IDeviceID;
import org.apache.tsfile.file.metadata.statistics.Statistics;
import org.apache.tsfile.read.TimeValuePair;
+import org.apache.tsfile.read.TsFileSequenceReader;
import org.apache.tsfile.read.common.Chunk;
import org.apache.tsfile.utils.Pair;
import org.apache.tsfile.utils.RamUsageEstimator;
@@ -71,10 +73,12 @@ import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.ArrayList;
import java.util.HashMap;
+import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
+import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
@@ -466,6 +470,37 @@ final class TsFileWriterManager {
// Consensus Lifecycle: PREPARE & LOAD
// -------------------------------------------------------------------------
+ private void validateStagedFiles(
+ final Map<TTimePartitionSlot, Long> timePartition2ExpectedDataEnd)
+ throws IOException, LoadFileException {
+ final Set<TTimePartitionSlot> partitions = new
HashSet<>(timePartition2Writer.keySet());
+ if (timePartition2ExpectedDataEnd != null) {
+ partitions.addAll(timePartition2ExpectedDataEnd.keySet());
+ }
+ for (final TTimePartitionSlot partition : partitions) {
+ final TsFilePrecalculatedChunkWriter writer =
timePartition2Writer.get(partition);
+ final File stagedFile =
+ writer == null
+ ? new File(taskDir, stagedFileName(partition) +
TsFileConstant.TSFILE_SUFFIX)
+ : writer.getFile();
+ final LoadTsFileProgress progress =
timePartition2Progress.get(partition);
+ final long fileLength = stagedFile.length();
+ final long expectedDataEnd =
expectedDataEndOf(timePartition2ExpectedDataEnd, partition);
+ if (writer == null
+ || progress == null
+ || !progress.isReady(fileLength)
+ || (expectedDataEnd >= 0L && fileLength < expectedDataEnd)) {
+ throw new LoadFileException(
+ String.format(
+
StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19,
+ stagedFile.getAbsolutePath(),
+ uuidDirName(stagedFile),
+ Math.max(progress == null ? 0L : progress.getTotalLength(),
expectedDataEnd),
+ fileLength));
+ }
+ }
+ }
+
/**
* Seals modification files, validates chunk continuity, and seals TsFiles
in the PREPARE round.
*/
@@ -477,12 +512,16 @@ final class TsFileWriterManager {
checkNotClosed();
touch();
- // 1. Close modification files first to freeze deletion views
+ // Validate every announced partition before sealing any file. A missing
partition or a file
+ // without durable chunk records must not turn into a successful partial
import.
+ validateStagedFiles(timePartition2ExpectedDataEnd);
+
+ // Close modification files to freeze deletion views.
for (final ModificationFile modFile :
timePartition2ModificationFile.values()) {
modFile.close();
}
- // 2. Validate continuity and seal writer metadata zone
+ // Seal writer metadata zones only after every partition has passed
validation.
for (final Map.Entry<TTimePartitionSlot, TsFilePrecalculatedChunkWriter>
entry :
timePartition2Writer.entrySet()) {
final TsFilePrecalculatedChunkWriter writer = entry.getValue();
@@ -490,23 +529,6 @@ final class TsFileWriterManager {
continue;
}
- final LoadTsFileProgress progress =
timePartition2Progress.get(entry.getKey());
- if (progress != null && progress.exists() && progress.getTotalLength() >
0) {
- final long fileLength = writer.getFile().length();
- final long expectedDataEnd =
- expectedDataEndOf(timePartition2ExpectedDataEnd, entry.getKey());
- if (!progress.isReady(fileLength)
- || (expectedDataEnd >= 0L && fileLength < expectedDataEnd)) {
- throw new LoadFileException(
- String.format(
- StorageEngineMessages
-
.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19,
- writer.getFile().getAbsolutePath(),
- uuidDirName(writer.getFile()),
- Math.max(progress.getTotalLength(), expectedDataEnd),
- fileLength));
- }
- }
// Rebuild from the staged devices so recovered writers also retain the
table schemas that
// Pipe needs to enumerate table data in the sealed file.
if (PathUtils.isTableModelDatabase(dataRegion.getDatabaseName())) {
@@ -721,7 +743,21 @@ final class TsFileWriterManager {
final TTimePartitionSlot timePartitionSlot =
new TTimePartitionSlot(parseTimePartitionStart(tsFile.getName()));
final long resumeOffset = progress.getTotalLength();
- if (tsFile.length() > resumeOffset) {
+ final TsFileResource resource = new TsFileResource(tsFile);
+ boolean sealed = false;
+ if (resource.resourceFileExists()) {
+ try (final TsFileSequenceReader reader =
+ new TsFileSequenceReader(tsFile.getAbsolutePath(), false)) {
+ sealed = TSFileConfig.MAGIC_STRING.equals(reader.readTailMagic());
+ }
+ if (sealed) {
+ resource.deserialize();
+ resource.setStatus(TsFileResourceStatus.NORMAL);
+ }
+ }
+ // A resource sidecar and a completed footer record a successful
PREPARE. Its coordinator
+ // may send COMMIT immediately after recovery, so preserve both the
sealed bytes and marks.
+ if (!sealed && tsFile.length() > resumeOffset) {
try (final FileChannel truncatingChannel =
FileChannel.open(tsFile.toPath(), StandardOpenOption.READ,
StandardOpenOption.WRITE)) {
truncatingChannel.truncate(resumeOffset);
@@ -745,15 +781,15 @@ final class TsFileWriterManager {
try {
channel =
FileChannel.open(tsFile.toPath(), StandardOpenOption.READ,
StandardOpenOption.WRITE);
- channel.position(resumeOffset);
+ channel.position(sealed ? tsFile.length() : resumeOffset);
final TsFilePrecalculatedChunkWriter writer =
- new TsFilePrecalculatedChunkWriter(tsFile, channel);
+ new TsFilePrecalculatedChunkWriter(tsFile, channel, sealed);
writer.setRejectHoles(LoadConsensusCompat.usesIoTConsensusV2());
writer.restoreChunkMetadata(restoreChunkMetadata(records));
publishedLengths.put(tsFile, tsFile.length());
timePartition2Writer.put(timePartitionSlot, writer);
- timePartition2Resource.put(timePartitionSlot, new
TsFileResource(tsFile));
+ timePartition2Resource.put(timePartitionSlot, resource);
timePartition2Progress.put(timePartitionSlot, progress);
final File modificationFile =
ModificationFile.getExclusiveMods(tsFile);
@@ -928,11 +964,15 @@ final class TsFileWriterManager {
/** The name a staged file of this region and time partition has. */
private String stagedFileName(final TTimePartitionSlot timePartitionSlot) {
+ return stagedFilePrefix() + timePartitionSlot.getStartTime();
+ }
+
+ private String stagedFilePrefix() {
return String.join(
IoTDBConstant.FILE_NAME_SEPARATOR,
dataRegion.getDatabaseName(),
dataRegion.getDataRegionIdString(),
- Long.toString(timePartitionSlot.getStartTime()));
+ "");
}
private static String uuidDirName(final File tsFile) {
@@ -940,10 +980,12 @@ final class TsFileWriterManager {
return parent == null ? tsFile.getName() : parent.getName();
}
- private static long parseTimePartitionStart(final String tsFileName) {
- final String baseName = tsFileName.substring(0,
tsFileName.lastIndexOf('.'));
- final int lastDash =
baseName.lastIndexOf(IoTDBConstant.FILE_NAME_SEPARATOR);
- return Long.parseLong(baseName.substring(lastDash + 1));
+ private long parseTimePartitionStart(final String tsFileName) {
+ // The partition may be negative: its sign is part of the number, not a
name separator.
+ return Long.parseLong(
+ tsFileName.substring(
+ stagedFilePrefix().length(),
+ tsFileName.length() - TsFileConstant.TSFILE_SUFFIX.length()));
}
/** The bytes this piece added to each staged file, as the refs a replica
reads them back with. */
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java
index 16109c64de2..9b17746f942 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java
@@ -24,6 +24,7 @@ import org.apache.iotdb.commons.utils.TimePartitionUtils;
import org.apache.iotdb.db.conf.IoTDBConfig;
import org.apache.iotdb.db.conf.IoTDBDescriptor;
import org.apache.iotdb.db.exception.load.LoadFileException;
+import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
import org.apache.iotdb.db.i18n.StorageEngineMessages;
import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry;
import
org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile;
@@ -455,7 +456,12 @@ public class TsFileSplitter {
private void handleModification(List<ModEntry> deletions) throws
LoadFileException {
for (final ModEntry mod : deletions) {
- consumer.apply(new DeletionData(mod));
+ if (!consumer.apply(new DeletionData(mod))) {
+ throw new IllegalStateException(
+ String.format(
+
DataNodeQueryMessages.DISPATCH_TSFILEDATA_ERROR_WHEN_PARSING_TSFILE_S,
+ tsFile.getPath()));
+ }
}
}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java
b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java
index b48b7d14cb1..7e375d9fe70 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java
@@ -109,12 +109,18 @@ public class TsFilePrecalculatedChunkWriter implements
AutoCloseable {
public TsFilePrecalculatedChunkWriter(final File file, final FileChannel
channel)
throws IOException {
+ this(file, channel, false);
+ }
+
+ /** Resumes a staged file, preserving a completed PREPARE when it is already
sealed. */
+ public TsFilePrecalculatedChunkWriter(
+ final File file, final FileChannel channel, final boolean sealed) throws
IOException {
this.file = file;
this.out =
new WritableFileChannelOutput(
Objects.requireNonNull(
channel,
StorageEngineMessages.EXCEPTION_CHANNEL_CANNOT_BE_NULL_0F79E0FB));
- this.sealed = false;
+ this.sealed = sealed;
this.maxPhysicalFileSize = channel.size();
LOGGER.info(
StorageEngineMessages
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcherTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcherTest.java
new file mode 100644
index 00000000000..b94ef5aa07b
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcherTest.java
@@ -0,0 +1,146 @@
+/*
+ * 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.queryengine.plan.scheduler.load;
+
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
+import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
+import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
+import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
+import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode;
+import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
+import org.apache.iotdb.db.storageengine.load.splitter.NonAlignedChunkData;
+
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.header.ChunkHeader;
+import org.apache.tsfile.file.metadata.StringArrayDeviceID;
+import org.apache.tsfile.file.metadata.enums.CompressionType;
+import org.apache.tsfile.file.metadata.enums.TSEncoding;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.io.File;
+import java.lang.reflect.Field;
+import java.util.Collection;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.when;
+
+public class PieceDispatcherTest {
+
+ @Test
+ public void testEvictionQueueStaysBoundedAcrossChunksAndFlushes() throws
Exception {
+ final MemoryBoundedBuffer buffer = Mockito.mock(MemoryBoundedBuffer.class);
+ final AtomicLong usedBytes = new AtomicLong();
+ final long budget = chunk(0L).getDataSize() * 5;
+ when(buffer.isMemoryEnough()).thenAnswer(ignored -> usedBytes.get() <=
budget);
+ Mockito.doAnswer(
+ invocation -> {
+ usedBytes.addAndGet(-invocation.<Long>getArgument(0));
+ return null;
+ })
+ .when(buffer)
+ .release(anyLong());
+ final AtomicInteger dispatchedChunks = new AtomicInteger();
+ final PieceDispatcher dispatcher =
+ new PieceDispatcher(
+ loadNode(),
+ buffer,
+ (piece, region) -> {
+ dispatchedChunks.addAndGet(piece.getAllTsFileData().size());
+ return true;
+ });
+ final TRegionReplicaSet replicaSet =
+ new TRegionReplicaSet()
+ .setRegionId(new TConsensusGroupId(TConsensusGroupType.DataRegion,
1));
+ final Field queueField =
PieceDispatcher.class.getDeclaredField("largestPieceRegions");
+ queueField.setAccessible(true);
+
+ for (int i = 0; i < 128; i++) {
+ final NonAlignedChunkData chunk = chunk(i % 2);
+ usedBytes.addAndGet(chunk.getDataSize());
+ dispatcher.offerChunk(chunk, replicaSet);
+ assertTrue(dispatcher.dispatchLargestUntilMemoryEnough());
+ assertTrue(
+ "Eviction tracking must be bounded by the two live partition
buffers",
+ ((Collection<?>) queueField.get(dispatcher)).size() <= 2);
+ }
+ assertTrue(dispatcher.flushAll());
+ assertEquals(128, dispatchedChunks.get());
+ assertEquals(0L, usedBytes.get());
+ assertTrue(
+ "Flushed pieces must leave no eviction entries",
+ ((Collection<?>) queueField.get(dispatcher)).isEmpty());
+ }
+
+ @Test
+ public void testFailedEvictionCanBeRetried() throws Exception {
+ final MemoryBoundedBuffer buffer = Mockito.mock(MemoryBoundedBuffer.class);
+ final AtomicBoolean memoryEnough = new AtomicBoolean(false);
+ when(buffer.isMemoryEnough()).thenAnswer(ignored -> memoryEnough.get());
+ Mockito.doAnswer(
+ ignored -> {
+ memoryEnough.set(true);
+ return null;
+ })
+ .when(buffer)
+ .release(anyLong());
+ final AtomicInteger attempts = new AtomicInteger();
+ final PieceDispatcher dispatcher =
+ new PieceDispatcher(loadNode(), buffer, (piece, region) ->
attempts.incrementAndGet() > 1);
+ final NonAlignedChunkData chunk = chunk(0L);
+ dispatcher.offerChunk(
+ chunk,
+ new TRegionReplicaSet()
+ .setRegionId(new TConsensusGroupId(TConsensusGroupType.DataRegion,
1)));
+
+ assertFalse(dispatcher.dispatchLargestUntilMemoryEnough());
+ assertTrue(dispatcher.dispatchLargestUntilMemoryEnough());
+ assertEquals(2, attempts.get());
+ Mockito.verify(buffer).release(chunk.getDataSize());
+ }
+
+ private static LoadSingleTsFileNode loadNode() {
+ final LoadSingleTsFileNode node = Mockito.mock(LoadSingleTsFileNode.class);
+ when(node.getPlanNodeId()).thenReturn(new PlanNodeId("bounded-queue"));
+ final TsFileResource resource = Mockito.mock(TsFileResource.class);
+ when(resource.getTsFile()).thenReturn(new File("bounded-queue.tsfile"));
+ when(node.getTsFileResource()).thenReturn(resource);
+ return node;
+ }
+
+ private static NonAlignedChunkData chunk(final long partition) {
+ final NonAlignedChunkData chunk =
+ new NonAlignedChunkData(
+ new StringArrayDeviceID("root", "test", "d1"),
+ new ChunkHeader(
+ "s1", 0, TSDataType.INT32, CompressionType.UNCOMPRESSED,
TSEncoding.PLAIN, 0),
+ new TTimePartitionSlot(partition));
+ chunk.writeDecodePage(new long[] {partition + 1}, new Object[] {1}, 1);
+ chunk.endChunk();
+ return chunk;
+ }
+}
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumerTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumerTest.java
new file mode 100644
index 00000000000..93bf7dc4c7b
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumerTest.java
@@ -0,0 +1,233 @@
+/*
+ * 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.queryengine.plan.scheduler.load;
+
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
+import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
+import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
+import org.apache.iotdb.commons.path.MeasurementPath;
+import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
+import org.apache.iotdb.db.conf.IoTDBConfig;
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode;
+import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
+import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry;
+import
org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile;
+import
org.apache.iotdb.db.storageengine.dataregion.modification.TreeDeletionEntry;
+import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
+import org.apache.iotdb.db.storageengine.load.LoadTsFileManager;
+import
org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock;
+import org.apache.iotdb.db.storageengine.load.splitter.DeletionData;
+import org.apache.iotdb.db.storageengine.load.splitter.NonAlignedChunkData;
+
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.header.ChunkHeader;
+import org.apache.tsfile.file.metadata.StringArrayDeviceID;
+import org.apache.tsfile.file.metadata.enums.CompressionType;
+import org.apache.tsfile.file.metadata.enums.TSEncoding;
+import org.apache.tsfile.utils.Pair;
+import org.junit.Assert;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.mockito.Mockito;
+
+import java.io.File;
+import java.lang.reflect.Constructor;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static org.mockito.ArgumentMatchers.anyList;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.when;
+
+public class TsFileSplitConsumerTest {
+
+ @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ @Test
+ public void testDeletionThatExceedsMemoryBudgetTriggersDispatch() throws
Exception {
+ final int memoryBudget =
IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize() >> 2;
+ final Constructor<LoadTsFileDataCacheMemoryBlock> constructor =
+
LoadTsFileDataCacheMemoryBlock.class.getDeclaredConstructor(long.class);
+ constructor.setAccessible(true);
+ final LoadTsFileDataCacheMemoryBlock memoryBlock =
constructor.newInstance(1024 * 1024L);
+
+ final TRegionReplicaSet replicaSet =
+ new TRegionReplicaSet()
+ .setRegionId(new TConsensusGroupId(TConsensusGroupType.DataRegion,
1));
+ final DataPartitionBatchFetcher partitionFetcher =
+ Mockito.mock(DataPartitionBatchFetcher.class);
+ when(partitionFetcher.queryDataPartition(anyList(), anyString()))
+ .thenReturn(List.of(replicaSet));
+
+ final LoadSingleTsFileNode loadNode =
Mockito.mock(LoadSingleTsFileNode.class);
+ when(loadNode.getPlanNodeId()).thenReturn(new
PlanNodeId("load-memory-budget"));
+ final TsFileResource resource = Mockito.mock(TsFileResource.class);
+ when(resource.getTsFile()).thenReturn(new
File("load-memory-budget.tsfile"));
+ when(loadNode.getTsFileResource()).thenReturn(resource);
+
+ final AtomicInteger dispatchedPieces = new AtomicInteger();
+ final TsFileSplitConsumer consumer =
+ new TsFileSplitConsumer(
+ loadNode,
+ memoryBlock,
+ new DataPartitionResolver(partitionFetcher, "root.test"),
+ null,
+ (pieceNode, targetReplicaSet) -> {
+ if
(pieceNode.getAllTsFileData().stream().anyMatch(DeletionData.class::isInstance))
{
+ dispatchedPieces.incrementAndGet();
+ }
+ return true;
+ });
+
+ final NonAlignedChunkData chunk =
+ new NonAlignedChunkData(
+ new StringArrayDeviceID("root", "test", "d1"),
+ new ChunkHeader(
+ "s1", 0, TSDataType.INT32, CompressionType.UNCOMPRESSED,
TSEncoding.PLAIN, 0),
+ new TTimePartitionSlot(0L));
+ chunk.writeDecodePage(new long[] {1L}, new Object[] {1}, 1);
+ chunk.endChunk();
+
+ final ModEntry deletion = Mockito.mock(ModEntry.class);
+ when(deletion.serializedSize()).thenReturn(memoryBudget + 1);
+
+ try {
+ Assert.assertTrue(consumer.apply(chunk));
+ Assert.assertEquals(0, dispatchedPieces.get());
+
+ Assert.assertTrue(consumer.apply(new DeletionData(deletion)));
+
+ Assert.assertEquals(
+ "A deletion must trigger largest-first dispatch when its buffered
piece exceeds budget",
+ 1,
+ dispatchedPieces.get());
+ } finally {
+ consumer.clear();
+ }
+ }
+
+ @Test
+ public void testDeletionsAreSentOncePerRegionAfterAllPartitions() throws
Exception {
+ final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig();
+ final String[] originalLoadDirs = config.getLoadTsFileDirs();
+ final File loadDir = temporaryFolder.newFolder("load");
+ config.setLoadTsFileDirs(new String[] {loadDir.getAbsolutePath()});
+ final LoadTsFileManager firstRegion = manager(1);
+ final LoadTsFileManager secondRegion = manager(2);
+ final Map<Integer, LoadTsFileManager> managers = Map.of(1, firstRegion, 2,
secondRegion);
+ final TRegionReplicaSet firstReplicaSet = replicaSet(1);
+ final TRegionReplicaSet secondReplicaSet = replicaSet(2);
+ final DataPartitionBatchFetcher partitionFetcher =
+ Mockito.mock(DataPartitionBatchFetcher.class);
+ when(partitionFetcher.queryDataPartition(anyList(), anyString()))
+ .thenAnswer(
+ invocation -> {
+ final List<Pair<org.apache.tsfile.file.metadata.IDeviceID,
TTimePartitionSlot>>
+ slots = invocation.getArgument(0);
+ return slots.stream()
+ .map(
+ slot ->
+ slot.right.getStartTime() < 2000L ? firstReplicaSet
: secondReplicaSet)
+ .toList();
+ });
+ final Constructor<LoadTsFileDataCacheMemoryBlock> constructor =
+
LoadTsFileDataCacheMemoryBlock.class.getDeclaredConstructor(long.class);
+ constructor.setAccessible(true);
+ final TsFileSplitConsumer consumer =
+ new TsFileSplitConsumer(
+ loadNode(),
+ constructor.newInstance(1024 * 1024L),
+ new DataPartitionResolver(partitionFetcher, "root"),
+ null,
+ (piece, target) -> {
+ try {
+ managers
+ .get(target.getRegionId().getId())
+ .writePiece("load-deletions", piece.getAllTsFileData());
+ return true;
+ } catch (Exception e) {
+ throw new AssertionError(e);
+ }
+ });
+ final ModEntry firstDeletion =
+ new TreeDeletionEntry(new MeasurementPath("root.test.d1.s1"), 0L, 10L);
+ final ModEntry secondDeletion =
+ new TreeDeletionEntry(new MeasurementPath("root.test.d1.s1"), 20L,
30L);
+ try {
+ for (long partition : new long[] {0L, 1000L, 2000L}) {
+ Assert.assertTrue(consumer.apply(chunk(partition)));
+ }
+ Assert.assertTrue(consumer.apply(new DeletionData(firstDeletion)));
+ Assert.assertTrue(consumer.apply(new DeletionData(secondDeletion)));
+ Assert.assertTrue(consumer.sendAllTsFileData());
+ for (long partition : new long[] {0L, 1000L, 2000L}) {
+ final int regionId = partition < 2000L ? 1 : 2;
+ final File taskDir = managers.get(regionId).getActiveTaskDirs().get(0);
+ final File staged =
+ new File(taskDir, "root.test-" + regionId + "-" + partition +
".tsfile");
+ Assert.assertEquals(
+ "Each partition must receive every deletion exactly once",
+ List.of(firstDeletion, secondDeletion),
+ ModificationFile.readAllModifications(staged, false));
+ }
+ } finally {
+ consumer.clear();
+ firstRegion.stop();
+ secondRegion.stop();
+ config.setLoadTsFileDirs(originalLoadDirs);
+ }
+ }
+
+ private static LoadTsFileManager manager(final int regionId) {
+ final DataRegion region = Mockito.mock(DataRegion.class);
+ when(region.getDatabaseName()).thenReturn("root.test");
+ when(region.getDataRegionIdString()).thenReturn(String.valueOf(regionId));
+ return new LoadTsFileManager(region);
+ }
+
+ private static TRegionReplicaSet replicaSet(final int regionId) {
+ return new TRegionReplicaSet()
+ .setRegionId(new TConsensusGroupId(TConsensusGroupType.DataRegion,
regionId));
+ }
+
+ private static LoadSingleTsFileNode loadNode() {
+ final LoadSingleTsFileNode node = Mockito.mock(LoadSingleTsFileNode.class);
+ when(node.getPlanNodeId()).thenReturn(new PlanNodeId("load-deletions"));
+ final TsFileResource resource = Mockito.mock(TsFileResource.class);
+ when(resource.getTsFile()).thenReturn(new File("load-deletions.tsfile"));
+ when(node.getTsFileResource()).thenReturn(resource);
+ return node;
+ }
+
+ private static NonAlignedChunkData chunk(final long partition) {
+ final NonAlignedChunkData chunk =
+ new NonAlignedChunkData(
+ new StringArrayDeviceID("root", "test", "d1"),
+ new ChunkHeader(
+ "s1", 0, TSDataType.INT32, CompressionType.UNCOMPRESSED,
TSEncoding.PLAIN, 0),
+ new TTimePartitionSlot(partition));
+ chunk.writeDecodePage(new long[] {partition + 1}, new Object[] {1}, 1);
+ chunk.endChunk();
+ return chunk;
+ }
+}
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java
index 1797c401160..4ccbc98aec6 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileSplitterTest.java
@@ -21,7 +21,10 @@ package org.apache.iotdb.db.storageengine.load.splitter;
import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
import org.apache.iotdb.commons.conf.CommonDescriptor;
+import org.apache.iotdb.commons.path.MeasurementPath;
import org.apache.iotdb.commons.utils.TimePartitionUtils;
+import
org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile;
+import
org.apache.iotdb.db.storageengine.dataregion.modification.TreeDeletionEntry;
import org.apache.tsfile.common.conf.TSFileConfig;
import org.apache.tsfile.enums.ColumnCategory;
@@ -61,9 +64,47 @@ import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
public class TsFileSplitterTest {
+ @Test
+ public void testRejectedDeletionStopsSplitting() throws Exception {
+ final File source = Files.createTempFile("rejected-load-deletion",
".tsfile").toFile();
+ final File mods = ModificationFile.getExclusiveMods(source);
+ try {
+ try (final TsFileIOWriter writer = new TsFileIOWriter(source)) {
+ writer.startChunkGroup(new StringArrayDeviceID("root", "sg", "d1"));
+ final ChunkWriterImpl chunkWriter =
+ new ChunkWriterImpl(new MeasurementSchema("s1", TSDataType.INT32));
+ chunkWriter.write(1L, 1);
+ chunkWriter.writeToFileWriter(writer);
+ writer.endChunkGroup();
+ writer.endFile();
+ }
+ try (final ModificationFile modificationFile = new
ModificationFile(mods, false)) {
+ modificationFile.write(new TreeDeletionEntry(new
MeasurementPath("root.sg.d1.s1"), 0L, 1L));
+ modificationFile.write(new TreeDeletionEntry(new
MeasurementPath("root.sg.d1.s1"), 2L, 3L));
+ }
+ final AtomicInteger consumedDeletions = new AtomicInteger();
+ final TsFileSplitter splitter =
+ new TsFileSplitter(
+ source,
+ data -> {
+ if (data instanceof DeletionData) {
+ consumedDeletions.incrementAndGet();
+ return false;
+ }
+ return true;
+ });
+ Assert.assertThrows(IllegalStateException.class,
splitter::splitTsFileByDataPartition);
+ Assert.assertEquals(1, consumedDeletions.get());
+ } finally {
+ Files.deleteIfExists(mods.toPath());
+ Files.deleteIfExists(source.toPath());
+ }
+ }
+
// Verify the splitter initializes the v3 deserialize configuration for a
valid v3 TsFile.
@Test
public void testSplitV3TsFile() throws Exception {
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManagerRecoveryTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManagerRecoveryTest.java
new file mode 100644
index 00000000000..9c76b75c6a1
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManagerRecoveryTest.java
@@ -0,0 +1,254 @@
+/*
+ * 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.storageengine.load;
+
+import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
+import org.apache.iotdb.commons.consensus.index.ProgressIndex;
+import org.apache.iotdb.commons.consensus.index.impl.SimpleProgressIndex;
+import org.apache.iotdb.db.exception.load.LoadFileException;
+import
org.apache.iotdb.db.queryengine.plan.scheduler.load.ChunkOffsetCalculator;
+import org.apache.iotdb.db.storageengine.dataregion.DataRegion;
+import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource;
+import org.apache.iotdb.db.storageengine.load.splitter.NonAlignedChunkData;
+
+import org.apache.tsfile.enums.TSDataType;
+import org.apache.tsfile.file.header.ChunkHeader;
+import org.apache.tsfile.file.metadata.IDeviceID;
+import org.apache.tsfile.file.metadata.StringArrayDeviceID;
+import org.apache.tsfile.file.metadata.enums.CompressionType;
+import org.apache.tsfile.file.metadata.enums.TSEncoding;
+import org.apache.tsfile.read.TsFileReader;
+import org.apache.tsfile.read.common.Path;
+import org.apache.tsfile.read.common.RowRecord;
+import org.apache.tsfile.read.expression.QueryExpression;
+import org.apache.tsfile.read.query.dataset.QueryDataSet;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+public class TsFileWriterManagerRecoveryTest {
+
+ private static final String DATABASE = "root.load_recovery_test";
+ private static final IDeviceID DEVICE =
+ new StringArrayDeviceID("root", "load_recovery_test", "d1");
+
+ @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+ private DataRegion dataRegion;
+
+ @Before
+ public void setUp() {
+ dataRegion = Mockito.mock(DataRegion.class);
+ Mockito.when(dataRegion.getDatabaseName()).thenReturn(DATABASE);
+ Mockito.when(dataRegion.getDataRegionIdString()).thenReturn("0");
+ }
+
+ @Test
+ public void testResumeNegativeTimePartition() throws Exception {
+ assertTimePartitionResumes(-604_800_000L);
+ }
+
+ @Test
+ public void testResumeMinimumTimePartition() throws Exception {
+ assertTimePartitionResumes(Long.MIN_VALUE);
+ }
+
+ @Test
+ public void testCommitPreparedFileAfterRecovery() throws Exception {
+ final File taskDir = temporaryFolder.newFolder("prepared-task");
+ final NonAlignedChunkData chunk = chunk(0L, 1);
+ new ChunkOffsetCalculator().assign(chunk);
+ final ProgressIndex progressIndex = new SimpleProgressIndex(1, 10L);
+ TsFileWriterManager manager = new TsFileWriterManager(dataRegion, taskDir);
+ try {
+ manager.writePiece(List.of(chunk));
+ manager.prepare(
+ true,
+ Map.of(chunk.getTimePartitionSlot(), progressIndex),
+ Map.of(chunk.getTimePartitionSlot(), chunkEnd(chunk)));
+ final File staged = stagedFile(taskDir, 0L);
+ final byte[] preparedBytes = Files.readAllBytes(staged.toPath());
+ manager.close(true);
+
+ manager = new TsFileWriterManager(dataRegion, taskDir).recoverFromDisk();
+
+ assertArrayEquals(
+ "Recovery must preserve a prepared TsFile",
+ preparedBytes,
+ Files.readAllBytes(staged.toPath()));
+ // The coordinator already received the PREPARE acknowledgement and only
sends COMMIT.
+ manager.loadAll(true, false);
+ final ArgumentCaptor<TsFileResource> imported =
ArgumentCaptor.forClass(TsFileResource.class);
+ Mockito.verify(dataRegion)
+ .loadNewTsFile(
+ imported.capture(),
+ Mockito.eq(false),
+ Mockito.eq(true),
+ Mockito.eq(false),
+ Mockito.any());
+ assertTrue(imported.getValue().isGeneratedByPipe());
+ assertEquals(progressIndex,
imported.getValue().getMaxProgressIndexAfterClose());
+ assertPoints(staged, 0L, 1);
+ } finally {
+ manager.close();
+ }
+ }
+
+ private void assertTimePartitionResumes(final long partitionStart) throws
Exception {
+ final File taskDir = temporaryFolder.newFolder("task");
+ final ChunkOffsetCalculator calculator = new ChunkOffsetCalculator();
+ final NonAlignedChunkData first = chunk(partitionStart, 1);
+ final NonAlignedChunkData second = chunk(partitionStart, 2);
+ calculator.assign(first);
+ calculator.assign(second);
+
+ TsFileWriterManager manager = new TsFileWriterManager(dataRegion, taskDir);
+ try {
+ manager.writePiece(List.of(first));
+ manager.close(true);
+ manager = new TsFileWriterManager(dataRegion, taskDir).recoverFromDisk();
+ manager.writePiece(List.of(second));
+ manager.prepare(
+ false, Collections.emptyMap(), Map.of(second.getTimePartitionSlot(),
chunkEnd(second)));
+
+ // Both the chunk staged before recovery and the one appended afterwards
must survive.
+ assertPoints(stagedFile(taskDir, partitionStart), partitionStart, 1, 2);
+ } finally {
+ manager.close();
+ }
+ }
+
+ @Test
+ public void testPrepareRejectsMissingPartition() throws Exception {
+ final File taskDir = temporaryFolder.newFolder("task");
+ final TsFileWriterManager manager = new TsFileWriterManager(dataRegion,
taskDir);
+ final NonAlignedChunkData first = chunk(0L, 1);
+ final NonAlignedChunkData missing = chunk(604_800_000L, 1);
+ new ChunkOffsetCalculator().assign(first);
+ new ChunkOffsetCalculator().assign(missing);
+ final Map<TTimePartitionSlot, Long> expectedDataEnd =
+ Map.of(
+ first.getTimePartitionSlot(), chunkEnd(first),
+ missing.getTimePartitionSlot(), chunkEnd(missing));
+ try {
+ manager.writePiece(List.of(first));
+ final long unsealedLength = stagedFile(taskDir, 0L).length();
+ assertPrepareFails(manager, expectedDataEnd, stagedFile(taskDir,
604_800_000L));
+ assertEquals(unsealedLength, stagedFile(taskDir, 0L).length());
+
+ manager.writePiece(List.of(missing));
+ manager.prepare(false, Collections.emptyMap(), expectedDataEnd);
+ assertPoints(stagedFile(taskDir, 0L), 0L, 1);
+ assertPoints(stagedFile(taskDir, 604_800_000L), 604_800_000L, 1);
+ } finally {
+ manager.close();
+ }
+ }
+
+ @Test
+ public void testPrepareRejectsPartitionWithoutRecordedChunks() throws
Exception {
+ final File taskDir = temporaryFolder.newFolder("task");
+ final TsFileWriterManager manager = new TsFileWriterManager(dataRegion,
taskDir);
+ final NonAlignedChunkData complete = chunk(0L, 1);
+ new ChunkOffsetCalculator().assign(complete);
+ final Map<TTimePartitionSlot, Long> expectedDataEnd =
+ Map.of(complete.getTimePartitionSlot(), chunkEnd(complete));
+ try {
+ // A failed piece may have opened the file without recording any durable
chunk.
+ try {
+ manager.writePiece(List.of(chunk(0L, 1)));
+ fail("A chunk without a layout must fail after opening its staged
file");
+ } catch (final IOException expected) {
+ assertTrue(stagedFile(taskDir, 0L).isFile());
+ }
+ assertPrepareFails(manager, expectedDataEnd, stagedFile(taskDir, 0L));
+
+ manager.writePiece(List.of(complete));
+ manager.prepare(false, Collections.emptyMap(), expectedDataEnd);
+ assertPoints(stagedFile(taskDir, 0L), 0L, 1);
+ } finally {
+ manager.close();
+ }
+ }
+
+ private static void assertPrepareFails(
+ final TsFileWriterManager manager,
+ final Map<TTimePartitionSlot, Long> expectedDataEnd,
+ final File missingDataFile)
+ throws Exception {
+ try {
+ manager.prepare(false, Collections.emptyMap(), expectedDataEnd);
+ fail("PREPARE must reject a partition whose announced data has not been
staged");
+ } catch (final LoadFileException e) {
+ assertTrue(e.getMessage(),
e.getMessage().contains(missingDataFile.getAbsolutePath()));
+ }
+ }
+
+ private static NonAlignedChunkData chunk(final long partitionStart, final
int value) {
+ final NonAlignedChunkData chunk =
+ new NonAlignedChunkData(
+ DEVICE,
+ new ChunkHeader(
+ "s1", 0, TSDataType.INT32, CompressionType.UNCOMPRESSED,
TSEncoding.PLAIN, 0),
+ new TTimePartitionSlot(partitionStart));
+ chunk.writeDecodePage(new long[] {partitionStart + value}, new Object[]
{value}, 1);
+ chunk.endChunk();
+ return chunk;
+ }
+
+ private static long chunkEnd(final NonAlignedChunkData chunk) {
+ return chunk.getChunkLayout().offset() + chunk.getChunkLayout().length();
+ }
+
+ private static File stagedFile(final File taskDir, final long
partitionStart) {
+ return new File(taskDir, DATABASE + "-0-" + partitionStart + ".tsfile");
+ }
+
+ private static void assertPoints(final File file, final long partitionStart,
final int... values)
+ throws Exception {
+ try (final TsFileReader reader = new TsFileReader(file)) {
+ final QueryDataSet result =
+ reader.query(QueryExpression.create(List.of(new Path(DEVICE, "s1",
false)), null));
+ for (final int value : values) {
+ assertTrue(result.hasNext());
+ final RowRecord row = result.next();
+ assertEquals(partitionStart + value, row.getTimestamp());
+ assertEquals(value, row.getField(0).getIntV());
+ }
+ assertFalse(result.hasNext());
+ }
+ }
+}
