[Fix] Harden split LOAD recovery, deletion dispatch and memory bounds

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

@@ -35,10 +35,14 @@
 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 @@ interface DispatchCallback {
   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 @@ void offerChunk(final ChunkData chunkData, final TRegionReplicaSet replicaSet)
         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 @@ void addDeletionToAll(final DeletionData deletionData) {
   /** 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 @@ boolean dispatchLargestUntilMemoryEnough() throws LoadFileException {
       final long pieceSize = pieceNode.getDataSize();
 
       if (!dispatchCallback.dispatch(pieceNode, pair.getLeft())) {
+        trackPieceSize(key, pieceSize);
         return false;
       }
 
@@ -192,10 +207,12 @@ boolean flushAll() throws LoadFileException {
   }
 
   // -------------------------------------------------------------------------
-  // 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 @@ private LoadTsFilePieceNode createPieceNode(final PartitionKey key) {
         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 @@ void clear() {
 
   private record PartitionKey(TConsensusGroupId regionId, long timePartitionStart) {}
 
-  private record HeapEntry(PartitionKey key, long dataSize) {}
+  private record PieceSize(PartitionKey key, long dataSize) {}
 }

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 @@ private boolean addOrSendChunkData(final ChunkData chunkData) throws LoadFileExc
   }
 
   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 @@ boolean sendAllTsFileData() throws LoadFileException {
   /** Drops all unrouted chunks, clears partition buffers, and releases cached memory accounting. */
   void clear() {
     nonDirectionalChunkData.clear();
+    isProcessingDeletions = false;
     dispatcher.clear();
     memoryBuffer.clear();
   }

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.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.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.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 @@ private void writeDeletion(final DeletionData deletionData) throws IOException {
   // 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,36 +512,23 @@ void prepare(
     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();
       if (writer.isSealed()) {
         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 @@ TsFileWriterManager recoverFromDisk() throws IOException {
       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 @@ TsFileWriterManager recoverFromDisk() throws IOException {
       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,22 +964,28 @@ private static void deleteIfExistsWithRetry(final Path path) throws IOException
 
   /** 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) {
     final File parent = tsFile.getParentFile();
     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. */

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

@@ -24,6 +24,7 @@
 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 @@ private void getChunkMetadata(
 
   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()));
+      }
     }
   }