From 396628989aa812fbea6438cca64c5837a9f91a2c Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Wed, 22 Jul 2026 16:16:58 +0800 Subject: [PATCH] Fix pipe tablet memory self-lock during batching (#18266) --- .../resource/memory/PipeMemoryManager.java | 16 ++- .../memory/PipeMemoryManagerResizeTest.java | 113 ++++++++++++++++++ 2 files changed, 128 insertions(+), 1 deletion(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java index 387a4228a2cd2..e21e66febeefa 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java @@ -209,6 +209,16 @@ < allowedMaxMemorySizeInBytesOfTabletsAndTsFiles() && (double) usedMemorySizeInBytesOfTsFiles < allowedMaxMemorySizeInBytesOfTsTiles(); } + private boolean isHardEnoughForResizing(final PipeMemoryBlock block) { + if (block instanceof PipeTabletMemoryBlock) { + return isHardEnough4TabletParsing(); + } + if (block instanceof PipeTsFileMemoryBlock) { + return isHardEnough4TsFileSlicing(); + } + return true; + } + public synchronized PipeMemoryBlock forceAllocate(long sizeInBytes) throws PipeRuntimeOutOfMemoryCriticalException { if (!PIPE_MEMORY_MANAGEMENT_ENABLED) { @@ -428,7 +438,11 @@ public synchronized void resize( long sizeInBytes = targetSize - oldSize; final int memoryAllocateMaxRetries = PipeConfig.getInstance().getPipeMemoryAllocateMaxRetries(); for (int i = 1; i <= memoryAllocateMaxRetries; i++) { - if (getTotalNonFloatingMemorySizeInBytes() - usedMemorySizeInBytes >= sizeInBytes) { + // Dynamically resized data-structure blocks must obey the same admission thresholds as + // blocks allocated with a non-zero initial size. Otherwise they can exhaust the pool and + // prevent downstream consumers from allocating the memory needed to release them. + if (isHardEnoughForResizing(block) + && getTotalNonFloatingMemorySizeInBytes() - usedMemorySizeInBytes >= sizeInBytes) { usedMemorySizeInBytes += sizeInBytes; if (oldSize == 0) { // If the memory block is not registered, we need to register it first. diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java new file mode 100644 index 0000000000000..be746addc9e39 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerResizeTest.java @@ -0,0 +1,113 @@ +/* + * 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.iotdb.commons.conf.CommonConfig; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; + +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +public class PipeMemoryManagerResizeTest { + + private final CommonConfig config = CommonDescriptor.getInstance().getConfig(); + + private boolean originalMemoryManagementEnabled; + private int originalAllocateMaxRetries; + private long originalAllocateRetryIntervalInMs; + private double originalFloatingMemoryProportion; + private double originalTabletRejectThreshold; + private double originalTsFileRejectThreshold; + + @Before + public void setUp() { + originalMemoryManagementEnabled = config.getPipeMemoryManagementEnabled(); + originalAllocateMaxRetries = config.getPipeMemoryAllocateMaxRetries(); + originalAllocateRetryIntervalInMs = config.getPipeMemoryAllocateRetryIntervalInMs(); + originalFloatingMemoryProportion = config.getPipeTotalFloatingMemoryProportion(); + originalTabletRejectThreshold = + config.getPipeDataStructureTabletMemoryBlockAllocationRejectThreshold(); + originalTsFileRejectThreshold = + config.getPipeDataStructureTsFileMemoryBlockAllocationRejectThreshold(); + + config.setPipeMemoryManagementEnabled(true); + config.setPipeMemoryAllocateMaxRetries(1); + config.setPipeMemoryAllocateRetryIntervalInMs(1); + config.setPipeTotalFloatingMemoryProportion(0.5); + config.setPipeDataStructureTabletMemoryBlockAllocationRejectThreshold(0.3); + config.setPipeDataStructureTsFileMemoryBlockAllocationRejectThreshold(0.3); + } + + @After + public void tearDown() { + config.setPipeMemoryManagementEnabled(originalMemoryManagementEnabled); + config.setPipeMemoryAllocateMaxRetries(originalAllocateMaxRetries); + config.setPipeMemoryAllocateRetryIntervalInMs(originalAllocateRetryIntervalInMs); + config.setPipeTotalFloatingMemoryProportion(originalFloatingMemoryProportion); + config.setPipeDataStructureTabletMemoryBlockAllocationRejectThreshold( + originalTabletRejectThreshold); + config.setPipeDataStructureTsFileMemoryBlockAllocationRejectThreshold( + originalTsFileRejectThreshold); + } + + @Test + public void testTabletResizeLeavesMemoryForSinkForwardProgress() { + final PipeMemoryManager manager = new PipeMemoryManager(); + final long totalNonFloatingMemorySizeInBytes = + PipeMemoryManager.getTotalNonFloatingMemorySizeInBytes(); + final long tabletMemorySizeInBytes = + (long) + (totalNonFloatingMemorySizeInBytes + * (config.getPipeDataStructureTabletMemoryBlockAllocationRejectThreshold() + + config.getPipeDataStructureTsFileMemoryBlockAllocationRejectThreshold() + / 2)) + + 1; + final long sinkMemorySizeInBytes = totalNonFloatingMemorySizeInBytes / 10; + final PipeTabletMemoryBlock retainedTablet = + manager.forceAllocateForTabletWithRetry(tabletMemorySizeInBytes); + final PipeTabletMemoryBlock pendingTablet = manager.forceAllocateForTabletWithRetry(0); + final PipeMemoryBlock sinkBatch = manager.forceAllocate(0); + + try { + Assert.assertThrows( + PipeRuntimeOutOfMemoryCriticalException.class, + () -> manager.forceResize(pendingTablet, 1)); + Assert.assertEquals(tabletMemorySizeInBytes, manager.getUsedMemorySizeInBytes()); + Assert.assertEquals(tabletMemorySizeInBytes, manager.getUsedMemorySizeInBytesOfTablets()); + + manager.forceResize(sinkBatch, sinkMemorySizeInBytes); + Assert.assertEquals( + tabletMemorySizeInBytes + sinkMemorySizeInBytes, manager.getUsedMemorySizeInBytes()); + + manager.release(retainedTablet); + manager.forceResize(pendingTablet, 1); + Assert.assertEquals(1, manager.getUsedMemorySizeInBytesOfTablets()); + } finally { + manager.release(retainedTablet); + manager.release(pendingTablet); + manager.release(sinkBatch); + } + + Assert.assertEquals(0, manager.getUsedMemorySizeInBytes()); + } +}