From 585ecae6b81ec31db785f5050d634b1afe7bfefa Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:01:47 +0800 Subject: [PATCH] Pipe: manage ConfigNode receiver temporary memory (#18144) (cherry picked from commit 33a38b6194de81658c3979aaffb2ca3b83d07858) --- .../runtime/PipeConfigNodeRuntimeAgent.java | 2 +- .../protocol/IoTDBConfigNodeReceiver.java | 39 +++++- .../PipeConfigNodeResourceManager.java | 19 +-- .../memory/PipeConfigNodeMemoryManager.java | 115 ++++++++++++++++++ .../conf/iotdb-system.properties.template | 3 +- .../request/PipeTransferCompressedReq.java | 21 ++++ 6 files changed, 177 insertions(+), 22 deletions(-) create mode 100644 iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/resource/memory/PipeConfigNodeMemoryManager.java diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigNodeRuntimeAgent.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigNodeRuntimeAgent.java index 10c710a64adca..900ccc8c10bbb 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/agent/runtime/PipeConfigNodeRuntimeAgent.java +++ b/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); } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/receiver/protocol/IoTDBConfigNodeReceiver.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/receiver/protocol/IoTDBConfigNodeReceiver.java index 64f4a26350c25..10999474e9efd 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/receiver/protocol/IoTDBConfigNodeReceiver.java +++ b/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: diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/resource/PipeConfigNodeResourceManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/resource/PipeConfigNodeResourceManager.java index 635e6b7b1f8c1..60fc2a9e6a06c 100644 --- a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/resource/PipeConfigNodeResourceManager.java +++ b/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(); } diff --git a/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/resource/memory/PipeConfigNodeMemoryManager.java b/iotdb-core/confignode/src/main/java/org/apache/iotdb/confignode/manager/pipe/resource/memory/PipeConfigNodeMemoryManager.java new file mode 100644 index 0000000000000..471a033671006 --- /dev/null +++ b/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(); + } +} diff --git a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template index b977ce9de8c1e..42c0aac0e1ab6 100644 --- a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template +++ b/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 diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java index 73bac4c05b401..7aa639140beca 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java +++ b/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 {