Skip to content

Commit 81b5ae4

Browse files
CaideyipiCaideyipi
committed
Fix LOAD recovery, deletion delivery and buffer eviction
1 parent b030a75 commit 81b5ae4

9 files changed

Lines changed: 548 additions & 37 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: 20 additions & 4 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;
@@ -741,7 +743,21 @@ TsFileWriterManager recoverFromDisk() throws IOException {
741743
final TTimePartitionSlot timePartitionSlot =
742744
new TTimePartitionSlot(parseTimePartitionStart(tsFile.getName()));
743745
final long resumeOffset = progress.getTotalLength();
744-
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) {
745761
try (final FileChannel truncatingChannel =
746762
FileChannel.open(tsFile.toPath(), StandardOpenOption.READ, StandardOpenOption.WRITE)) {
747763
truncatingChannel.truncate(resumeOffset);
@@ -765,15 +781,15 @@ TsFileWriterManager recoverFromDisk() throws IOException {
765781
try {
766782
channel =
767783
FileChannel.open(tsFile.toPath(), StandardOpenOption.READ, StandardOpenOption.WRITE);
768-
channel.position(resumeOffset);
784+
channel.position(sealed ? tsFile.length() : resumeOffset);
769785
final TsFilePrecalculatedChunkWriter writer =
770-
new TsFilePrecalculatedChunkWriter(tsFile, channel);
786+
new TsFilePrecalculatedChunkWriter(tsFile, channel, sealed);
771787
writer.setRejectHoles(LoadConsensusCompat.usesIoTConsensusV2());
772788
writer.restoreChunkMetadata(restoreChunkMetadata(records));
773789
publishedLengths.put(tsFile, tsFile.length());
774790

775791
timePartition2Writer.put(timePartitionSlot, writer);
776-
timePartition2Resource.put(timePartitionSlot, new TsFileResource(tsFile));
792+
timePartition2Resource.put(timePartitionSlot, resource);
777793
timePartition2Progress.put(timePartitionSlot, progress);
778794

779795
final File modificationFile = ModificationFile.getExclusiveMods(tsFile);

‎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

‎iotdb-core/datanode/src/main/java/org/apache/tsfile/write/TsFilePrecalculatedChunkWriter.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -109,12 +109,18 @@ public TsFilePrecalculatedChunkWriter(final File file) throws IOException {
109109

110110
public TsFilePrecalculatedChunkWriter(final File file, final FileChannel channel)
111111
throws IOException {
112+
this(file, channel, false);
113+
}
114+
115+
/** Resumes a staged file, preserving a completed PREPARE when it is already sealed. */
116+
public TsFilePrecalculatedChunkWriter(
117+
final File file, final FileChannel channel, final boolean sealed) throws IOException {
112118
this.file = file;
113119
this.out =
114120
new WritableFileChannelOutput(
115121
Objects.requireNonNull(
116122
channel, StorageEngineMessages.EXCEPTION_CHANNEL_CANNOT_BE_NULL_0F79E0FB));
117-
this.sealed = false;
123+
this.sealed = sealed;
118124
this.maxPhysicalFileSize = channel.size();
119125
LOGGER.info(
120126
StorageEngineMessages

0 commit comments

Comments
 (0)