Pipe: Optimize Batch and WAL memory allocation algorithms

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/connector/payload/evolvable/batch/PipeTabletEventBatch.java

@@ -22,7 +22,9 @@
 import org.apache.iotdb.commons.pipe.event.EnrichedEvent;
 import org.apache.iotdb.db.pipe.connector.protocol.thrift.async.IoTDBDataRegionAsyncConnector;
 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.PipeDynamicMemoryBlock;
+import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlockType;
+import org.apache.iotdb.db.pipe.resource.memory.PipeModelFixedMemoryBlock;
 import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALPipeException;
 import org.apache.iotdb.pipe.api.event.Event;
 import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent;
@@ -38,14 +40,19 @@
 public abstract class PipeTabletEventBatch implements AutoCloseable {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(PipeTabletEventBatch.class);
+  private static final PipeModelFixedMemoryBlock PIPE_MODEL_FIXED_MEMORY_BLOCK =
+      PipeDataNodeResourceManager.memory()
+          .forceAllocateForModelFixedMemoryBlock(
+              PipeDataNodeResourceManager.memory().getAllocatedMemorySizeInBytesOfBatch(),
+              PipeMemoryBlockType.BATCH);
 
   protected final List<EnrichedEvent> events = new ArrayList<>();
 
   private final int maxDelayInMs;
   private long firstEventProcessingTime = Long.MIN_VALUE;
 
   protected long totalBufferSize = 0;
-  private final PipeMemoryBlock allocatedMemoryBlock;
+  private final PipeDynamicMemoryBlock allocatedMemoryBlock;
 
   protected volatile boolean isClosed = false;
 
@@ -54,19 +61,8 @@ protected PipeTabletEventBatch(final int maxDelayInMs, final long requestMaxBatc
 
     // limit in buffer size
     this.allocatedMemoryBlock =
-        PipeDataNodeResourceManager.memory()
-            .tryAllocate(requestMaxBatchSizeInBytes)
-            .setShrinkMethod(oldMemory -> Math.max(oldMemory / 2, 0))
-            .setShrinkCallback(
-                (oldMemory, newMemory) ->
-                    LOGGER.info(
-                        "The batch size limit has shrunk from {} to {}.", oldMemory, newMemory))
-            .setExpandMethod(
-                oldMemory -> Math.min(Math.max(oldMemory, 1) * 2, requestMaxBatchSizeInBytes))
-            .setExpandCallback(
-                (oldMemory, newMemory) ->
-                    LOGGER.info(
-                        "The batch size limit has expanded from {} to {}.", oldMemory, newMemory));
+        PIPE_MODEL_FIXED_MEMORY_BLOCK.registerPipeBatchMemoryBlock(requestMaxBatchSizeInBytes);
+    allocatedMemoryBlock.setExpandable(false);
 
     if (getMaxBatchSizeInBytes() != requestMaxBatchSizeInBytes) {
       LOGGER.info(
@@ -131,8 +127,12 @@ protected abstract boolean constructBatch(final TabletInsertionEvent event)
       throws WALPipeException, IOException;
 
   public boolean shouldEmit() {
-    return totalBufferSize >= getMaxBatchSizeInBytes()
-        || System.currentTimeMillis() - firstEventProcessingTime >= maxDelayInMs;
+    final long diff = System.currentTimeMillis() - firstEventProcessingTime;
+    if (totalBufferSize >= getMaxBatchSizeInBytes() || diff >= maxDelayInMs) {
+      allocatedMemoryBlock.updateCurrentMemoryEfficiencyAdjustMem((double) diff / maxDelayInMs);
+      return true;
+    }
+    return false;
   }
 
   private long getMaxBatchSizeInBytes() {

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

@@ -0,0 +1,156 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.pipe.resource.memory;
+
+import org.apache.tsfile.utils.Pair;
+
+import javax.validation.constraints.NotNull;
+
+import java.util.function.Consumer;
+import java.util.stream.Stream;
+
+public class PipeDynamicMemoryBlock {
+
+  private final PipeModelFixedMemoryBlock fixedMemoryBlock;
+
+  private boolean isExpandable = true;
+
+  private Consumer<PipeDynamicMemoryBlock> expand = null;
+
+  private volatile boolean released = false;
+
+  private volatile long memoryUsageInBytes;
+
+  private volatile double historyMemoryEfficiency;
+
+  private volatile double currentMemoryEfficiency;
+
+  PipeDynamicMemoryBlock(
+      final @NotNull PipeModelFixedMemoryBlock fixedMemoryBlock, final long memoryUsageInBytes) {
+    this.memoryUsageInBytes = Math.min(memoryUsageInBytes, 0);
+    this.fixedMemoryBlock = fixedMemoryBlock;
+  }
+
+  public long getMemoryUsageInBytes() {
+    return memoryUsageInBytes;
+  }
+
+  public void setMemoryUsageInBytes(final long memoryUsageInBytes) {
+    this.memoryUsageInBytes = memoryUsageInBytes;
+  }
+
+  public Pair<Double, Double> getMemoryEfficiency() {
+    synchronized (fixedMemoryBlock) {
+      return new Pair<>(historyMemoryEfficiency, currentMemoryEfficiency);
+    }
+  }
+
+  public void setExpandable(boolean expandable) {
+    isExpandable = expandable;
+  }
+
+  public void setExpand(Consumer<PipeDynamicMemoryBlock> expand) {
+    this.expand = expand;
+  }
+
+  public double getMemoryBlockUsageRatio() {
+    return (double) memoryUsageInBytes / fixedMemoryBlock.getMemoryUsageInBytes();
+  }
+
+  public double getFixedMemoryBlockUsageRatio() {
+    return (double) fixedMemoryBlock.getMemoryAllocatedInBytes()
+        / fixedMemoryBlock.getMemoryUsageInBytes();
+  }
+
+  public long canAllocateMemorySize() {
+    return fixedMemoryBlock.getMemoryUsageInBytes() - fixedMemoryBlock.getMemoryAllocatedInBytes();
+  }
+
+  public synchronized long getExpectedAverageAllocatedMemorySize() {
+    return fixedMemoryBlock.getMemoryUsageInBytes() / fixedMemoryBlock.getMemoryBlocks().size();
+  }
+
+  public void updateCurrentMemoryEfficiencyAdjustMem(double currentMemoryEfficiency) {
+    synchronized (fixedMemoryBlock) {
+      this.historyMemoryEfficiency = this.currentMemoryEfficiency;
+      if (Double.isNaN(currentMemoryEfficiency)
+          || Double.isInfinite(currentMemoryEfficiency)
+          || currentMemoryEfficiency < 0.0) {
+        currentMemoryEfficiency = 0.0;
+      }
+      this.currentMemoryEfficiency = Math.min(currentMemoryEfficiency, 1.0);
+      fixedMemoryBlock.dynamicallyAdjustMemory(this);
+    }
+  }
+
+  public long getFixedMemoryCapacity() {
+    return fixedMemoryBlock.getMemoryUsageInBytes();
+  }
+
+  public void updateMemoryEfficiency(
+      double currentMemoryEfficiency, double historyMemoryEfficiency) {
+    synchronized (fixedMemoryBlock) {
+      if (Double.isNaN(currentMemoryEfficiency)
+          || Double.isInfinite(currentMemoryEfficiency)
+          || currentMemoryEfficiency < 0.0) {
+        currentMemoryEfficiency = 0.0;
+      }
+
+      if (Double.isNaN(historyMemoryEfficiency)
+          || Double.isInfinite(historyMemoryEfficiency)
+          || historyMemoryEfficiency < 0.0) {
+        currentMemoryEfficiency = 0.0;
+      }
+
+      this.historyMemoryEfficiency = Math.min(historyMemoryEfficiency, 1.0);
+      this.currentMemoryEfficiency = Math.min(currentMemoryEfficiency, 1.0);
+    }
+  }
+
+  public Stream<PipeDynamicMemoryBlock> getMemoryBlocks() {
+    return fixedMemoryBlock.getMemoryBlocksStream();
+  }
+
+  public void applyForDynamicMemory(final long memoryUsageInBytes) {
+    fixedMemoryBlock.resetMemoryBlockSize(this, memoryUsageInBytes);
+  }
+
+  public boolean isReleased() {
+    return released;
+  }
+
+  public void close() {
+    if (released) {
+      return;
+    }
+    synchronized (fixedMemoryBlock) {
+      if (!released) {
+        fixedMemoryBlock.releaseMemory(this);
+        released = true;
+      }
+    }
+  }
+
+  void doExpand() {
+    if (isExpandable && expand != null) {
+      expand.accept(this);
+    }
+  }
+}

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

@@ -23,4 +23,6 @@ public enum PipeMemoryBlockType {
   NORMAL,
   TABLET,
   TS_FILE,
+  BATCH,
+  WAL
 }

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

@@ -25,6 +25,7 @@
 import org.apache.iotdb.commons.pipe.config.PipeConfig;
 import org.apache.iotdb.db.conf.IoTDBDescriptor;
 import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent;
+import org.apache.iotdb.db.pipe.resource.memory.strategy.ThresholdAllocationStrategy;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -100,6 +101,18 @@ private double allowedMaxMemorySizeInBytesOfTsTiles() {
         * getTotalNonFloatingMemorySizeInBytes();
   }
 
+  public long getAllocatedMemorySizeInBytesOfWAL() {
+    return (long)
+        (PIPE_CONFIG.getPipeDataStructureWalMemoryProportion()
+            * getTotalNonFloatingMemorySizeInBytes());
+  }
+
+  public long getAllocatedMemorySizeInBytesOfBatch() {
+    return (long)
+        (PIPE_CONFIG.getPipeDataStructureBatchMemoryProportion()
+            * getTotalNonFloatingMemorySizeInBytes());
+  }
+
   public boolean isEnough4TabletParsing() {
     return (double) usedMemorySizeInBytesOfTablets + (double) usedMemorySizeInBytesOfTsFiles
             < EXCEED_PROTECT_THRESHOLD * allowedMaxMemorySizeInBytesOfTabletsAndTsFiles()
@@ -225,6 +238,39 @@ public PipeTsFileMemoryBlock forceAllocateForTsFileWithRetry(long tsFileSizeInBy
     }
   }
 
+  public PipeModelFixedMemoryBlock forceAllocateForModelFixedMemoryBlock(
+      long fixedSizeInBytes, PipeMemoryBlockType type)
+      throws PipeRuntimeOutOfMemoryCriticalException {
+    if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
+      return new PipeModelFixedMemoryBlock(Long.MAX_VALUE, new ThresholdAllocationStrategy());
+    }
+
+    if (fixedSizeInBytes == 0) {
+      return (PipeModelFixedMemoryBlock) registerMemoryBlock(0, type);
+    }
+
+    for (int i = 1, size = PIPE_CONFIG.getPipeMemoryAllocateMaxRetries(); i <= size; i++) {
+      if (getFreeMemorySizeInBytes() >= fixedSizeInBytes) {
+        break;
+      }
+
+      try {
+        Thread.sleep(PIPE_CONFIG.getPipeMemoryAllocateRetryIntervalInMs());
+      } catch (InterruptedException ex) {
+        Thread.currentThread().interrupt();
+        LOGGER.warn("forceAllocateWithRetry: interrupted while waiting for available memory", ex);
+      }
+    }
+
+    if (getFreeMemorySizeInBytes() < fixedSizeInBytes) {
+      return (PipeModelFixedMemoryBlock) forceAllocateWithRetry(getFreeMemorySizeInBytes(), type);
+    }
+
+    synchronized (this) {
+      return (PipeModelFixedMemoryBlock) forceAllocateWithRetry(fixedSizeInBytes, type);
+    }
+  }
+
   private PipeMemoryBlock forceAllocateWithRetry(long sizeInBytes, PipeMemoryBlockType type)
       throws PipeRuntimeOutOfMemoryCriticalException {
     if (!PIPE_MEMORY_MANAGEMENT_ENABLED) {
@@ -233,6 +279,9 @@ private PipeMemoryBlock forceAllocateWithRetry(long sizeInBytes, PipeMemoryBlock
           return new PipeTabletMemoryBlock(sizeInBytes);
         case TS_FILE:
           return new PipeTsFileMemoryBlock(sizeInBytes);
+        case BATCH:
+        case WAL:
+          return new PipeModelFixedMemoryBlock(sizeInBytes, new ThresholdAllocationStrategy());
         default:
           return new PipeMemoryBlock(sizeInBytes);
       }
@@ -466,6 +515,11 @@ private PipeMemoryBlock registerMemoryBlock(long sizeInBytes, PipeMemoryBlockTyp
       case TS_FILE:
         returnedMemoryBlock = new PipeTsFileMemoryBlock(sizeInBytes);
         break;
+      case BATCH:
+      case WAL:
+        returnedMemoryBlock =
+            new PipeModelFixedMemoryBlock(sizeInBytes, new ThresholdAllocationStrategy());
+        break;
       default:
         returnedMemoryBlock = new PipeMemoryBlock(sizeInBytes);
         break;