Pipe: throttle TsFile parser retries under memory pressure

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

@@ -295,6 +295,9 @@ public final class DataNodePipeMessages {
   public static final String FAILED_TO_ALLOCATE_MEMORY_FOR_PARSING_TSFILE =
       "{}: failed to allocate memory for parsing TsFile {}, tablet event no. {}, "
           + "will release parser memory and retry the TsFile event later.";
+  public static final String FAILED_TO_CONSUME_PARSED_TABLET_FROM_TSFILE_KEEP_PARSER =
+      "{}: failed to consume parsed tablet from TsFile {}, tablet event no. {}, retry count is {}, "
+          + "will keep parser and retry locally for a short time.";
   public static final String FAILED_TO_BUILD_TABLET = "Failed to build tablet";
   public static final String FAILED_TO_CHECK_NEXT = "Failed to check next";
   public static final String FAILED_TO_CLOSE_TSFILEREADER = "Failed to close TsFileReader";

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

@@ -281,6 +281,9 @@ public final class DataNodePipeMessages {
   public static final String FAILED_TO_ALLOCATE_MEMORY_FOR_PARSING_TSFILE =
       "{}：为解析 TsFile {} 分配内存失败，tablet 事件编号 {}，"
           + "将释放解析器内存并稍后重试该 TsFile 事件。";
+  public static final String FAILED_TO_CONSUME_PARSED_TABLET_FROM_TSFILE_KEEP_PARSER =
+      "{}：消费 TsFile {} 解析出的 tablet 失败，tablet 事件编号 {}，重试次数 {}，"
+          + "将暂时保留解析器并在本地短暂重试。";
   public static final String FAILED_TO_BUILD_TABLET = "构建 tablet 失败";
   public static final String FAILED_TO_CHECK_NEXT = "check next 失败";
   public static final String FAILED_TO_CLOSE_TSFILEREADER = "关闭 TsFileReader 失败";

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

@@ -97,6 +97,7 @@ public class PipeTsFileInsertionEvent extends PipeInsertionEvent
   private final boolean isTsFileSealed;
   private final AtomicBoolean isClosed;
   private final AtomicReference<TsFileInsertionEventParser> eventParser;
+  private final AtomicBoolean isTsFileParserMemoryReserved = new AtomicBoolean(false);
 
   // The point count of the TsFile. Used for metrics on IoTConsensusV2' receiver side.
   // May be updated after it is flushed. Should be negative if not set.
@@ -733,12 +734,8 @@ public void consumeTabletInsertionEventsWithRetry(
       while (iterator.hasNext()) {
         final TabletInsertionEvent parsedEvent = iterator.next();
         tabletEventCount++;
-        try {
-          consumer.consume((PipeRawTabletInsertionEvent) parsedEvent);
-        } catch (final PipeRuntimeOutOfMemoryCriticalException e) {
-          releaseParsedTabletEvent(parsedEvent);
-          throw e;
-        }
+        consumeParsedTabletInsertionEventWithRetry(
+            consumer, callerName, tabletEventCount, parsedEvent);
       }
     } catch (final PipeRuntimeOutOfMemoryCriticalException e) {
       close();
@@ -752,6 +749,57 @@ public void consumeTabletInsertionEventsWithRetry(
     }
   }
 
+  private void consumeParsedTabletInsertionEventWithRetry(
+      final TabletInsertionEventConsumer consumer,
+      final String callerName,
+      final int tabletEventCount,
+      final TabletInsertionEvent parsedEvent)
+      throws Exception {
+    final PipeMemoryManager memoryManager = PipeDataNodeResourceManager.memory();
+    long firstOutOfMemoryTimeInMs = Long.MIN_VALUE;
+    int retryCount = 0;
+    while (true) {
+      try {
+        consumer.consume((PipeRawTabletInsertionEvent) parsedEvent);
+        return;
+      } catch (final PipeRuntimeOutOfMemoryCriticalException e) {
+        if (firstOutOfMemoryTimeInMs == Long.MIN_VALUE) {
+          firstOutOfMemoryTimeInMs = System.currentTimeMillis();
+        }
+        if (memoryManager.shouldReleaseTsFileParserOnOutOfMemory(
+            firstOutOfMemoryTimeInMs, ++retryCount)) {
+          releaseParsedTabletEvent(parsedEvent);
+          throw e;
+        }
+        logParserRetryOnOutOfMemory(callerName, tabletEventCount, retryCount, e);
+        try {
+          Thread.sleep(PipeConfig.getInstance().getPipeMemoryAllocateRetryIntervalInMs());
+        } catch (final InterruptedException interruptedException) {
+          Thread.currentThread().interrupt();
+          releaseParsedTabletEvent(parsedEvent);
+          throw e;
+        }
+      }
+    }
+  }
+
+  private void logParserRetryOnOutOfMemory(
+      final String callerName,
+      final int tabletEventCount,
+      final int retryCount,
+      final PipeRuntimeOutOfMemoryCriticalException e) {
+    if (retryCount != 1 && retryCount % 10 != 0) {
+      return;
+    }
+    LOGGER.warn(
+        DataNodePipeMessages.FAILED_TO_CONSUME_PARSED_TABLET_FROM_TSFILE_KEEP_PARSER,
+        callerName,
+        getTsFile(),
+        tabletEventCount,
+        retryCount,
+        e);
+  }
+
   private void releaseParsedTabletEvent(final TabletInsertionEvent parsedEvent) {
     if (parsedEvent instanceof PipeRawTabletInsertionEvent
         && ((PipeRawTabletInsertionEvent) parsedEvent).getReferenceCount() == 0
@@ -801,7 +849,7 @@ public Iterable<TabletInsertionEvent> toTabletInsertionEvents(final long timeout
 
   private void waitForResourceEnough4Parsing(final long timeoutMs) throws InterruptedException {
     final PipeMemoryManager memoryManager = PipeDataNodeResourceManager.memory();
-    if (memoryManager.isEnough4TabletParsing()) {
+    if (tryReserveTsFileParserMemory(memoryManager)) {
       return;
     }
 
@@ -810,7 +858,7 @@ private void waitForResourceEnough4Parsing(final long timeoutMs) throws Interrup
 
     final long memoryCheckIntervalMs =
         PipeConfig.getInstance().getPipeCheckMemoryEnoughIntervalMs();
-    while (!memoryManager.isEnough4TabletParsing()) {
+    while (!tryReserveTsFileParserMemory(memoryManager)) {
       Thread.sleep(memoryCheckIntervalMs);
 
       final long currentTime = System.currentTimeMillis();
@@ -847,6 +895,29 @@ private void waitForResourceEnough4Parsing(final long timeoutMs) throws Interrup
         waitTimeSeconds);
   }
 
+  private boolean tryReserveTsFileParserMemory(final PipeMemoryManager memoryManager) {
+    synchronized (isTsFileParserMemoryReserved) {
+      if (isTsFileParserMemoryReserved.get()) {
+        return true;
+      }
+
+      if (!memoryManager.tryReserveTsFileParserMemory()) {
+        return false;
+      }
+
+      isTsFileParserMemoryReserved.set(true);
+      return true;
+    }
+  }
+
+  private void releaseTsFileParserMemoryIfReserved() {
+    synchronized (isTsFileParserMemoryReserved) {
+      if (isTsFileParserMemoryReserved.compareAndSet(true, false)) {
+        PipeDataNodeResourceManager.memory().releaseTsFileParserMemory();
+      }
+    }
+  }
+
   /** The method is used to prevent circular replication in IoTConsensusV2 */
   public boolean isGeneratedByIoTConsensusV2() {
     return isGeneratedByIoTConsensusV2;
@@ -922,6 +993,7 @@ public void close() {
           }
           return null;
         });
+    releaseTsFileParserMemoryIfReserved();
   }
 
   /////////////////////////// Object ///////////////////////////
@@ -962,7 +1034,8 @@ public PipeEventResource eventResourceBuilder() {
         this.isWithMod,
         this.modFile,
         this.sharedModFile,
-        this.eventParser);
+        this.eventParser,
+        this.isTsFileParserMemoryReserved);
   }
 
   private static class PipeTsFileInsertionEventResource extends PipeEventResource {
@@ -974,6 +1047,7 @@ private static class PipeTsFileInsertionEventResource extends PipeEventResource
     private final AtomicReference<TsFileInsertionEventParser> eventParser;
     private final String pipeName;
     private final long creationTime;
+    private final AtomicBoolean isTsFileParserMemoryReserved;
 
     private PipeTsFileInsertionEventResource(
         final AtomicBoolean isReleased,
@@ -984,7 +1058,8 @@ private PipeTsFileInsertionEventResource(
         final boolean isWithMod,
         final File modFile,
         final File sharedModFile,
-        final AtomicReference<TsFileInsertionEventParser> eventParser) {
+        final AtomicReference<TsFileInsertionEventParser> eventParser,
+        final AtomicBoolean isTsFileParserMemoryReserved) {
       super(isReleased, referenceCount);
       this.pipeName = pipeName;
       this.creationTime = creationTime;
@@ -993,6 +1068,7 @@ private PipeTsFileInsertionEventResource(
       this.modFile = modFile;
       this.sharedModFile = sharedModFile;
       this.eventParser = eventParser;
+      this.isTsFileParserMemoryReserved = isTsFileParserMemoryReserved;
     }
 
     @Override
@@ -1016,6 +1092,11 @@ protected void finalizeResource() {
               }
               return null;
             });
+        synchronized (isTsFileParserMemoryReserved) {
+          if (isTsFileParserMemoryReserved.compareAndSet(true, false)) {
+            PipeDataNodeResourceManager.memory().releaseTsFileParserMemory();
+          }
+        }
       } catch (final Exception e) {
         LOGGER.warn(
             DataNodePipeMessages.DECREASE_REFERENCE_COUNT_FOR_TSFILE_ERROR, tsFile.getPath(), e);

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java

@@ -60,6 +60,8 @@ public class PipeMemoryManager {
 
   private volatile long usedMemorySizeInBytesOfTsFiles;
 
+  private volatile long reservedTsFileParserCount;
+
   // Only non-zero memory blocks will be added to this set.
   private final Set<PipeMemoryBlock> allocatedBlocks = new HashSet<>();
   private final Set<PipeMemoryBlock> shrinkableBlocks = new HashSet<>();
@@ -101,6 +103,34 @@ private double allowedMaxMemorySizeInBytesOfTsTiles() {
         * getTotalNonFloatingMemorySizeInBytes();
   }
 
+  private static long getTsFileParserMemorySizeInBytes() {
+    return Math.max(
+        PIPE_CONFIG.getTsFileParserMemory(), PIPE_CONFIG.getPipeMemoryAllocateMinSizeInBytes());
+  }
+
+  private long getReservedTsFileParserMemorySizeInBytes() {
+    return reservedTsFileParserCount * getTsFileParserMemorySizeInBytes();
+  }
+
+  private boolean isEnough4TabletParsingWithReservedParserMemory(final long extraMemoryInBytes) {
+    final double tabletMemoryWithParserMemory =
+        (double) usedMemorySizeInBytesOfTablets
+            + getReservedTsFileParserMemorySizeInBytes()
+            + extraMemoryInBytes;
+    return tabletMemoryWithParserMemory + (double) usedMemorySizeInBytesOfTsFiles
+            < EXCEED_PROTECT_THRESHOLD * allowedMaxMemorySizeInBytesOfTabletsAndTsFiles()
+        && tabletMemoryWithParserMemory
+            < EXCEED_PROTECT_THRESHOLD * allowedMaxMemorySizeInBytesOfTablets();
+  }
+
+  private boolean isHardEnough4TabletParsingWithReservedParserMemory() {
+    final double tabletMemoryWithParserMemory =
+        (double) usedMemorySizeInBytesOfTablets + getReservedTsFileParserMemorySizeInBytes();
+    return tabletMemoryWithParserMemory + (double) usedMemorySizeInBytesOfTsFiles
+            < allowedMaxMemorySizeInBytesOfTabletsAndTsFiles()
+        && tabletMemoryWithParserMemory < allowedMaxMemorySizeInBytesOfTablets();
+  }
+
   public boolean isEnough4TabletParsing() {
     return (double) usedMemorySizeInBytesOfTablets + (double) usedMemorySizeInBytesOfTsFiles
             < EXCEED_PROTECT_THRESHOLD * allowedMaxMemorySizeInBytesOfTabletsAndTsFiles()
@@ -114,6 +144,54 @@ < allowedMaxMemorySizeInBytesOfTabletsAndTsFiles()
         && (double) usedMemorySizeInBytesOfTablets < allowedMaxMemorySizeInBytesOfTablets();
   }
 
+  public synchronized boolean tryReserveTsFileParserMemory() {
+    if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
+      return true;
+    }
+
+    final long parserMemorySizeInBytes = getTsFileParserMemorySizeInBytes();
+    if (isEnough4TabletParsingWithReservedParserMemory(parserMemorySizeInBytes)) {
+      reservedTsFileParserCount++;
+      return true;
+    }
+
+    return false;
+  }
+
+  public synchronized void releaseTsFileParserMemory() {
+    if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
+      return;
+    }
+
+    reservedTsFileParserCount = Math.max(0, reservedTsFileParserCount - 1);
+    this.notifyAll();
+  }
+
+  public boolean shouldReleaseTsFileParserOnOutOfMemory(
+      final long firstOutOfMemoryTimeInMs, final int retryCount) {
+    final long retryIntervalInMs = PIPE_CONFIG.getPipeMemoryAllocateRetryIntervalInMs();
+    final long minRetryTimeInMs = Math.max(retryIntervalInMs * 2, 1);
+    final long maxRetryTimeInMs =
+        Math.max(
+            minRetryTimeInMs, retryIntervalInMs * PIPE_CONFIG.getPipeMemoryAllocateMaxRetries());
+
+    final long elapsedTimeInMs = System.currentTimeMillis() - firstOutOfMemoryTimeInMs;
+    if (elapsedTimeInMs < minRetryTimeInMs) {
+      return false;
+    }
+
+    if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
+      return elapsedTimeInMs >= maxRetryTimeInMs;
+    }
+
+    if (!isHardEnough4TabletParsingWithReservedParserMemory()) {
+      return true;
+    }
+
+    return retryCount >= PIPE_CONFIG.getPipeMemoryAllocateMaxRetries()
+        || elapsedTimeInMs >= maxRetryTimeInMs;
+  }
+
   public boolean isEnough4TsFileSlicing() {
     return (double) usedMemorySizeInBytesOfTablets + (double) usedMemorySizeInBytesOfTsFiles
             < EXCEED_PROTECT_THRESHOLD * allowedMaxMemorySizeInBytesOfTabletsAndTsFiles()

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionEventParserTest.java

@@ -237,6 +237,62 @@ public void testConsumeTabletInsertionEventsWithRetryReleasesParserOnOutOfMemory
     Assert.assertNull(getEventParser(event).get());
   }
 
+  @Test
+  public void testConsumeTabletInsertionEventsWithRetryKeepsParserForTransientOutOfMemory()
+      throws Exception {
+    nonalignedTsFile =
+        TsFileGeneratorUtils.generateNonAlignedTsFile(
+            "nonaligned-consume-transient-oom.tsfile", 1, 1, 10, 0, 100, 10, 10);
+    resource = new TsFileResource(nonalignedTsFile);
+    resource.setStatusForTest(TsFileResourceStatus.NORMAL);
+
+    final IDeviceID deviceID = IDeviceID.Factory.DEFAULT_FACTORY.create("root.testsg.d0");
+    resource.updateStartTime(deviceID, 0);
+    resource.updateEndTime(deviceID, 9);
+
+    final PipeTsFileInsertionEvent event =
+        new PipeTsFileInsertionEvent(
+            false,
+            "root",
+            resource,
+            null,
+            false,
+            false,
+            false,
+            null,
+            null,
+            0,
+            null,
+            new PrefixTreePattern("root"),
+            null,
+            null,
+            null,
+            null,
+            true,
+            Long.MIN_VALUE,
+            Long.MAX_VALUE);
+    final AtomicInteger retryCount = new AtomicInteger(0);
+    final AtomicReference<PipeRawTabletInsertionEvent> parsedEventReference =
+        new AtomicReference<>();
+
+    event.consumeTabletInsertionEventsWithRetry(
+        parsedEvent -> {
+          parsedEventReference.set(parsedEvent);
+          if (retryCount.getAndIncrement() == 0) {
+            throw new PipeRuntimeOutOfMemoryCriticalException("transient oom");
+          }
+          parsedEvent.clearReferenceCount(getClass().getName());
+        },
+        "test");
+
+    Assert.assertEquals(2, retryCount.get());
+    Assert.assertNotNull(parsedEventReference.get());
+    Assert.assertTrue(parsedEventReference.get().isReleased());
+    Assert.assertNotNull(getEventParser(event).get());
+
+    event.close();
+  }
+
   @Test
   public void testScanParserSplitNonAlignedSinglePageChunkByEstimatedPageMemory() throws Exception {
     final long originalPipeMaxReaderChunkSize =