From a982a26d1fa99d33453040964f1796c4a23667c6 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:28:46 +0800 Subject: [PATCH] [Pipe] Unify TsFile parser allocation and resize quotas --- .../iotdb/db/i18n/DataNodePipeMessages.java | 4 + .../iotdb/db/i18n/DataNodePipeMessages.java | 4 + ...FileInsertionEventParserMemoryManager.java | 23 +-- .../memory/PipeMemoryBlockCategory.java | 1 + .../resource/memory/PipeMemoryBlockType.java | 1 + .../resource/memory/PipeMemoryManager.java | 82 +++++++- .../memory/PipeTsFileParserMemoryBlock.java | 33 ++++ .../memory/PipeMemoryManagerResizeTest.java | 180 ++++++++++++++++++ 8 files changed, 305 insertions(+), 23 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeTsFileParserMemoryBlock.java diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java index 8eb2b8a6116d0..8569f76908df6 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java @@ -2403,6 +2403,10 @@ private DataNodePipeMessages() {} public static final String PIPE_EXCEPTION_FORCERESIZE_FAILED_TO_ALLOCATE_MEMORY_AFTER_D_RETRIES_TOTAL_8C6948BC = "forceResize: failed to allocate memory after %d retries, total memory size %d bytes, used " + "memory size %d bytes, requested memory size %d bytes"; + public static final String EXCEPTION_FORCERESIZE_FAILED_AFTER_ARG_RETRIES_TOTAL_ARG_BYTES_USED_ARG_BYTES_USEDTABLETS_ARG_BYTES_USEDTSFILES_ARG_BYTES_TABLETLIMIT_ARG_BYTES_TSFILELIMIT_ARG_BYTES_TABLETANDTSFILELIMIT_ARG_BYTES_PROTECTEDPARSERLIMIT_ARG_BYTES_REQUESTED_ARG_BYTES_311C8C7E = + "forceResize failed after %d retries: total=%d bytes, used=%d bytes, " + + "usedTablets=%d bytes, usedTsFiles=%d bytes, tabletLimit=%d bytes, tsFileLimit=%d bytes, " + + "tabletAndTsFileLimit=%d bytes, protectedParserLimit=%d bytes, requested=%d bytes"; public static final String EXCEPTION_UNSUPPORTED_BATCH_TYPE_ARG_WHEN_TRANSFERRING_TABLET_INSERTION_EVENT_66153E12 = "Unsupported batch type %s when transferring tablet insertion event."; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java index ee86a424c299f..57f98bad5f06c 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java @@ -2233,6 +2233,10 @@ private DataNodePipeMessages() {} "forceAllocate：重试 %d 次后仍无法分配内存，总内存大小 %d bytes，已用内存大小 %d bytes，请求内存大小 %d bytes"; public static final String PIPE_EXCEPTION_FORCERESIZE_FAILED_TO_ALLOCATE_MEMORY_AFTER_D_RETRIES_TOTAL_8C6948BC = "forceResize：重试 %d 次后仍无法分配内存，总内存大小 %d bytes，已用内存大小 %d bytes，请求内存大小 %d bytes"; + public static final String EXCEPTION_FORCERESIZE_FAILED_AFTER_ARG_RETRIES_TOTAL_ARG_BYTES_USED_ARG_BYTES_USEDTABLETS_ARG_BYTES_USEDTSFILES_ARG_BYTES_TABLETLIMIT_ARG_BYTES_TSFILELIMIT_ARG_BYTES_TABLETANDTSFILELIMIT_ARG_BYTES_PROTECTEDPARSERLIMIT_ARG_BYTES_REQUESTED_ARG_BYTES_311C8C7E = + "forceResize 重试 %d 次后仍无法分配内存：total=%d bytes，used=%d bytes，" + + "usedTablets=%d bytes，usedTsFiles=%d bytes，tabletLimit=%d bytes，tsFileLimit=%d bytes，" + + "tabletAndTsFileLimit=%d bytes，protectedParserLimit=%d bytes，requested=%d bytes"; public static final String EXCEPTION_UNSUPPORTED_BATCH_TYPE_ARG_WHEN_TRANSFERRING_TABLET_INSERTION_EVENT_66153E12 = "传输 tablet insertion event 时不支持 batch 类型 %s。"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java index 36b522f199995..8e7048d16d7a0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java @@ -20,10 +20,10 @@ package org.apache.iotdb.db.pipe.event.common.tsfile.parser; import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; -import org.apache.iotdb.db.i18n.DataNodePipeMessages; 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.PipeMemoryBlockCategory; +import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryManager; /** Allocates parser working memory from the pool owned by the caller. */ public interface TsFileInsertionEventParserMemoryManager { @@ -63,11 +63,8 @@ public TsFileInsertionEventParserMemoryBlock forceAllocateForTabletWithRetry( final String name, final long sizeInBytes) { return new PipeBlock( PipeDataNodeResourceManager.memory() - .forceAllocateForTabletWithRetry( - name, - sizeInBytes, - PipeMemoryBlockCategory.TABLET, - TsFileInsertionEventParser.class.getSimpleName())); + .forceAllocateForTsFileParserWithRetry( + name, sizeInBytes, TsFileInsertionEventParser.class.getSimpleName())); } @Override @@ -90,8 +87,7 @@ public TsFileInsertionEventParserMemoryBlock forceAllocateForTabletWithRetry( final Object assigner) { return new PipeBlock( PipeDataNodeResourceManager.memory() - .forceAllocateForTabletWithRetry( - name, sizeInBytes, PipeMemoryBlockCategory.TABLET, assigner, parent)); + .forceAllocateForTsFileParserWithRetry(name, sizeInBytes, assigner, parent)); } @Override @@ -123,15 +119,10 @@ public long getMemoryUsageInBytes() { @Override public void forceResize(final long newSizeInBytes) { final long oldSize = delegate.getMemoryUsageInBytes(); - if (!PipeDataNodeResourceManager.memory().tryResize(delegate, newSizeInBytes)) { + final PipeMemoryManager memoryManager = PipeDataNodeResourceManager.memory(); + if (!memoryManager.tryResize(delegate, newSizeInBytes)) { throw new PipeRuntimeOutOfMemoryCriticalException( - String.format( - DataNodePipeMessages - .PIPE_EXCEPTION_FORCERESIZE_FAILED_TO_ALLOCATE_MEMORY_AFTER_D_RETRIES_TOTAL_8C6948BC, - 0, - PipeDataNodeResourceManager.memory().getTotalNonFloatingMemorySizeInBytes(), - PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes(), - newSizeInBytes - oldSize)); + memoryManager.getResizeFailureMessage(0, newSizeInBytes - oldSize)); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryBlockCategory.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryBlockCategory.java index 6b3fce264f320..0427e226d181d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryBlockCategory.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryBlockCategory.java @@ -48,6 +48,7 @@ public static PipeMemoryBlockCategory fromType(final PipeMemoryBlockType type) { return switch (type) { case NORMAL -> OTHER; case TABLET -> TABLET; + case TS_FILE_PARSER -> PARSER; case TS_FILE -> TS_FILE; case BATCH -> BATCH; case WAL -> WAL; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryBlockType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryBlockType.java index 846fc7dd1ce17..78fd9aff6704a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryBlockType.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryBlockType.java @@ -22,6 +22,7 @@ public enum PipeMemoryBlockType { NORMAL, TABLET, + TS_FILE_PARSER, TS_FILE, BATCH, WAL 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 b86c4ab5954fa..3d51448df10c9 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 @@ -152,6 +152,35 @@ private double allowedMaxMemorySizeInBytesOfTsTiles( * totalNonFloatingMemorySizeInBytes; } + private boolean isEnoughForTsFileParserMemory( + final long extraMemoryInBytes, final long totalNonFloatingMemorySizeInBytes) { + return (double) usedMemorySizeInBytesOfTablets + + (double) usedMemorySizeInBytesOfTsFiles + + (double) extraMemoryInBytes + < EXCEED_PROTECT_THRESHOLD + * allowedMaxMemorySizeInBytesOfTabletsAndTsFiles(totalNonFloatingMemorySizeInBytes); + } + + /** Reports both global availability and the category limits that can reject a resize. */ + public synchronized String getResizeFailureMessage(final int retries, final long requestedSize) { + final long totalNonFloatingMemorySizeInBytes = getTotalNonFloatingMemorySizeInBytes(); + final double combinedLimit = + allowedMaxMemorySizeInBytesOfTabletsAndTsFiles(totalNonFloatingMemorySizeInBytes); + return String.format( + DataNodePipeMessages + .EXCEPTION_FORCERESIZE_FAILED_AFTER_ARG_RETRIES_TOTAL_ARG_BYTES_USED_ARG_BYTES_USEDTABLETS_ARG_BYTES_USEDTSFILES_ARG_BYTES_TABLETLIMIT_ARG_BYTES_TSFILELIMIT_ARG_BYTES_TABLETANDTSFILELIMIT_ARG_BYTES_PROTECTEDPARSERLIMIT_ARG_BYTES_REQUESTED_ARG_BYTES_311C8C7E, + retries, + totalNonFloatingMemorySizeInBytes, + memoryBlock.getUsedMemoryInBytes(), + usedMemorySizeInBytesOfTablets, + usedMemorySizeInBytesOfTsFiles, + Math.round(allowedMaxMemorySizeInBytesOfTablets(totalNonFloatingMemorySizeInBytes)), + Math.round(allowedMaxMemorySizeInBytesOfTsTiles(totalNonFloatingMemorySizeInBytes)), + Math.round(combinedLimit), + Math.round(EXCEED_PROTECT_THRESHOLD * combinedLimit), + requestedSize); + } + private static long getTsFileParserMemorySizeInBytes() { return Math.max( PIPE_CONFIG.getTsFileParserMemory(), PIPE_CONFIG.getPipeMemoryAllocateMinSizeInBytes()); @@ -556,6 +585,10 @@ < allowedMaxMemorySizeInBytesOfTabletsAndTsFiles(totalNonFloatingMemorySizeInByt private boolean isHardEnoughForResizing( final PipeMemoryBlock block, final long extraMemoryInBytes) { + if (block instanceof PipeTsFileParserMemoryBlock) { + return isEnoughForTsFileParserMemory( + extraMemoryInBytes, getTotalNonFloatingMemorySizeInBytes()); + } if (!(block instanceof PipeTabletMemoryBlock) && !(block instanceof PipeTsFileMemoryBlock)) { return true; } @@ -757,6 +790,39 @@ public PipeTabletMemoryBlock forceAllocateForTabletWithRetry( name, tabletSizeInBytes, PipeMemoryBlockCategory.EVENT_CHILD, null, parent); } + /** Allocates TsFile parser memory against the protected tablet + TsFile combined quota. */ + public PipeTsFileParserMemoryBlock forceAllocateForTsFileParserWithRetry( + final String name, final long parserMemorySizeInBytes, final Object assigner) { + return forceAllocateForTsFileParserWithRetry(name, parserMemorySizeInBytes, assigner, null); + } + + /** Allocates TsFile parser memory as an event child against the combined quota. */ + public synchronized PipeTsFileParserMemoryBlock forceAllocateForTsFileParserWithRetry( + final String name, + final long parserMemorySizeInBytes, + final Object assigner, + final PipeMemoryBlock parent) { + final PipeMemoryBlockCategory category = + parent == null ? PipeMemoryBlockCategory.PARSER : PipeMemoryBlockCategory.EVENT_CHILD; + final PipeTsFileParserMemoryBlock block = + (PipeTsFileParserMemoryBlock) + registerMemoryBlock( + name, 0, PipeMemoryBlockType.TS_FILE_PARSER, category, assigner, parent); + if (!PIPE_MEMORY_MANAGEMENT_ENABLED || parserMemorySizeInBytes == 0) { + return block; + } + + try { + // Initial allocation and subsequent growth must check the same post-allocation usage + // under the manager lock, including after a wait releases that lock. + forceResize(block, parserMemorySizeInBytes); + return block; + } catch (final RuntimeException e) { + release(block); + throw e; + } + } + public PipeTsFileMemoryBlock forceAllocateForTsFileWithRetry( final String name, final long tsFileSizeInBytes) throws PipeRuntimeOutOfMemoryCriticalException { @@ -1001,17 +1067,14 @@ public synchronized void resize( if (force) { throw new PipeRuntimeOutOfMemoryCriticalException( - String.format( - DataNodePipeMessages - .PIPE_EXCEPTION_FORCERESIZE_FAILED_TO_ALLOCATE_MEMORY_AFTER_D_RETRIES_TOTAL_8C6948BC, - memoryAllocateMaxRetries, - getTotalNonFloatingMemorySizeInBytes(), - memoryBlock.getUsedMemoryInBytes(), - sizeInBytes)); + getResizeFailureMessage(memoryAllocateMaxRetries, sizeInBytes)); } } private boolean tryResizeInternal(final PipeMemoryBlock block, final long targetSize) { + if (block.isReleased()) { + return false; + } // Parent blocks expose an aggregate usage. Do not let a direct resize reduce that aggregate // below the bytes still owned by live children. final long normalizedTargetSize = @@ -1249,6 +1312,11 @@ private synchronized PipeMemoryBlock registerMemoryBlock( returnedMemoryBlock = new PipeTabletMemoryBlock(this, name, 0, normalizedCategory, null, normalizedParent); break; + case TS_FILE_PARSER: + returnedMemoryBlock = + new PipeTsFileParserMemoryBlock( + this, name, 0, normalizedCategory, null, normalizedParent); + break; case TS_FILE: returnedMemoryBlock = new PipeTsFileMemoryBlock(this, name, 0, normalizedCategory, null, normalizedParent); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeTsFileParserMemoryBlock.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeTsFileParserMemoryBlock.java new file mode 100644 index 0000000000000..999610b7517da --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeTsFileParserMemoryBlock.java @@ -0,0 +1,33 @@ +/* + * 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; + +/** Parser working memory charged to TsFiles and admitted against the protected combined quota. */ +public class PipeTsFileParserMemoryBlock extends PipeTsFileMemoryBlock { + + PipeTsFileParserMemoryBlock( + final PipeMemoryManager pipeMemoryManager, + final String name, + final long memoryUsageInBytes, + final PipeMemoryBlockCategory category, + final String assigner, + final PipeMemoryBlock parent) { + super(pipeMemoryManager, name, memoryUsageInBytes, category, assigner, parent); + } +} 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 index 9fd6820feceb3..85cfb650d9c9e 100644 --- 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 @@ -41,10 +41,16 @@ import org.mockito.Mockito; import java.lang.reflect.Field; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; import java.util.stream.Collectors; @@ -117,6 +123,180 @@ public void testTabletResizeCannotCrossTabletHardLimit() { } } + @Test + public void testTsFileParserAllocationAndResizeShareProtectedCombinedLimit() { + final AtomicLong floatingMemoryUsageInBytes = new AtomicLong(0); + final PipeMemoryManager manager = + new PipeMemoryManager( + new AtomicLongMemoryBlock( + "PipeMemoryManagerResizeTest", + null, + TOTAL_MEMORY_SIZE_IN_BYTES, + MemoryBlockType.DYNAMIC), + floatingMemoryUsageInBytes::get); + final PipeTsFileParserMemoryBlock parser = + manager.forceAllocateForTsFileParserWithRetry("parser", 1000, "test"); + + try { + Assert.assertTrue(manager.tryResize(parser, 1100)); + Assert.assertFalse(manager.tryResize(parser, 1140)); + Assert.assertEquals(1100, parser.getMemoryUsageInBytes()); + Assert.assertEquals(0, manager.getUsedMemorySizeInBytesOfTablets()); + Assert.assertEquals(1100, manager.getUsedMemorySizeInBytesOfTsFiles()); + Assert.assertEquals(1100, manager.getUsedMemorySizeInBytes()); + + final PipeRuntimeOutOfMemoryCriticalException allocationException = + Assert.assertThrows( + PipeRuntimeOutOfMemoryCriticalException.class, + () -> manager.forceAllocateForTsFileParserWithRetry("parser-2", 40, "test")); + Assert.assertTrue( + allocationException.getMessage(), + allocationException.getMessage().contains("usedTsFiles=1100 bytes")); + Assert.assertTrue( + allocationException.getMessage(), + allocationException.getMessage().contains("tabletLimit=900 bytes")); + Assert.assertTrue( + allocationException.getMessage().contains("tabletAndTsFileLimit=1200 bytes")); + Assert.assertTrue( + allocationException.getMessage().contains("protectedParserLimit=1140 bytes")); + Assert.assertEquals(2, manager.getPipeMemoryBlockInfoList().size()); + + // Parser working memory must not consume the tablet-only quota: the generated tablet + // still needs to be retained and sent while the parser owns its decoded pages/chunks. + final PipeTabletMemoryBlock output = manager.forceAllocateForTabletWithRetry("output", 50); + try { + Assert.assertTrue(manager.tryResize(output, 90)); + Assert.assertEquals(90, manager.getUsedMemorySizeInBytesOfTablets()); + Assert.assertEquals(1100, manager.getUsedMemorySizeInBytesOfTsFiles()); + } finally { + manager.release(output); + } + } finally { + manager.release(parser); + } + + Assert.assertEquals(0, manager.getUsedMemorySizeInBytes()); + } + + @Test + public void testZeroSizedTsFileParserCanGrowBeyondTabletLimitAndRelease() { + final PipeMemoryManager manager = + new PipeMemoryManager( + new AtomicLongMemoryBlock( + "PipeMemoryManagerResizeTest", + null, + TOTAL_MEMORY_SIZE_IN_BYTES, + MemoryBlockType.DYNAMIC), + () -> 0); + final PipeMemoryBlock event = manager.forceAllocate("event", 0); + final PipeTsFileParserMemoryBlock parser = + manager.forceAllocateForTsFileParserWithRetry("parser", 0, "test", event); + try { + Assert.assertTrue(manager.tryResize(parser, 1000)); + Assert.assertEquals(1000, event.getMemoryUsageInBytes()); + Assert.assertEquals(1000, manager.getUsedMemorySizeInBytesOfTsFiles()); + Assert.assertEquals(1000, manager.getUsedMemorySizeInBytes()); + + Assert.assertTrue(manager.tryResize(parser, 100)); + Assert.assertEquals(100, event.getMemoryUsageInBytes()); + Assert.assertEquals(100, manager.getUsedMemorySizeInBytesOfTsFiles()); + manager.release(event); + Assert.assertTrue(parser.isReleased()); + Assert.assertEquals(0, manager.getUsedMemorySizeInBytesOfTsFiles()); + Assert.assertEquals(0, manager.getUsedMemorySizeInBytesOfTablets()); + Assert.assertEquals(0, manager.getUsedMemorySizeInBytes()); + } finally { + manager.release(event); + } + } + + @Test + public void testTsFileParserSharesQuotaWithTabletsAndTsFilesAndHonorsGlobalLimit() { + final PipeMemoryManager manager = + new PipeMemoryManager( + new AtomicLongMemoryBlock( + "PipeMemoryManagerResizeTest", + null, + TOTAL_MEMORY_SIZE_IN_BYTES, + MemoryBlockType.DYNAMIC), + () -> 0); + final PipeTabletMemoryBlock tablet = manager.forceAllocateForTabletWithRetry("tablet", 200); + final PipeTsFileMemoryBlock tsFile = manager.forceAllocateForTsFileWithRetry("tsFile", 500); + final PipeTsFileParserMemoryBlock parser = + manager.forceAllocateForTsFileParserWithRetry("parser", 400, "test"); + final PipeMemoryBlock sink = manager.forceAllocate("sink", 900); + try { + Assert.assertFalse(manager.tryResize(parser, 440)); + Assert.assertFalse(manager.tryResize(parser, 401)); + Assert.assertEquals(200, manager.getUsedMemorySizeInBytesOfTablets()); + Assert.assertEquals(900, manager.getUsedMemorySizeInBytesOfTsFiles()); + Assert.assertEquals(2000, manager.getUsedMemorySizeInBytes()); + + manager.release(sink); + Assert.assertTrue(manager.tryResize(parser, 439)); + Assert.assertFalse(manager.tryResize(parser, 440)); + Assert.assertEquals(200, manager.getUsedMemorySizeInBytesOfTablets()); + Assert.assertEquals(939, manager.getUsedMemorySizeInBytesOfTsFiles()); + Assert.assertEquals(1139, manager.getUsedMemorySizeInBytes()); + } finally { + manager.release(tablet); + manager.release(tsFile); + manager.release(parser); + manager.release(sink); + } + Assert.assertEquals(0, manager.getUsedMemorySizeInBytesOfTablets()); + Assert.assertEquals(0, manager.getUsedMemorySizeInBytesOfTsFiles()); + Assert.assertEquals(0, manager.getUsedMemorySizeInBytes()); + } + + @Test + public void testConcurrentTsFileParserAllocationsCannotOvershootQuota() throws Exception { + final PipeMemoryManager manager = + new PipeMemoryManager( + new AtomicLongMemoryBlock( + "PipeMemoryManagerResizeTest", + null, + TOTAL_MEMORY_SIZE_IN_BYTES, + MemoryBlockType.DYNAMIC), + () -> 0); + final ExecutorService executor = Executors.newFixedThreadPool(4); + final CountDownLatch start = new CountDownLatch(1); + final List> futures = new ArrayList<>(); + final List allocated = new ArrayList<>(); + try { + for (int i = 0; i < 20; i++) { + futures.add( + executor.submit( + () -> { + start.await(); + try { + return manager.forceAllocateForTsFileParserWithRetry("parser", 300, "test"); + } catch (final PipeRuntimeOutOfMemoryCriticalException e) { + return null; + } + })); + } + start.countDown(); + for (final Future future : futures) { + final PipeTsFileParserMemoryBlock block = future.get(10, TimeUnit.SECONDS); + if (block != null) { + allocated.add(block); + } + } + Assert.assertEquals(3, allocated.size()); + Assert.assertEquals(900, manager.getUsedMemorySizeInBytes()); + Assert.assertEquals(0, manager.getUsedMemorySizeInBytesOfTablets()); + Assert.assertEquals(900, manager.getUsedMemorySizeInBytesOfTsFiles()); + Assert.assertEquals(4, manager.getPipeMemoryBlockInfoList().size()); + } finally { + executor.shutdownNow(); + allocated.forEach(manager::release); + } + Assert.assertEquals(0, manager.getUsedMemorySizeInBytes()); + Assert.assertEquals(0, manager.getUsedMemorySizeInBytesOfTsFiles()); + Assert.assertEquals(0, manager.getUsedMemorySizeInBytesOfTablets()); + } + @Test public void testTabletResizeLeavesMemoryForSinkForwardProgress() { final PipeMemoryManager manager =