From 17f452c80bd71f1f9e183e04336bd55e01293a76 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 7 Aug 2026 13:04:26 +0800 Subject: [PATCH] Pipe: dynamically share idle floating memory --- .../resource/memory/PipeMemoryManager.java | 38 +++++++++++++++---- .../memory/PipeMemoryManagerResizeTest.java | 38 ++++++++++++++++++- .../iotdb/commons/conf/CommonConfig.java | 3 ++ .../iotdb/commons/pipe/config/PipeConfig.java | 1 + 4 files changed, 71 insertions(+), 9 deletions(-) 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 dd0485992ee45..90a45c1542ec3 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 @@ -41,6 +41,7 @@ import java.util.Map; import java.util.Objects; import java.util.Set; +import java.util.function.LongSupplier; import java.util.function.LongUnaryOperator; public class PipeMemoryManager { @@ -55,6 +56,8 @@ public class PipeMemoryManager { // TODO @spricoder: consider combine memory block and used MemorySizeInBytes private final IMemoryBlock memoryBlock; + private final LongSupplier floatingMemoryUsageSupplier; + private static final double EXCEED_PROTECT_THRESHOLD = 0.95; private volatile long usedMemorySizeInBytesOfTablets; @@ -83,7 +86,8 @@ public PipeMemoryManager() { IoTDBDescriptor.getInstance() .getMemoryConfig() .getPipeMemoryManager() - .exactAllocate("Stream", MemoryBlockType.DYNAMIC)); + .exactAllocate("Stream", MemoryBlockType.DYNAMIC), + () -> PipeDataNodeAgent.task().getAllFloatingMemoryUsageInByte()); PipeDataNodeAgent.runtime() .registerPeriodicalJob( "PipeMemoryManager#tryExpandAll()", @@ -92,7 +96,13 @@ public PipeMemoryManager() { } PipeMemoryManager(final IMemoryBlock memoryBlock) { + this(memoryBlock, () -> PipeDataNodeAgent.task().getAllFloatingMemoryUsageInByte()); + } + + PipeMemoryManager( + final IMemoryBlock memoryBlock, final LongSupplier floatingMemoryUsageSupplier) { this.memoryBlock = memoryBlock; + this.floatingMemoryUsageSupplier = floatingMemoryUsageSupplier; } // NOTE: Here we unify the memory threshold judgment for tablet and tsfile memory block, because @@ -1038,19 +1048,31 @@ public long getUsedMemorySizeInBytesOfTsFiles() { } public long getFreeMemorySizeInBytes() { - return memoryBlock.getFreeMemoryInBytes(); + return Math.max(0, getTotalNonFloatingMemorySizeInBytes() - memoryBlock.getUsedMemoryInBytes()); } public long getTotalNonFloatingMemorySizeInBytes() { - return (long) - (memoryBlock.getTotalMemorySizeInBytes() - * (1 - PipeConfig.getInstance().getPipeTotalFloatingMemoryProportion())); + // Floating memory is an upper limit for retained InsertNodes instead of a statically reserved + // partition. Non-floating allocations can borrow all floating memory that is not actually in + // use, which is especially important for TsFile-only pipes. + return Math.max( + 0, memoryBlock.getTotalMemorySizeInBytes() - getUsedFloatingMemorySizeInBytes()); } public long getTotalFloatingMemorySizeInBytes() { - return (long) - (memoryBlock.getTotalMemorySizeInBytes() - * PipeConfig.getInstance().getPipeTotalFloatingMemoryProportion()); + final long configuredUpperLimit = + Math.max( + 0, + (long) + (memoryBlock.getTotalMemorySizeInBytes() + * PipeConfig.getInstance().getPipeTotalFloatingMemoryProportion())); + final long memoryNotUsedByNonFloatingAllocations = + Math.max(0, memoryBlock.getTotalMemorySizeInBytes() - memoryBlock.getUsedMemoryInBytes()); + return Math.min(configuredUpperLimit, memoryNotUsedByNonFloatingAllocations); + } + + private long getUsedFloatingMemorySizeInBytes() { + return Math.max(0, floatingMemoryUsageSupplier.getAsLong()); } public long getTotalMemorySizeInBytes() { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java index c151857e8cbc7..d63fe1c6e8309 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java @@ -30,10 +30,12 @@ import org.junit.Before; import org.junit.Test; +import java.util.concurrent.atomic.AtomicLong; + public class PipeMemoryManagerResizeTest { private static final long TOTAL_MEMORY_SIZE_IN_BYTES = 2000; - private static final long TABLET_MEMORY_SIZE_IN_BYTES = 451; + private static final long TABLET_MEMORY_SIZE_IN_BYTES = 901; private static final long SINK_MEMORY_SIZE_IN_BYTES = 100; private final CommonConfig config = CommonDescriptor.getInstance().getConfig(); @@ -136,4 +138,38 @@ public void testTabletResizeLeavesMemoryForSinkForwardProgress() { Assert.assertEquals(0, manager.getUsedMemorySizeInBytes()); } + + @Test + public void testFloatingAndNonFloatingMemoryShareTheSamePool() { + final AtomicLong floatingMemoryUsageInBytes = new AtomicLong(0); + final PipeMemoryManager manager = + new PipeMemoryManager( + new AtomicLongMemoryBlock( + "PipeMemoryManagerResizeTest", + null, + TOTAL_MEMORY_SIZE_IN_BYTES, + MemoryBlockType.DYNAMIC), + floatingMemoryUsageInBytes::get); + + Assert.assertEquals(TOTAL_MEMORY_SIZE_IN_BYTES, manager.getTotalNonFloatingMemorySizeInBytes()); + Assert.assertEquals( + TOTAL_MEMORY_SIZE_IN_BYTES / 2, manager.getTotalFloatingMemorySizeInBytes()); + + final PipeTsFileMemoryBlock nonFloatingMemory = manager.forceAllocateForTsFileWithRetry(1200); + try { + // Non-floating memory can borrow the unused half that was previously reserved for InsertNode + // queues. Its usage also reduces the current floating-memory limit symmetrically. + Assert.assertEquals(1200, manager.getUsedMemorySizeInBytes()); + Assert.assertEquals(800, manager.getTotalFloatingMemorySizeInBytes()); + + floatingMemoryUsageInBytes.set(500); + Assert.assertEquals(1500, manager.getTotalNonFloatingMemorySizeInBytes()); + Assert.assertEquals(300, manager.getFreeMemorySizeInBytes()); + + Assert.assertThrows( + PipeRuntimeOutOfMemoryCriticalException.class, () -> manager.forceAllocate(301)); + } finally { + manager.release(nonFloatingMemory); + } + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java index 73f5bf52fc33e..4adacb0d7aa1c 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java @@ -235,6 +235,9 @@ public class CommonConfig { private int pipeDataStructureTabletSizeInBytes = 16 * 1024 * 1024; private double pipeDataStructureTabletMemoryBlockAllocationRejectThreshold = 0.3; private double pipeDataStructureTsFileMemoryBlockAllocationRejectThreshold = 0.3; + + // Maximum proportion for floating memory retained by InsertNode queues. Unused floating memory + // can be borrowed by non-floating Pipe allocations. private volatile double pipeTotalFloatingMemoryProportion = 0.5; // Check if memory check is enabled for Pipe diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeConfig.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeConfig.java index 561d2923c07e7..e04237129e867 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeConfig.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeConfig.java @@ -67,6 +67,7 @@ public double getPipeDataStructureTsFileMemoryBlockAllocationRejectThreshold() { } public double getPipeTotalFloatingMemoryProportion() { + // This is the upper limit of floating memory, not a statically reserved partition. return COMMON_CONFIG.getPipeTotalFloatingMemoryProportion(); }