From 4c46c7c5c621b046f43eecbe536930b8df6ad223 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:18:12 +0800 Subject: [PATCH] Pipe: throttle TsFile parser retries under memory pressure (#18141) (cherry picked from commit 82690873b64ffb19d23cbd688cd86610b70ad231) --- .../tsfile/PipeTsFileInsertionEvent.java | 101 ++++++++++++++++-- .../resource/memory/PipeMemoryManager.java | 80 ++++++++++++++ .../TsFileInsertionDataContainerTest.java | 62 +++++++++++ 3 files changed, 233 insertions(+), 10 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java index 2416e6caf82d..98c660ea3513 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java +++ b/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 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 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 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 dataContainer) { + final AtomicReference 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); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java index afdbafad1a7c..387a4228a2cd 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java +++ b/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 allocatedBlocks = new HashSet<>(); private final Set 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() diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionDataContainerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionDataContainerTest.java index fa82d15db61e..862e2b553fe4 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionDataContainerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionDataContainerTest.java @@ -105,11 +105,19 @@ 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 @@ -117,6 +125,12 @@ 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 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 =