[To dev/1.3] Pipe: throttle TsFile parser retries under memory pressure

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

@@ -85,6 +85,7 @@ public class PipeTsFileInsertionEvent extends EnrichedEvent
 
   protected final AtomicBoolean isClosed;
   protected final AtomicReference<TsFileInsertionDataContainer> dataContainer;
+  private final AtomicBoolean isTsFileParserMemoryReserved = new AtomicBoolean(false);
 
   // The point count of the TsFile. Used for metrics on PipeConsensus' receiver side.
   // May be updated after it is flushed. Should be negative if not set.
@@ -487,12 +488,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();
@@ -506,6 +503,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(
+        "{}: failed to consume parsed tablet from TsFile {}, tablet event no. {}, retry count is {}, will keep parser and retry locally for a short time.",
+        callerName,
+        getTsFile(),
+        tabletEventCount,
+        retryCount,
+        e);
+  }
+
   private void releaseParsedTabletEvent(final TabletInsertionEvent parsedEvent) {
     if (parsedEvent instanceof PipeRawTabletInsertionEvent
         && ((PipeRawTabletInsertionEvent) parsedEvent).getReferenceCount() == 0
@@ -556,7 +604,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;
     }
 
@@ -565,7 +613,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();
@@ -600,6 +648,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 PipeConsensus */
   public boolean isGeneratedByPipeConsensus() {
     return isGeneratedByPipeConsensus;
@@ -668,6 +739,7 @@ public void close() {
           }
           return null;
         });
+    releaseTsFileParserMemoryIfReserved();
   }
 
   /////////////////////////// Object ///////////////////////////
@@ -706,7 +778,8 @@ public PipeEventResource eventResourceBuilder() {
         this.tsFile,
         this.isWithMod,
         this.modFile,
-        this.dataContainer);
+        this.dataContainer,
+        this.isTsFileParserMemoryReserved);
   }
 
   private static class PipeTsFileInsertionEventResource extends PipeEventResource {
@@ -716,6 +789,7 @@ private static class PipeTsFileInsertionEventResource extends PipeEventResource
     private final File modFile;
     private final AtomicReference<TsFileInsertionDataContainer> dataContainer;
     private final String pipeName;
+    private final AtomicBoolean isTsFileParserMemoryReserved;
 
     private PipeTsFileInsertionEventResource(
         final AtomicBoolean isReleased,
@@ -724,13 +798,15 @@ private PipeTsFileInsertionEventResource(
         final File tsFile,
         final boolean isWithMod,
         final File modFile,
-        final AtomicReference<TsFileInsertionDataContainer> dataContainer) {
+        final AtomicReference<TsFileInsertionDataContainer> dataContainer,
+        final AtomicBoolean isTsFileParserMemoryReserved) {
       super(isReleased, referenceCount);
       this.pipeName = pipeName;
       this.tsFile = tsFile;
       this.isWithMod = isWithMod;
       this.modFile = modFile;
       this.dataContainer = dataContainer;
+      this.isTsFileParserMemoryReserved = isTsFileParserMemoryReserved;
     }
 
     @Override
@@ -750,6 +826,11 @@ protected void finalizeResource() {
               }
               return null;
             });
+        synchronized (isTsFileParserMemoryReserved) {
+          if (isTsFileParserMemoryReserved.compareAndSet(true, false)) {
+            PipeDataNodeResourceManager.memory().releaseTsFileParserMemory();
+          }
+        }
       } catch (final Exception e) {
         LOGGER.warn("Decrease reference count for TsFile {} error.", tsFile.getPath(), e);
       }

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

@@ -54,6 +54,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<>();
@@ -103,6 +105,34 @@ private static double allowedMaxMemorySizeInBytesOfTsTiles() {
         * getTotalNonFloatingMemorySizeInBytes();
   }
 
+  private static long getTsFileParserMemorySizeInBytes() {
+    return Math.max(
+        PipeConfig.getInstance().getTsFileParserMemory(), MEMORY_ALLOCATE_MIN_SIZE_IN_BYTES);
+  }
+
+  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()
@@ -116,6 +146,56 @@ < 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 =
+        PipeConfig.getInstance().getPipeMemoryAllocateRetryIntervalInMs();
+    final long minRetryTimeInMs = Math.max(retryIntervalInMs * 2, 1);
+    final long maxRetryTimeInMs =
+        Math.max(
+            minRetryTimeInMs,
+            retryIntervalInMs * PipeConfig.getInstance().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 >= PipeConfig.getInstance().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/TsFileInsertionDataContainerTest.java

@@ -105,18 +105,32 @@ public class TsFileInsertionDataContainerTest {
   private File nonalignedTsFile;
   private TsFileResource resource;
   private boolean isPipeMemoryManagementEnabled;
+  private long originalPipeMemoryAllocateRetryIntervalInMs;
+  private int originalPipeMemoryAllocateMaxRetries;
 
   @Before
   public void setUp() throws Exception {
     isPipeMemoryManagementEnabled = PipeConfig.getInstance().getPipeMemoryManagementEnabled();
+    originalPipeMemoryAllocateRetryIntervalInMs =
+        PipeConfig.getInstance().getPipeMemoryAllocateRetryIntervalInMs();
+    originalPipeMemoryAllocateMaxRetries =
+        PipeConfig.getInstance().getPipeMemoryAllocateMaxRetries();
     CommonDescriptor.getInstance().getConfig().setPipeMemoryManagementEnabled(false);
+    CommonDescriptor.getInstance().getConfig().setPipeMemoryAllocateRetryIntervalInMs(1);
+    CommonDescriptor.getInstance().getConfig().setPipeMemoryAllocateMaxRetries(2);
   }
 
   @After
   public void tearDown() throws Exception {
     CommonDescriptor.getInstance()
         .getConfig()
         .setPipeMemoryManagementEnabled(isPipeMemoryManagementEnabled);
+    CommonDescriptor.getInstance()
+        .getConfig()
+        .setPipeMemoryAllocateRetryIntervalInMs(originalPipeMemoryAllocateRetryIntervalInMs);
+    CommonDescriptor.getInstance()
+        .getConfig()
+        .setPipeMemoryAllocateMaxRetries(originalPipeMemoryAllocateMaxRetries);
     if (alignedTsFile != null) {
       alignedTsFile.delete();
     }
@@ -216,6 +230,54 @@ public void testConsumeTabletInsertionEventsWithRetryReleasesContainerOnOutOfMem
     Assert.assertNull(getDataContainer(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 = new PlainDeviceID("root.testsg.d0");
+    resource.updateStartTime(deviceID, 0);
+    resource.updateEndTime(deviceID, 9);
+
+    final PipeTsFileInsertionEvent event =
+        new PipeTsFileInsertionEvent(
+            resource,
+            null,
+            false,
+            false,
+            false,
+            null,
+            0,
+            null,
+            new PrefixPipePattern("root"),
+            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(getDataContainer(event).get());
+
+    event.close();
+  }
+
   @Test
   public void testScanParserSplitNonAlignedSinglePageChunkByEstimatedPageMemory() throws Exception {
     final long originalPipeMaxReaderChunkSize =