Pipe: dynamically share idle floating memory (#18422) (#18437)

Author: Caideyipi

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

@@ -28,7 +28,6 @@
 import org.apache.iotdb.db.pipe.event.common.tsfile.container.scan.TsFileInsertionScanDataContainer;
 import org.apache.iotdb.db.pipe.metric.overview.PipeTsFileToTabletsMetrics;
 import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager;
-import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager;
 import org.apache.iotdb.db.pipe.resource.tsfile.PipeTsFilePublicResource;
 
 import org.apache.tsfile.file.metadata.IDeviceID;
@@ -133,7 +132,7 @@ public TsFileInsertionDataContainer provide(final boolean isWithMod) throws IOEx
 
     // Use scan container to save memory
     if ((double) PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes()
-            / PipeMemoryManager.getTotalNonFloatingMemorySizeInBytes()
+            / PipeDataNodeResourceManager.memory().getTotalNonFloatingMemorySizeInBytes()
         > PipeTsFilePublicResource.MEMORY_SUFFICIENT_THRESHOLD) {
       return new TsFileInsertionScanDataContainer(
           pipeName,

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/metric/overview/PipeResourceMetrics.java

@@ -75,14 +75,14 @@ public void bindTo(final AbstractMetricService metricService) {
         Metric.PIPE_MEM.toString(),
         MetricLevel.IMPORTANT,
         PipeDataNodeResourceManager.memory(),
-        o -> PipeMemoryManager.getTotalNonFloatingMemorySizeInBytes(),
+        PipeMemoryManager::getTotalNonFloatingMemorySizeInBytes,
         Tag.NAME.toString(),
         PIPE_TOTAL_MEMORY);
     metricService.createAutoGauge(
         Metric.PIPE_MEM.toString(),
         MetricLevel.IMPORTANT,
         PipeDataNodeResourceManager.memory(),
-        o -> PipeMemoryManager.getTotalFloatingMemorySizeInBytes(),
+        PipeMemoryManager::getTotalFloatingMemorySizeInBytes,
         Tag.NAME.toString(),
         PIPE_FLOATING_MEMORY);
     metricService.createAutoGauge(

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

@@ -38,6 +38,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 {
@@ -53,6 +54,9 @@ public class PipeMemoryManager {
   private static final long MEMORY_ALLOCATE_MIN_SIZE_IN_BYTES =
       PipeConfig.getInstance().getPipeMemoryAllocateMinSizeInBytes();
 
+  private final long totalMemorySizeInBytes;
+  private final LongSupplier floatingMemoryUsageSupplier;
+
   private long usedMemorySizeInBytes;
 
   private static final double EXCEED_PROTECT_THRESHOLD = 0.95;
@@ -79,13 +83,22 @@ public class PipeMemoryManager {
   private final Set<PipeMemoryBlock> expandableBlocks = new HashSet<>();
 
   public PipeMemoryManager() {
+    this(
+        TOTAL_MEMORY_SIZE_IN_BYTES,
+        () -> PipeDataNodeAgent.task().getAllFloatingMemoryUsageInByte());
     PipeDataNodeAgent.runtime()
         .registerPeriodicalJob(
             "PipeMemoryManager#tryExpandAll()",
             this::tryExpandAllAndCheckConsistency,
             PipeConfig.getInstance().getPipeMemoryExpanderIntervalSeconds());
   }
 
+  PipeMemoryManager(
+      final long totalMemorySizeInBytes, final LongSupplier floatingMemoryUsageSupplier) {
+    this.totalMemorySizeInBytes = totalMemorySizeInBytes;
+    this.floatingMemoryUsageSupplier = floatingMemoryUsageSupplier;
+  }
+
   // NOTE: Here we unify the memory threshold judgment for tablet and tsfile memory block, because
   // introducing too many heuristic rules not conducive to flexible dynamic adjustment of memory
   // configuration:
@@ -96,15 +109,15 @@ public PipeMemoryManager() {
   // 3. The sum of the memory proportion occupied by the tablet memory block and the tsfile memory
   // block does not exceed TABLET_MEMORY_REJECT_THRESHOLD + TS_FILE_MEMORY_REJECT_THRESHOLD
 
-  private static double allowedMaxMemorySizeInBytesOfTabletsAndTsFiles() {
+  private double allowedMaxMemorySizeInBytesOfTabletsAndTsFiles() {
     return (PipeConfig.getInstance()
                 .getPipeDataStructureTabletMemoryBlockAllocationRejectThreshold()
             + PipeConfig.getInstance()
                 .getPipeDataStructureTsFileMemoryBlockAllocationRejectThreshold())
         * getTotalNonFloatingMemorySizeInBytes();
   }
 
-  private static double allowedMaxMemorySizeInBytesOfTablets() {
+  private double allowedMaxMemorySizeInBytesOfTablets() {
     return (PipeConfig.getInstance()
                 .getPipeDataStructureTabletMemoryBlockAllocationRejectThreshold()
             + PipeConfig.getInstance()
@@ -113,7 +126,7 @@ private static double allowedMaxMemorySizeInBytesOfTablets() {
         * getTotalNonFloatingMemorySizeInBytes();
   }
 
-  private static double allowedMaxMemorySizeInBytesOfTsTiles() {
+  private double allowedMaxMemorySizeInBytesOfTsTiles() {
     return (PipeConfig.getInstance()
                 .getPipeDataStructureTsFileMemoryBlockAllocationRejectThreshold()
             + PipeConfig.getInstance()
@@ -1037,19 +1050,30 @@ public long getUsedMemorySizeInBytesOfTsFiles() {
   }
 
   public long getFreeMemorySizeInBytes() {
-    return TOTAL_MEMORY_SIZE_IN_BYTES - usedMemorySizeInBytes;
+    return Math.max(0, getTotalNonFloatingMemorySizeInBytes() - usedMemorySizeInBytes);
   }
 
-  public static long getTotalNonFloatingMemorySizeInBytes() {
-    return (long)
-        (TOTAL_MEMORY_SIZE_IN_BYTES
-            * (1 - PipeConfig.getInstance().getPipeTotalFloatingMemoryProportion()));
+  public long getTotalNonFloatingMemorySizeInBytes() {
+    // 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, totalMemorySizeInBytes - getUsedFloatingMemorySizeInBytes());
+  }
+
+  public long getTotalFloatingMemorySizeInBytes() {
+    final long configuredUpperLimit =
+        Math.max(
+            0,
+            (long)
+                (totalMemorySizeInBytes
+                    * PipeConfig.getInstance().getPipeTotalFloatingMemoryProportion()));
+    final long memoryNotUsedByNonFloatingAllocations =
+        Math.max(0, totalMemorySizeInBytes - usedMemorySizeInBytes);
+    return Math.min(configuredUpperLimit, memoryNotUsedByNonFloatingAllocations);
   }
 
-  public static long getTotalFloatingMemorySizeInBytes() {
-    return (long)
-        (TOTAL_MEMORY_SIZE_IN_BYTES
-            * PipeConfig.getInstance().getPipeTotalFloatingMemoryProportion());
+  private long getUsedFloatingMemorySizeInBytes() {
+    return Math.max(0, floatingMemoryUsageSupplier.getAsLong());
   }
 
   public static long getTotalMemorySizeInBytes() {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/source/dataregion/realtime/PipeRealtimeDataRegionHybridSource.java

@@ -29,7 +29,7 @@
 import org.apache.iotdb.db.pipe.event.realtime.PipeRealtimeEvent;
 import org.apache.iotdb.db.pipe.metric.overview.PipeDataNodeRemainingEventAndTimeOperator;
 import org.apache.iotdb.db.pipe.metric.overview.PipeDataNodeSinglePipeMetrics;
-import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager;
+import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager;
 import org.apache.iotdb.db.pipe.source.dataregion.realtime.assigner.PipeTsFileEpochProgressIndexKeeper;
 import org.apache.iotdb.db.pipe.source.dataregion.realtime.epoch.TsFileEpoch;
 import org.apache.iotdb.pipe.api.event.Event;
@@ -178,7 +178,8 @@ private boolean canNotUseTabletAnymore(final PipeRealtimeEvent event) {
     final long floatingMemoryUsageInByte =
         PipeDataNodeAgent.task().getFloatingMemoryUsageInByte(pipeName);
     final long pipeCount = PipeDataNodeAgent.task().getPipeCount();
-    long totalFloatingMemorySizeInBytes = PipeMemoryManager.getTotalFloatingMemorySizeInBytes();
+    long totalFloatingMemorySizeInBytes =
+        PipeDataNodeResourceManager.memory().getTotalFloatingMemorySizeInBytes();
     // If the occupied memory has reached the max, it may cause a large latency to the receiver due
     // to queuing. To reduce the latency, we lower the memory limit forcibly in the single tsFile
     // since the tsFile is doomed to be transferred, then more downgrading will just cause more

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/PipeProcessorSubtaskExecutorTest.java

@@ -70,7 +70,7 @@ public void testTsFileInsertionEventPreservesOutOfMemoryCause() {
     try {
       memoryBlock =
           memoryManager.forceAllocateForTabletWithRetry(
-              PipeMemoryManager.getTotalNonFloatingMemorySizeInBytes());
+              memoryManager.getTotalNonFloatingMemorySizeInBytes());
       Assert.assertFalse(memoryManager.isEnough4TabletParsing());
 
       final File tsFile =

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java

@@ -28,8 +28,11 @@
 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 final CommonConfig config = CommonDescriptor.getInstance().getConfig();
 
   private boolean originalMemoryManagementEnabled;
@@ -76,7 +79,7 @@ public void testTabletResizeCannotCrossTabletHardLimit() {
     final PipeTabletMemoryBlock tablet = manager.forceAllocateForTabletWithRetry(0);
     final long tabletMemorySizeInBytes =
         (long)
-                (PipeMemoryManager.getTotalNonFloatingMemorySizeInBytes()
+                (manager.getTotalNonFloatingMemorySizeInBytes()
                     * (config.getPipeDataStructureTabletMemoryBlockAllocationRejectThreshold()
                         + config.getPipeDataStructureTsFileMemoryBlockAllocationRejectThreshold()
                             / 2))
@@ -97,8 +100,7 @@ public void testTabletResizeCannotCrossTabletHardLimit() {
   @Test
   public void testTabletResizeLeavesMemoryForSinkForwardProgress() {
     final PipeMemoryManager manager = new PipeMemoryManager();
-    final long totalNonFloatingMemorySizeInBytes =
-        PipeMemoryManager.getTotalNonFloatingMemorySizeInBytes();
+    final long totalNonFloatingMemorySizeInBytes = manager.getTotalNonFloatingMemorySizeInBytes();
     final long tabletMemorySizeInBytes =
         (long)
                 (totalNonFloatingMemorySizeInBytes
@@ -134,4 +136,32 @@ public void testTabletResizeLeavesMemoryForSinkForwardProgress() {
 
     Assert.assertEquals(0, manager.getUsedMemorySizeInBytes());
   }
+
+  @Test
+  public void testFloatingAndNonFloatingMemoryShareTheSamePool() {
+    final AtomicLong floatingMemoryUsageInBytes = new AtomicLong(0);
+    final PipeMemoryManager manager =
+        new PipeMemoryManager(TOTAL_MEMORY_SIZE_IN_BYTES, 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

@@ -224,6 +224,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

@@ -66,6 +66,7 @@ public double getPipeDataStructureTsFileMemoryBlockAllocationRejectThreshold() {
   }
 
   public double getPipeTotalFloatingMemoryProportion() {
+    // This is the upper limit of floating memory, not a statically reserved partition.
     return COMMON_CONFIG.getPipeTotalFloatingMemoryProportion();
   }