Skip to content

Commit 576d6a7

Browse files
authored
[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
1 parent 86ae957 commit 576d6a7

9 files changed

Lines changed: 805 additions & 61 deletions

File tree

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java‎

Lines changed: 35 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -35,10 +35,14 @@
3535
import org.slf4j.LoggerFactory;
3636

3737
import java.io.IOException;
38+
import java.util.Comparator;
3839
import java.util.HashMap;
40+
import java.util.HashSet;
3941
import java.util.Map;
42+
import java.util.NavigableSet;
4043
import java.util.Objects;
41-
import java.util.PriorityQueue;
44+
import java.util.Set;
45+
import java.util.TreeSet;
4246

4347
/**
4448
* Buffers chunks and deletions of one source TsFile per target region partition, managing
@@ -63,9 +67,13 @@ interface DispatchCallback {
6367
private final Map<PartitionKey, ChunkOffsetCalculator> partition2OffsetCalculator =
6468
new HashMap<>();
6569

66-
/** Max-heap tracking (partitionKey, bufferedSize) for memory budget eviction. */
67-
private final PriorityQueue<HeapEntry> largestPieceRegions =
68-
new PriorityQueue<>((a, b) -> Long.compare(b.dataSize(), a.dataSize()));
70+
/** One ordered entry per non-empty partition, so tracking stays bounded as chunks stream in. */
71+
private final NavigableSet<PieceSize> largestPieceRegions =
72+
new TreeSet<>(
73+
Comparator.comparingLong(PieceSize::dataSize)
74+
.reversed()
75+
.thenComparing(entry -> entry.key().regionId())
76+
.thenComparingLong(entry -> entry.key().timePartitionStart()));
6977

7078
PieceDispatcher(
7179
final LoadSingleTsFileNode singleTsFileNode,
@@ -116,23 +124,29 @@ void offerChunk(final ChunkData chunkData, final TRegionReplicaSet replicaSet)
116124
partition2ReplicaSetAndNode.computeIfAbsent(
117125
key, k -> new Pair<>(replicaSet, createPieceNode(k)));
118126

127+
largestPieceRegions.remove(new PieceSize(key, targetPair.getRight().getDataSize()));
119128
targetPair.getRight().addTsFileData(chunkData);
120-
trackHeapEntry(key, targetPair.getRight().getDataSize());
129+
trackPieceSize(key, targetPair.getRight().getDataSize());
121130
}
122131

123-
/** Broadcasts deletion data across every active partition buffer. */
124-
void addDeletionToAll(final DeletionData deletionData) {
132+
/** Buffers one copy per region; its receiver applies the deletion to all staged partitions. */
133+
void addDeletionToAllRegions(final DeletionData deletionData) {
125134
if (deletionData == null || partition2ReplicaSetAndNode.isEmpty()) {
126135
return;
127136
}
128137

129138
final long deletionSize = deletionData.getDataSize();
139+
final Set<TConsensusGroupId> regions = new HashSet<>();
130140
for (final Map.Entry<PartitionKey, Pair<TRegionReplicaSet, LoadTsFilePieceNode>> entry :
131141
partition2ReplicaSetAndNode.entrySet()) {
142+
if (!regions.add(entry.getKey().regionId())) {
143+
continue;
144+
}
132145
memoryBuffer.add(deletionSize);
133146
final LoadTsFilePieceNode pieceNode = entry.getValue().getRight();
147+
largestPieceRegions.remove(new PieceSize(entry.getKey(), pieceNode.getDataSize()));
134148
pieceNode.addTsFileData(deletionData);
135-
trackHeapEntry(entry.getKey(), pieceNode.getDataSize());
149+
trackPieceSize(entry.getKey(), pieceNode.getDataSize());
136150
}
137151
}
138152

@@ -143,7 +157,7 @@ void addDeletionToAll(final DeletionData deletionData) {
143157
/** Progressively dispatches largest buffered pieces until memory usage is back within budget. */
144158
boolean dispatchLargestUntilMemoryEnough() throws LoadFileException {
145159
while (!memoryBuffer.isMemoryEnough()) {
146-
final PartitionKey key = pollLargestValidPartition();
160+
final PartitionKey key = pollLargestPartition();
147161
if (key == null) {
148162
break;
149163
}
@@ -158,6 +172,7 @@ boolean dispatchLargestUntilMemoryEnough() throws LoadFileException {
158172
final long pieceSize = pieceNode.getDataSize();
159173

160174
if (!dispatchCallback.dispatch(pieceNode, pair.getLeft())) {
175+
trackPieceSize(key, pieceSize);
161176
return false;
162177
}
163178

@@ -192,10 +207,12 @@ boolean flushAll() throws LoadFileException {
192207
}
193208

194209
// -------------------------------------------------------------------------
195-
// Helpers & Eviction Heap Management
210+
// Helpers & Eviction Tracking
196211
// -------------------------------------------------------------------------
197212

198213
private void replacePieceNode(final PartitionKey key, final TRegionReplicaSet replicaSet) {
214+
largestPieceRegions.remove(
215+
new PieceSize(key, partition2ReplicaSetAndNode.get(key).getRight().getDataSize()));
199216
partition2ReplicaSetAndNode.put(key, new Pair<>(replicaSet, createPieceNode(key)));
200217
}
201218

@@ -208,25 +225,15 @@ private LoadTsFilePieceNode createPieceNode(final PartitionKey key) {
208225
pieceIndex);
209226
}
210227

211-
private void trackHeapEntry(final PartitionKey key, final long size) {
212-
largestPieceRegions.offer(new HeapEntry(key, size));
228+
private void trackPieceSize(final PartitionKey key, final long size) {
229+
if (size > 0) {
230+
largestPieceRegions.add(new PieceSize(key, size));
231+
}
213232
}
214233

215-
private PartitionKey pollLargestValidPartition() {
216-
while (!largestPieceRegions.isEmpty()) {
217-
final HeapEntry entry = largestPieceRegions.poll();
218-
final Pair<TRegionReplicaSet, LoadTsFilePieceNode> pair =
219-
partition2ReplicaSetAndNode.get(entry.key());
220-
if (pair == null) {
221-
continue;
222-
}
223-
224-
final long actualSize = pair.getRight().getDataSize();
225-
if (entry.dataSize() == actualSize && actualSize > 0) {
226-
return entry.key();
227-
}
228-
}
229-
return null;
234+
private PartitionKey pollLargestPartition() {
235+
final PieceSize entry = largestPieceRegions.pollFirst();
236+
return entry == null ? null : entry.key();
230237
}
231238

232239
void clear() {
@@ -242,5 +249,5 @@ void clear() {
242249

243250
private record PartitionKey(TConsensusGroupId regionId, long timePartitionStart) {}
244251

245-
private record HeapEntry(PartitionKey key, long dataSize) {}
252+
private record PieceSize(PartitionKey key, long dataSize) {}
246253
}

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ public class TsFileSplitConsumer implements TsFileSplitter.TsFileDataConsumer {
4949
private final Consumer<TTimePartitionSlot> progressIndexCallback;
5050

5151
private final List<ChunkData> nonDirectionalChunkData = new ArrayList<>();
52+
private boolean isProcessingDeletions;
5253

5354
public TsFileSplitConsumer(
5455
final LoadSingleTsFileNode singleTsFileNode,
@@ -111,10 +112,17 @@ private boolean addOrSendChunkData(final ChunkData chunkData) throws LoadFileExc
111112
}
112113

113114
private boolean addOrSendDeletionData(final DeletionData deletionData) throws LoadFileException {
114-
// Route pending chunks first to ensure deletions never precede the chunks they apply to
115+
// The receiver broadcasts each deletion to all staged partitions of its region. Stage all
116+
// chunks first, then send each deletion once per region, including partitions already evicted.
115117
routeChunkData();
116-
dispatcher.addDeletionToAll(deletionData);
117-
return true;
118+
if (!isProcessingDeletions) {
119+
if (!dispatcher.flushAll()) {
120+
return false;
121+
}
122+
isProcessingDeletions = true;
123+
}
124+
dispatcher.addDeletionToAllRegions(deletionData);
125+
return dispatcher.dispatchLargestUntilMemoryEnough();
118126
}
119127

120128
// -------------------------------------------------------------------------
@@ -157,6 +165,7 @@ boolean sendAllTsFileData() throws LoadFileException {
157165
/** Drops all unrouted chunks, clears partition buffers, and releases cached memory accounting. */
158166
void clear() {
159167
nonDirectionalChunkData.clear();
168+
isProcessingDeletions = false;
160169
dispatcher.clear();
161170
memoryBuffer.clear();
162171
}

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java‎

Lines changed: 70 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import org.apache.iotdb.db.storageengine.load.splitter.DeletionData;
4646
import org.apache.iotdb.db.storageengine.load.splitter.TsFileData;
4747

48+
import org.apache.tsfile.common.conf.TSFileConfig;
4849
import org.apache.tsfile.common.constant.TsFileConstant;
4950
import org.apache.tsfile.enums.TSDataType;
5051
import org.apache.tsfile.exception.write.PageException;
@@ -54,6 +55,7 @@
5455
import org.apache.tsfile.file.metadata.IDeviceID;
5556
import org.apache.tsfile.file.metadata.statistics.Statistics;
5657
import org.apache.tsfile.read.TimeValuePair;
58+
import org.apache.tsfile.read.TsFileSequenceReader;
5759
import org.apache.tsfile.read.common.Chunk;
5860
import org.apache.tsfile.utils.Pair;
5961
import org.apache.tsfile.utils.RamUsageEstimator;
@@ -71,10 +73,12 @@
7173
import java.nio.file.StandardOpenOption;
7274
import java.util.ArrayList;
7375
import java.util.HashMap;
76+
import java.util.HashSet;
7477
import java.util.List;
7578
import java.util.Map;
7679
import java.util.Objects;
7780
import java.util.Optional;
81+
import java.util.Set;
7882
import java.util.concurrent.ConcurrentHashMap;
7983
import java.util.concurrent.atomic.AtomicLong;
8084
import java.util.concurrent.atomic.AtomicReference;
@@ -466,6 +470,37 @@ private void writeDeletion(final DeletionData deletionData) throws IOException {
466470
// Consensus Lifecycle: PREPARE & LOAD
467471
// -------------------------------------------------------------------------
468472

473+
private void validateStagedFiles(
474+
final Map<TTimePartitionSlot, Long> timePartition2ExpectedDataEnd)
475+
throws IOException, LoadFileException {
476+
final Set<TTimePartitionSlot> partitions = new HashSet<>(timePartition2Writer.keySet());
477+
if (timePartition2ExpectedDataEnd != null) {
478+
partitions.addAll(timePartition2ExpectedDataEnd.keySet());
479+
}
480+
for (final TTimePartitionSlot partition : partitions) {
481+
final TsFilePrecalculatedChunkWriter writer = timePartition2Writer.get(partition);
482+
final File stagedFile =
483+
writer == null
484+
? new File(taskDir, stagedFileName(partition) + TsFileConstant.TSFILE_SUFFIX)
485+
: writer.getFile();
486+
final LoadTsFileProgress progress = timePartition2Progress.get(partition);
487+
final long fileLength = stagedFile.length();
488+
final long expectedDataEnd = expectedDataEndOf(timePartition2ExpectedDataEnd, partition);
489+
if (writer == null
490+
|| progress == null
491+
|| !progress.isReady(fileLength)
492+
|| (expectedDataEnd >= 0L && fileLength < expectedDataEnd)) {
493+
throw new LoadFileException(
494+
String.format(
495+
StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19,
496+
stagedFile.getAbsolutePath(),
497+
uuidDirName(stagedFile),
498+
Math.max(progress == null ? 0L : progress.getTotalLength(), expectedDataEnd),
499+
fileLength));
500+
}
501+
}
502+
}
503+
469504
/**
470505
* Seals modification files, validates chunk continuity, and seals TsFiles in the PREPARE round.
471506
*/
@@ -477,36 +512,23 @@ void prepare(
477512
checkNotClosed();
478513
touch();
479514

480-
// 1. Close modification files first to freeze deletion views
515+
// Validate every announced partition before sealing any file. A missing partition or a file
516+
// without durable chunk records must not turn into a successful partial import.
517+
validateStagedFiles(timePartition2ExpectedDataEnd);
518+
519+
// Close modification files to freeze deletion views.
481520
for (final ModificationFile modFile : timePartition2ModificationFile.values()) {
482521
modFile.close();
483522
}
484523

485-
// 2. Validate continuity and seal writer metadata zone
524+
// Seal writer metadata zones only after every partition has passed validation.
486525
for (final Map.Entry<TTimePartitionSlot, TsFilePrecalculatedChunkWriter> entry :
487526
timePartition2Writer.entrySet()) {
488527
final TsFilePrecalculatedChunkWriter writer = entry.getValue();
489528
if (writer.isSealed()) {
490529
continue;
491530
}
492531

493-
final LoadTsFileProgress progress = timePartition2Progress.get(entry.getKey());
494-
if (progress != null && progress.exists() && progress.getTotalLength() > 0) {
495-
final long fileLength = writer.getFile().length();
496-
final long expectedDataEnd =
497-
expectedDataEndOf(timePartition2ExpectedDataEnd, entry.getKey());
498-
if (!progress.isReady(fileLength)
499-
|| (expectedDataEnd >= 0L && fileLength < expectedDataEnd)) {
500-
throw new LoadFileException(
501-
String.format(
502-
StorageEngineMessages
503-
.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19,
504-
writer.getFile().getAbsolutePath(),
505-
uuidDirName(writer.getFile()),
506-
Math.max(progress.getTotalLength(), expectedDataEnd),
507-
fileLength));
508-
}
509-
}
510532
// Rebuild from the staged devices so recovered writers also retain the table schemas that
511533
// Pipe needs to enumerate table data in the sealed file.
512534
if (PathUtils.isTableModelDatabase(dataRegion.getDatabaseName())) {
@@ -721,7 +743,21 @@ TsFileWriterManager recoverFromDisk() throws IOException {
721743
final TTimePartitionSlot timePartitionSlot =
722744
new TTimePartitionSlot(parseTimePartitionStart(tsFile.getName()));
723745
final long resumeOffset = progress.getTotalLength();
724-
if (tsFile.length() > resumeOffset) {
746+
final TsFileResource resource = new TsFileResource(tsFile);
747+
boolean sealed = false;
748+
if (resource.resourceFileExists()) {
749+
try (final TsFileSequenceReader reader =
750+
new TsFileSequenceReader(tsFile.getAbsolutePath(), false)) {
751+
sealed = TSFileConfig.MAGIC_STRING.equals(reader.readTailMagic());
752+
}
753+
if (sealed) {
754+
resource.deserialize();
755+
resource.setStatus(TsFileResourceStatus.NORMAL);
756+
}
757+
}
758+
// A resource sidecar and a completed footer record a successful PREPARE. Its coordinator
759+
// may send COMMIT immediately after recovery, so preserve both the sealed bytes and marks.
760+
if (!sealed && tsFile.length() > resumeOffset) {
725761
try (final FileChannel truncatingChannel =
726762
FileChannel.open(tsFile.toPath(), StandardOpenOption.READ, StandardOpenOption.WRITE)) {
727763
truncatingChannel.truncate(resumeOffset);
@@ -745,15 +781,15 @@ TsFileWriterManager recoverFromDisk() throws IOException {
745781
try {
746782
channel =
747783
FileChannel.open(tsFile.toPath(), StandardOpenOption.READ, StandardOpenOption.WRITE);
748-
channel.position(resumeOffset);
784+
channel.position(sealed ? tsFile.length() : resumeOffset);
749785
final TsFilePrecalculatedChunkWriter writer =
750-
new TsFilePrecalculatedChunkWriter(tsFile, channel);
786+
new TsFilePrecalculatedChunkWriter(tsFile, channel, sealed);
751787
writer.setRejectHoles(LoadConsensusCompat.usesIoTConsensusV2());
752788
writer.restoreChunkMetadata(restoreChunkMetadata(records));
753789
publishedLengths.put(tsFile, tsFile.length());
754790

755791
timePartition2Writer.put(timePartitionSlot, writer);
756-
timePartition2Resource.put(timePartitionSlot, new TsFileResource(tsFile));
792+
timePartition2Resource.put(timePartitionSlot, resource);
757793
timePartition2Progress.put(timePartitionSlot, progress);
758794

759795
final File modificationFile = ModificationFile.getExclusiveMods(tsFile);
@@ -928,22 +964,28 @@ private static void deleteIfExistsWithRetry(final Path path) throws IOException
928964

929965
/** The name a staged file of this region and time partition has. */
930966
private String stagedFileName(final TTimePartitionSlot timePartitionSlot) {
967+
return stagedFilePrefix() + timePartitionSlot.getStartTime();
968+
}
969+
970+
private String stagedFilePrefix() {
931971
return String.join(
932972
IoTDBConstant.FILE_NAME_SEPARATOR,
933973
dataRegion.getDatabaseName(),
934974
dataRegion.getDataRegionIdString(),
935-
Long.toString(timePartitionSlot.getStartTime()));
975+
"");
936976
}
937977

938978
private static String uuidDirName(final File tsFile) {
939979
final File parent = tsFile.getParentFile();
940980
return parent == null ? tsFile.getName() : parent.getName();
941981
}
942982

943-
private static long parseTimePartitionStart(final String tsFileName) {
944-
final String baseName = tsFileName.substring(0, tsFileName.lastIndexOf('.'));
945-
final int lastDash = baseName.lastIndexOf(IoTDBConstant.FILE_NAME_SEPARATOR);
946-
return Long.parseLong(baseName.substring(lastDash + 1));
983+
private long parseTimePartitionStart(final String tsFileName) {
984+
// The partition may be negative: its sign is part of the number, not a name separator.
985+
return Long.parseLong(
986+
tsFileName.substring(
987+
stagedFilePrefix().length(),
988+
tsFileName.length() - TsFileConstant.TSFILE_SUFFIX.length()));
947989
}
948990

949991
/** The bytes this piece added to each staged file, as the refs a replica reads them back with. */

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.apache.iotdb.db.conf.IoTDBConfig;
2525
import org.apache.iotdb.db.conf.IoTDBDescriptor;
2626
import org.apache.iotdb.db.exception.load.LoadFileException;
27+
import org.apache.iotdb.db.i18n.DataNodeQueryMessages;
2728
import org.apache.iotdb.db.i18n.StorageEngineMessages;
2829
import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry;
2930
import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile;
@@ -455,7 +456,12 @@ private void getChunkMetadata(
455456

456457
private void handleModification(List<ModEntry> deletions) throws LoadFileException {
457458
for (final ModEntry mod : deletions) {
458-
consumer.apply(new DeletionData(mod));
459+
if (!consumer.apply(new DeletionData(mod))) {
460+
throw new IllegalStateException(
461+
String.format(
462+
DataNodeQueryMessages.DISPATCH_TSFILEDATA_ERROR_WHEN_PARSING_TSFILE_S,
463+
tsFile.getPath()));
464+
}
459465
}
460466
}
461467

0 commit comments

Comments
 (0)