From 807c7e605a0622eb91ae9368c563fc27c37eac39 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Tue, 15 Sep 2026 18:16:13 +0800 Subject: [PATCH] [Pipe] Restore sink progress under memory pressure --- .../evolvable/batch/PipeTabletEventBatch.java | 7 +++--- .../PipeTransferBatchReqBuilderTest.java | 25 +++++++++++++++++++ .../task/subtask/PipeAbstractSinkSubtask.java | 3 +-- .../pipe/task/PipeSleepIntervalTest.java | 17 +++++++++++++ 4 files changed, 46 insertions(+), 6 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventBatch.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventBatch.java index b84985675facf..4cf05ae619620 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventBatch.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTabletEventBatch.java @@ -24,8 +24,8 @@ import org.apache.iotdb.commons.pipe.event.EnrichedEvent; import org.apache.iotdb.db.i18n.DataNodePipeMessages; import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; +import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlock; import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager; -import org.apache.iotdb.db.pipe.resource.memory.PipeTabletMemoryBlock; import org.apache.iotdb.db.pipe.sink.protocol.thrift.async.IoTDBDataRegionAsyncSink; import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALPipeException; import org.apache.iotdb.pipe.api.event.Event; @@ -51,7 +51,7 @@ public abstract class PipeTabletEventBatch implements AutoCloseable { private long firstEventProcessingTime = Long.MIN_VALUE; protected long totalBufferSize = 0; - private final PipeTabletMemoryBlock allocatedMemoryBlock; + private final PipeMemoryBlock allocatedMemoryBlock; private boolean shouldEmitOnMemoryPressure = false; protected volatile boolean isClosed = false; @@ -64,8 +64,7 @@ protected PipeTabletEventBatch( // limit in buffer size this.maxBatchSizeInBytes = requestMaxBatchSizeInBytes; - this.allocatedMemoryBlock = - PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(0); + this.allocatedMemoryBlock = PipeDataNodeResourceManager.memory().forceAllocate(0); if (recordMetric != null) { this.recordMetric = recordMetric; } else { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java index b366c37ecb42b..d3710ae0e3849 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/payload/evolvable/batch/PipeTransferBatchReqBuilderTest.java @@ -22,6 +22,8 @@ import org.apache.iotdb.common.rpc.thrift.TEndPoint; import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; +import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; +import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager; import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters; import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent; @@ -94,6 +96,29 @@ protected boolean constructBatch(final TabletInsertionEvent event) { } } + @Test + public void testBatchMemoryIsNotCountedAsTabletMemory() throws Exception { + final PipeMemoryManager memoryManager = PipeDataNodeResourceManager.memory(); + final PipeRawTabletInsertionEvent event = createEvent(1); + + try { + Assert.assertTrue(event.increaseReferenceCount(getClass().getName())); + final long tabletMemoryBeforeBatch = memoryManager.getUsedMemorySizeInBytesOfTablets(); + final long totalMemoryBeforeBatch = memoryManager.getUsedMemorySizeInBytes(); + + try (final PipeTabletEventBatch batch = + new PipeTabletEventPlainBatch(Integer.MAX_VALUE, Long.MAX_VALUE, null)) { + batch.onEvent(event); + + Assert.assertEquals( + tabletMemoryBeforeBatch, memoryManager.getUsedMemorySizeInBytesOfTablets()); + Assert.assertTrue(memoryManager.getUsedMemorySizeInBytes() > totalMemoryBeforeBatch); + } + } finally { + event.clearReferenceCount(getClass().getName()); + } + } + private static PipeRawTabletInsertionEvent createEvent(final int value) { final Tablet tablet = new Tablet( diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/subtask/PipeAbstractSinkSubtask.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/subtask/PipeAbstractSinkSubtask.java index 7a07345c4ae2e..6dc6e28d8a179 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/subtask/PipeAbstractSinkSubtask.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/agent/task/subtask/PipeAbstractSinkSubtask.java @@ -361,9 +361,8 @@ protected void handleException(final Event event, final Exception e) { if (failureType == PipeResourceFailureType.MEMORY_TIMEOUT) { PipeLogger.log(LOGGER::info, e, PipeMessages.TEMPORARILY_OUT_OF_MEMORY); - } else { - sleep4NonReportException(); } + sleep4NonReportException(); } else if (e instanceof PipeRuntimeSinkNonReportTimeConfigurableException) { if (lastExceptionTime == Long.MAX_VALUE) { lastExceptionTime = System.currentTimeMillis(); diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/task/PipeSleepIntervalTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/task/PipeSleepIntervalTest.java index a9b72391b2b75..015279acc9866 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/task/PipeSleepIntervalTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/task/PipeSleepIntervalTest.java @@ -22,9 +22,11 @@ import org.apache.iotdb.commons.conf.CommonConfig; import org.apache.iotdb.commons.conf.CommonDescriptor; import org.apache.iotdb.commons.exception.pipe.PipeRuntimeException; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.commons.pipe.agent.task.subtask.PipeAbstractSinkSubtask; import org.apache.iotdb.commons.pipe.config.PipeConfig; import org.apache.iotdb.commons.pipe.event.EnrichedEvent; +import org.apache.iotdb.pipe.api.event.Event; import org.apache.iotdb.pipe.api.exception.PipeConnectionException; import org.junit.After; @@ -65,6 +67,10 @@ boolean isAuthenticationFailureException(final Throwable throwable) { void sleepWithoutHighPriorityTask(final long sleepMillis) throws InterruptedException { sleepIfNoHighPriorityTask(sleepMillis); } + + void handle(final Event event, final Exception exception) { + handleException(event, exception); + } } private long oldPipeSinkSubtaskSleepIntervalInitMs; @@ -136,4 +142,15 @@ public void testSleepIfNoHighPriorityTaskWaits() throws Exception { Assert.assertTrue(System.currentTimeMillis() - startTime >= 15L); } } + + @Test + public void testMemoryTimeoutRetryWaits() { + try (final TestSinkSubtask subtask = new TestSinkSubtask()) { + final long startTime = System.currentTimeMillis(); + subtask.handle(null, new PipeRuntimeOutOfMemoryCriticalException("memory unavailable")); + Assert.assertTrue( + System.currentTimeMillis() - startTime + >= PipeConfig.getInstance().getPipeSinkSubtaskSleepIntervalInitMs()); + } + } }