[IOTDB-6240] Pipe: Fix bug with PipeMemoryManager allocating memory f…

…or the DisruptorQueue (#11485) Allocating memory for the DisruptorQueue requires ensuring that the memory size is a power of 2, so the Memory Framework provides `Function` method for users to customize their own reallocate strategy.
apache · Nov 7, 2023 · 256194a · 256194a
1 parent 4e686b9
commit 256194a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/...de/src/main/java/org/apache/iotdb/db/pipe/extractor/realtime/assigner/DisruptorQueue.java b/...de/src/main/java/org/apache/iotdb/db/pipe/extractor/realtime/assigner/DisruptorQueue.java
@@ -54,7 +54,9 @@ public DisruptorQueue(EventHandler<PipeRealtimeEvent> eventHandler) {
         config.getPipeExtractorAssignerDisruptorRingBufferEntrySizeInBytes();
 
     allocatedMemoryBlock =
-        PipeResourceManager.memory().tryAllocate(ringBufferSize * ringBufferEntrySizeInBytes);
+        PipeResourceManager.memory()
+            .tryAllocate(
+                ringBufferSize * ringBufferEntrySizeInBytes, (currentSize) -> currentSize / 2);
 
     disruptor =
         new Disruptor<>(

diff --git a/...re/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java b/...re/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java
@@ -31,6 +31,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.List;
+import java.util.function.Function;
 
 public class PipeMemoryManager {
 
@@ -145,6 +146,11 @@ private long calculateTabletSizeInBytes(Tablet tablet) {
   }
 
   public synchronized PipeMemoryBlock tryAllocate(long sizeInBytes) {
+    return tryAllocate(sizeInBytes, (currentSize) -> currentSize * 2 / 3);
+  }
+
+  public synchronized PipeMemoryBlock tryAllocate(
+      long sizeInBytes, Function<Long, Long> customAllocateStrategy) {
     if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
       return new PipeMemoryBlock(sizeInBytes);
     }
@@ -171,7 +177,9 @@ public synchronized PipeMemoryBlock tryAllocate(long sizeInBytes) {
       }
 
       sizeToAllocateInBytes =
-          Math.max(sizeToAllocateInBytes * 2 / 3, MEMORY_ALLOCATE_MIN_SIZE_IN_BYTES);
+          Math.max(
+              customAllocateStrategy.apply(sizeToAllocateInBytes),
+              MEMORY_ALLOCATE_MIN_SIZE_IN_BYTES);
     }
 
     LOGGER.warn(