[Pipe] Support receiver-side async takeover for TsFile conversion

iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java

@@ -539,6 +539,10 @@ private StorageEngineMessages() {}
   public static final String ACTIVE_LOAD_METRIC_COLLECTOR_REGISTERED = "Active load metric collector periodical jobs registered";
   public static final String DATABASE_NAME_MUST_NOT_BE_EMPTY = "Database name must not be empty.";
   public static final String USER_NAME_MUST_NOT_BE_EMPTY = "User name must not be empty.";
+  public static final String EXCEPTION_CONVERSION_TASK_ID_MUST_NOT_BE_EMPTY_411D064E =
+      "Conversion task ID must not be empty.";
+  public static final String LOG_FAILED_TO_CLOSE_PIPE_TSFILE_CONVERSION_CONTEXT_8E4D886B =
+      "Failed to close Pipe TsFile conversion context.";
   public static final String ERROR_EXECUTING_ACTIVE_LOAD_JOB = "Error occurred when executing active load periodical job.";
   public static final String ACTIVE_LOAD_EXECUTOR_STARTED = "Active load periodical jobs executor is started successfully.";
   public static final String ACTIVE_LOAD_EXECUTOR_STOPPED = "Active load periodical jobs executor is stopped successfully.";

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java

@@ -539,6 +539,10 @@ private StorageEngineMessages() {}
   public static final String ACTIVE_LOAD_METRIC_COLLECTOR_REGISTERED = "Active 加载指标收集定期任务已注册";
   public static final String DATABASE_NAME_MUST_NOT_BE_EMPTY = "数据库名称不能为空。";
   public static final String USER_NAME_MUST_NOT_BE_EMPTY = "用户名不能为空。";
+  public static final String EXCEPTION_CONVERSION_TASK_ID_MUST_NOT_BE_EMPTY_411D064E =
+      "转换任务 ID 不能为空。";
+  public static final String LOG_FAILED_TO_CLOSE_PIPE_TSFILE_CONVERSION_CONTEXT_8E4D886B =
+      "关闭 Pipe TsFile 转换上下文失败。";
   public static final String ERROR_EXECUTING_ACTIVE_LOAD_JOB = "执行 Active 加载定期任务时发生错误。";
   public static final String ACTIVE_LOAD_EXECUTOR_STARTED = "Active 加载定期任务执行器已成功启动。";
   public static final String ACTIVE_LOAD_EXECUTOR_STOPPED = "Active 加载定期任务执行器已成功停止。";

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParser.java

@@ -30,6 +30,7 @@
 import org.apache.iotdb.commons.pipe.config.PipeConfig;
 import org.apache.iotdb.commons.pipe.datastructure.pattern.TreePattern;
 import org.apache.iotdb.db.auth.AuthorityChecker;
+import org.apache.iotdb.db.exception.load.LoadRuntimeOutOfMemoryException;
 import org.apache.iotdb.db.i18n.DataNodePipeMessages;
 import org.apache.iotdb.db.pipe.event.common.PipeInsertionEvent;
 import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent;
