[to dev/1.3] Pipe: manage ConfigNode receiver temporary memory (#18144)

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

@@ -97,7 +97,7 @@ public synchronized void stop() {
 
   private void initPipePeriodicalLogReducer() {
     PipePeriodicalLogReducer.setMemoryResizeFunction(
-        PipeConfigNodeResourceManager::resizeLogReducerMemory);
+        PipeConfigNodeResourceManager.memory()::resizeLogReducerMemory);
     PipeLogger.setLogger(PipePeriodicalLogReducer::log);
   }
 

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

@@ -23,6 +23,7 @@
 import org.apache.iotdb.commons.auth.entity.PrivilegeType;
 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.pipe.datastructure.pattern.IoTDBPipePattern;
@@ -57,6 +58,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;
@@ -147,14 +149,21 @@ public TPipeTransferResp receive(final TPipeTransferReq req) {
                 .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(
@@ -163,7 +172,14 @@ public TPipeTransferResp receive(final TPipeTransferReq req) {
                 .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;
         }
@@ -216,6 +232,17 @@ private TPipeTransferResp checkPipeTransferAuthenticated(final PipeRequestType t
     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

@@ -22,16 +22,14 @@
 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;
 
-import java.util.concurrent.atomic.AtomicLong;
-
 public class PipeConfigNodeResourceManager {
 
   private final PipeSnapshotResourceManager pipeSnapshotResourceManager;
-  private final AtomicLong pipeLogReducerMemoryUsageInBytes = new AtomicLong(0);
+  private final PipeConfigNodeMemoryManager pipeMemoryManager;
   private final PipeLogManager pipeLogManager;
   private final PipePhantomReferenceManager pipePhantomReferenceManager;
 
@@ -40,8 +38,8 @@ public static PipeSnapshotResourceManager snapshot() {
         .pipeSnapshotResourceManager;
   }
 
-  public static long resizeLogReducerMemory(final long targetSizeInBytes) {
-    return PipeResourceManagerHolder.INSTANCE.resizePipeLogReducerMemory(targetSizeInBytes);
+  public static PipeConfigNodeMemoryManager memory() {
+    return PipeConfigNodeResourceManager.PipeResourceManagerHolder.INSTANCE.pipeMemoryManager;
   }
 
   public static PipeLogManager log() {
@@ -54,16 +52,9 @@ public static PipePhantomReferenceManager ref() {
 
   ///////////////////////////// SINGLETON /////////////////////////////
 
-  private long resizePipeLogReducerMemory(final long targetSizeInBytes) {
-    final long pipeMemorySizeInBytes =
-        ConfigNodeDescriptor.getInstance().getMemoryConfig().getPipeMemorySizeInBytes();
-    final long resizedSizeInBytes = Math.min(Math.max(0, targetSizeInBytes), pipeMemorySizeInBytes);
-    pipeLogReducerMemoryUsageInBytes.set(resizedSizeInBytes);
-    return pipeLogReducerMemoryUsageInBytes.get();
-  }
-
   private PipeConfigNodeResourceManager() {
     pipeSnapshotResourceManager = new PipeConfigNodeSnapshotResourceManager();
+    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.pipe.config.PipeConfig;
+import org.apache.iotdb.confignode.conf.ConfigNodeDescriptor;
+
+public class PipeConfigNodeMemoryManager {
+
+  private static final PipeConfig PIPE_CONFIG = PipeConfig.getInstance();
+
+  private long usedMemorySizeInBytes;
+
+  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 (tryAllocate(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 (tryAllocate(sizeInBytes)) {
+        return;
+      }
+
+      try {
+        final long retryIntervalInMs = PIPE_CONFIG.getPipeMemoryAllocateRetryIntervalInMs();
+        if (retryIntervalInMs > 0) {
+          wait(retryIntervalInMs);
+        }
+      } catch (final InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new PipeRuntimeOutOfMemoryCriticalException("Temporarily out of memory.");
+      }
+    }
+
+    throw new PipeRuntimeOutOfMemoryCriticalException(
+        String.format(
+            "forceAllocate: failed to allocate memory after %d retries, "
+                + "total memory size %d bytes, used memory size %d bytes, "
+                + "requested memory size %d bytes",
+            memoryAllocateMaxRetries,
+            getTotalMemorySizeInBytes(),
+            usedMemorySizeInBytes,
+            sizeInBytes));
+  }
+
+  private boolean tryAllocate(final long sizeInBytes) {
+    if (getFreeMemorySizeInBytes() < sizeInBytes) {
+      return false;
+    }
+
+    usedMemorySizeInBytes += sizeInBytes;
+    return true;
+  }
+
+  private synchronized void releaseMemory(final long sizeInBytes) {
+    usedMemorySizeInBytes = Math.max(0, usedMemorySizeInBytes - sizeInBytes);
+    notifyAll();
+  }
+
+  public synchronized long getUsedMemorySizeInBytes() {
+    return usedMemorySizeInBytes;
+  }
+
+  public synchronized long getFreeMemorySizeInBytes() {
+    return Math.max(0, getTotalMemorySizeInBytes() - usedMemorySizeInBytes);
+  }
+
+  public long getTotalMemorySizeInBytes() {
+    return ConfigNodeDescriptor.getInstance().getMemoryConfig().getPipeMemorySizeInBytes();
+  }
+}

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

@@ -731,7 +731,8 @@ disk_space_warning_threshold=0.05
 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
 

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java

@@ -112,6 +112,27 @@ public static TPipeTransferReq fromTPipeTransferReq(final TPipeTransferReq trans
     return decompressedReq;
   }
 
+  /** Get the largest intermediate decompressed body size without consuming the request body. */
+  public static int getMaxDecompressedLengthInBytes(final TPipeTransferReq transferReq) {
+    final ByteBuffer compressedBuffer = transferReq.body.duplicate();
+
+    int maxDecompressedLength = compressedBuffer.remaining();
+    final int compressorsSize = ReadWriteIOUtils.readByte(compressedBuffer);
+    for (int i = 0; i < compressorsSize; ++i) {
+      ReadWriteIOUtils.readByte(compressedBuffer);
+      final int decompressedLength = ReadWriteIOUtils.readInt(compressedBuffer);
+      checkDecompressedLength(decompressedLength);
+      maxDecompressedLength = Math.max(maxDecompressedLength, decompressedLength);
+    }
+    return maxDecompressedLength;
+  }
+
+  /** Get the largest additional decompressed body size beyond the current transfer frame. */
+  public static int getMaxAdditionalDecompressedLengthInBytes(final TPipeTransferReq transferReq) {
+    final int transferFrameBodySize = transferReq.body.duplicate().remaining();
+    return Math.max(0, getMaxDecompressedLengthInBytes(transferReq) - transferFrameBodySize);
+  }
+
   /** This method is used to prevent decompression bomb attacks. */
   private static void checkDecompressedLength(final int decompressedLength)
       throws IllegalArgumentException {