Pipe: dynamically share idle floating memory with TsFile sync

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 @@
   // 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 @@
         IoTDBDescriptor.getInstance()
             .getMemoryConfig()
             .getPipeMemoryManager()
-            .exactAllocate("Stream", MemoryBlockType.DYNAMIC));
+            .exactAllocate("Stream", MemoryBlockType.DYNAMIC),
+        () -> PipeDataNodeAgent.task().getAllFloatingMemoryUsageInByte());
     PipeDataNodeAgent.runtime()
         .registerPeriodicalJob(
             "PipeMemoryManager#tryExpandAll()",
@@ -92,7 +96,13 @@
   }
 
   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
@@ -477,7 +487,7 @@
     if (block instanceof PipeTabletMemoryBlock) {
       return (double) usedMemorySizeInBytesOfTablets
                   + (double) extraMemoryInBytes
                   + (double) usedMemorySizeInBytesOfTsFiles
               < allowedMaxMemorySizeInBytesOfTabletsAndTsFiles()
           && (double) usedMemorySizeInBytesOfTablets + (double) extraMemoryInBytes
               < allowedMaxMemorySizeInBytesOfTablets();
@@ -485,7 +495,7 @@
     if (block instanceof PipeTsFileMemoryBlock) {
       return (double) usedMemorySizeInBytesOfTablets
                   + (double) usedMemorySizeInBytesOfTsFiles
                   + (double) extraMemoryInBytes
               < allowedMaxMemorySizeInBytesOfTabletsAndTsFiles()
           && (double) usedMemorySizeInBytesOfTsFiles + (double) extraMemoryInBytes
               < allowedMaxMemorySizeInBytesOfTsTiles();
@@ -1038,19 +1048,31 @@
   }
 
   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() {

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);
+    }
+  }
 }

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

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();
   }