@@ -93,6 +94,8 @@ public class TsFileInsertionEventScanParser extends TsFileInsertionEventParser {
 
   private IChunkReader chunkReader;
   private BatchData data;
+  private BatchData pendingPageDataAfterMemoryPressure;
+  private Tablet pendingTabletAfterMemoryPressure;
   private final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForBatchData;
   private final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForChunk;
   private TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForTsFileInput;
@@ -115,6 +118,7 @@ public class TsFileInsertionEventScanParser extends TsFileInsertionEventParser {
   private CachedAlignedValueChunk cachedAlignedValueChunk;
 
   private byte lastMarker = Byte.MIN_VALUE;
+  private boolean shouldRetryChunkHeaderAfterMemoryPressure;
 
   public TsFileInsertionEventScanParser(
       final String pipeName,
@@ -271,6 +275,7 @@ public Iterable<TabletInsertionEvent> toTabletInsertionEvents() {
                 @Override
                 public boolean hasNext() {
                   throwIfDeferredException();
+                  retryChunkHeaderAfterMemoryPressureIfNecessary();
                   final boolean hasNext = Objects.nonNull(chunkReader);
                   if (hasNext && !parseStartTimeRecorded) {
                     // Record start time on first hasNext() that returns true
@@ -290,8 +295,11 @@ public TabletInsertionEvent next() {
                   }
 
                   // Release the previous parser-owned tablet buffer before allocating the next
-                  // tablet.
-                  releaseTabletMemoryBlock();
+                  // tablet. A tablet retained after memory pressure still owns this block and must
+                  // keep it until the retry has successfully returned the pending event.
+                  if (pendingTabletAfterMemoryPressure == null) {
+                    releaseTabletMemoryBlock();
+                  }
                   // currentIsAligned is initialized when TsFileInsertionEventScanParser is
                   // constructed.
                   // When the getNextTablet function is called, currentIsAligned may be updated,
@@ -350,6 +358,7 @@ public Iterable<Pair<Tablet, Boolean>> toTabletWithIsAligneds() {
           @Override
           public boolean hasNext() {
             throwIfDeferredException();
+            retryChunkHeaderAfterMemoryPressureIfNecessary();
             return Objects.nonNull(chunkReader);
           }
 
@@ -394,15 +403,25 @@ public List<String> getCurrentMeasurements() {
   }
 
   private Tablet getNextTablet() {
+    Tablet tablet = null;
+    boolean tabletMemoryReserved = false;
     try {
-      Tablet tablet = null;
+      if (data != null && !data.hasCurrent() && pendingPageDataAfterMemoryPressure != null) {
+        data = nextPageData();
+      }
+      tablet = pendingTabletAfterMemoryPressure;
+      tabletMemoryReserved = tablet != null;
+      pendingTabletAfterMemoryPressure = null;
 
       if (!data.hasCurrent()) {
-        tablet = new Tablet(currentDeviceString, currentMeasurements, 1);
-        return tablet;
+        if (tablet != null) {
+          PipeTabletUtils.compactBitMaps(tablet);
+          return tablet;
+        }
+        return new Tablet(currentDeviceString, currentMeasurements, 1);
       }
 
-      boolean isFirstRow = true;
+      boolean isFirstRow = tablet == null;
       while (data.hasCurrent()) {
         if (currentIsMultiPage
             || data.currentTime() >= startTime && data.currentTime() <= endTime) {
@@ -417,6 +436,8 @@ private Tablet getNextTablet() {
                 < rowCountAndMemorySize.getRight()) {
               allocatedMemoryBlockForTablet.forceResize(rowCountAndMemorySize.getRight());
             }
+            tabletMemoryReserved = true;
+            pendingTabletAfterMemoryPressure = tablet;
             isFirstRow = false;
           }
 
@@ -450,8 +471,10 @@ private Tablet getNextTablet() {
         }
       }
       PipeTabletUtils.compactBitMaps(tablet);
+      pendingTabletAfterMemoryPressure = null;
       return tablet;
-    } catch (final PipeRuntimeOutOfMemoryCriticalException e) {
+    } catch (final PipeRuntimeOutOfMemoryCriticalException | LoadRuntimeOutOfMemoryException e) {
+      pendingTabletAfterMemoryPressure = tabletMemoryReserved ? tablet : null;
       // Keep the parser state so the caller can yield its parser slot and retry from the same
       // unconsumed data after memory is available again.
       throw e;
@@ -473,6 +496,21 @@ private void throwIfDeferredException() {
         exception);
   }
 
+  private void retryChunkHeaderAfterMemoryPressureIfNecessary() {
+    if (!shouldRetryChunkHeaderAfterMemoryPressure) {
+      return;
+    }
+    try {
+      prepareData();
+      shouldRetryChunkHeaderAfterMemoryPressure = false;
+    } catch (final PipeRuntimeOutOfMemoryCriticalException | LoadRuntimeOutOfMemoryException e) {
+      throw e;
+    } catch (final Exception e) {
+      close();
+      throw new PipeException(DataNodePipeMessages.FAILED_TO_GET_NEXT_TABLET_INSERTION_EVENT, e);
+    }
+  }
+
   private boolean isLastTabletWithoutDeferredException() {
     return Objects.isNull(deferredException) && Objects.isNull(chunkReader);
   }
@@ -501,9 +539,24 @@ private void prepareData() throws IOException, IllegalPathException {
   }
 
   private BatchData nextPageData() throws IOException {
+    if (pendingPageDataAfterMemoryPressure != null) {
+      resizePageDataMemoryIfNeeded(
+          PipeMemoryWeightUtil.calculateBatchDataRamBytesUsed(pendingPageDataAfterMemoryPressure));
+      final BatchData pendingPageData = pendingPageDataAfterMemoryPressure;
+      pendingPageDataAfterMemoryPressure = null;
+      return pendingPageData;
+    }
+
     resizePageDataMemoryForCurrentPageIfNeeded();
     final BatchData nextData = chunkReader.nextPageData();
-    resizePageDataMemoryIfNeeded(PipeMemoryWeightUtil.calculateBatchDataRamBytesUsed(nextData));
+    try {
+      resizePageDataMemoryIfNeeded(PipeMemoryWeightUtil.calculateBatchDataRamBytesUsed(nextData));
+    } catch (final PipeRuntimeOutOfMemoryCriticalException | LoadRuntimeOutOfMemoryException e) {
+      // The chunk reader has already advanced. Retain the decoded page until its memory can be
+      // accounted for instead of advancing to the following page on retry.
+      pendingPageDataAfterMemoryPressure = nextData;
+      throw e;
+    }
     return nextData;
   }
 
@@ -699,8 +752,20 @@ private void moveToNextChunkReader()
               break;
             }
 
-            if (chunkHeader.getDataSize() > allocatedMemoryBlockForChunk.getMemoryUsageInBytes()) {
-              allocatedMemoryBlockForChunk.forceResize(chunkHeader.getDataSize());
+            try {
+              if (chunkHeader.getDataSize()
+                  > allocatedMemoryBlockForChunk.getMemoryUsageInBytes()) {
+                allocatedMemoryBlockForChunk.forceResize(chunkHeader.getDataSize());
+              }
+            } catch (final PipeRuntimeOutOfMemoryCriticalException
+                | LoadRuntimeOutOfMemoryException e) {
+              // The marker has already been consumed. Rewind to the start of the header and retain
+              // the marker so a retry does not interpret header bytes as a marker.
+              tsFileSequenceReader.position(currentChunkHeaderOffset + 1);
+              lastMarker = marker;
+              chunkReader = null;
+              shouldRetryChunkHeaderAfterMemoryPressure = true;
+              throw e;
             }
 
             Chunk chunk =
@@ -762,10 +827,19 @@ private void moveToNextChunkReader()
               cachedAlignedValueChunk = null;
             }
 
-            if (returnPendingAlignedChunkBeforeCaching(valueChunk)) {
-              return;
+            try {
+              if (returnPendingAlignedChunkBeforeCaching(valueChunk)) {
+                return;
+              }
+              cacheAlignedValueChunk(valueChunk);
+            } catch (final PipeRuntimeOutOfMemoryCriticalException
+                | LoadRuntimeOutOfMemoryException e) {
+              // The value chunk has already been read from the file. Keep it as the next logical
+              // input so a retry neither skips it nor reads it twice.
+              cachedAlignedValueChunk = valueChunk;
+              shouldRetryChunkHeaderAfterMemoryPressure = true;
+              throw e;
             }
-            cacheAlignedValueChunk(valueChunk);
             break;
           }
         case MetaMarker.CHUNK_GROUP_HEADER:
@@ -882,56 +956,70 @@ private boolean filterChunk(
   }
 
   private boolean useNextPendingAlignedChunk(final byte marker) throws IOException {
-    while (!pendingAlignedChunkGroups.isEmpty()) {
-      final PendingAlignedChunkGroup pendingAlignedChunkGroup = pendingAlignedChunkGroups.remove(0);
-      pendingAlignedChunkSize =
-          Math.max(0, pendingAlignedChunkSize - pendingAlignedChunkGroup.chunkSize);
-
-      if (pendingAlignedChunkGroup.valueChunkList.isEmpty()) {
-        continue;
-      }
+    try {
+      while (!pendingAlignedChunkGroups.isEmpty()) {
+        final PendingAlignedChunkGroup pendingAlignedChunkGroup = pendingAlignedChunkGroups.get(0);
 
-      final Chunk timeChunk = timeChunkList.get(pendingAlignedChunkGroup.timeChunkIndex);
-      timeChunk.getData().rewind();
-      for (final Chunk valueChunk : pendingAlignedChunkGroup.valueChunkList) {
-        valueChunk.getData().rewind();
-      }
+        if (pendingAlignedChunkGroup.valueChunkList.isEmpty()) {
+          pendingAlignedChunkGroups.remove(0);
+          pendingAlignedChunkSize =
+              Math.max(0, pendingAlignedChunkSize - pendingAlignedChunkGroup.chunkSize);
+          continue;
+        }
 
-      currentMeasurements.clear();
-      currentMeasurements.addAll(pendingAlignedChunkGroup.measurements);
-      modsInfos.clear();
-      modsInfos.addAll(pendingAlignedChunkGroup.modsInfos);
+        final Chunk timeChunk = timeChunkList.get(pendingAlignedChunkGroup.timeChunkIndex);
+        timeChunk.getData().rewind();
+        for (final Chunk valueChunk : pendingAlignedChunkGroup.valueChunkList) {
+          valueChunk.getData().rewind();
+        }
 
-      currentIsMultiPage = isMultiPageList.get(pendingAlignedChunkGroup.timeChunkIndex);
-      if (!currentIsMultiPage) {
-        resizePageDataMemoryIfNeeded(
-            AlignedSinglePageWholeChunkReader.calculatePageEstimatedMemoryUsageInBytes(
-                timeChunk, pendingAlignedChunkGroup.valueChunkList));
-      }
-      final List<Long> pageEstimatedMemoryUsageInBytesList =
-          currentIsMultiPage
-              ? AlignedSinglePageWholeChunkReader
-                  .calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(
-                      timeChunk, pendingAlignedChunkGroup.valueChunkList)
-              : Collections.emptyList();
-      final long maxPageEstimatedMemoryUsageInBytes =
-          pageEstimatedMemoryUsageInBytesList.isEmpty()
-              ? 0
-              : pageEstimatedMemoryUsageInBytesList.get(0);
-      resizePageDataMemoryIfNeeded(maxPageEstimatedMemoryUsageInBytes);
-      chunkReader =
-          currentIsMultiPage
-              ? new MemoryControlledChunkReader(
-                  new AlignedChunkReader(
-                      timeChunk, pendingAlignedChunkGroup.valueChunkList, filter),
-                  pageEstimatedMemoryUsageInBytesList)
-              : new AlignedSinglePageWholeChunkReader(
-                  timeChunk, pendingAlignedChunkGroup.valueChunkList, null);
-      currentIsAligned = true;
-      if (marker != Byte.MIN_VALUE) {
-        lastMarker = marker;
+        final boolean nextIsMultiPage =
+            isMultiPageList.get(pendingAlignedChunkGroup.timeChunkIndex);
+        if (!nextIsMultiPage) {
+          resizePageDataMemoryIfNeeded(
+              AlignedSinglePageWholeChunkReader.calculatePageEstimatedMemoryUsageInBytes(
+                  timeChunk, pendingAlignedChunkGroup.valueChunkList));
+        }
+        final List<Long> pageEstimatedMemoryUsageInBytesList =
+            nextIsMultiPage
+                ? AlignedSinglePageWholeChunkReader
+                    .calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(
+                        timeChunk, pendingAlignedChunkGroup.valueChunkList)
+                : Collections.emptyList();
+        final long maxPageEstimatedMemoryUsageInBytes =
+            pageEstimatedMemoryUsageInBytesList.isEmpty()
+                ? 0
+                : pageEstimatedMemoryUsageInBytesList.get(0);
+        resizePageDataMemoryIfNeeded(maxPageEstimatedMemoryUsageInBytes);
+        final IChunkReader nextChunkReader =
+            nextIsMultiPage
+                ? new MemoryControlledChunkReader(
+                    new AlignedChunkReader(
+                        timeChunk, pendingAlignedChunkGroup.valueChunkList, filter),
+                    pageEstimatedMemoryUsageInBytesList)
+                : new AlignedSinglePageWholeChunkReader(
+                    timeChunk, pendingAlignedChunkGroup.valueChunkList, null);
+
+        // Publish the state transition only after all memory reservations and reader construction
+        // succeed. Otherwise a retry would lose the already-read aligned chunk group.
+        pendingAlignedChunkGroups.remove(0);
+        pendingAlignedChunkSize =
+            Math.max(0, pendingAlignedChunkSize - pendingAlignedChunkGroup.chunkSize);
+        currentMeasurements.clear();
+        currentMeasurements.addAll(pendingAlignedChunkGroup.measurements);
+        modsInfos.clear();
+        modsInfos.addAll(pendingAlignedChunkGroup.modsInfos);
+        currentIsMultiPage = nextIsMultiPage;
+        chunkReader = nextChunkReader;
+        currentIsAligned = true;
+        if (marker != Byte.MIN_VALUE) {
+          lastMarker = marker;
+        }
+        return true;
       }
-      return true;
+    } catch (final PipeRuntimeOutOfMemoryCriticalException | LoadRuntimeOutOfMemoryException e) {
+      shouldRetryChunkHeaderAfterMemoryPressure = true;
+      throw e;
     }
     return false;
   }
@@ -1158,6 +1246,11 @@ private byte toValueChunkMarker(final ChunkHeader chunkHeader) {
   @Override
   public void close() {
     super.close();
+    pendingPageDataAfterMemoryPressure = null;
+    pendingTabletAfterMemoryPressure = null;
+    cachedAlignedValueChunk = null;
+    pendingAlignedChunkGroups.clear();
+    pendingAlignedChunkSize = 0;
 
     if (allocatedMemoryBlockForBatchData != null) {
       allocatedMemoryBlockForBatchData.close();