From a35e8a42a6e324fad9111a327cf4bda52655eec6 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:53:10 +0800 Subject: [PATCH 1/2] fix(load): isolate reused tsfile parser memory from pipe pool --- .../parser/TsFileInsertionEventParser.java | 44 ++++++- ...TsFileInsertionEventParserMemoryBlock.java | 37 ++++++ ...FileInsertionEventParserMemoryManager.java | 75 +++++++++++ .../TsFileInsertionEventQueryParser.java | 69 ++++++++-- ...sertionEventQueryParserTabletIterator.java | 10 +- .../scan/TsFileInsertionEventScanParser.java | 92 ++++++++++---- .../TsFileInsertionEventTableParser.java | 92 +++++++++++--- ...sertionEventTableParserTabletIterator.java | 40 +++--- ...tementDataTypeConvertExecutionVisitor.java | 4 +- .../LoadTreeTsFileTabletIterator.java | 23 +++- .../memory/LoadTsFileParserMemoryManager.java | 120 ++++++++++++++++++ .../event/TsFileInsertionEventParserTest.java | 56 +++++--- ...ntDataTypeConvertExecutionVisitorTest.java | 25 ++++ ...adTsFileParserPipeMemoryIsolationTest.java | 66 ++++++++++ .../memory/LoadTsFileMemoryManagerTest.java | 17 +++ 15 files changed, 665 insertions(+), 105 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryBlock.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileParserMemoryManager.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileParserPipeMemoryIsolationTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParser.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParser.java index da5cc6b88dba1..eefcafec9c53b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParser.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParser.java @@ -28,8 +28,6 @@ import org.apache.iotdb.db.pipe.event.common.PipeInsertionEvent; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.table.TsFileInsertionEventTableParser; import org.apache.iotdb.db.pipe.metric.overview.PipeTsFileToTabletsMetrics; -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.PipeMemoryWeightUtil; import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry; import org.apache.iotdb.db.utils.datastructure.PatternTreeMapFactory; @@ -64,14 +62,16 @@ public abstract class TsFileInsertionEventParser implements AutoCloseable { protected final PipeInsertionEvent sourceEvent; // used to report progress // mods entry - protected PipeMemoryBlock allocatedMemoryBlockForModifications; + protected TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForModifications; protected PatternTreeMap currentModifications; protected long parseStartTimeNano = -1; protected boolean parseStartTimeRecorded = false; protected boolean parseEndTimeRecorded = false; - protected final PipeMemoryBlock allocatedMemoryBlockForTablet; + protected final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForTablet; + + protected final TsFileInsertionEventParserMemoryManager memoryManager; protected TsFileSequenceReader tsFileSequenceReader; @@ -90,6 +90,36 @@ protected TsFileInsertionEventParser( final boolean skipIfNoPrivileges, final PipeInsertionEvent sourceEvent, final boolean isWithMod) { + this( + tsFile, + pipeName, + creationTime, + treePattern, + tablePattern, + startTime, + endTime, + pipeTaskMeta, + entity, + skipIfNoPrivileges, + sourceEvent, + isWithMod, + TsFileInsertionEventParserMemoryManager.pipe()); + } + + protected TsFileInsertionEventParser( + final File tsFile, + final String pipeName, + final long creationTime, + final TreePattern treePattern, + final TablePattern tablePattern, + final long startTime, + final long endTime, + final PipeTaskMeta pipeTaskMeta, + final IAuditEntity entity, + final boolean skipIfNoPrivileges, + final PipeInsertionEvent sourceEvent, + final boolean isWithMod, + final TsFileInsertionEventParserMemoryManager memoryManager) { this.pipeName = pipeName; this.creationTime = creationTime; this.entity = entity; @@ -106,9 +136,9 @@ protected TsFileInsertionEventParser( this.pipeTaskMeta = pipeTaskMeta; this.sourceEvent = sourceEvent; + this.memoryManager = memoryManager; - this.allocatedMemoryBlockForTablet = - PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(0); + this.allocatedMemoryBlockForTablet = memoryManager.forceAllocateForTabletWithRetry(0); LOGGER.debug( DataNodePipeMessages.TSFILE_HAS_INITIALIZED_PIPENAME_CREATION_TIME_PATTERN, @@ -180,7 +210,7 @@ protected void recordTabletMetrics(final Tablet tablet) { protected void releaseTabletMemoryBlock() { if (allocatedMemoryBlockForTablet != null && allocatedMemoryBlockForTablet.getMemoryUsageInBytes() > 0) { - PipeDataNodeResourceManager.memory().forceResize(allocatedMemoryBlockForTablet, 0); + allocatedMemoryBlockForTablet.forceResize(0); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryBlock.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryBlock.java new file mode 100644 index 0000000000000..84a458028e0ec --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryBlock.java @@ -0,0 +1,37 @@ +/* + * 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.event.common.tsfile.parser; + +/** + * Memory block used by a tsfile parser. + * + *

The parser is shared by Pipe and Load. Keeping the block behind this small interface allows + * the same parsing code to use the owning subsystem's memory pool instead of hard-coding the Pipe + * pool. + */ +public interface TsFileInsertionEventParserMemoryBlock extends AutoCloseable { + + long getMemoryUsageInBytes(); + + void forceResize(long newSizeInBytes); + + @Override + void close(); +} 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 new file mode 100644 index 0000000000000..68669ce4b55d8 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/TsFileInsertionEventParserMemoryManager.java @@ -0,0 +1,75 @@ +/* + * 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.event.common.tsfile.parser; + +import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; +import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlock; + +/** Allocates parser working memory from the pool owned by the caller. */ +public interface TsFileInsertionEventParserMemoryManager { + + TsFileInsertionEventParserMemoryBlock forceAllocateForTabletWithRetry(long sizeInBytes); + + TsFileInsertionEventParserMemoryBlock forceAllocate(long sizeInBytes); + + static TsFileInsertionEventParserMemoryManager pipe() { + return PipeHolder.INSTANCE; + } + + final class PipeHolder { + private static final TsFileInsertionEventParserMemoryManager INSTANCE = + new TsFileInsertionEventParserMemoryManager() { + @Override + public TsFileInsertionEventParserMemoryBlock forceAllocateForTabletWithRetry( + final long sizeInBytes) { + return new PipeBlock( + PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(sizeInBytes)); + } + + @Override + public TsFileInsertionEventParserMemoryBlock forceAllocate(final long sizeInBytes) { + return new PipeBlock(PipeDataNodeResourceManager.memory().forceAllocate(sizeInBytes)); + } + }; + } + + final class PipeBlock implements TsFileInsertionEventParserMemoryBlock { + private final PipeMemoryBlock delegate; + + private PipeBlock(final PipeMemoryBlock delegate) { + this.delegate = delegate; + } + + @Override + public long getMemoryUsageInBytes() { + return delegate.getMemoryUsageInBytes(); + } + + @Override + public void forceResize(final long newSizeInBytes) { + PipeDataNodeResourceManager.memory().forceResize(delegate, newSizeInBytes); + } + + @Override + public void close() { + delegate.close(); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/query/TsFileInsertionEventQueryParser.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/query/TsFileInsertionEventQueryParser.java index ba840b11c32dc..c26910b426f18 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/query/TsFileInsertionEventQueryParser.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/query/TsFileInsertionEventQueryParser.java @@ -35,9 +35,10 @@ import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; import org.apache.iotdb.db.pipe.event.common.tablet.PipeTabletUtils.TabletStringInternPool; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParser; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryBlock; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryManager; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.util.ModsOperationUtil; 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.PipeMemoryWeightUtil; import org.apache.iotdb.db.pipe.resource.tsfile.PipeTsFileResourceManager; import org.apache.iotdb.db.utils.datastructure.PatternTreeMapFactory; @@ -74,7 +75,7 @@ public class TsFileInsertionEventQueryParser extends TsFileInsertionEventParser private static final Logger LOGGER = LoggerFactory.getLogger(TsFileInsertionEventQueryParser.class); - private final PipeMemoryBlock allocatedMemoryBlock; + private final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlock; private final TsFileReader tsFileReader; private final Iterator>> deviceMeasurementsMapIterator; @@ -166,6 +167,39 @@ public TsFileInsertionEventQueryParser( final Map> deviceMeasurementsMapOverride, final boolean isWithMod) throws IOException, IllegalPathException { + this( + pipeName, + creationTime, + tsFile, + pattern, + startTime, + endTime, + pipeTaskMeta, + sourceEvent, + entity, + skipIfNoPrivileges, + deviceIsAlignedMap, + deviceMeasurementsMapOverride, + isWithMod, + TsFileInsertionEventParserMemoryManager.pipe()); + } + + public TsFileInsertionEventQueryParser( + final String pipeName, + final long creationTime, + final File tsFile, + final TreePattern pattern, + final long startTime, + final long endTime, + final PipeTaskMeta pipeTaskMeta, + final PipeInsertionEvent sourceEvent, + final IAuditEntity entity, + final boolean skipIfNoPrivileges, + final Map deviceIsAlignedMap, + final Map> deviceMeasurementsMapOverride, + final boolean isWithMod, + final TsFileInsertionEventParserMemoryManager memoryManager) + throws IOException, IllegalPathException { super( tsFile, pipeName, @@ -178,7 +212,8 @@ public TsFileInsertionEventQueryParser( entity, skipIfNoPrivileges, sourceEvent, - isWithMod); + isWithMod, + memoryManager); try { currentModifications = @@ -186,8 +221,7 @@ public TsFileInsertionEventQueryParser( ? ModsOperationUtil.loadModificationsFromTsFile(tsFile) : PatternTreeMapFactory.getModsPatternTreeMap(); allocatedMemoryBlockForModifications = - PipeDataNodeResourceManager.memory() - .forceAllocateForTabletWithRetry(currentModifications.ramBytesUsed()); + memoryManager.forceAllocateForTabletWithRetry(currentModifications.ramBytesUsed()); final PipeTsFileResourceManager tsFileResourceManager = PipeDataNodeResourceManager.tsfile(); final Map> deviceMeasurementsMap; @@ -248,8 +282,7 @@ public TsFileInsertionEventQueryParser( memoryRequiredInBytes += PipeMemoryWeightUtil.memoryOfIDeviceID2StrList(deviceMeasurementsMap); } - allocatedMemoryBlock = - PipeDataNodeResourceManager.memory().forceAllocate(memoryRequiredInBytes); + allocatedMemoryBlock = memoryManager.forceAllocate(memoryRequiredInBytes); final Iterator>> iterator = deviceMeasurementsMap.entrySet().iterator(); @@ -325,6 +358,25 @@ public TsFileInsertionEventQueryParser( final Map> deviceMeasurementsMapOverride, final boolean isWithMod) throws IOException, IllegalPathException { + this( + tsFile, + pattern, + startTime, + endTime, + deviceMeasurementsMapOverride, + isWithMod, + TsFileInsertionEventParserMemoryManager.pipe()); + } + + public TsFileInsertionEventQueryParser( + final File tsFile, + final TreePattern pattern, + final long startTime, + final long endTime, + final Map> deviceMeasurementsMapOverride, + final boolean isWithMod, + final TsFileInsertionEventParserMemoryManager memoryManager) + throws IOException, IllegalPathException { this( null, 0, @@ -338,7 +390,8 @@ public TsFileInsertionEventQueryParser( false, null, deviceMeasurementsMapOverride, - isWithMod); + isWithMod, + memoryManager); } private Map> filterDeviceMeasurementsMapByPattern( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/query/TsFileInsertionEventQueryParserTabletIterator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/query/TsFileInsertionEventQueryParserTabletIterator.java index 5c7c06089fe9b..281247e1529de 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/query/TsFileInsertionEventQueryParserTabletIterator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/query/TsFileInsertionEventQueryParserTabletIterator.java @@ -23,9 +23,8 @@ import org.apache.iotdb.db.i18n.DataNodePipeMessages; import org.apache.iotdb.db.pipe.event.common.tablet.PipeTabletUtils; import org.apache.iotdb.db.pipe.event.common.tablet.PipeTabletUtils.TabletStringInternPool; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryBlock; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.util.ModsOperationUtil; -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.PipeMemoryWeightUtil; import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry; import org.apache.iotdb.db.utils.datastructure.PatternTreeMapFactory; @@ -70,7 +69,7 @@ public class TsFileInsertionEventQueryParserTabletIterator implements Iterator measurementModsList; @@ -83,7 +82,7 @@ public class TsFileInsertionEventQueryParserTabletIterator implements Iterator measurements, final IExpression timeFilterExpression, - final PipeMemoryBlock allocatedBlockForTablet, + final TsFileInsertionEventParserMemoryBlock allocatedBlockForTablet, final PatternTreeMap currentModifications, final TabletStringInternPool tabletStringInternPool) throws IOException { @@ -172,8 +171,7 @@ private Tablet buildNextTablet() throws IOException { // Used for tree model deviceIdString, schemas, rowCountAndMemorySize.getLeft()); if (allocatedBlockForTablet.getMemoryUsageInBytes() < rowCountAndMemorySize.getRight()) { - PipeDataNodeResourceManager.memory() - .forceResize(allocatedBlockForTablet, rowCountAndMemorySize.getRight()); + allocatedBlockForTablet.forceResize(rowCountAndMemorySize.getRight()); } this.rowRecord = null; // Clear the saved first row isFirstRow = false; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParser.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParser.java index 79abbd65fae6c..3d963c27bb356 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParser.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParser.java @@ -36,9 +36,9 @@ import org.apache.iotdb.db.pipe.event.common.tablet.PipeTabletUtils; import org.apache.iotdb.db.pipe.event.common.tablet.PipeTabletUtils.TabletStringInternPool; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParser; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryBlock; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryManager; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.util.ModsOperationUtil; -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.PipeMemoryWeightUtil; import org.apache.iotdb.db.utils.datastructure.PatternTreeMapFactory; import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent; @@ -93,9 +93,9 @@ public class TsFileInsertionEventScanParser extends TsFileInsertionEventParser { private IChunkReader chunkReader; private BatchData data; - private final PipeMemoryBlock allocatedMemoryBlockForBatchData; - private final PipeMemoryBlock allocatedMemoryBlockForChunk; - private PipeMemoryBlock allocatedMemoryBlockForTsFileInput; + private final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForBatchData; + private final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForChunk; + private TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForTsFileInput; private boolean currentIsMultiPage; private IDeviceID currentDevice; @@ -129,6 +129,35 @@ public TsFileInsertionEventScanParser( final PipeInsertionEvent sourceEvent, final boolean isWithMod) throws IOException, IllegalPathException { + this( + pipeName, + creationTime, + tsFile, + pattern, + startTime, + endTime, + pipeTaskMeta, + entity, + skipIfNoPrivileges, + sourceEvent, + isWithMod, + TsFileInsertionEventParserMemoryManager.pipe()); + } + + public TsFileInsertionEventScanParser( + final String pipeName, + final long creationTime, + final File tsFile, + final TreePattern pattern, + final long startTime, + final long endTime, + final PipeTaskMeta pipeTaskMeta, + final IAuditEntity entity, + final boolean skipIfNoPrivileges, + final PipeInsertionEvent sourceEvent, + final boolean isWithMod, + final TsFileInsertionEventParserMemoryManager memoryManager) + throws IOException, IllegalPathException { super( tsFile, pipeName, @@ -141,16 +170,15 @@ public TsFileInsertionEventScanParser( entity, skipIfNoPrivileges, sourceEvent, - isWithMod); + isWithMod, + memoryManager); this.startTime = startTime; this.endTime = endTime; filter = Objects.nonNull(timeFilterExpression) ? timeFilterExpression.getFilter() : null; - this.allocatedMemoryBlockForBatchData = - PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(0); - this.allocatedMemoryBlockForChunk = - PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(0); + this.allocatedMemoryBlockForBatchData = memoryManager.forceAllocateForTabletWithRetry(0); + this.allocatedMemoryBlockForChunk = memoryManager.forceAllocateForTabletWithRetry(0); try { currentModifications = @@ -158,8 +186,7 @@ public TsFileInsertionEventScanParser( ? ModsOperationUtil.loadModificationsFromTsFile(tsFile) : PatternTreeMapFactory.getModsPatternTreeMap(); allocatedMemoryBlockForModifications = - PipeDataNodeResourceManager.memory() - .forceAllocateForTabletWithRetry(currentModifications.ramBytesUsed()); + memoryManager.forceAllocateForTabletWithRetry(currentModifications.ramBytesUsed()); tsFileSequenceReader = createTsFileSequenceReader(tsFile, !currentModifications.isEmpty()); tsFileSequenceReader.position((long) TSFileConfig.MAGIC_STRING.getBytes().length + 1); @@ -183,6 +210,27 @@ public TsFileInsertionEventScanParser( final PipeInsertionEvent sourceEvent, final boolean isWithMod) throws IOException, IllegalPathException { + this( + tsFile, + pattern, + startTime, + endTime, + pipeTaskMeta, + sourceEvent, + isWithMod, + TsFileInsertionEventParserMemoryManager.pipe()); + } + + public TsFileInsertionEventScanParser( + final File tsFile, + final TreePattern pattern, + final long startTime, + final long endTime, + final PipeTaskMeta pipeTaskMeta, + final PipeInsertionEvent sourceEvent, + final boolean isWithMod, + final TsFileInsertionEventParserMemoryManager memoryManager) + throws IOException, IllegalPathException { this( null, 0, @@ -194,7 +242,8 @@ public TsFileInsertionEventScanParser( null, false, sourceEvent, - isWithMod); + isWithMod, + memoryManager); } private TsFileSequenceReader createTsFileSequenceReader( @@ -204,8 +253,7 @@ private TsFileSequenceReader createTsFileSequenceReader( } allocatedMemoryBlockForTsFileInput = - PipeDataNodeResourceManager.memory() - .forceAllocateForTabletWithRetry(TS_FILE_INPUT_BUFFER_SIZE_IN_BYTES); + memoryManager.forceAllocateForTabletWithRetry(TS_FILE_INPUT_BUFFER_SIZE_IN_BYTES); return new TsFileSequenceReader( new BufferedTsFileInput(tsFile.toPath(), TS_FILE_INPUT_BUFFER_SIZE_IN_BYTES), false, @@ -367,8 +415,7 @@ private Tablet getNextTablet() { currentDeviceString, currentMeasurements, rowCountAndMemorySize.getLeft()); if (allocatedMemoryBlockForTablet.getMemoryUsageInBytes() < rowCountAndMemorySize.getRight()) { - PipeDataNodeResourceManager.memory() - .forceResize(allocatedMemoryBlockForTablet, rowCountAndMemorySize.getRight()); + allocatedMemoryBlockForTablet.forceResize(rowCountAndMemorySize.getRight()); } isFirstRow = false; } @@ -472,8 +519,7 @@ private void resizePageDataMemoryForCurrentPageIfNeeded() { private void resizePageDataMemoryIfNeeded(final long estimatedMemoryUsageInBytes) { if (allocatedMemoryBlockForBatchData.getMemoryUsageInBytes() < estimatedMemoryUsageInBytes) { - PipeDataNodeResourceManager.memory() - .forceResize(allocatedMemoryBlockForBatchData, estimatedMemoryUsageInBytes); + allocatedMemoryBlockForBatchData.forceResize(estimatedMemoryUsageInBytes); } } @@ -654,8 +700,7 @@ private void moveToNextChunkReader() } if (chunkHeader.getDataSize() > allocatedMemoryBlockForChunk.getMemoryUsageInBytes()) { - PipeDataNodeResourceManager.memory() - .forceResize(allocatedMemoryBlockForChunk, chunkHeader.getDataSize()); + allocatedMemoryBlockForChunk.forceResize(chunkHeader.getDataSize()); } Chunk chunk = @@ -1030,7 +1075,7 @@ private void resizeChunkMemoryBlockIfFirstValueChunkExceedsLimit( final long chunkSize = pendingAlignedChunkGroup.chunkSize + valueChunk.valueChunkSize; if (chunkSize > allocatedMemoryBlockForChunk.getMemoryUsageInBytes()) { - PipeDataNodeResourceManager.memory().forceResize(allocatedMemoryBlockForChunk, chunkSize); + allocatedMemoryBlockForChunk.forceResize(chunkSize); } } @@ -1046,8 +1091,7 @@ private void resizePageDataMemoryBlockIfFirstValueChunkExceedsLimit( calculateMaxAlignedPageMemorySizeWithBatchData( pendingAlignedChunkGroup.timeChunkIndex, pendingAlignedChunkGroup, valueChunk); if (pageMemorySize > getPageDataMemoryLimitInBytes()) { - PipeDataNodeResourceManager.memory() - .forceResize(allocatedMemoryBlockForBatchData, pageMemorySize); + allocatedMemoryBlockForBatchData.forceResize(pageMemorySize); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/table/TsFileInsertionEventTableParser.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/table/TsFileInsertionEventTableParser.java index 187d55a22a245..67e0beac96722 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/table/TsFileInsertionEventTableParser.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/table/TsFileInsertionEventTableParser.java @@ -29,9 +29,9 @@ import org.apache.iotdb.db.pipe.event.common.PipeInsertionEvent; import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParser; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryBlock; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryManager; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.util.ModsOperationUtil; -import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; -import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlock; import org.apache.iotdb.db.utils.datastructure.PatternTreeMapFactory; import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent; import org.apache.iotdb.pipe.api.exception.PipeException; @@ -52,10 +52,10 @@ public class TsFileInsertionEventTableParser extends TsFileInsertionEventParser private final TablePattern tablePattern; private final boolean isWithMod; - private final PipeMemoryBlock allocatedMemoryBlockForBatchData; - private final PipeMemoryBlock allocatedMemoryBlockForChunk; - private final PipeMemoryBlock allocatedMemoryBlockForChunkMeta; - private final PipeMemoryBlock allocatedMemoryBlockForTableSchemas; + private final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForBatchData; + private final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForChunk; + private final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForChunkMeta; + private final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForTableSchemas; public TsFileInsertionEventTableParser( final String pipeName, @@ -69,6 +69,33 @@ public TsFileInsertionEventTableParser( final PipeInsertionEvent sourceEvent, final boolean isWithMod) throws IOException { + this( + pipeName, + creationTime, + tsFile, + pattern, + startTime, + endTime, + pipeTaskMeta, + entity, + sourceEvent, + isWithMod, + TsFileInsertionEventParserMemoryManager.pipe()); + } + + public TsFileInsertionEventTableParser( + final String pipeName, + final long creationTime, + final File tsFile, + final TablePattern pattern, + final long startTime, + final long endTime, + final PipeTaskMeta pipeTaskMeta, + final IAuditEntity entity, + final PipeInsertionEvent sourceEvent, + final boolean isWithMod, + final TsFileInsertionEventParserMemoryManager memoryManager) + throws IOException { super( tsFile, pipeName, @@ -81,7 +108,8 @@ public TsFileInsertionEventTableParser( entity, true, sourceEvent, - isWithMod); + isWithMod, + memoryManager); this.isWithMod = isWithMod; try { @@ -90,16 +118,11 @@ public TsFileInsertionEventTableParser( ? ModsOperationUtil.loadModificationsFromTsFile(tsFile) : PatternTreeMapFactory.getModsPatternTreeMap(); allocatedMemoryBlockForModifications = - PipeDataNodeResourceManager.memory() - .forceAllocateForTabletWithRetry(currentModifications.ramBytesUsed()); - this.allocatedMemoryBlockForChunk = - PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(0); - this.allocatedMemoryBlockForBatchData = - PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(0); - this.allocatedMemoryBlockForChunkMeta = - PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(0); - this.allocatedMemoryBlockForTableSchemas = - PipeDataNodeResourceManager.memory().forceAllocateForTabletWithRetry(0); + memoryManager.forceAllocateForTabletWithRetry(currentModifications.ramBytesUsed()); + this.allocatedMemoryBlockForChunk = memoryManager.forceAllocateForTabletWithRetry(0); + this.allocatedMemoryBlockForBatchData = memoryManager.forceAllocateForTabletWithRetry(0); + this.allocatedMemoryBlockForChunkMeta = memoryManager.forceAllocateForTabletWithRetry(0); + this.allocatedMemoryBlockForTableSchemas = memoryManager.forceAllocateForTabletWithRetry(0); this.startTime = startTime; this.endTime = endTime; @@ -124,7 +147,40 @@ public TsFileInsertionEventTableParser( final boolean isWithMod) throws IOException { this( - null, 0, tsFile, pattern, startTime, endTime, pipeTaskMeta, entity, sourceEvent, isWithMod); + tsFile, + pattern, + startTime, + endTime, + pipeTaskMeta, + entity, + sourceEvent, + isWithMod, + TsFileInsertionEventParserMemoryManager.pipe()); + } + + public TsFileInsertionEventTableParser( + final File tsFile, + final TablePattern pattern, + final long startTime, + final long endTime, + final PipeTaskMeta pipeTaskMeta, + final IAuditEntity entity, + final PipeInsertionEvent sourceEvent, + final boolean isWithMod, + final TsFileInsertionEventParserMemoryManager memoryManager) + throws IOException { + this( + null, + 0, + tsFile, + pattern, + startTime, + endTime, + pipeTaskMeta, + entity, + sourceEvent, + isWithMod, + memoryManager); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/table/TsFileInsertionEventTableParserTabletIterator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/table/TsFileInsertionEventTableParserTabletIterator.java index 384f475a3592b..ab591070d6b62 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/table/TsFileInsertionEventTableParserTabletIterator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/table/TsFileInsertionEventTableParserTabletIterator.java @@ -23,9 +23,8 @@ import org.apache.iotdb.db.i18n.DataNodePipeMessages; import org.apache.iotdb.db.pipe.event.common.tablet.PipeTabletUtils; import org.apache.iotdb.db.pipe.event.common.tablet.PipeTabletUtils.TabletStringInternPool; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryBlock; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.util.ModsOperationUtil; -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.PipeMemoryWeightUtil; import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry; import org.apache.iotdb.db.utils.datastructure.PatternTreeMapFactory; @@ -78,11 +77,11 @@ public class TsFileInsertionEventTableParserTabletIterator implements Iterator modifications; @@ -119,11 +118,11 @@ public class TsFileInsertionEventTableParserTabletIterator implements Iterator> predicate, - final PipeMemoryBlock allocatedMemoryBlockForTablet, - final PipeMemoryBlock allocatedMemoryBlockForBatchData, - final PipeMemoryBlock allocatedMemoryBlockForChunk, - final PipeMemoryBlock allocatedMemoryBlockForChunkMeta, - final PipeMemoryBlock allocatedMemoryBlockForTableSchema, + final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForTablet, + final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForBatchData, + final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForChunk, + final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForChunkMeta, + final TsFileInsertionEventParserMemoryBlock allocatedMemoryBlockForTableSchema, final PatternTreeMap modifications, final long startTime, final long endTime) @@ -156,8 +155,7 @@ public TsFileInsertionEventTableParserTabletIterator( tableSchemaEntry.getKey().length() + PipeMemoryWeightUtil.calculateTableSchemaBytesUsed(tableSchemaEntry.getValue()); if (tableSchemaSize > allocatedMemoryBlockForTableSchema.getMemoryUsageInBytes()) { - PipeDataNodeResourceManager.memory() - .forceResize(this.allocatedMemoryBlockForTableSchema, tableSchemaSize); + this.allocatedMemoryBlockForTableSchema.forceResize(tableSchemaSize); } } @@ -179,8 +177,7 @@ public boolean hasNext() { batchData = chunkReader.nextPageData(); final long size = PipeMemoryWeightUtil.calculateBatchDataRamBytesUsed(batchData); if (allocatedMemoryBlockForBatchData.getMemoryUsageInBytes() < size) { - PipeDataNodeResourceManager.memory() - .forceResize(allocatedMemoryBlockForBatchData, size); + allocatedMemoryBlockForBatchData.forceResize(size); } state = State.CHECK_DATA; break; @@ -230,8 +227,7 @@ public boolean hasNext() { size += PipeMemoryWeightUtil.calculateAlignedChunkMetaBytesUsed(alignedChunkMetadata); if (allocatedMemoryBlockForChunkMeta.getMemoryUsageInBytes() < size) { - PipeDataNodeResourceManager.memory() - .forceResize(allocatedMemoryBlockForChunkMeta, size); + allocatedMemoryBlockForChunkMeta.forceResize(size); } } @@ -317,8 +313,7 @@ private Tablet buildNextTablet() { PipeMemoryWeightUtil.calculateTabletRowCountAndMemory(batchData); if (allocatedMemoryBlockForTablet.getMemoryUsageInBytes() < rowCountAndMemorySize.getRight()) { - PipeDataNodeResourceManager.memory() - .forceResize(allocatedMemoryBlockForTablet, rowCountAndMemorySize.getRight()); + allocatedMemoryBlockForTablet.forceResize(rowCountAndMemorySize.getRight()); } tablet = @@ -360,8 +355,7 @@ private void initChunkReader(final AbstractAlignedChunkMetadata alignedChunkMeta timeChunk = reader.readMemChunk((ChunkMetadata) alignedChunkMetadata.getTimeChunkMetadata()); timeChunkSize = PipeMemoryWeightUtil.calculateChunkRamBytesUsed(timeChunk); if (allocatedMemoryBlockForChunk.getMemoryUsageInBytes() < timeChunkSize) { - PipeDataNodeResourceManager.memory() - .forceResize(allocatedMemoryBlockForChunk, timeChunkSize); + allocatedMemoryBlockForChunk.forceResize(timeChunkSize); } } timeChunk.getData().rewind(); @@ -418,7 +412,7 @@ private void initChunkReader(final AbstractAlignedChunkMetadata alignedChunkMeta if (!hasSelectedNonNullChunk) { // If the first chunk exceeds the memory limit, we need to allocate more memory size = newSize; - PipeDataNodeResourceManager.memory().forceResize(allocatedMemoryBlockForChunk, size); + allocatedMemoryBlockForChunk.forceResize(size); } else { break; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTableStatementDataTypeConvertExecutionVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTableStatementDataTypeConvertExecutionVisitor.java index ce3ddccde92c2..b4abe2d6ebd40 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTableStatementDataTypeConvertExecutionVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTableStatementDataTypeConvertExecutionVisitor.java @@ -28,6 +28,7 @@ import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.AstVisitor; import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LoadTsFile; import org.apache.iotdb.db.queryengine.plan.statement.Statement; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileParserMemoryManager; import org.apache.iotdb.db.storageengine.load.util.LoadUtil; import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent; import org.apache.iotdb.rpc.TSStatusCode; @@ -85,7 +86,8 @@ public Optional visitLoadTsFile( null, null, null, - true)) { + true, + LoadTsFileParserMemoryManager.getInstance())) { for (final TabletInsertionEvent tabletInsertionEvent : parser.toTabletInsertionEvents()) { if (!(tabletInsertionEvent instanceof PipeRawTabletInsertionEvent)) { continue; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTreeTsFileTabletIterator.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTreeTsFileTabletIterator.java index 31ecbc2402070..3257bdc0cdd3f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTreeTsFileTabletIterator.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/converter/LoadTreeTsFileTabletIterator.java @@ -22,10 +22,12 @@ import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; import org.apache.iotdb.commons.pipe.datastructure.pattern.IoTDBTreePattern; import org.apache.iotdb.commons.pipe.datastructure.pattern.TreePattern; +import org.apache.iotdb.db.exception.load.LoadRuntimeOutOfMemoryException; import org.apache.iotdb.db.i18n.StorageEngineMessages; import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.query.TsFileInsertionEventQueryParser; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.scan.TsFileInsertionEventScanParser; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileParserMemoryManager; import org.apache.iotdb.pipe.api.event.dml.insertion.TabletInsertionEvent; import org.apache.tsfile.file.metadata.IDeviceID; @@ -153,10 +155,20 @@ private void ensureActiveIterator() throws Exception { try { scanParser = new TsFileInsertionEventScanParser( - file, LOAD_TREE_PATTERN, Long.MIN_VALUE, Long.MAX_VALUE, null, null, isWithMod); + file, + LOAD_TREE_PATTERN, + Long.MIN_VALUE, + Long.MAX_VALUE, + null, + null, + isWithMod, + LoadTsFileParserMemoryManager.getInstance()); activeIterator = scanParser.toTabletWithIsAligneds().iterator(); return; } catch (final Exception e) { + if (shouldRethrow(e)) { + throw toRuntimeException(e); + } if (!switchFromScanToQuery(e)) { throw toRuntimeException(e); } @@ -324,7 +336,8 @@ private boolean activateNextQueryParser() { activeQueryTask.startTime, activeQueryTask.endTime, activeQueryTask.toDeviceMeasurementsMap(), - isWithMod); + isWithMod, + LoadTsFileParserMemoryManager.getInstance()); final Iterator tabletIterator = activeQueryParser.toTabletInsertionEvents().iterator(); activeIterator = @@ -352,6 +365,9 @@ public Pair next() { }; return true; } catch (final Exception e) { + if (shouldRethrow(e)) { + throw toRuntimeException(e); + } LOGGER.warn( StorageEngineMessages .MESSAGE_LOAD_FAILED_TO_INITIALIZE_QUERY_FALLBACK_FOR_DEVICE_ARG_MEASUREMENTS_ARG_IN_TSFILE_ARG_SPLIT_OR_SKIP_THIS_QUERY_TASK_AND_CONTINUE_C6F69685, @@ -389,7 +405,8 @@ private boolean shouldRethrow(final Exception e) { Throwable current = e; while (Objects.nonNull(current)) { if (current instanceof InterruptedException - || current instanceof PipeRuntimeOutOfMemoryCriticalException) { + || current instanceof PipeRuntimeOutOfMemoryCriticalException + || current instanceof LoadRuntimeOutOfMemoryException) { return true; } current = current.getCause(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileParserMemoryManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileParserMemoryManager.java new file mode 100644 index 0000000000000..a7430a3f819fc --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileParserMemoryManager.java @@ -0,0 +1,120 @@ +/* + * 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.storageengine.load.memory; + +import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryBlock; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryManager; + +/** + * Allocates the working memory of TsFile parsers reused by Load from the query engine memory pool. + */ +public class LoadTsFileParserMemoryManager implements TsFileInsertionEventParserMemoryManager { + + private static final LoadTsFileMemoryManager LOAD_MEMORY_MANAGER = + LoadTsFileMemoryManager.getInstance(); + + private LoadTsFileParserMemoryManager() {} + + public static LoadTsFileParserMemoryManager getInstance() { + return LoadTsFileParserMemoryManagerHolder.INSTANCE; + } + + @Override + public TsFileInsertionEventParserMemoryBlock forceAllocateForTabletWithRetry( + final long sizeInBytes) { + return new LoadParserMemoryBlock(sizeInBytes); + } + + @Override + public TsFileInsertionEventParserMemoryBlock forceAllocate(final long sizeInBytes) { + return new LoadParserMemoryBlock(sizeInBytes); + } + + private static class LoadParserMemoryBlock implements TsFileInsertionEventParserMemoryBlock { + + private LoadTsFileMemoryBlock delegate; + private long memoryUsageInBytes; + private boolean isClosed; + + private LoadParserMemoryBlock(final long sizeInBytes) { + checkNonNegative(sizeInBytes); + if (sizeInBytes > 0) { + delegate = LOAD_MEMORY_MANAGER.allocateMemoryBlock(sizeInBytes); + } + memoryUsageInBytes = sizeInBytes; + } + + @Override + public synchronized long getMemoryUsageInBytes() { + return memoryUsageInBytes; + } + + @Override + public synchronized void forceResize(final long newSizeInBytes) { + checkNonNegative(newSizeInBytes); + if (isClosed || memoryUsageInBytes == newSizeInBytes) { + return; + } + + resizeDelegate(newSizeInBytes); + memoryUsageInBytes = newSizeInBytes; + } + + private void resizeDelegate(final long newSizeInBytes) { + if (newSizeInBytes == 0) { + delegate.close(); + delegate = null; + } else if (delegate == null) { + delegate = LOAD_MEMORY_MANAGER.allocateMemoryBlock(newSizeInBytes); + } else { + delegate.forceResize(newSizeInBytes); + } + } + + @Override + public synchronized void close() { + if (isClosed) { + return; + } + isClosed = true; + memoryUsageInBytes = 0; + if (delegate != null) { + delegate.close(); + delegate = null; + } + } + + private static void checkNonNegative(final long sizeInBytes) { + if (sizeInBytes < 0) { + throw new IllegalArgumentException( + String.format( + StorageEngineMessages + .STORAGE_EXCEPTION_LOAD_INVALID_MEMORY_SIZE_D_BYTES_MUST_BE_NON_NEGATIVE_A0146353, + sizeInBytes)); + } + } + } + + private static class LoadTsFileParserMemoryManagerHolder { + private static final LoadTsFileParserMemoryManager INSTANCE = + new LoadTsFileParserMemoryManager(); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionEventParserTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionEventParserTest.java index bdc7c1e4e38bf..f12664ff8b599 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionEventParserTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/TsFileInsertionEventParserTest.java @@ -30,12 +30,12 @@ import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; import org.apache.iotdb.db.pipe.event.common.tsfile.PipeTsFileInsertionEvent; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParser; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryBlock; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.query.TsFileInsertionEventQueryParser; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.scan.AlignedSinglePageWholeChunkReader; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.scan.SinglePageWholeChunkReader; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.scan.TsFileInsertionEventScanParser; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.table.TsFileInsertionEventTableParser; -import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryBlock; import org.apache.iotdb.db.pipe.resource.memory.PipeMemoryWeightUtil; import org.apache.iotdb.db.storageengine.dataregion.compaction.io.CompactionTsFileWriter; import org.apache.iotdb.db.storageengine.dataregion.compaction.schedule.constant.CompactionType; @@ -174,11 +174,22 @@ public void testScanParserReleasesTabletMemoryAfterRawTabletGenerated() throws E false)) { replaceAllocatedTabletMemory( parser, - new PipeMemoryBlock(0) { + new TsFileInsertionEventParserMemoryBlock() { + private long memoryUsageInBytes; + + @Override + public long getMemoryUsageInBytes() { + return memoryUsageInBytes; + } + + @Override + public void forceResize(final long newSizeInBytes) { + memoryUsageInBytes = newSizeInBytes; + } + @Override public void close() { - Assert.assertEquals(0, getMemoryUsageInBytes()); - super.close(); + Assert.assertEquals(0, memoryUsageInBytes); } }); @@ -211,13 +222,25 @@ public void testScanParserKeepsIteratorOnOutOfMemory() throws Exception { final AtomicInteger memoryUsageReadCount = new AtomicInteger(0); replaceAllocatedTabletMemory( parser, - new PipeMemoryBlock(0) { + new TsFileInsertionEventParserMemoryBlock() { + private long memoryUsageInBytes; + @Override public long getMemoryUsageInBytes() { if (memoryUsageReadCount.incrementAndGet() == 2) { throw new PipeRuntimeOutOfMemoryCriticalException("expected oom"); } - return super.getMemoryUsageInBytes(); + return memoryUsageInBytes; + } + + @Override + public void forceResize(final long newSizeInBytes) { + memoryUsageInBytes = newSizeInBytes; + } + + @Override + public void close() { + memoryUsageInBytes = 0; } }); @@ -2182,37 +2205,40 @@ private int getNonNullSize(final Tablet tablet) { return count; } - private PipeMemoryBlock getAllocatedChunkMemory(final TsFileInsertionEventScanParser parser) + private TsFileInsertionEventParserMemoryBlock getAllocatedChunkMemory( + final TsFileInsertionEventScanParser parser) throws NoSuchFieldException, IllegalAccessException { final Field field = TsFileInsertionEventScanParser.class.getDeclaredField("allocatedMemoryBlockForChunk"); field.setAccessible(true); - return (PipeMemoryBlock) field.get(parser); + return (TsFileInsertionEventParserMemoryBlock) field.get(parser); } - private PipeMemoryBlock getAllocatedBatchDataMemory(final TsFileInsertionEventScanParser parser) + private TsFileInsertionEventParserMemoryBlock getAllocatedBatchDataMemory( + final TsFileInsertionEventScanParser parser) throws NoSuchFieldException, IllegalAccessException { final Field field = TsFileInsertionEventScanParser.class.getDeclaredField("allocatedMemoryBlockForBatchData"); field.setAccessible(true); - return (PipeMemoryBlock) field.get(parser); + return (TsFileInsertionEventParserMemoryBlock) field.get(parser); } - private PipeMemoryBlock getAllocatedTabletMemory(final TsFileInsertionEventParser parser) - throws NoSuchFieldException, IllegalAccessException { + private TsFileInsertionEventParserMemoryBlock getAllocatedTabletMemory( + final TsFileInsertionEventParser parser) throws NoSuchFieldException, IllegalAccessException { final Field field = TsFileInsertionEventParser.class.getDeclaredField("allocatedMemoryBlockForTablet"); field.setAccessible(true); - return (PipeMemoryBlock) field.get(parser); + return (TsFileInsertionEventParserMemoryBlock) field.get(parser); } private void replaceAllocatedTabletMemory( - final TsFileInsertionEventParser parser, final PipeMemoryBlock replacement) + final TsFileInsertionEventParser parser, + final TsFileInsertionEventParserMemoryBlock replacement) throws NoSuchFieldException, IllegalAccessException { final Field field = TsFileInsertionEventParser.class.getDeclaredField("allocatedMemoryBlockForTablet"); field.setAccessible(true); - ((PipeMemoryBlock) field.get(parser)).close(); + ((TsFileInsertionEventParserMemoryBlock) field.get(parser)).close(); field.set(parser, replacement); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/converter/LoadTreeStatementDataTypeConvertExecutionVisitorTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/converter/LoadTreeStatementDataTypeConvertExecutionVisitorTest.java index 29766f5e5ebd5..c82b9bf7a2fd7 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/converter/LoadTreeStatementDataTypeConvertExecutionVisitorTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/converter/LoadTreeStatementDataTypeConvertExecutionVisitorTest.java @@ -25,10 +25,12 @@ import org.apache.iotdb.commons.path.PartialPath; import org.apache.iotdb.commons.pipe.datastructure.pattern.IoTDBTreePattern; import org.apache.iotdb.db.pipe.event.common.tsfile.parser.scan.TsFileInsertionEventScanParser; +import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; import org.apache.iotdb.db.queryengine.plan.statement.Statement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertMultiTabletsStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.InsertTabletStatement; import org.apache.iotdb.db.queryengine.plan.statement.crud.LoadTsFileStatement; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileMemoryManager; import org.apache.iotdb.rpc.TSStatusCode; import org.apache.tsfile.enums.TSDataType; @@ -111,6 +113,29 @@ public void testFallbackToQueryForRemainingDevicesWhenScanParserHitsCorruption() Assert.assertEquals(loadedPointCountBeforeCorruption, loadedPointCountAfterFallback); } + @Test + public void testLoadScanParserUsesQueryMemoryPoolInsteadOfPipeMemory() throws Exception { + tsFile = new File("load-tree-parser-query-memory.tsfile"); + writeTsFile(tsFile); + + final long pipeMemoryBefore = PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes(); + final LoadTsFileMemoryManager loadMemoryManager = LoadTsFileMemoryManager.getInstance(); + final long loadMemoryBefore = loadMemoryManager.getUsedMemorySizeInBytes(); + + try (final LoadTreeTsFileTabletIterator tabletIterator = + new LoadTreeTsFileTabletIterator(tsFile, true)) { + Assert.assertTrue(tabletIterator.hasNext()); + Assert.assertNotNull(tabletIterator.next()); + Assert.assertTrue(loadMemoryManager.getUsedMemorySizeInBytes() > loadMemoryBefore); + Assert.assertEquals( + pipeMemoryBefore, PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes()); + } + + Assert.assertEquals(loadMemoryBefore, loadMemoryManager.getUsedMemorySizeInBytes()); + Assert.assertEquals( + pipeMemoryBefore, PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes()); + } + @Test public void testFallbackToQueryWhenFirstNonAlignedDeviceIsCorrupted() throws Exception { tsFile = new File("load-tree-query-fallback-corrupted-first-non-aligned-device.tsfile"); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileParserPipeMemoryIsolationTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileParserPipeMemoryIsolationTest.java new file mode 100644 index 0000000000000..2b403b079c3f9 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/converter/LoadTsFileParserPipeMemoryIsolationTest.java @@ -0,0 +1,66 @@ +/* + * 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.storageengine.load.converter; + +import org.apache.iotdb.db.pipe.resource.PipeDataNodeResourceManager; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileMemoryManager; + +import org.apache.tsfile.utils.TsFileGeneratorUtils; +import org.junit.Assert; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.powermock.api.mockito.PowerMockito; +import org.powermock.core.classloader.annotations.PowerMockIgnore; +import org.powermock.core.classloader.annotations.PrepareForTest; +import org.powermock.modules.junit4.PowerMockRunner; + +import java.io.File; + +@PowerMockIgnore({"com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*", "javax.management.*"}) +@RunWith(PowerMockRunner.class) +@PrepareForTest(PipeDataNodeResourceManager.class) +public class LoadTsFileParserPipeMemoryIsolationTest { + + @Test + public void testLoadParserDoesNotAccessPipeMemoryPool() throws Exception { + final File tsFile = new File("load-parser-pipe-memory-isolation.tsfile"); + try { + TsFileGeneratorUtils.generateNonAlignedTsFile(tsFile.getPath(), 1, 1, 10, 0, 100, 10, 10); + + PowerMockito.mockStatic(PipeDataNodeResourceManager.class); + PowerMockito.when(PipeDataNodeResourceManager.memory()) + .thenThrow(new AssertionError("Load parser must not access Pipe memory")); + + final LoadTsFileMemoryManager loadMemoryManager = LoadTsFileMemoryManager.getInstance(); + final long loadMemoryBefore = loadMemoryManager.getUsedMemorySizeInBytes(); + try (final LoadTreeTsFileTabletIterator tabletIterator = + new LoadTreeTsFileTabletIterator(tsFile, true)) { + Assert.assertTrue(tabletIterator.hasNext()); + Assert.assertTrue(loadMemoryManager.getUsedMemorySizeInBytes() > loadMemoryBefore); + Assert.assertNotNull(tabletIterator.next()); + } + Assert.assertEquals(loadMemoryBefore, loadMemoryManager.getUsedMemorySizeInBytes()); + } finally { + if (tsFile.exists()) { + Assert.assertTrue(tsFile.delete()); + } + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileMemoryManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileMemoryManagerTest.java index f3a1bf7e4116f..ebbc666533018 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileMemoryManagerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileMemoryManagerTest.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.storageengine.load.memory; import org.apache.iotdb.db.exception.load.LoadRuntimeOutOfMemoryException; +import org.apache.iotdb.db.pipe.event.common.tsfile.parser.TsFileInsertionEventParserMemoryBlock; import org.junit.Assert; import org.junit.Test; @@ -83,6 +84,22 @@ public void testAllocateDataCacheMemoryBlockRollsBackPartialAllocationOnFailure( } } + @Test + public void testParserMemoryBlockGrowsAndReleasesFromQueryPool() throws Exception { + final LoadTsFileMemoryManager manager = LoadTsFileMemoryManager.getInstance(); + final long usedMemoryBefore = manager.getUsedMemorySizeInBytes(); + final TsFileInsertionEventParserMemoryBlock block = + LoadTsFileParserMemoryManager.getInstance().forceAllocate(0); + + Assert.assertEquals(0L, block.getMemoryUsageInBytes()); + block.forceResize(1024); + Assert.assertEquals(usedMemoryBefore + 1024, manager.getUsedMemorySizeInBytes()); + block.forceResize(0); + Assert.assertEquals(usedMemoryBefore, manager.getUsedMemorySizeInBytes()); + block.close(); + Assert.assertEquals(usedMemoryBefore, manager.getUsedMemorySizeInBytes()); + } + private static LoadTsFileMemoryManager newMemoryManager() throws Exception { final Constructor constructor = LoadTsFileMemoryManager.class.getDeclaredConstructor(); From 277e61e1ac30074277c1d808f4522ec86580794d Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:10:33 +0800 Subject: [PATCH 2/2] test(load): cover repeated zero-size parser resize --- .../storageengine/load/memory/LoadTsFileMemoryManagerTest.java | 3 +++ 1 file changed, 3 insertions(+) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileMemoryManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileMemoryManagerTest.java index ebbc666533018..f0ffe1d5f9985 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileMemoryManagerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/memory/LoadTsFileMemoryManagerTest.java @@ -96,6 +96,9 @@ public void testParserMemoryBlockGrowsAndReleasesFromQueryPool() throws Exceptio Assert.assertEquals(usedMemoryBefore + 1024, manager.getUsedMemorySizeInBytes()); block.forceResize(0); Assert.assertEquals(usedMemoryBefore, manager.getUsedMemorySizeInBytes()); + block.forceResize(0); + Assert.assertEquals(0L, block.getMemoryUsageInBytes()); + Assert.assertEquals(usedMemoryBefore, manager.getUsedMemorySizeInBytes()); block.close(); Assert.assertEquals(usedMemoryBefore, manager.getUsedMemorySizeInBytes()); }