Pipe: manage ConfigNode receiver temporary memory

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/conf/ConfigNodeMemoryConfig.java

@@ -70,8 +70,8 @@ public void init(final TrimProperties properties) {
         MemoryConfig.global().getOrCreateMemoryManager("ConfigNodeOnHeap", maxMemoryAvailable);
     pipeMemoryManager =
         onHeapMemoryManager.getOrCreateMemoryManager(PIPE_MEMORY_MANAGER_NAME, pipeMemorySize);
-    // Keep the rest of ConfigNode heap unconnected for now. The memory framework currently only
-    // serves PipePeriodicalLogReducer on ConfigNode.
+    // Keep the rest of ConfigNode heap unconnected for now. The memory framework currently
+    // serves ConfigNode Pipe memory control.
 
     LOGGER.info(
         ConfigNodeMessages.INITIAL_CONFIGNODE_ALLOCATE_MEMORY_FOR_PIPE,

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigNodeRuntimeAgent.java

@@ -21,7 +21,6 @@
 
 import org.apache.iotdb.commons.exception.IllegalPathException;
 import org.apache.iotdb.commons.exception.pipe.PipeRuntimeException;
-import org.apache.iotdb.commons.memory.IMemoryBlock;
 import org.apache.iotdb.commons.pipe.agent.runtime.PipePeriodicalJobExecutor;
 import org.apache.iotdb.commons.pipe.agent.runtime.PipePeriodicalPhantomReferenceCleaner;
 import org.apache.iotdb.commons.pipe.agent.task.meta.PipeTaskMeta;
@@ -98,18 +97,8 @@ public synchronized void stop() {
   }
 
   private void initPipePeriodicalLogReducer() {
-    final IMemoryBlock pipeLogReducerMemoryBlock = PipeConfigNodeResourceManager.logReducerMemory();
     PipePeriodicalLogReducer.setMemoryResizeFunction(
-        targetSizeInBytes -> {
-          final long nonNegativeTargetSizeInBytes = Math.max(0, targetSizeInBytes);
-          final long oldSizeInBytes = pipeLogReducerMemoryBlock.getUsedMemoryInBytes();
-          if (oldSizeInBytes < nonNegativeTargetSizeInBytes) {
-            pipeLogReducerMemoryBlock.allocate(nonNegativeTargetSizeInBytes - oldSizeInBytes);
-          } else if (oldSizeInBytes > nonNegativeTargetSizeInBytes) {
-            pipeLogReducerMemoryBlock.release(oldSizeInBytes - nonNegativeTargetSizeInBytes);
-          }
-          return pipeLogReducerMemoryBlock.getUsedMemoryInBytes();
-        });
+        PipeConfigNodeResourceManager.memory()::resizeLogReducerMemory);
     PipeLogger.setLogger(PipePeriodicalLogReducer::log);
   }
 

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/receiver/protocol/IoTDBConfigNodeReceiver.java

@@ -27,6 +27,7 @@
 import org.apache.iotdb.commons.auth.entity.PrivilegeUnion;
 import org.apache.iotdb.commons.conf.CommonDescriptor;
 import org.apache.iotdb.commons.conf.IoTDBConstant;
+import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException;
 import org.apache.iotdb.commons.path.PartialPath;
 import org.apache.iotdb.commons.path.PathPatternTree;
 import org.apache.iotdb.commons.path.PathPatternTreeUtils;
@@ -96,6 +97,7 @@
 import org.apache.iotdb.confignode.manager.pipe.metric.receiver.PipeConfigNodeReceiverMetrics;
 import org.apache.iotdb.confignode.manager.pipe.receiver.visitor.PipeConfigPhysicalPlanExceptionVisitor;
 import org.apache.iotdb.confignode.manager.pipe.receiver.visitor.PipeConfigPhysicalPlanTSStatusVisitor;
+import org.apache.iotdb.confignode.manager.pipe.resource.PipeConfigNodeResourceManager;
 import org.apache.iotdb.confignode.manager.pipe.sink.payload.PipeTransferConfigNodeHandshakeV2Req;
 import org.apache.iotdb.confignode.manager.pipe.sink.payload.PipeTransferConfigPlanReq;
 import org.apache.iotdb.confignode.manager.pipe.sink.payload.PipeTransferConfigSnapshotPieceReq;
@@ -174,7 +176,7 @@
   private final CNAuditLogger auditLogger = configManager.getAuditLogger();
 
   @Override
   public TPipeTransferResp receive(final TPipeTransferReq req) {
     try {
       final short rawRequestType = req.getType();
       if (PipeRequestType.isValidatedRequestType(rawRequestType)) {
@@ -214,15 +216,22 @@
             PipeConfigNodeReceiverMetrics.getInstance()
                 .recordTransferConfigPlanTimer(System.nanoTime() - startTime);
             return resp;
           case TRANSFER_CONFIG_SNAPSHOT_PIECE:
-            resp =
-                handleTransferFilePiece(
+            try {
+              try (final AutoCloseable ignored =
+                  PipeConfigNodeResourceManager.memory()
+                      .tryAllocateReceiverMemory(getRequestBodySizeInBytes(req))) {
+                return handleTransferFilePiece(
                     PipeTransferConfigSnapshotPieceReq.fromTPipeTransferReq(req),
                     req instanceof AirGapPseudoTPipeTransferRequest,
                     false);
-            PipeConfigNodeReceiverMetrics.getInstance()
-                .recordTransferConfigSnapshotPieceTimer(System.nanoTime() - startTime);
-            return resp;
+              } catch (final PipeRuntimeOutOfMemoryCriticalException e) {
+                return getReceiverTemporaryUnavailableResp(e);
+              }
+            } finally {
+              PipeConfigNodeReceiverMetrics.getInstance()
+                  .recordTransferConfigSnapshotPieceTimer(System.nanoTime() - startTime);
+            }
           case TRANSFER_CONFIG_SNAPSHOT_SEAL:
             resp =
                 handleTransferFileSealV2(
@@ -231,7 +240,14 @@
                 .recordTransferConfigSnapshotSealTimer(System.nanoTime() - startTime);
             return resp;
           case TRANSFER_COMPRESSED:
-            return receive(PipeTransferCompressedReq.fromTPipeTransferReq(req));
+            try (final AutoCloseable ignored =
+                PipeConfigNodeResourceManager.memory()
+                    .tryAllocateReceiverMemory(
+                        PipeTransferCompressedReq.getMaxAdditionalDecompressedLengthInBytes(req))) {
+              return receive(PipeTransferCompressedReq.fromTPipeTransferReq(req));
+            } catch (final PipeRuntimeOutOfMemoryCriticalException e) {
+              return getReceiverTemporaryUnavailableResp(e);
+            }
           default:
             break;
         }
@@ -287,6 +303,17 @@
     return new TPipeTransferResp(getNotLoggedInStatus());
   }
 
+  private static long getRequestBodySizeInBytes(final TPipeTransferReq req) {
+    return req.getBody() == null ? 0 : req.getBody().length;
+  }
+
+  private static TPipeTransferResp getReceiverTemporaryUnavailableResp(
+      final PipeRuntimeOutOfMemoryCriticalException e) {
+    return new TPipeTransferResp(
+        new TSStatus(TSStatusCode.PIPE_RECEIVER_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode())
+            .setMessage(e.getMessage()));
+  }
+
   private static boolean requiresAuthentication(final PipeRequestType type) {
     switch (type) {
       case TRANSFER_CONFIG_PLAN:

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/resource/PipeConfigNodeResourceManager.java

@@ -19,19 +19,17 @@
 
 package org.apache.iotdb.confignode.manager.pipe.resource;
 
-import org.apache.iotdb.commons.memory.IMemoryBlock;
-import org.apache.iotdb.commons.memory.MemoryBlockType;
 import org.apache.iotdb.commons.pipe.resource.log.PipeLogManager;
 import org.apache.iotdb.commons.pipe.resource.ref.PipePhantomReferenceManager;
 import org.apache.iotdb.commons.pipe.resource.snapshot.PipeSnapshotResourceManager;
-import org.apache.iotdb.confignode.conf.ConfigNodeDescriptor;
+import org.apache.iotdb.confignode.manager.pipe.resource.memory.PipeConfigNodeMemoryManager;
 import org.apache.iotdb.confignode.manager.pipe.resource.ref.PipeConfigNodePhantomReferenceManager;
 import org.apache.iotdb.confignode.manager.pipe.resource.snapshot.PipeConfigNodeSnapshotResourceManager;
 
 public class PipeConfigNodeResourceManager {
 
   private final PipeSnapshotResourceManager pipeSnapshotResourceManager;
-  private final IMemoryBlock pipeLogReducerMemoryBlock;
+  private final PipeConfigNodeMemoryManager pipeMemoryManager;
   private final PipeLogManager pipeLogManager;
   private final PipePhantomReferenceManager pipePhantomReferenceManager;
 
@@ -40,9 +38,8 @@ public static PipeSnapshotResourceManager snapshot() {
         .pipeSnapshotResourceManager;
   }
 
-  public static IMemoryBlock logReducerMemory() {
-    return PipeConfigNodeResourceManager.PipeResourceManagerHolder.INSTANCE
-        .pipeLogReducerMemoryBlock;
+  public static PipeConfigNodeMemoryManager memory() {
+    return PipeConfigNodeResourceManager.PipeResourceManagerHolder.INSTANCE.pipeMemoryManager;
   }
 
   public static PipeLogManager log() {
@@ -57,11 +54,7 @@ public static PipePhantomReferenceManager ref() {
 
   private PipeConfigNodeResourceManager() {
     pipeSnapshotResourceManager = new PipeConfigNodeSnapshotResourceManager();
-    pipeLogReducerMemoryBlock =
-        ConfigNodeDescriptor.getInstance()
-            .getMemoryConfig()
-            .getPipeMemoryManager()
-            .exactAllocate("PipePeriodicalLogReducer", MemoryBlockType.DYNAMIC);
+    pipeMemoryManager = new PipeConfigNodeMemoryManager();
     pipeLogManager = new PipeLogManager();
     pipePhantomReferenceManager = new PipeConfigNodePhantomReferenceManager();
   }

iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/resource/memory/PipeConfigNodeMemoryManager.java

@@ -0,0 +1,115 @@
+/*
+ * 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.confignode.manager.pipe.resource.memory;
+
+import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException;
+import org.apache.iotdb.commons.i18n.CommonMessages;
+import org.apache.iotdb.commons.i18n.PipeMessages;
+import org.apache.iotdb.commons.memory.IMemoryBlock;
+import org.apache.iotdb.commons.memory.MemoryBlockType;
+import org.apache.iotdb.commons.pipe.config.PipeConfig;
+import org.apache.iotdb.confignode.conf.ConfigNodeDescriptor;
+
+public class PipeConfigNodeMemoryManager {
+
+  private static final PipeConfig PIPE_CONFIG = PipeConfig.getInstance();
+
+  private final IMemoryBlock memoryBlock =
+      ConfigNodeDescriptor.getInstance()
+          .getMemoryConfig()
+          .getPipeMemoryManager()
+          .exactAllocate("Stream", MemoryBlockType.DYNAMIC);
+
+  private long pipeLogReducerMemoryUsageInBytes;
+
+  public synchronized long resizeLogReducerMemory(final long targetSizeInBytes) {
+    final long nonNegativeTargetSizeInBytes = Math.max(0, targetSizeInBytes);
+    if (pipeLogReducerMemoryUsageInBytes < nonNegativeTargetSizeInBytes) {
+      final long deltaSizeInBytes = nonNegativeTargetSizeInBytes - pipeLogReducerMemoryUsageInBytes;
+      if (memoryBlock.allocate(deltaSizeInBytes)) {
+        pipeLogReducerMemoryUsageInBytes = nonNegativeTargetSizeInBytes;
+      }
+    } else if (pipeLogReducerMemoryUsageInBytes > nonNegativeTargetSizeInBytes) {
+      final long deltaSizeInBytes = pipeLogReducerMemoryUsageInBytes - nonNegativeTargetSizeInBytes;
+      releaseMemory(deltaSizeInBytes);
+      pipeLogReducerMemoryUsageInBytes = nonNegativeTargetSizeInBytes;
+    }
+    return pipeLogReducerMemoryUsageInBytes;
+  }
+
+  public AutoCloseable tryAllocateReceiverMemory(final long requestedMemorySizeInBytes)
+      throws PipeRuntimeOutOfMemoryCriticalException {
+    final long nonNegativeRequestedMemorySizeInBytes = Math.max(0, requestedMemorySizeInBytes);
+    if (nonNegativeRequestedMemorySizeInBytes == 0) {
+      return () -> {};
+    }
+
+    forceAllocateWithRetry(nonNegativeRequestedMemorySizeInBytes);
+    return () -> releaseMemory(nonNegativeRequestedMemorySizeInBytes);
+  }
+
+  private synchronized void forceAllocateWithRetry(final long sizeInBytes)
+      throws PipeRuntimeOutOfMemoryCriticalException {
+    final int memoryAllocateMaxRetries = PIPE_CONFIG.getPipeMemoryAllocateMaxRetries();
+    for (int i = 0; i < memoryAllocateMaxRetries; ++i) {
+      if (memoryBlock.allocate(sizeInBytes)) {
+        return;
+      }
+
+      try {
+        final long retryIntervalInMs = PIPE_CONFIG.getPipeMemoryAllocateRetryIntervalInMs();
+        if (retryIntervalInMs > 0) {
+          wait(retryIntervalInMs);
+        }
+      } catch (final InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new PipeRuntimeOutOfMemoryCriticalException(PipeMessages.TEMPORARILY_OUT_OF_MEMORY);
+      }
+    }
+
+    throw new PipeRuntimeOutOfMemoryCriticalException(
+        String.format(
+            CommonMessages.EXCEPTION_EXACTALLOCATE_FAILED_ALLOCATE_MEMORY_AFTER_ARG_RETRIES_957A647B
+                + CommonMessages
+                    .EXCEPTION_TOTAL_MEMORY_SIZE_ARG_BYTES_USED_MEMORY_SIZE_ARG_BYTES_9FC9A9C6
+                + CommonMessages.EXCEPTION_REQUESTED_MEMORY_SIZE_ARG_BYTES_E6340842,
+            memoryAllocateMaxRetries,
+            memoryBlock.getTotalMemorySizeInBytes(),
+            memoryBlock.getUsedMemoryInBytes(),
+            sizeInBytes));
+  }
+
+  private synchronized void releaseMemory(final long sizeInBytes) {
+    memoryBlock.release(sizeInBytes);
+    notifyAll();
+  }
+
+  public long getUsedMemorySizeInBytes() {
+    return memoryBlock.getUsedMemoryInBytes();
+  }
+
+  public long getFreeMemorySizeInBytes() {
+    return memoryBlock.getFreeMemoryInBytes();
+  }
+
+  public long getTotalMemorySizeInBytes() {
+    return memoryBlock.getTotalMemorySizeInBytes();
+  }
+}

iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template

@@ -808,7 +808,8 @@ partition_table_recover_max_read_megabytes_per_second=10
 datanode_memory_proportion=3:3:1:1:1:1
 
 # ConfigNode Memory Allocation Ratio: Pipe and Free Memory.
-# The parameter form is a:b, where a and b are integers. Currently, only PipePeriodicalLogReducer is connected to Pipe memory on ConfigNode.
+# The parameter form is a:b, where a and b are integers. Currently, ConfigNode Pipe uses
+# this memory for log reducer and receiver temporary buffers.
 # effectiveMode: restart
 confignode_memory_proportion=1:9