From be2cff36f567f9ebe7ce84e96b5168977e418cb0 Mon Sep 17 00:00:00 2001 From: HTHou Date: Fri, 23 Oct 2026 16:00:19 +0800 Subject: [PATCH 1/6] Adapt file-scoped TsFile contexts and resource lifecycles --- .../iotdb/db/i18n/DataNodeQueryMessages.java | 4 + .../iotdb/db/i18n/StorageEngineMessages.java | 6 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 4 + .../iotdb/db/i18n/StorageEngineMessages.java | 6 + .../AlignedSinglePageWholeChunkReader.java | 64 ++- .../scan/SinglePageWholeChunkReader.java | 39 +- .../scan/TsFileInsertionEventScanParser.java | 61 ++- .../TsFileInsertionEventTableParser.java | 41 +- ...sertionEventTableParserTabletIterator.java | 63 ++- .../PipeTableModelTsFileBuilderV2.java | 8 +- .../builder/PipeTreeModelTsFileBuilderV2.java | 7 +- .../fragment/FragmentInstanceContext.java | 6 + .../execution/fragment/QueryChunkReaders.java | 94 ++++ .../execution/fragment/QueryContext.java | 5 + .../operator/source/FileLoaderUtils.java | 11 +- .../operator/source/SeriesScanUtil.java | 9 + .../scheduler/load/LoadFallbackHandler.java | 45 +- .../scheduler/load/LoadTsFileScheduler.java | 4 + .../load/TabletRewriteLoadStrategy.java | 66 +++ .../db/storageengine/buffer/ChunkCache.java | 24 +- .../impl/FastCompactionPerformer.java | 8 + .../impl/ReadChunkCompactionPerformer.java | 22 +- .../impl/ReadPointCompactionPerformer.java | 16 +- .../subtask/ReadPointPerformerSubTask.java | 9 +- ...dChunkAlignedSeriesCompactionExecutor.java | 6 + ...dChunkAlignedSeriesCompactionExecutor.java | 12 +- .../SingleSeriesCompactionExecutor.java | 7 + .../writer/AbstractCompactionWriter.java | 33 ++ .../writer/AbstractCrossCompactionWriter.java | 12 +- .../writer/AbstractInnerCompactionWriter.java | 5 +- .../writer/FastInnerCompactionWriter.java | 23 + .../compaction/io/CompactionTsFileReader.java | 4 +- .../compaction/io/CompactionTsFileWriter.java | 14 +- .../repair/RepairDataFileScanUtil.java | 163 ++++--- .../dataregion/flush/MemTableFlushTask.java | 4 + .../dataregion/memtable/AbstractMemTable.java | 28 +- .../AlignedWritableMemChunkGroup.java | 6 + .../dataregion/memtable/IMemTable.java | 5 + .../memtable/PrimitiveMemTable.java | 6 +- .../dataregion/memtable/TsFileProcessor.java | 9 +- .../impl/DiskAlignedChunkHandleImpl.java | 12 +- .../filescan/impl/DiskChunkHandleImpl.java | 21 +- .../reader/chunk/DiskAlignedChunkLoader.java | 85 ++-- .../read/reader/chunk/DiskChunkLoader.java | 4 + .../utils/SharedTimeDataBuffer.java | 21 +- .../dataregion/utils/TsFileResourceUtils.java | 6 +- .../load/TsFileWriterManager.java | 6 + .../load/splitter/AlignedChunkData.java | 4 +- .../load/splitter/ChunkData.java | 12 + .../load/splitter/NonAlignedChunkData.java | 2 +- .../load/splitter/TsFileSplitter.java | 28 +- .../db/tools/TsFileSplitByPartitionTool.java | 13 +- .../iotdb/db/tools/TsFileSplitTool.java | 8 +- .../db/tools/utils/TsFileSequenceScan.java | 11 +- .../apache/iotdb/db/utils/EncryptDBUtils.java | 84 ++++ .../FileEncryptionParameterProvider.java | 39 ++ .../fragment/QueryChunkReadersTest.java | 88 ++++ .../load/TabletRewriteLoadStrategyTest.java | 177 ++++++++ .../writer/AbstractCompactionWriterTest.java | 35 +- .../db/utils/FileEncryptionContextTest.java | 408 ++++++++++++++++++ .../TestFileEncryptionParameterProvider.java | 181 ++++++++ ...b.db.utils.FileEncryptionParameterProvider | 20 + pom.xml | 2 +- 63 files changed, 1981 insertions(+), 245 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryChunkReaders.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TabletRewriteLoadStrategy.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/FileEncryptionParameterProvider.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryChunkReadersTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TabletRewriteLoadStrategyTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/FileEncryptionContextTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/TestFileEncryptionParameterProvider.java create mode 100644 iotdb-core/datanode/src/test/resources/META-INF/services/org.apache.iotdb.db.utils.FileEncryptionParameterProvider diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index dabf70dd732ce..694cc55a7601d 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -20,6 +20,10 @@ package org.apache.iotdb.db.i18n; public final class DataNodeQueryMessages { + public static final String EXCEPTION_QUERY_CHUNK_RESOURCES_ARE_ALREADY_CLOSED_CA6A0559 = + "Query chunk resources are already closed"; + public static final String LOG_FAILED_TO_RELEASE_QUERY_CHUNK_RESOURCES_91573FE5 = + "Failed to release query chunk resources"; public static final String EXCEPTION_FAILED_TO_CONVERT_TSFILE_ARG_TO_TABLETS_STATUS_ARG_E7C9F477 = "Failed to convert TsFile %s to tablets. Status: %s"; public static final String EXCEPTION_FAILED_TO_CONVERT_TSFILE_ARG_TO_TABLETS_ARG_96358FDD = diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java index 25b73c3015274..9563ec23f1920 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -20,6 +20,12 @@ package org.apache.iotdb.db.i18n; public final class StorageEngineMessages { + public static final String EXCEPTION_FILE_ENCRYPTION_PARAMETER_PROVIDER_RETURNED_NULL_15C41723 = + "File encryption parameter provider returned null"; + public static final String EXCEPTION_PRECALCULATED_CHUNK_STAGING_IS_NOT_SUPPORTED_BY_THE_FILE_ENCRYPTION_POLICY_13DA8C32 = + "Precalculated chunk staging is not supported by the file encryption policy"; + public static final String EXCEPTION_MULTIPLE_FILE_ENCRYPTION_PARAMETER_PROVIDERS_ARE_CONFIGURED_5E7DDA5B = + "Multiple file encryption parameter providers are configured"; public static final String LOG_REBUILT_UNREADABLE_WAL_FOOTER_FILE_ARG_VERSION_ARG_ORIGINALFILEBYTES_ARG_WRITERPROGRESS_ARG_1DBF836E = "Rebuilt unreadable WAL footer: file={}, version={}, originalFileBytes={}, writerProgress={}."; diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java index 9addf5f954abb..61ba5e23056c7 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java @@ -20,6 +20,10 @@ package org.apache.iotdb.db.i18n; public final class DataNodeQueryMessages { + public static final String EXCEPTION_QUERY_CHUNK_RESOURCES_ARE_ALREADY_CLOSED_CA6A0559 = + "查询 Chunk 资源已关闭"; + public static final String LOG_FAILED_TO_RELEASE_QUERY_CHUNK_RESOURCES_91573FE5 = + "释放查询 Chunk 资源失败"; public static final String EXCEPTION_FAILED_TO_CONVERT_TSFILE_ARG_TO_TABLETS_STATUS_ARG_E7C9F477 = "将 TsFile %s 转换为 tablet 失败。状态：%s"; public static final String EXCEPTION_FAILED_TO_CONVERT_TSFILE_ARG_TO_TABLETS_ARG_96358FDD = diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java index 93f9b29b6f8d9..6a4b3aaf90de6 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -20,6 +20,12 @@ package org.apache.iotdb.db.i18n; public final class StorageEngineMessages { + public static final String EXCEPTION_FILE_ENCRYPTION_PARAMETER_PROVIDER_RETURNED_NULL_15C41723 = + "文件加密参数 Provider 返回 null"; + public static final String EXCEPTION_PRECALCULATED_CHUNK_STAGING_IS_NOT_SUPPORTED_BY_THE_FILE_ENCRYPTION_POLICY_13DA8C32 = + "文件加密策略不支持预计算 Chunk 暂存"; + public static final String EXCEPTION_MULTIPLE_FILE_ENCRYPTION_PARAMETER_PROVIDERS_ARE_CONFIGURED_5E7DDA5B = + "配置了多个文件加密参数 Provider"; public static final String LOG_REBUILT_UNREADABLE_WAL_FOOTER_FILE_ARG_VERSION_ARG_ORIGINALFILEBYTES_ARG_WRITERPROGRESS_ARG_1DBF836E = "已重建不可读的 WAL footer：file={}，version={}，originalFileBytes={}，writerProgress={}。"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedSinglePageWholeChunkReader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedSinglePageWholeChunkReader.java index c1ea72da8cf8e..376e275888036 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedSinglePageWholeChunkReader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/AlignedSinglePageWholeChunkReader.java @@ -63,37 +63,61 @@ public class AlignedSinglePageWholeChunkReader extends AbstractChunkReader // deleted intervals of all the sub sensors private final List> valueDeleteIntervalsList = new ArrayList<>(); private final long pageEstimatedMemoryUsageInBytes; + private final Chunk ownedTimeChunk; + private final List ownedValueChunks; public AlignedSinglePageWholeChunkReader( Chunk timeChunk, List valueChunkList, LongConsumer filteredRowsRecord) throws IOException { super(Long.MIN_VALUE, null, filteredRowsRecord); + ownedTimeChunk = timeChunk; + ownedValueChunks = new ArrayList<>(valueChunkList); this.timeChunkHeader = timeChunk.getHeader(); this.timeChunkDataBuffer = timeChunk.getData(); this.encryptParam = timeChunk.getEncryptParam(); - this.pageEstimatedMemoryUsageInBytes = - calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(timeChunk, valueChunkList); - - valueChunkList.forEach( - chunk -> { - this.valueChunkHeaderList.add(chunk == null ? null : chunk.getHeader()); - this.valueChunkDataBufferList.add(chunk == null ? null : chunk.getData()); - this.valueDeleteIntervalsList.add(chunk == null ? null : chunk.getDeleteIntervalList()); - }); + try { + this.pageEstimatedMemoryUsageInBytes = + calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(timeChunk, valueChunkList); + + valueChunkList.forEach( + chunk -> { + this.valueChunkHeaderList.add(chunk == null ? null : chunk.getHeader()); + this.valueChunkDataBufferList.add(chunk == null ? null : chunk.getData()); + this.valueDeleteIntervalsList.add(chunk == null ? null : chunk.getDeleteIntervalList()); + }); + + initAllPageReaders(); + } catch (IOException | RuntimeException | Error e) { + try { + close(); + } catch (RuntimeException failure) { + e.addSuppressed(failure); + } + throw e; + } + } - initAllPageReaders(); + @Override + public void close() { + ownedTimeChunk.close(); + for (Chunk chunk : ownedValueChunks) { + if (chunk != null) { + chunk.close(); + } + } } private void initAllPageReaders() throws IOException { + int pageIndex = 0; while (timeChunkDataBuffer.remaining() > 0) { - AlignedPageReader alignedPageReader = deserializeFromSinglePageChunk(); + AlignedPageReader alignedPageReader = deserializeFromSinglePageChunk(pageIndex++); if (alignedPageReader != null) { pageReaderList.add(alignedPageReader); } } } - private AlignedPageReader deserializeFromSinglePageChunk() throws IOException { + private AlignedPageReader deserializeFromSinglePageChunk(int pageIndex) throws IOException { PageHeader timePageHeader = PageHeader.deserializeFrom(timeChunkDataBuffer, (Statistics) null); List valuePageHeaderList = new ArrayList<>(); @@ -115,7 +139,7 @@ private AlignedPageReader deserializeFromSinglePageChunk() throws IOException { skipCurrentPage(timePageHeader, valuePageHeaderList); return null; } - return constructAlignedPageReader(timePageHeader, valuePageHeaderList); + return constructAlignedPageReader(timePageHeader, valuePageHeaderList, pageIndex); } private void skipCurrentPage(PageHeader timePageHeader, List valuePageHeader) { @@ -133,11 +157,17 @@ private void skipCurrentPage(PageHeader timePageHeader, List valuePa } private AlignedPageReader constructAlignedPageReader( - PageHeader timePageHeader, List rawValuePageHeaderList) throws IOException { + PageHeader timePageHeader, List rawValuePageHeaderList, int pageIndex) + throws IOException { IDecryptor decryptor = IDecryptor.getDecryptor(encryptParam); ByteBuffer timePageData = ChunkReader.deserializePageData( - timePageHeader, timeChunkDataBuffer, timeChunkHeader, decryptor); + timePageHeader, + timeChunkDataBuffer, + timeChunkHeader, + decryptor, + encryptParam, + pageIndex); List valuePageHeaderList = new ArrayList<>(); LazyLoadPageData[] valuePageDataArray = new LazyLoadPageData[rawValuePageHeaderList.size()]; @@ -166,7 +196,9 @@ private AlignedPageReader constructAlignedPageReader( valueChunkDataBufferList.get(i).array(), currentPagePosition, IUnCompressor.getUnCompressor(valueChunkHeader.getCompressionType()), - encryptParam); + encryptParam, + pageIndex, + valueChunkHeader.getChunkOrdinal()); valueDataTypeList.add(valueChunkHeader.getDataType()); valueDecoderList.add( Decoder.getDecoderByType( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/SinglePageWholeChunkReader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/SinglePageWholeChunkReader.java index 7bf4889b20e5c..1c1a30f4d0e0e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/SinglePageWholeChunkReader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/SinglePageWholeChunkReader.java @@ -35,6 +35,7 @@ import org.apache.tsfile.read.common.Chunk; import org.apache.tsfile.read.common.type.Type; import org.apache.tsfile.read.reader.chunk.AbstractChunkReader; +import org.apache.tsfile.read.reader.chunk.ChunkReader; import org.apache.tsfile.read.reader.page.LazyLoadPageData; import org.apache.tsfile.read.reader.page.PageReader; import org.apache.tsfile.utils.RamUsageEstimator; @@ -55,29 +56,47 @@ public class SinglePageWholeChunkReader extends AbstractChunkReader private final ByteBuffer chunkDataBuffer; private final EncryptParameter encryptParam; private final long pageEstimatedMemoryUsageInBytes; + private final Chunk ownedChunk; public SinglePageWholeChunkReader(Chunk chunk) throws IOException { super(Long.MIN_VALUE, null, null); + ownedChunk = chunk; this.chunkHeader = chunk.getHeader(); this.chunkDataBuffer = chunk.getData(); this.encryptParam = chunk.getEncryptParam(); - this.pageEstimatedMemoryUsageInBytes = - calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(chunk); - initAllPageReaders(); + try { + this.pageEstimatedMemoryUsageInBytes = + calculateMaxPageEstimatedMemoryUsageInBytesWithBatchData(chunk); + initAllPageReaders(); + } catch (IOException | RuntimeException | Error e) { + try { + chunk.close(); + } catch (RuntimeException failure) { + e.addSuppressed(failure); + } + throw e; + } + } + + @Override + public void close() { + ownedChunk.close(); } private void initAllPageReaders() throws IOException { + int pageIndex = 0; // construct next satisfied page header while (chunkDataBuffer.remaining() > 0) { pageReaderList.add( constructPageReader( PageHeader.deserializeFrom( - chunkDataBuffer, (Statistics) null))); + chunkDataBuffer, (Statistics) null), + pageIndex++)); } } - private PageReader constructPageReader(PageHeader pageHeader) throws IOException { + private PageReader constructPageReader(PageHeader pageHeader, int pageIndex) throws IOException { final int currentPagePosition = chunkDataBuffer.position(); chunkDataBuffer.position(currentPagePosition + pageHeader.getCompressedSize()); return new PageReader( @@ -86,7 +105,9 @@ private PageReader constructPageReader(PageHeader pageHeader) throws IOException chunkDataBuffer.array(), currentPagePosition, IUnCompressor.getUnCompressor(chunkHeader.getCompressionType()), - encryptParam), + encryptParam, + pageIndex, + chunkHeader.getChunkOrdinal()), chunkHeader.getDataType(), Decoder.getDecoderByType(chunkHeader.getEncodingType(), chunkHeader.getDataType()), getConfiguredTimeDecoder(), @@ -194,11 +215,13 @@ private static int countSinglePageTimeValues(final Chunk chunk) throws IOExcepti chunkDataBuffer.rewind(); final PageHeader pageHeader = deserializePageHeader(chunkDataBuffer, chunk.getHeader()); final ByteBuffer pageData = - deserializePageData( + ChunkReader.deserializePageData( pageHeader, chunkDataBuffer, chunk.getHeader(), - IDecryptor.getDecryptor(chunk.getEncryptParam())); + IDecryptor.getDecryptor(chunk.getEncryptParam()), + chunk.getEncryptParam(), + 0); final Decoder decoder = Decoder.getDecoderByType(chunk.getHeader().getEncodingType(), TSDataType.INT64); 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 c368b0087a6f0..f3748abc9ad37 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 @@ -32,6 +32,7 @@ import org.apache.iotdb.db.auth.AuthorityChecker; import org.apache.iotdb.db.exception.load.LoadRuntimeOutOfMemoryException; import org.apache.iotdb.db.i18n.DataNodePipeMessages; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; 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.tablet.PipeTabletUtils; @@ -47,7 +48,6 @@ import org.apache.iotdb.pipe.api.exception.PipeException; import org.apache.iotdb.rpc.TSStatusCode; -import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.common.constant.TsFileConstant; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.MetaMarker; @@ -71,6 +71,8 @@ import org.apache.tsfile.write.record.Tablet; import org.apache.tsfile.write.schema.IMeasurementSchema; import org.apache.tsfile.write.schema.MeasurementSchema; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import java.io.File; import java.io.IOException; @@ -84,6 +86,8 @@ import java.util.Objects; public class TsFileInsertionEventScanParser extends TsFileInsertionEventParser { + private static final Logger LOGGER = + LoggerFactory.getLogger(TsFileInsertionEventScanParser.class); private static final int TS_FILE_INPUT_BUFFER_SIZE_IN_BYTES = 8 * 1024; @@ -197,7 +201,7 @@ public TsFileInsertionEventScanParser( currentModifications.ramBytesUsed()); tsFileSequenceReader = createTsFileSequenceReader(tsFile, !currentModifications.isEmpty()); - tsFileSequenceReader.position((long) TSFileConfig.MAGIC_STRING.getBytes().length + 1); + tsFileSequenceReader.position(tsFileSequenceReader.getDataStartOffset()); prepareData(); if (Objects.isNull(chunkReader)) { @@ -621,6 +625,7 @@ private boolean putValueToColumns(final BatchData data, final Tablet tablet, fin private void moveToNextChunkReader() throws IOException, IllegalStateException, IllegalPathException { + releaseChunkReader(); ChunkHeader chunkHeader; currentMeasurements.clear(); modsInfos.clear(); @@ -676,7 +681,9 @@ private void moveToNextChunkReader() Chunk chunk = new Chunk( - chunkHeader, tsFileSequenceReader.readChunk(-1, chunkHeader.getDataSize())); + chunkHeader, + tsFileSequenceReader.readChunk(-1, chunkHeader.getDataSize()), + tsFileSequenceReader.getEncryptParam()); final List pageEstimatedMemoryUsageInBytesList = SinglePageWholeChunkReader .calculatePageEstimatedMemoryUsageInBytesWithBatchDataList(chunk); @@ -729,7 +736,9 @@ private void moveToNextChunkReader() } final Chunk chunk = new Chunk( - chunkHeader, tsFileSequenceReader.readChunk(-1, chunkHeader.getDataSize())); + chunkHeader, + tsFileSequenceReader.readChunk(-1, chunkHeader.getDataSize()), + tsFileSequenceReader.getEncryptParam()); valueChunk = new CachedAlignedValueChunk(valueIndex, chunk, chunkHeader.getDataSize()); } else { @@ -808,7 +817,10 @@ private boolean filterChunk( if ((chunkHeader.getChunkType() & TsFileConstant.TIME_COLUMN_MASK) == TsFileConstant.TIME_COLUMN_MASK) { final Chunk timeChunk = - new Chunk(chunkHeader, tsFileSequenceReader.readChunk(-1, chunkHeader.getDataSize())); + new Chunk( + chunkHeader, + tsFileSequenceReader.readChunk(-1, chunkHeader.getDataSize()), + tsFileSequenceReader.getEncryptParam()); timeChunkList.add(timeChunk); final boolean isMultiPage = marker == MetaMarker.TIME_CHUNK_HEADER; isMultiPageList.add(isMultiPage); @@ -904,10 +916,10 @@ private boolean useNextPendingAlignedChunk(final byte marker) throws IOException nextIsMultiPage ? new MemoryControlledChunkReader( new AlignedChunkReader( - timeChunk, pendingAlignedChunkGroup.valueChunkList, filter), + retainChunk(timeChunk), pendingAlignedChunkGroup.valueChunkList, filter), pageEstimatedMemoryUsageInBytesList) : new AlignedSinglePageWholeChunkReader( - timeChunk, pendingAlignedChunkGroup.valueChunkList, null); + retainChunk(timeChunk), pendingAlignedChunkGroup.valueChunkList, null); // Publish the state transition only after all memory reservations and reader construction // succeed. Otherwise a retry would lose the already-read aligned chunk group. @@ -1058,6 +1070,18 @@ private void validateAlignedValueChunkTimeIndex(final int timeChunkIndex) throws } private void clearCachedAlignedChunkData() { + for (PendingAlignedChunkGroup group : pendingAlignedChunkGroups) { + group.valueChunkList.forEach( + chunk -> { + if (chunk != null) { + chunk.close(); + } + }); + } + if (cachedAlignedValueChunk != null) { + cachedAlignedValueChunk.chunk.close(); + } + timeChunkList.forEach(Chunk::close); pendingAlignedChunkGroups.clear(); pendingAlignedChunkSize = 0; cachedAlignedValueChunk = null; @@ -1120,6 +1144,12 @@ private byte toValueChunkMarker(final ChunkHeader chunkHeader) { @Override public void close() { + try { + releaseChunkReader(); + } catch (IOException e) { + LOGGER.warn(DataNodeQueryMessages.LOG_FAILED_TO_RELEASE_QUERY_CHUNK_RESOURCES_91573FE5, e); + } + clearCachedAlignedChunkData(); super.close(); pendingPageDataAfterMemoryPressure = null; pendingTabletAfterMemoryPressure = null; @@ -1140,6 +1170,23 @@ public void close() { } } + private void releaseChunkReader() throws IOException { + IChunkReader current = chunkReader; + chunkReader = null; + if (current != null) { + current.close(); + } + } + + private static Chunk retainChunk(Chunk chunk) { + return new Chunk( + chunk.getHeader(), + chunk.getData().duplicate(), + chunk.getDeleteIntervalList(), + chunk.getChunkStatistic(), + chunk.getEncryptParam()); + } + private Statistics findAlignedChunkStatistics( List metadataList, long currentChunkHeaderOffset) { for (IChunkMetadata metadata : metadataList) { 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 1f32af7dd78dd..7563cecbc6a15 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 @@ -47,6 +47,7 @@ import java.util.Objects; public class TsFileInsertionEventTableParser extends TsFileInsertionEventParser { + private TsFileInsertionEventTableParserTabletIterator activeTabletIterator; private final long startTime; private final long endTime; @@ -291,6 +292,7 @@ && hasTablePrivilege(entry.getKey()), currentModifications, startTime, endTime); + activeTabletIterator = tabletIterator; } while (tabletIterator.hasNext()) { @@ -374,22 +376,29 @@ private static boolean isMemoryPressureException(final Throwable throwable) { @Override public void close() { - super.close(); - - if (allocatedMemoryBlockForBatchData != null) { - allocatedMemoryBlockForBatchData.close(); - } - - if (allocatedMemoryBlockForChunk != null) { - allocatedMemoryBlockForChunk.close(); - } - - if (allocatedMemoryBlockForChunkMeta != null) { - allocatedMemoryBlockForChunkMeta.close(); - } - - if (allocatedMemoryBlockForTableSchemas != null) { - allocatedMemoryBlockForTableSchemas.close(); + try { + if (activeTabletIterator != null) { + activeTabletIterator.close(); + activeTabletIterator = null; + } + } finally { + super.close(); + + if (allocatedMemoryBlockForBatchData != null) { + allocatedMemoryBlockForBatchData.close(); + } + + if (allocatedMemoryBlockForChunk != null) { + allocatedMemoryBlockForChunk.close(); + } + + if (allocatedMemoryBlockForChunkMeta != null) { + allocatedMemoryBlockForChunkMeta.close(); + } + + if (allocatedMemoryBlockForTableSchemas != null) { + allocatedMemoryBlockForTableSchemas.close(); + } } } } 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 91b613d79875c..f8affa9da9f5d 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 @@ -58,6 +58,7 @@ import org.apache.tsfile.write.schema.MeasurementSchema; import java.io.IOException; +import java.io.UncheckedIOException; import java.util.ArrayList; import java.util.HashMap; import java.util.Iterator; @@ -66,7 +67,8 @@ import java.util.Objects; import java.util.function.Predicate; -public class TsFileInsertionEventTableParserTabletIterator implements Iterator { +public class TsFileInsertionEventTableParserTabletIterator + implements Iterator, AutoCloseable { private final long startTime; private final long endTime; @@ -188,11 +190,15 @@ public boolean hasNext() { break; } case INIT_CHUNK_READER: + closeChunkReader(); if (currentChunkMetadata != null || (chunkMetadataList != null && chunkMetadataList.hasNext())) { if (currentChunkMetadata == null) { currentChunkMetadata = chunkMetadataList.next(); - timeChunk = null; + if (timeChunk != null) { + timeChunk.close(); + timeChunk = null; + } offset = 0; } initChunkReader(currentChunkMetadata); @@ -285,6 +291,7 @@ public boolean hasNext() { state = State.INIT_CHUNK_METADATA; break; } + close(); return false; } } @@ -422,6 +429,7 @@ private void initChunkReader(final AbstractAlignedChunkMetadata alignedChunkMeta Chunk chunk = null; if (metadata != null) { chunk = reader.readMemChunk((ChunkMetadata) metadata); + valueChunkList.add(chunk); final long newSize = size + PipeMemoryWeightUtil.calculateChunkRamBytesUsed(chunk); if (newSize > allocatedMemoryBlockForChunk.getMemoryUsageInBytes()) { if (!hasSelectedNonNullChunk) { @@ -429,6 +437,8 @@ private void initChunkReader(final AbstractAlignedChunkMetadata alignedChunkMeta size = newSize; allocatedMemoryBlockForChunk.forceResize(size); } else { + valueChunkList.remove(valueChunkList.size() - 1); + chunk.close(); break; } } else { @@ -439,10 +449,18 @@ private void initChunkReader(final AbstractAlignedChunkMetadata alignedChunkMeta columnTypes.add(ColumnCategory.FIELD); measurementList.add(measurementName); dataTypeList.add(schema.getType()); - valueChunkList.add(chunk); + if (chunk == null) { + valueChunkList.add(null); + } hasSelectedField = true; } } catch (final PipeRuntimeOutOfMemoryCriticalException | LoadRuntimeOutOfMemoryException e) { + valueChunkList.forEach( + chunk -> { + if (chunk != null) { + chunk.close(); + } + }); // The current field subset has not been published to a chunk reader yet. Restart that subset // from the same offset after memory pressure is relieved. Remove fields appended before the // failed reservation as well; otherwise the retry would expose duplicate columns in the @@ -452,6 +470,14 @@ private void initChunkReader(final AbstractAlignedChunkMetadata alignedChunkMeta dataTypeList.subList(deviceIdSize, dataTypeList.size()).clear(); offset = initialOffset; throw e; + } catch (IOException | RuntimeException | Error failure) { + valueChunkList.forEach( + chunk -> { + if (chunk != null) { + chunk.close(); + } + }); + throw failure; } if (offset >= fieldSchemaList.size()) { @@ -464,11 +490,40 @@ private void initChunkReader(final AbstractAlignedChunkMetadata alignedChunkMeta return; } - this.chunkReader = new TableChunkReader(timeChunk, valueChunkList, null); + Chunk retainedTimeChunk = + new Chunk( + timeChunk.getHeader(), + timeChunk.getData().duplicate(), + timeChunk.getDeleteIntervalList(), + timeChunk.getChunkStatistic(), + timeChunk.getEncryptParam()); + this.chunkReader = new TableChunkReader(retainedTimeChunk, valueChunkList, null); this.modsInfoList = ModsOperationUtil.initializeMeasurementMods(deviceID, measurementList, modifications, true); } + private void closeChunkReader() throws IOException { + IChunkReader previous = chunkReader; + chunkReader = null; + if (previous != null) { + previous.close(); + } + } + + @Override + public void close() { + try { + closeChunkReader(); + } catch (IOException e) { + throw new UncheckedIOException(e); + } finally { + if (timeChunk != null) { + timeChunk.close(); + timeChunk = null; + } + } + } + private boolean areAllFieldsDeletedByMods( final IDeviceID currentDeviceID, final AbstractAlignedChunkMetadata alignedChunkMetadata) { if (modifications.isEmpty() || fieldSchemaList.isEmpty()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/builder/PipeTableModelTsFileBuilderV2.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/builder/PipeTableModelTsFileBuilderV2.java index dc1a58ff0121e..ad11596a8d566 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/builder/PipeTableModelTsFileBuilderV2.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/builder/PipeTableModelTsFileBuilderV2.java @@ -28,7 +28,9 @@ import org.apache.iotdb.db.storageengine.dataregion.flush.MemTableFlushTask; import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable; import org.apache.iotdb.db.storageengine.dataregion.memtable.PrimitiveMemTable; +import org.apache.iotdb.db.utils.EncryptDBUtils; +import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.enums.ColumnCategory; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.write.WriteProcessException; @@ -182,7 +184,11 @@ private List> writeTabletsToTsFiles(final String dataBase) File file = null; try { file = createFile(); - try (final RestorableTsFileIOWriter writer = new RestorableTsFileIOWriter(file)) { + try (final RestorableTsFileIOWriter writer = + EncryptDBUtils.createRestorableWriterForNewFile( + file, dataBase, EncryptDBUtils.getFirstEncryptParamFromDatabase(null))) { + memTable.setEncryptParameter( + EncryptUtils.getEncryptParameter(writer.getEncryptParameter())); writeTabletsIntoOneFile(dataBase, memTable, writer); sealedFiles.add(new Pair<>(dataBase, writer.getFile())); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/builder/PipeTreeModelTsFileBuilderV2.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/builder/PipeTreeModelTsFileBuilderV2.java index 21dd3bb9cb052..9cb8849764bc6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/builder/PipeTreeModelTsFileBuilderV2.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/util/builder/PipeTreeModelTsFileBuilderV2.java @@ -27,7 +27,9 @@ import org.apache.iotdb.db.storageengine.dataregion.flush.MemTableFlushTask; import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable; import org.apache.iotdb.db.storageengine.dataregion.memtable.PrimitiveMemTable; +import org.apache.iotdb.db.utils.EncryptDBUtils; +import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.write.WriteProcessException; import org.apache.tsfile.utils.BitMap; @@ -161,7 +163,10 @@ private List> writeTabletsToTsFiles() throws WriteProcessExce File file = null; try { file = createFile(); - try (final RestorableTsFileIOWriter writer = new RestorableTsFileIOWriter(file)) { + try (final RestorableTsFileIOWriter writer = + EncryptDBUtils.createRestorableWriterForNewFile(file, null)) { + memTable.setEncryptParameter( + EncryptUtils.getEncryptParameter(writer.getEncryptParameter())); writeTabletsIntoOneFile(memTable, writer); sealedFiles.add(new Pair<>(null, writer.getFile())); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java index 16fecf6bc4983..a060bd5b91b51 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/FragmentInstanceContext.java @@ -71,6 +71,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.io.IOException; import java.time.ZoneId; import java.time.format.DateTimeParseException; import java.util.ArrayList; @@ -1347,6 +1348,11 @@ private void releaseTVListOwnedByQuery() { * be decreased. */ public synchronized void releaseResource() { + try { + getQueryChunkReaders().close(); + } catch (IOException e) { + LOGGER.warn(DataNodeQueryMessages.LOG_FAILED_TO_RELEASE_QUERY_CHUNK_RESOURCES_91573FE5, e); + } batchQueryDataSourceLeases.forEach(QueryDataSourceLease::close); // For schema related query FI, closedFilePaths and unClosedFilePaths will be null if (closedFilePaths != null) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryChunkReaders.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryChunkReaders.java new file mode 100644 index 0000000000000..0fac0fb717c39 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryChunkReaders.java @@ -0,0 +1,94 @@ +/* + * 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.queryengine.execution.fragment; + +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; + +import org.apache.tsfile.read.reader.IChunkReader; +import org.apache.tsfile.read.reader.IPageReader; + +import java.io.IOException; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; + +/** Keeps lazy pages' file contexts alive, releasing each owner once its last page is consumed. */ +public final class QueryChunkReaders implements AutoCloseable { + private final Map pages = new IdentityHashMap<>(); + private boolean closed; + + public synchronized void retain(IChunkReader reader, List pageReaders) + throws IOException { + if (closed) { + reader.close(); + throw new IOException( + DataNodeQueryMessages.EXCEPTION_QUERY_CHUNK_RESOURCES_ARE_ALREADY_CLOSED_CA6A0559); + } + if (pageReaders.isEmpty()) { + reader.close(); + return; + } + Lease lease = new Lease(reader, pageReaders.size()); + for (IPageReader page : pageReaders) { + pages.put(page, lease); + } + } + + public synchronized void release(IPageReader page) throws IOException { + Lease lease = pages.remove(page); + if (lease != null && --lease.remaining == 0) { + lease.reader.close(); + } + } + + @Override + public synchronized void close() throws IOException { + closed = true; + Map owners = new IdentityHashMap<>(); + for (Lease lease : pages.values()) { + owners.put(lease, Boolean.TRUE); + } + pages.clear(); + IOException failure = null; + for (Lease lease : owners.keySet()) { + try { + lease.reader.close(); + } catch (IOException | RuntimeException e) { + if (failure == null) { + failure = e instanceof IOException ? (IOException) e : new IOException(e); + } else { + failure.addSuppressed(e); + } + } + } + if (failure != null) { + throw failure; + } + } + + private static final class Lease { + private final IChunkReader reader; + private int remaining; + + private Lease(IChunkReader reader, int remaining) { + this.reader = reader; + this.remaining = remaining; + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryContext.java index b25dc7a083ee2..978b66a35a302 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryContext.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryContext.java @@ -48,6 +48,11 @@ /** QueryContext contains the shared information with in a query. */ public class QueryContext { + private final QueryChunkReaders queryChunkReaders = new QueryChunkReaders(); + + public QueryChunkReaders getQueryChunkReaders() { + return queryChunkReaders; + } private static final Logger LOGGER = LoggerFactory.getLogger(QueryContext.class); private QueryStatistics queryStatistics = new QueryStatistics(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/FileLoaderUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/FileLoaderUtils.java index ab034a46806f2..4374e974db5bc 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/FileLoaderUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/FileLoaderUtils.java @@ -552,7 +552,16 @@ public static List loadPageReaderList( if (chunkMetaData.isDataTypeModifiedAndCannotUseStatistics()) { chunkReader.markDataTypeModifiedAndCannotUseStatistics(); } - return chunkReader.loadPageReaderList(); + try { + return chunkReader.loadPageReaderList(); + } catch (IOException | RuntimeException | Error failure) { + try { + chunkReader.close(); + } catch (IOException | RuntimeException | Error cleanup) { + failure.addSuppressed(cleanup); + } + throw failure; + } } /** diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtil.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtil.java index 1582cacabc6cd..e790ac611f416 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtil.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/operator/source/SeriesScanUtil.java @@ -81,6 +81,7 @@ import java.io.File; import java.io.IOException; import java.io.Serializable; +import java.io.UncheckedIOException; import java.util.Collections; import java.util.Comparator; import java.util.LinkedList; @@ -887,6 +888,13 @@ public Statistics currentPageStatistics(int index) { } public void skipCurrentPage() { + if (firstPageReader != null) { + try { + context.getQueryChunkReaders().release(firstPageReader.getPageReader()); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } firstPageReader = null; } @@ -1740,6 +1748,7 @@ public TsBlock getAllSatisfiedPageData(boolean ascending) throws IOException { } throw e; } finally { + context.getQueryChunkReaders().release(data); long time = System.nanoTime() - startTime; if (isAligned) { if (isMem) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadFallbackHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadFallbackHandler.java index 27839d168b4ed..4d0752d3f8639 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadFallbackHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadFallbackHandler.java @@ -132,9 +132,8 @@ private void convertAndRetry() { try { final TSStatus status = - failedNode.isTableModel() - ? executeTableModelConversion(converter, failedNode, filePath) - : executeTreeModelConversion(converter, failedNode, filePath); + convertTsFileToTablets( + converter, failedNode, isGeneratedByPipe, failedNode.isDeleteAfterLoad()); if (converter.isSuccessful(status)) { iterator.remove(); @@ -173,29 +172,48 @@ private void convertAndRetry() { resolveFinalState(failureDetails); } - private TSStatus executeTableModelConversion( + static TSStatus convertTsFileToTablets( final LoadTsFileDataTypeConverter converter, final LoadSingleTsFileNode node, - final String filePath) { + final boolean isGeneratedByPipe, + final boolean deleteAfterLoad) + throws FileNotFoundException { + final String filePath = node.getTsFileResource().getTsFilePath(); + return node.isTableModel() + ? executeTableModelConversion(converter, node, filePath, isGeneratedByPipe, deleteAfterLoad) + : executeTreeModelConversion(converter, node, filePath, isGeneratedByPipe, deleteAfterLoad); + } + + private static TSStatus executeTableModelConversion( + final LoadTsFileDataTypeConverter converter, + final LoadSingleTsFileNode node, + final String filePath, + final boolean isGeneratedByPipe, + final boolean deleteAfterLoad) { final LoadTsFile statement = (isGeneratedByPipe ? LoadTsFile.createForPipe(null, filePath, Collections.emptyMap()) : LoadTsFile.createUnchecked(null, filePath, Collections.emptyMap())) .setDatabase(node.getDatabase()) - .setDeleteAfterLoad(node.isDeleteAfterLoad()) + .setDeleteAfterLoad(deleteAfterLoad) .setConvertOnTypeMismatch(true); + if (isGeneratedByPipe) { + statement.markIsGeneratedByPipe(); + } return converter.convertForTableModel(statement).orElse(null); } - private TSStatus executeTreeModelConversion( + private static TSStatus executeTreeModelConversion( final LoadTsFileDataTypeConverter converter, final LoadSingleTsFileNode node, - final String filePath) + final String filePath, + final boolean isGeneratedByPipe, + final boolean deleteAfterLoad) throws FileNotFoundException { final String database = LoadTsFileScheduler.getPartitionQueryDatabase(node, isGeneratedByPipe); final LoadTsFileStatement statement = - buildRetryTreeLoadStatement(filePath, node.isDeleteAfterLoad(), database); + buildRetryTreeLoadStatement(filePath, deleteAfterLoad, database, isGeneratedByPipe); return converter.convertForTreeModel(statement).orElse(null); } @@ -230,6 +248,15 @@ private String getFailedFilePathsString() { private LoadTsFileStatement buildRetryTreeLoadStatement( final String filePath, final boolean deleteAfterLoad, final String database) throws FileNotFoundException { + return buildRetryTreeLoadStatement(filePath, deleteAfterLoad, database, isGeneratedByPipe); + } + + private static LoadTsFileStatement buildRetryTreeLoadStatement( + final String filePath, + final boolean deleteAfterLoad, + final String database, + final boolean isGeneratedByPipe) + throws FileNotFoundException { final LoadTsFileStatement statement = (isGeneratedByPipe ? LoadTsFileStatement.createForPipe(filePath) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java index ea4bf8679187f..18678a043eeea 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileScheduler.java @@ -36,6 +36,7 @@ import org.apache.iotdb.db.service.RegionMigrateService; import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileMemoryManager; +import org.apache.iotdb.db.utils.EncryptDBUtils; import io.airlift.units.Duration; import org.slf4j.Logger; @@ -458,6 +459,9 @@ private boolean processSingleNode( */ private TsFileLoadStrategy resolveStrategy(final LoadSingleTsFileNode node, final String userName) throws Exception { + if (EncryptDBUtils.requiresFileRewrite()) { + return new TabletRewriteLoadStrategy(queryContext, isGeneratedByPipe); + } // The judgement resolves the partitions of the devices anyway, so the routing reuses them. final DataPartitionResolver partitionResolver = new DataPartitionResolver(partitionFetcher, userName); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TabletRewriteLoadStrategy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TabletRewriteLoadStrategy.java new file mode 100644 index 0000000000000..07089dab53ae9 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TabletRewriteLoadStrategy.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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.storageengine.load.converter.LoadTsFileDataTypeConverter; +import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; + +/** Rewrites files through the regular tablet write pipeline without staging plain chunks. */ +public class TabletRewriteLoadStrategy implements TsFileLoadStrategy { + + private final LoadTsFileDataTypeConverter converter; + private final boolean isGeneratedByPipe; + private final LoadTsFileCostMetricsSet costMetrics; + + public TabletRewriteLoadStrategy( + final MPPQueryContext queryContext, final boolean isGeneratedByPipe) { + this( + new LoadTsFileDataTypeConverter(queryContext, isGeneratedByPipe), + isGeneratedByPipe, + LoadTsFileCostMetricsSet.getInstance()); + } + + TabletRewriteLoadStrategy( + final LoadTsFileDataTypeConverter converter, + final boolean isGeneratedByPipe, + final LoadTsFileCostMetricsSet costMetrics) { + this.converter = converter; + this.isGeneratedByPipe = isGeneratedByPipe; + this.costMetrics = costMetrics; + } + + @Override + public boolean execute(final LoadSingleTsFileNode node) throws LoadFileException { + final long startTime = System.nanoTime(); + try { + // Only the scheduler may delete the source, after the entire strategy has succeeded. + return converter.isSuccessful( + LoadFallbackHandler.convertTsFileToTablets(converter, node, isGeneratedByPipe, false)); + } catch (final Exception e) { + throw new LoadFileException(e); + } finally { + costMetrics.recordPhaseTimeCost( + LoadTsFileCostMetricsSet.SCHEDULER_CAST_TABLETS, System.nanoTime() - startTime); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/buffer/ChunkCache.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/buffer/ChunkCache.java index dc569acfb415a..9a8192bc45a89 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/buffer/ChunkCache.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/buffer/ChunkCache.java @@ -35,6 +35,7 @@ import com.github.benmanes.caffeine.cache.Cache; import com.github.benmanes.caffeine.cache.Caffeine; +import com.github.benmanes.caffeine.cache.RemovalCause; import com.github.benmanes.caffeine.cache.Weigher; import org.apache.tsfile.file.metadata.statistics.Statistics; import org.apache.tsfile.read.TsFileSequenceReader; @@ -95,6 +96,14 @@ private ChunkCache() { (key, chunk) -> (int) (key.getRetainedSizeInBytes() + chunk.getRetainedSizeInBytes())) .recordStats() + .removalListener( + (ChunkCacheKey key, Chunk chunk, RemovalCause cause) -> { + if (chunk != null) { + synchronized (chunk) { + chunk.close(); + } + } + }) .build(); // add metrics @@ -162,8 +171,9 @@ private Chunk get( ChunkLoader chunkLoader = new ChunkLoader(ioSizeRecorder, externalTsFile); try { if (!CACHE_ENABLE || externalTsFile) { - Chunk chunk = chunkLoader.apply(chunkCacheKey); - return constructChunk(chunk, timeRangeList, chunkStatistic); + try (Chunk chunk = chunkLoader.apply(chunkCacheKey)) { + return constructChunk(chunk, timeRangeList, chunkStatistic); + } } Chunk chunk = lruCache.get(chunkCacheKey, chunkLoader); @@ -172,7 +182,15 @@ private Chunk get( DEBUG_LOGGER.info(StorageEngineMessages.GET_CHUNK_FROM_CACHE, chunkCacheKey); } - return constructChunk(chunk, timeRangeList, chunkStatistic); + synchronized (chunk) { + if (chunk.getEncryptParam() == null || !chunk.getEncryptParam().isDestroyed()) { + return constructChunk(chunk, timeRangeList, chunkStatistic); + } + } + // Eviction can race with a cache hit. Reload instead of retaining a closed parameter. + try (Chunk reloaded = chunkLoader.apply(chunkCacheKey)) { + return constructChunk(reloaded, timeRangeList, chunkStatistic); + } } catch (IoTDBIORuntimeException e) { throw e.getCause(); } finally { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/FastCompactionPerformer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/FastCompactionPerformer.java index 4af319684885a..7a86bf60701f6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/FastCompactionPerformer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/FastCompactionPerformer.java @@ -49,6 +49,7 @@ import org.apache.iotdb.db.storageengine.dataregion.compaction.selector.estimator.FastCrossSpaceCompactionEstimator; import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.utils.EncryptDBUtils; import org.apache.iotdb.db.utils.datastructure.PatternTreeMapFactory; import org.apache.tsfile.common.conf.TSFileDescriptor; @@ -157,6 +158,13 @@ public FastCompactionPerformer(boolean isCrossCompaction, EncryptParameter encry @Override public void perform() throws Exception { + if (EncryptDBUtils.requiresFileRewrite()) { + ReadPointCompactionPerformer rewriting = + new ReadPointCompactionPerformer(seqFiles, unseqFiles, targetFiles, encryptParameter); + rewriting.setSummary(subTaskSummary); + rewriting.perform(); + return; + } this.subTaskSummary.setTemporalFileNum(targetFiles.size()); try (MultiTsFileDeviceIterator deviceIterator = new MultiTsFileDeviceIterator(seqFiles, unseqFiles, readerCacheMap); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadChunkCompactionPerformer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadChunkCompactionPerformer.java index 82bdbd3e6496d..63cb9916838e8 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadChunkCompactionPerformer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadChunkCompactionPerformer.java @@ -137,6 +137,26 @@ public void perform() InterruptedException, StorageEngineException, PageException { + if (EncryptDBUtils.requiresFileRewrite()) { + ReadPointCompactionPerformer rewriting = + new ReadPointCompactionPerformer( + seqFiles, Collections.emptyList(), targetResources, firstEncryptParameter); + rewriting.setSummary(summary); + try { + rewriting.perform(); + } catch (IOException + | MetadataException + | InterruptedException + | StorageEngineException + | PageException e) { + throw e; + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new IOException(e); + } + return; + } try (MultiTsFileDeviceIterator deviceIterator = new MultiTsFileDeviceIterator(seqFiles)) { schema = CompactionTableSchemaCollector.collectSchema( @@ -213,7 +233,7 @@ private void rollCompactionFileWriter() throws IOException { private void useNewWriter() throws IOException { currentWriter = - new CompactionTsFileWriter( + CompactionTsFileWriter.create( targetResources.get(currentTargetFileIndex).getTsFile(), memoryBudgetForFileWriter, CompactionType.INNER_SEQ_COMPACTION, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadPointCompactionPerformer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadPointCompactionPerformer.java index 517f88cda77fd..216ca3b2fb8d1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadPointCompactionPerformer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadPointCompactionPerformer.java @@ -183,7 +183,11 @@ public void perform() throws Exception { CompactionUtils.updatePlanIndexes(targetFiles, seqFiles, unseqFiles); } finally { - QueryResourceManager.getInstance().endQuery(queryId); + try { + fragmentInstanceContext.getQueryChunkReaders().close(); + } finally { + QueryResourceManager.getInstance().endQuery(queryId); + } } } @@ -237,13 +241,13 @@ private void compactAlignedSeries( if (dataBlockReader.hasNextBatch()) { // chunkgroup is serialized only when at least one timeseries under this device has data compactionWriter.startChunkGroup(device, true); - measurementSchemas.add(0, timeSchema); compactionWriter.startMeasurement( TsFileConstant.TIME_COLUMN_ID, - new AlignedChunkWriterImpl( - measurementSchemas.remove(0), - measurementSchemas, - EncryptUtils.getEncryptParameter(getEncryptParameter())), + encryptParameter -> + new AlignedChunkWriterImpl( + timeSchema, + measurementSchemas, + EncryptUtils.getEncryptParameter(encryptParameter)), 0); writeWithReader(compactionWriter, dataBlockReader, device, 0, true); compactionWriter.endMeasurement(0); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/task/subtask/ReadPointPerformerSubTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/task/subtask/ReadPointPerformerSubTask.java index 74f7259074508..6554e9dcdcfdb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/task/subtask/ReadPointPerformerSubTask.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/task/subtask/ReadPointPerformerSubTask.java @@ -93,10 +93,11 @@ public Void call() throws Exception { if (dataBlockReader.hasNextBatch()) { compactionWriter.startMeasurement( measurement, - new ChunkWriterImpl( - measurementSchemas.get(0), - true, - EncryptUtils.getEncryptParameter(compactionWriter.getEncryptParameter())), + encryptParameter -> + new ChunkWriterImpl( + measurementSchemas.get(0), + true, + EncryptUtils.getEncryptParameter(encryptParameter)), taskId); ReadPointCompactionPerformer.writeWithReader( compactionWriter, dataBlockReader, device, taskId, false); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/batch/BatchedReadChunkAlignedSeriesCompactionExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/batch/BatchedReadChunkAlignedSeriesCompactionExecutor.java index c2bfe18e12ca1..f50b5d168b513 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/batch/BatchedReadChunkAlignedSeriesCompactionExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/batch/BatchedReadChunkAlignedSeriesCompactionExecutor.java @@ -430,6 +430,9 @@ public FollowingBatchReadChunkAlignedSeriesCompactionFlushController() { @Override public boolean canCompactCurrentChunkByDirectlyFlush( ChunkLoader timeChunk, List valueChunks) throws IOException { + if (requiresPageRewrite()) { + return false; + } CompactChunkPlan compactChunkPlan = batchCompactionPlan.getCompactChunkPlan(currentCompactChunk); boolean isCurrentChunkCompactedByDirectlyFlush = @@ -452,6 +455,9 @@ protected boolean canFlushCurrentChunkWriter() { @Override protected boolean canCompactCurrentPageByDirectlyFlush( PageLoader timePage, List valuePages) { + if (requiresPageRewrite()) { + return false; + } int currentPage = ((FollowingBatchCompactionAlignedChunkWriter) chunkWriter).getCurrentPage(); for (int i = 0; i < valuePages.size(); i++) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/ReadChunkAlignedSeriesCompactionExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/ReadChunkAlignedSeriesCompactionExecutor.java index 90181e111a2ca..9b8fe593842dc 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/ReadChunkAlignedSeriesCompactionExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/ReadChunkAlignedSeriesCompactionExecutor.java @@ -38,6 +38,7 @@ import org.apache.iotdb.db.storageengine.dataregion.compaction.io.CompactionTsFileReader; import org.apache.iotdb.db.storageengine.dataregion.compaction.io.CompactionTsFileWriter; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.utils.EncryptDBUtils; import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.common.conf.TSFileDescriptor; @@ -571,7 +572,9 @@ public ReadChunkAlignedSeriesCompactionFlushController(int compactionFileLevel) public boolean canCompactCurrentChunkByDirectlyFlush( ChunkLoader timeChunk, List valueChunks) throws IOException { - return canFlushCurrentChunkWriter() && canFlushChunk(timeChunk, valueChunks); + return !requiresPageRewrite() + && canFlushCurrentChunkWriter() + && canFlushChunk(timeChunk, valueChunks); } protected boolean canFlushCurrentChunkWriter() { @@ -612,6 +615,9 @@ private boolean canFlushChunk(ChunkLoader timeChunk, List valueChun protected boolean canCompactCurrentPageByDirectlyFlush( PageLoader timePage, List valuePages) { + if (requiresPageRewrite()) { + return false; + } boolean isHighLevelCompaction = compactionTargetFileLevel > largeFileLevelSeparator; if (isHighLevelCompaction) { return canFlushPage(timePage, valuePages); @@ -624,6 +630,10 @@ private boolean canSealCurrentPageWriter() { return chunkWriter.checkIsUnsealedPageOverThreshold(targetPageSize, targetPagePointNum, true); } + protected boolean requiresPageRewrite() { + return !EncryptDBUtils.canCopyChunk(null, writer.getEncryptParameter()); + } + private boolean canFlushPage(PageLoader timePage, List valuePages) { if (timePage.getModifiedStatus() == ModifiedStatus.PARTIAL_DELETED) { return false; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/SingleSeriesCompactionExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/SingleSeriesCompactionExecutor.java index 44f8427a6f905..62fa4e00d31d9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/SingleSeriesCompactionExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/SingleSeriesCompactionExecutor.java @@ -25,6 +25,7 @@ import org.apache.iotdb.db.storageengine.dataregion.compaction.execute.task.CompactionTaskSummary; import org.apache.iotdb.db.storageengine.dataregion.compaction.io.CompactionTsFileWriter; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.utils.EncryptDBUtils; import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.file.MetaMarker; @@ -155,6 +156,12 @@ public void execute() throws IOException { constructChunkWriterFromReadChunk(currentChunk); } + if (!EncryptDBUtils.canCopyChunk( + currentChunk.getEncryptParam(), fileWriter.getEncryptParameter())) { + processModifiedChunk(currentChunk); + continue; + } + // if this chunk is modified, deserialize it into points if (chunkMetadata.getDeleteIntervalList() != null) { processModifiedChunk(currentChunk); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriter.java index 2b99ca67552d7..02ad8a6f0de9b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriter.java @@ -53,6 +53,7 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.util.List; +import java.util.function.Function; public abstract class AbstractCompactionWriter implements AutoCloseable { protected int subTaskNum = IoTDBDescriptor.getInstance().getConfig().getSubCompactionTaskNum(); @@ -65,6 +66,10 @@ public abstract class AbstractCompactionWriter implements AutoCloseable { // The index of the array corresponds to subTaskId. protected IChunkWriter[] chunkWriters = new IChunkWriter[subTaskNum]; + @SuppressWarnings("unchecked") + private final Function[] chunkWriterFactories = + new Function[subTaskNum]; + // Each sub task has point count in current measurment, which is used to check size. // The index of the array corresponds to subTaskId. protected int[] chunkPointNumArray = new int[subTaskNum]; @@ -137,6 +142,19 @@ public ModEntry getTTLLowerBoundForCurrentDevice() { } public void startMeasurement(String measurement, IChunkWriter chunkWriter, int subTaskId) { + chunkWriterFactories[subTaskId] = null; + initializeChunkWriter(measurement, chunkWriter, subTaskId); + } + + public void startMeasurement( + String measurement, + Function chunkWriterFactory, + int subTaskId) { + chunkWriterFactories[subTaskId] = chunkWriterFactory; + initializeChunkWriter(measurement, chunkWriterFactory.apply(getEncryptParameter()), subTaskId); + } + + private void initializeChunkWriter(String measurement, IChunkWriter chunkWriter, int subTaskId) { resetChunkWriterStatistics(subTaskId); lastTimeSet[subTaskId] = false; chunkWriters[subTaskId] = chunkWriter; @@ -144,6 +162,21 @@ public void startMeasurement(String measurement, IChunkWriter chunkWriter, int s hasVariableLengthTypeArray[subTaskId] = containsVariableLengthType(chunkWriter); } + protected void switchChunkWriterEncryptionParameter( + EncryptParameter encryptParameter, int subTaskId) { + Function chunkWriterFactory = chunkWriterFactories[subTaskId]; + if (chunkWriterFactory == null) { + return; + } + IChunkWriter chunkWriter = chunkWriterFactory.apply(encryptParameter); + chunkWriters[subTaskId] = chunkWriter; + hasVariableLengthTypeArray[subTaskId] = containsVariableLengthType(chunkWriter); + } + + protected void clearChunkWriterFactory(int subTaskId) { + chunkWriterFactories[subTaskId] = null; + } + private boolean containsVariableLengthType(IChunkWriter chunkWriter) { if (chunkWriter instanceof ChunkWriterImpl) { return ((ChunkWriterImpl) chunkWriter).getDataType().isBinary(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCrossCompactionWriter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCrossCompactionWriter.java index ac4120614de8b..8daa119a3c55c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCrossCompactionWriter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCrossCompactionWriter.java @@ -98,7 +98,7 @@ protected AbstractCrossCompactionWriter( / targetResources.size()); for (int i = 0; i < targetResources.size(); i++) { this.targetFileWriters.add( - new CompactionTsFileWriter( + CompactionTsFileWriter.create( targetResources.get(i).getTsFile(), memorySizeForEachWriter, CompactionType.CROSS_COMPACTION, @@ -141,6 +141,7 @@ public void endMeasurement(int subTaskId) throws IOException { sealChunk( targetFileWriters.get(seqFileIndexArray[subTaskId]), chunkWriters[subTaskId], subTaskId); seqFileIndexArray[subTaskId] = 0; + clearChunkWriterFactory(subTaskId); } @Override @@ -211,6 +212,7 @@ protected void checkTimeAndMayFlushChunkToCurrentFile(long timestamp, int subTas checkPreviousTimestamp(timestamp, subTaskId); int fileIndex = seqFileIndexArray[subTaskId]; + int previousFileIndex = fileIndex; boolean hasFlushedCurrentChunk = false; // If timestamp is later than the current source seq tsfile, then flush chunk writer and move to // next file @@ -224,6 +226,10 @@ protected void checkTimeAndMayFlushChunkToCurrentFile(long timestamp, int subTas } seqFileIndexArray[subTaskId] = ++fileIndex; } + if (fileIndex != previousFileIndex) { + switchChunkWriterEncryptionParameter( + targetFileWriters.get(fileIndex).getEncryptParameter(), subTaskId); + } } private void checkIsDeviceExistAndGetDeviceEndTime() throws IOException { @@ -278,6 +284,8 @@ public void setSchemaForAllTargetFile(List schemas) { @Override public EncryptParameter getEncryptParameter() { - return encryptParameter; + return targetFileWriters.isEmpty() + ? encryptParameter + : targetFileWriters.get(0).getEncryptParameter(); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractInnerCompactionWriter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractInnerCompactionWriter.java index c59c658708cb4..7e193485d95ec 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractInnerCompactionWriter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractInnerCompactionWriter.java @@ -114,7 +114,7 @@ private void rollCompactionFileWriter() throws IOException { private void useNewWriter() throws IOException { fileWriter = - new CompactionTsFileWriter( + CompactionTsFileWriter.create( targetResources.get(currentFileIndex).getTsFile(), memoryBudgetForFileWriter, targetResources.get(currentFileIndex).isSeq() @@ -137,6 +137,7 @@ public void endChunkGroup() throws IOException { @Override public void endMeasurement(int subTaskId) throws IOException { sealChunk(fileWriter, chunkWriters[subTaskId], subTaskId); + clearChunkWriterFactory(subTaskId); } @Override @@ -189,6 +190,6 @@ public long getWriterSize() throws IOException { @Override public EncryptParameter getEncryptParameter() { - return encryptParameter; + return fileWriter == null ? encryptParameter : fileWriter.getEncryptParameter(); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/FastInnerCompactionWriter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/FastInnerCompactionWriter.java index ccac7b4c63516..2615cfea74273 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/FastInnerCompactionWriter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/FastInnerCompactionWriter.java @@ -78,6 +78,9 @@ public void write(TsBlock tsBlock, int subTaskId) throws IOException { @Override public boolean flushNonAlignedChunk(Chunk chunk, ChunkMetadata chunkMetadata, int subTaskId) throws IOException { + if (requiresPageRewrite()) { + return false; + } checkPreviousTimestamp(chunkMetadata.getStartTime(), subTaskId); if (chunkPointNumArray[subTaskId] != 0 && chunkWriters[subTaskId].checkIsChunkSizeOverThreshold( @@ -107,6 +110,9 @@ public boolean flushNonAlignedChunk(Chunk chunk, ChunkMetadata chunkMetadata, in @Override public boolean flushAlignedChunk(ChunkMetadataElement chunkMetadataElement, int subTaskId) throws IOException { + if (requiresPageRewrite()) { + return false; + } AbstractAlignedChunkMetadata alignedChunkMetadata = (AbstractAlignedChunkMetadata) chunkMetadataElement.chunkMetadata; IChunkMetadata timeChunkMetadata = alignedChunkMetadata.getTimeChunkMetadata(); @@ -140,6 +146,9 @@ public boolean flushBatchedValueChunk( int subTaskId, AbstractCompactionFlushController flushController) throws IOException { + if (requiresPageRewrite()) { + return false; + } AbstractAlignedChunkMetadata alignedChunkMetadata = (AbstractAlignedChunkMetadata) chunkMetadataElement.chunkMetadata; IChunkMetadata timeChunkMetadata = alignedChunkMetadata.getTimeChunkMetadata(); @@ -176,6 +185,9 @@ public boolean flushBatchedValueChunk( */ public boolean flushAlignedPage(AlignedPageElement alignedPageElement, int subTaskId) throws IOException, PageException { + if (requiresPageRewrite()) { + return false; + } PageHeader timePageHeader = alignedPageElement.getTimePageHeader(); List valuePageHeaders = alignedPageElement.getValuePageHeaders(); ByteBuffer compressedTimePageData = alignedPageElement.getTimePageData(); @@ -212,6 +224,9 @@ public boolean flushBatchedValuePage( int subTaskId, AbstractCompactionFlushController flushController) throws PageException, IOException { + if (requiresPageRewrite()) { + return false; + } PageHeader timePageHeader = alignedPageElement.getTimePageHeader(); List valuePageHeaders = alignedPageElement.getValuePageHeaders(); ByteBuffer compressedTimePageData = alignedPageElement.getTimePageData(); @@ -255,6 +270,9 @@ public boolean flushBatchedValuePage( public boolean flushNonAlignedPage( ByteBuffer compressedPageData, PageHeader pageHeader, int subTaskId) throws IOException, PageException { + if (requiresPageRewrite()) { + return false; + } checkPreviousTimestamp(pageHeader.getStartTime(), subTaskId); boolean isUnsealedPageOverThreshold = chunkWriters[subTaskId].checkIsUnsealedPageOverThreshold( @@ -314,4 +332,9 @@ private boolean checkIsPageLargeEnough(PageHeader pageHeader) { return pageHeader.getStatistics().getCount() >= pagePointNumLowerBoundInCompaction || pageHeader.getSerializedPageSize() >= pageSizeLowerBoundInCompaction; } + + private boolean requiresPageRewrite() { + return fileWriter != null + && !EncryptDBUtils.canCopyChunk(null, fileWriter.getEncryptParameter()); + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/io/CompactionTsFileReader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/io/CompactionTsFileReader.java index 142d17341fa3e..63617d9dbf78a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/io/CompactionTsFileReader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/io/CompactionTsFileReader.java @@ -81,9 +81,9 @@ public void markEndOfAlignedSeries() { ((CompactionTsFileInput) tsFileInput).markEndOfAlignedSeries(); } - @SuppressWarnings("java:S2177") + @Override public ChunkHeader readChunkHeader(long position) throws IOException { - return ChunkHeader.deserializeFrom(tsFileInput, position); + return super.readChunkHeader(position); } public Map>> diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/io/CompactionTsFileWriter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/io/CompactionTsFileWriter.java index 1f822e3b9345d..bc22353c020b6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/io/CompactionTsFileWriter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/io/CompactionTsFileWriter.java @@ -67,7 +67,7 @@ public CompactionTsFileWriter( File file, long maxMetadataSize, CompactionType type, EncryptParameter encryptParameter) throws IOException { super(file, maxMetadataSize, encryptParameter); - this.firstEncryptParameter = encryptParameter; + this.firstEncryptParameter = super.getEncryptParameter(); this.type = type; super.out = new CompactionTsFileOutput( @@ -78,6 +78,18 @@ public EncryptParameter getEncryptParameter() { return firstEncryptParameter; } + public static CompactionTsFileWriter create( + File file, long maxMetadataSize, CompactionType type, EncryptParameter legacyParameter) + throws IOException { + EncryptParameter parameter = + EncryptDBUtils.createEncryptParameterForNewFile(file, null, legacyParameter); + try { + return new CompactionTsFileWriter(file, maxMetadataSize, type, parameter); + } finally { + EncryptDBUtils.closeFileEncryptionParameter(parameter); + } + } + public void markStartingWritingAligned() { isWritingAligned = true; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/repair/RepairDataFileScanUtil.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/repair/RepairDataFileScanUtil.java index 97381eb29d01c..0c45770242597 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/repair/RepairDataFileScanUtil.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/repair/RepairDataFileScanUtil.java @@ -60,6 +60,7 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -243,41 +244,42 @@ private void checkAlignedDeviceSeries( private void checkTimeChunkInAlignedSeries( TsFileSequenceReader reader, IDeviceID device, ChunkMetadata timeChunkMetadata) throws IOException { - Chunk timeChunk = reader.readMemChunk(timeChunkMetadata); - - CompactionChunkReader chunkReader = new CompactionChunkReader(timeChunk); - ByteBuffer chunkDataBuffer = timeChunk.getData(); - ChunkHeader chunkHeader = timeChunk.getHeader(); - long actualChunkStartTime = Long.MAX_VALUE; - long actualChunkEndTime = Long.MIN_VALUE; - while (chunkDataBuffer.hasRemaining()) { - // deserialize a PageHeader from chunkDataBuffer - PageHeader pageHeader = null; - if (((byte) (chunkHeader.getChunkType() & 0x3F)) == MetaMarker.ONLY_ONE_PAGE_CHUNK_HEADER) { - pageHeader = PageHeader.deserializeFrom(chunkDataBuffer, timeChunk.getChunkStatistic()); - } else { - pageHeader = PageHeader.deserializeFrom(chunkDataBuffer, chunkHeader.getDataType()); + try (Chunk timeChunk = reader.readMemChunk(timeChunkMetadata)) { + CompactionChunkReader chunkReader = new CompactionChunkReader(timeChunk); + ByteBuffer chunkDataBuffer = timeChunk.getData(); + ChunkHeader chunkHeader = timeChunk.getHeader(); + long actualChunkStartTime = Long.MAX_VALUE; + long actualChunkEndTime = Long.MIN_VALUE; + int pageIndex = 0; + while (chunkDataBuffer.hasRemaining()) { + // deserialize a PageHeader from chunkDataBuffer + PageHeader pageHeader = null; + if (((byte) (chunkHeader.getChunkType() & 0x3F)) == MetaMarker.ONLY_ONE_PAGE_CHUNK_HEADER) { + pageHeader = PageHeader.deserializeFrom(chunkDataBuffer, timeChunk.getChunkStatistic()); + } else { + pageHeader = PageHeader.deserializeFrom(chunkDataBuffer, chunkHeader.getDataType()); + } + actualChunkStartTime = Math.min(actualChunkStartTime, pageHeader.getStartTime()); + actualChunkEndTime = Math.max(actualChunkEndTime, pageHeader.getEndTime()); + ByteBuffer pageData = chunkReader.readPageDataWithoutUncompressing(pageHeader); + IDecryptor decryptor = IDecryptor.getDecryptor(timeChunk.getEncryptParam()); + ByteBuffer uncompressedPageData = + decodePage(timeChunk, pageHeader, pageData, decryptor, pageIndex++); + try { + validateTimeData(device, uncompressedPageData, pageHeader); + } finally { + wipePage(timeChunk, uncompressedPageData); + } + } + if (actualChunkStartTime > actualChunkEndTime) { + return; + } + TimeRange actualChunkTimeRange = new TimeRange(actualChunkStartTime, actualChunkEndTime); + if (!actualChunkTimeRange.equals( + new TimeRange(timeChunkMetadata.getStartTime(), timeChunkMetadata.getEndTime()))) { + throw new CompactionStatisticsCheckFailedException( + device, timeChunkMetadata, actualChunkTimeRange); } - actualChunkStartTime = Math.min(actualChunkStartTime, pageHeader.getStartTime()); - actualChunkEndTime = Math.max(actualChunkEndTime, pageHeader.getEndTime()); - ByteBuffer pageData = chunkReader.readPageDataWithoutUncompressing(pageHeader); - IDecryptor decryptor = IDecryptor.getDecryptor(timeChunk.getEncryptParam()); - ByteBuffer uncompressedPageData = - decryptAndUncompressPageData( - pageHeader, - IUnCompressor.getUnCompressor(chunkHeader.getCompressionType()), - pageData, - decryptor); - validateTimeData(device, uncompressedPageData, pageHeader); - } - if (actualChunkStartTime > actualChunkEndTime) { - return; - } - TimeRange actualChunkTimeRange = new TimeRange(actualChunkStartTime, actualChunkEndTime); - if (!actualChunkTimeRange.equals( - new TimeRange(timeChunkMetadata.getStartTime(), timeChunkMetadata.getEndTime()))) { - throw new CompactionStatisticsCheckFailedException( - device, timeChunkMetadata, actualChunkTimeRange); } } @@ -358,41 +360,68 @@ private void checkSingleNonAlignedSeries( private void checkChunkOfNonAlignedSeries( TsFileSequenceReader reader, IDeviceID deviceID, ChunkMetadata chunkMetadata) throws IOException { - Chunk chunk = reader.readMemChunk(chunkMetadata); - ChunkHeader chunkHeader = chunk.getHeader(); - CompactionChunkReader chunkReader = new CompactionChunkReader(chunk); - ByteBuffer chunkDataBuffer = chunk.getData(); - long actualChunkStartTime = Long.MAX_VALUE; - long actualChunkEndTime = Long.MIN_VALUE; - while (chunkDataBuffer.hasRemaining()) { - // deserialize a PageHeader from chunkDataBuffer - PageHeader pageHeader = null; - if (((byte) (chunkHeader.getChunkType() & 0x3F)) == MetaMarker.ONLY_ONE_PAGE_CHUNK_HEADER) { - pageHeader = PageHeader.deserializeFrom(chunkDataBuffer, chunk.getChunkStatistic()); - } else { - pageHeader = PageHeader.deserializeFrom(chunkDataBuffer, chunkHeader.getDataType()); + try (Chunk chunk = reader.readMemChunk(chunkMetadata)) { + ChunkHeader chunkHeader = chunk.getHeader(); + CompactionChunkReader chunkReader = new CompactionChunkReader(chunk); + ByteBuffer chunkDataBuffer = chunk.getData(); + long actualChunkStartTime = Long.MAX_VALUE; + long actualChunkEndTime = Long.MIN_VALUE; + int pageIndex = 0; + while (chunkDataBuffer.hasRemaining()) { + // deserialize a PageHeader from chunkDataBuffer + PageHeader pageHeader = null; + if (((byte) (chunkHeader.getChunkType() & 0x3F)) == MetaMarker.ONLY_ONE_PAGE_CHUNK_HEADER) { + pageHeader = PageHeader.deserializeFrom(chunkDataBuffer, chunk.getChunkStatistic()); + } else { + pageHeader = PageHeader.deserializeFrom(chunkDataBuffer, chunkHeader.getDataType()); + } + actualChunkStartTime = Math.min(actualChunkStartTime, pageHeader.getStartTime()); + actualChunkEndTime = Math.max(actualChunkEndTime, pageHeader.getEndTime()); + ByteBuffer pageData = chunkReader.readPageDataWithoutUncompressing(pageHeader); + IDecryptor decryptor = IDecryptor.getDecryptor(chunk.getEncryptParam()); + ByteBuffer uncompressedPageData = + decodePage(chunk, pageHeader, pageData, decryptor, pageIndex++); + try { + ByteBuffer timeBuffer = getTimeBufferFromNonAlignedPage(uncompressedPageData); + validateTimeData(deviceID, timeBuffer, pageHeader); + } finally { + wipePage(chunk, uncompressedPageData); + } + } + if (actualChunkStartTime > actualChunkEndTime) { + return; + } + TimeRange actualChunkTimeRange = new TimeRange(actualChunkStartTime, actualChunkEndTime); + if (!actualChunkTimeRange.equals( + new TimeRange(chunkMetadata.getStartTime(), chunkMetadata.getEndTime()))) { + throw new CompactionStatisticsCheckFailedException( + deviceID, chunkMetadata, actualChunkTimeRange); } - actualChunkStartTime = Math.min(actualChunkStartTime, pageHeader.getStartTime()); - actualChunkEndTime = Math.max(actualChunkEndTime, pageHeader.getEndTime()); - ByteBuffer pageData = chunkReader.readPageDataWithoutUncompressing(pageHeader); - IDecryptor decryptor = IDecryptor.getDecryptor(chunk.getEncryptParam()); - ByteBuffer uncompressedPageData = - decryptAndUncompressPageData( - pageHeader, - IUnCompressor.getUnCompressor(chunkHeader.getCompressionType()), - pageData, - decryptor); - ByteBuffer timeBuffer = getTimeBufferFromNonAlignedPage(uncompressedPageData); - validateTimeData(deviceID, timeBuffer, pageHeader); - } - if (actualChunkStartTime > actualChunkEndTime) { - return; } - TimeRange actualChunkTimeRange = new TimeRange(actualChunkStartTime, actualChunkEndTime); - if (!actualChunkTimeRange.equals( - new TimeRange(chunkMetadata.getStartTime(), chunkMetadata.getEndTime()))) { - throw new CompactionStatisticsCheckFailedException( - deviceID, chunkMetadata, actualChunkTimeRange); + } + + private static ByteBuffer decodePage( + Chunk chunk, PageHeader page, ByteBuffer body, IDecryptor decryptor, int pageIndex) + throws IOException { + IUnCompressor uncompressor = + IUnCompressor.getUnCompressor(chunk.getHeader().getCompressionType()); + return chunk.getEncryptParam() != null && chunk.getEncryptParam().isTdePageAead() + ? decryptAndUncompressPageData( + page, + uncompressor, + body, + decryptor, + chunk.getEncryptParam(), + pageIndex, + chunk.getHeader().getChunkOrdinal()) + : decryptAndUncompressPageData(page, uncompressor, body, decryptor); + } + + private static void wipePage(Chunk chunk, ByteBuffer page) { + if (chunk.getEncryptParam() != null + && chunk.getEncryptParam().isTdePageAead() + && page.hasArray()) { + Arrays.fill(page.array(), (byte) 0); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/flush/MemTableFlushTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/flush/MemTableFlushTask.java index 57b1f73de7c1c..7777e48be3a41 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/flush/MemTableFlushTask.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/flush/MemTableFlushTask.java @@ -245,6 +245,10 @@ private void doSyncFlushMemTable() throws ExecutionException, InterruptedExcepti try { long writePlanIndicesStartTime = System.currentTimeMillis(); writer.writePlanIndices(); + if (writer.getEncryptParameter() != null && writer.getEncryptParameter().isTdePageAead()) { + // The file header and data must be durable before the WAL checkpoint can be released. + writer.getTsFileOutput().force(); + } WRITING_METRICS.recordFlushCost( WritingMetrics.WRITE_PLAN_INDICES, System.currentTimeMillis() - writePlanIndicesStartTime); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractMemTable.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractMemTable.java index f51f2c6810564..9924a174253b2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractMemTable.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AbstractMemTable.java @@ -51,6 +51,7 @@ import org.apache.iotdb.db.utils.ModificationUtils; import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.ChunkMetadata; import org.apache.tsfile.file.metadata.IChunkMetadata; @@ -139,17 +140,20 @@ public abstract class AbstractMemTable implements IMemTable { private String database; private String dataRegionId; + private EncryptParameter encryptParameter; protected AbstractMemTable() { this.database = null; this.dataRegionId = null; this.memTableMap = new HashMap<>(); + this.encryptParameter = EncryptDBUtils.getSecondEncryptParamFromDatabase(null); } protected AbstractMemTable(String database, String dataRegionId) { this.database = database; this.dataRegionId = dataRegionId; this.memTableMap = new HashMap<>(); + this.encryptParameter = EncryptDBUtils.getSecondEncryptParamFromDatabase(database); } protected AbstractMemTable( @@ -157,6 +161,7 @@ protected AbstractMemTable( this.database = database; this.dataRegionId = dataRegionId; this.memTableMap = memTableMap; + this.encryptParameter = EncryptDBUtils.getSecondEncryptParamFromDatabase(database); } @Override @@ -174,11 +179,7 @@ public Map getMemTableMap() { private IWritableMemChunkGroup createMemChunkGroupIfNotExistAndGet( IDeviceID deviceId, List schemaList) { IWritableMemChunkGroup memChunkGroup = - memTableMap.computeIfAbsent( - deviceId, - k -> - new WritableMemChunkGroup( - EncryptDBUtils.getSecondEncryptParamFromDatabase(database))); + memTableMap.computeIfAbsent(deviceId, k -> new WritableMemChunkGroup(encryptParameter)); for (IMeasurementSchema schema : schemaList) { if (schema != null && !memChunkGroup.contains(schema.getMeasurementName())) { seriesNumber++; @@ -196,7 +197,7 @@ private IWritableMemChunkGroup createAlignedMemChunkGroupIfNotExistAndGet( deviceId, k -> { seriesNumber += filteredSchemaList.size(); - return new AlignedWritableMemChunkGroup(filteredSchemaList); + return new AlignedWritableMemChunkGroup(filteredSchemaList, encryptParameter); }); for (IMeasurementSchema schema : schemaList) { if (schema != null && !memChunkGroup.contains(schema.getMeasurementName())) { @@ -1267,8 +1268,21 @@ public String getDataRegionId() { public void setDatabaseAndDataRegionId(String database, String dataRegionId) { this.database = database; this.dataRegionId = dataRegionId; + if (encryptParameter == null || !encryptParameter.isTdePageAead()) { + setEncryptParameter(EncryptDBUtils.getSecondEncryptParamFromDatabase(database)); + } + } + + @Override + public EncryptParameter getEncryptParameter() { + return encryptParameter; + } + + @Override + public void setEncryptParameter(EncryptParameter encryptParameter) { + this.encryptParameter = encryptParameter; for (IWritableMemChunkGroup memChunkGroup : memTableMap.values()) { - memChunkGroup.setEncryptParameter(EncryptDBUtils.getSecondEncryptParamFromDatabase(database)); + memChunkGroup.setEncryptParameter(encryptParameter); } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedWritableMemChunkGroup.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedWritableMemChunkGroup.java index e1235ca935b0d..9071a96b277dc 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedWritableMemChunkGroup.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/AlignedWritableMemChunkGroup.java @@ -50,6 +50,12 @@ public AlignedWritableMemChunkGroup(List schemaList) { memChunk = new AlignedWritableMemChunk(schemaList); } + public AlignedWritableMemChunkGroup( + List schemaList, EncryptParameter encryptParameter) { + this.encryptParameter = encryptParameter; + memChunk = new AlignedWritableMemChunk(schemaList, encryptParameter); + } + @TestOnly public AlignedWritableMemChunkGroup( AlignedWritableMemChunk memChunk, List schemaList) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/IMemTable.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/IMemTable.java index 066b8931be47d..e31dfc27b7f00 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/IMemTable.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/IMemTable.java @@ -35,6 +35,7 @@ import org.apache.iotdb.db.storageengine.rescon.memory.WriteMemoryBlock; import org.apache.iotdb.db.storageengine.rescon.memory.WriteMemoryManager; +import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.file.metadata.IChunkMetadata; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.read.filter.basic.Filter; @@ -234,5 +235,9 @@ void queryForDeviceRegionScan( void setDatabaseAndDataRegionId(String database, String dataRegionId); + EncryptParameter getEncryptParameter(); + + void setEncryptParameter(EncryptParameter encryptParameter); + void checkDataType(InsertNode node) throws DataTypeInconsistentException; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/PrimitiveMemTable.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/PrimitiveMemTable.java index fd8ecb0671e44..7e6faf7d79f9a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/PrimitiveMemTable.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/PrimitiveMemTable.java @@ -42,8 +42,10 @@ public PrimitiveMemTable( @Override public IMemTable copy() { Map newMap = new HashMap<>(getMemTableMap()); - - return new PrimitiveMemTable(getDatabase(), getDataRegionId(), newMap); + PrimitiveMemTable copiedMemTable = + new PrimitiveMemTable(getDatabase(), getDataRegionId(), newMap); + copiedMemTable.setEncryptParameter(getEncryptParameter()); + return copiedMemTable; } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessor.java index f790116011736..6ca52288e8249 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/memtable/TsFileProcessor.java @@ -86,6 +86,7 @@ import org.apache.iotdb.rpc.RpcUtils; import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.file.metadata.AbstractAlignedChunkMetadata; import org.apache.tsfile.file.metadata.AlignedChunkMetadata; @@ -243,10 +244,8 @@ public TsFileProcessor( this.tsFileResource = new TsFileResource(tsfile, this); this.dataRegionInfo = dataRegionInfo; this.writer = - new RestorableTsFileIOWriter( - tsfile, - EncryptDBUtils.getFirstEncryptParamFromDatabase( - dataRegionInfo.getDataRegion().getDatabaseName())); + EncryptDBUtils.createRestorableWriterForNewFile( + tsfile, dataRegionInfo.getDataRegion().getDatabaseName()); this.updateLatestFlushTimeCallback = updateLatestFlushTimeCallback; this.walNode = WALManager.getInstance() @@ -544,6 +543,8 @@ private void createNewWorkingMemTable() { dataRegionInfo.getDataRegion().getDatabaseName(), dataRegionInfo.getDataRegion().getDataRegionIdString()); workMemTable.setWriteMemoryManager(dataRegionInfo.getWriteMemoryManager()); + workMemTable.setEncryptParameter( + EncryptUtils.getEncryptParameter(writer.getEncryptParameter())); walNode.onMemTableCreated(workMemTable, tsFileResource.getTsFilePath()); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/filescan/impl/DiskAlignedChunkHandleImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/filescan/impl/DiskAlignedChunkHandleImpl.java index be62a3b2e7f1e..f9673d16fc446 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/filescan/impl/DiskAlignedChunkHandleImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/filescan/impl/DiskAlignedChunkHandleImpl.java @@ -23,7 +23,6 @@ import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileID; import org.apache.iotdb.db.storageengine.dataregion.utils.SharedTimeDataBuffer; -import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.encrypt.IDecryptor; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.file.metadata.statistics.Statistics; @@ -39,9 +38,6 @@ public class DiskAlignedChunkHandleImpl extends DiskChunkHandleImpl { private static final int MASK = 0x80; private final SharedTimeDataBuffer sharedTimeDataBuffer; - private int pageIndex = 0; - - private EncryptParameter encryptParam; public DiskAlignedChunkHandleImpl( IDeviceID deviceID, @@ -60,7 +56,6 @@ public DiskAlignedChunkHandleImpl( protected void init(TsFileSequenceReader reader) throws IOException { sharedTimeDataBuffer.init(reader); super.init(reader); - this.encryptParam = reader.getEncryptParam(); } @Override @@ -71,12 +66,14 @@ public long[] getDataTime() throws IOException { this.currentPageHeader, this.currentChunkDataBuffer, this.currentChunkHeader, - decryptor); + decryptor, + encryptParam, + cryptoPageIndex); int size = ReadWriteIOUtils.readInt(currentPageDataBuffer); byte[] bitmap = new byte[(size + 7) / 8]; currentPageDataBuffer.get(bitmap); - long[] timeData = sharedTimeDataBuffer.getPageTime(pageIndex); + long[] timeData = sharedTimeDataBuffer.getPageTime(cryptoPageIndex); if (timeData.length != size) { throw new UnsupportedOperationException(StorageEngineMessages.TIME_DATA_SIZE_NOT_MATCH); } @@ -90,7 +87,6 @@ public long[] getDataTime() throws IOException { validTimeList[index++] = timeData[i]; } - pageIndex++; return validTimeList; } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/filescan/impl/DiskChunkHandleImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/filescan/impl/DiskChunkHandleImpl.java index 36a87dbfe1d47..d7b5d4c082dc6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/filescan/impl/DiskChunkHandleImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/filescan/impl/DiskChunkHandleImpl.java @@ -50,7 +50,8 @@ public class DiskChunkHandleImpl implements IChunkHandle { private final String measurement; private final String filePath; private final TsFileID tsFileID; - private EncryptParameter encryptParam; + protected EncryptParameter encryptParam; + protected int cryptoPageIndex = -1; protected ChunkHeader currentChunkHeader; protected PageHeader currentPageHeader; protected ByteBuffer currentChunkDataBuffer; @@ -85,10 +86,12 @@ protected void init(TsFileSequenceReader reader) throws IOException { if (currentChunkDataBuffer != null) { return; } - Chunk chunk = reader.readMemChunk(offset); - this.currentChunkDataBuffer = chunk.getData(); - this.currentChunkHeader = chunk.getHeader(); - this.encryptParam = chunk.getEncryptParam(); + try (Chunk chunk = reader.readMemChunk(offset)) { + this.currentChunkDataBuffer = chunk.getData(); + this.currentChunkHeader = chunk.getHeader(); + // The query's FileReaderManager reference owns this file-level context. + this.encryptParam = reader.getEncryptParam(); + } } // Check if there is more pages to be scanned in Chunk. @@ -107,6 +110,7 @@ public void nextPage() throws IOException { init(reader); } if (currentChunkDataBuffer.hasRemaining()) { + cryptoPageIndex++; // If there is only one page, page statistics is not stored in the chunk header, which is the // same as chunkStatistics if ((byte) (this.currentChunkHeader.getChunkType() & 0x3F) @@ -137,7 +141,12 @@ public long[] getDataTime() throws IOException { IDecryptor decryptor = IDecryptor.getDecryptor(encryptParam); ByteBuffer currentPageDataBuffer = ChunkReader.deserializePageData( - currentPageHeader, this.currentChunkDataBuffer, this.currentChunkHeader, decryptor); + currentPageHeader, + this.currentChunkDataBuffer, + this.currentChunkHeader, + decryptor, + encryptParam, + cryptoPageIndex); int timeBufferLength = ReadWriteForEncodingUtils.readUnsignedVarInt(currentPageDataBuffer); ByteBuffer timeBuffer = currentPageDataBuffer.slice(); timeBuffer.limit(timeBufferLength); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/reader/chunk/DiskAlignedChunkLoader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/reader/chunk/DiskAlignedChunkLoader.java index 5af552502857a..06c43821c2b0d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/reader/chunk/DiskAlignedChunkLoader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/reader/chunk/DiskAlignedChunkLoader.java @@ -95,44 +95,57 @@ public IChunkReader getChunkReader(IChunkMetadata chunkMetaData, Filter globalTi timeChunkMetadata.getStatistics(), context); List valueChunkList = new ArrayList<>(); - for (IChunkMetadata valueChunkMetadata : alignedChunkMetadata.getValueChunkMetadataList()) { - Chunk chunk = - valueChunkMetadata == null - ? null - : ChunkCache.getInstance() - .get( - new ChunkCache.ChunkCacheKey( - resource.getTsFilePath(), - resource.getTsFileID(), - valueChunkMetadata.getOffsetOfChunkHeader(), - resource.isClosed()), - valueChunkMetadata.getDeleteIntervalList(), - valueChunkMetadata.getStatistics(), - context); - final TsFileID tsFileID = getTsFileID(); - if (chunk != null - && tsFileID.regionId > 0 - && chunkMetaData.getDataType() == TSDataType.OBJECT) { - chunk - .getHeader() - .setReplaceDecoder( - decoder -> ObjectTypeUtils.getReplaceDecoder(decoder, tsFileID.regionId)); + try { + for (IChunkMetadata valueChunkMetadata : alignedChunkMetadata.getValueChunkMetadataList()) { + Chunk chunk = + valueChunkMetadata == null + ? null + : ChunkCache.getInstance() + .get( + new ChunkCache.ChunkCacheKey( + resource.getTsFilePath(), + resource.getTsFileID(), + valueChunkMetadata.getOffsetOfChunkHeader(), + resource.isClosed()), + valueChunkMetadata.getDeleteIntervalList(), + valueChunkMetadata.getStatistics(), + context); + final TsFileID tsFileID = getTsFileID(); + if (chunk != null + && tsFileID.regionId > 0 + && chunkMetaData.getDataType() == TSDataType.OBJECT) { + chunk + .getHeader() + .setReplaceDecoder( + decoder -> ObjectTypeUtils.getReplaceDecoder(decoder, tsFileID.regionId)); + } + valueChunkList.add(chunk); } - valueChunkList.add(chunk); + LongConsumer filterRowsRecorder = + this.context.getQueryStatistics()::addFilteredRowsOfPageLevel; + long t2 = System.nanoTime(); + IChunkReader chunkReader = + ignoreAllNullRows + ? new AlignedChunkReader( + timeChunk, valueChunkList, globalTimeFilter, filterRowsRecorder) + : new TableChunkReader( + timeChunk, valueChunkList, globalTimeFilter, filterRowsRecorder); + if (timeChunk.getEncryptParam() != null && timeChunk.getEncryptParam().isTdePageAead()) { + context.getQueryChunkReaders().retain(chunkReader, chunkReader.loadPageReaderList()); + } + SERIES_SCAN_COST_METRIC_SET.recordSeriesScanCost( + INIT_CHUNK_READER_ALIGNED_DISK, System.nanoTime() - t2); + + return chunkReader; + } catch (IOException | RuntimeException | Error failure) { + timeChunk.close(); + for (Chunk valueChunk : valueChunkList) { + if (valueChunk != null) { + valueChunk.close(); + } + } + throw failure; } - LongConsumer filterRowsRecorder = - this.context.getQueryStatistics()::addFilteredRowsOfPageLevel; - long t2 = System.nanoTime(); - IChunkReader chunkReader = - ignoreAllNullRows - ? new AlignedChunkReader( - timeChunk, valueChunkList, globalTimeFilter, filterRowsRecorder) - : new TableChunkReader( - timeChunk, valueChunkList, globalTimeFilter, filterRowsRecorder); - SERIES_SCAN_COST_METRIC_SET.recordSeriesScanCost( - INIT_CHUNK_READER_ALIGNED_DISK, System.nanoTime() - t2); - - return chunkReader; } finally { long time = System.nanoTime() - t1; context.getQueryStatistics().getConstructAlignedChunkReadersDiskCount().getAndAdd(1); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/reader/chunk/DiskChunkLoader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/reader/chunk/DiskChunkLoader.java index eb7abb7ed07f6..77b46a3dcae1d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/reader/chunk/DiskChunkLoader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/read/reader/chunk/DiskChunkLoader.java @@ -92,6 +92,7 @@ public IChunkReader getChunkReader(IChunkMetadata chunkMetaData, Filter globalTi byte chunkType = chunk.getHeader().getChunkType(); if (chunkType != MetaMarker.CHUNK_HEADER && chunkType != MetaMarker.ONLY_ONE_PAGE_CHUNK_HEADER) { + chunk.close(); throw new ChunkTypeInconsistentException(); } @@ -109,6 +110,9 @@ public IChunkReader getChunkReader(IChunkMetadata chunkMetaData, Filter globalTi chunk, globalTimeFilter, this.context.getQueryStatistics()::addFilteredRowsOfPageLevel); + if (chunk.getEncryptParam() != null && chunk.getEncryptParam().isTdePageAead()) { + context.getQueryChunkReaders().retain(chunkReader, chunkReader.loadPageReaderList()); + } SeriesScanCostMetricSet.getInstance() .recordSeriesScanCost(INIT_CHUNK_READER_NONALIGNED_DISK, System.nanoTime() - t2); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/utils/SharedTimeDataBuffer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/utils/SharedTimeDataBuffer.java index 184d939da8293..3199bff11d17a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/utils/SharedTimeDataBuffer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/utils/SharedTimeDataBuffer.java @@ -61,18 +61,20 @@ public void init(TsFileSequenceReader reader) throws IOException { if (timeBuffer != null) { return; } - Chunk timeChunk = reader.readMemChunk(timeChunkMetaData.getOffsetOfChunkHeader()); - timeChunkHeader = timeChunk.getHeader(); - timeBuffer = timeChunk.getData(); - encryptParam = timeChunk.getEncryptParam(); + try (Chunk timeChunk = reader.readMemChunk(timeChunkMetaData.getOffsetOfChunkHeader())) { + timeChunkHeader = timeChunk.getHeader(); + timeBuffer = timeChunk.getData(); + encryptParam = reader.getEncryptParam(); + } } public long[] getPageTime(int pageId) throws IOException { - int size = timeData.size(); - if (pageId < size) { + if (pageId >= 0 && pageId < timeData.size()) { return timeData.get(pageId); - } else if (pageId == size) { - loadPageData(); + } else if (pageId >= timeData.size()) { + while (timeData.size() <= pageId) { + loadPageData(); + } return timeData.get(pageId); } else { throw new UnsupportedOperationException( @@ -92,7 +94,8 @@ private void loadPageData() throws IOException { : PageHeader.deserializeFrom(timeBuffer, timeChunkHeader.getDataType()); IDecryptor decryptor = IDecryptor.getDecryptor(encryptParam); ByteBuffer timePageData = - ChunkReader.deserializePageData(timePageHeader, timeBuffer, timeChunkHeader, decryptor); + ChunkReader.deserializePageData( + timePageHeader, timeBuffer, timeChunkHeader, decryptor, encryptParam, timeData.size()); long[] pageData = new long[(int) timePageHeader.getNumOfValues()]; int index = 0; while (defaultTimeDecoder.hasNext(timePageData)) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/utils/TsFileResourceUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/utils/TsFileResourceUtils.java index 1a77855631262..fa28541a6247c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/utils/TsFileResourceUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/utils/TsFileResourceUtils.java @@ -158,7 +158,7 @@ public static boolean validateTsFileDataCorrectness(TsFileResource resource) { List> alignedTimeBatches = new ArrayList<>(); Map valueColumn2TimeBatchIndex = new HashMap<>(); - reader.position((long) TSFileConfig.MAGIC_STRING.getBytes().length + 1); + reader.position(reader.getDataStartOffset()); int pageIndex = 0; byte marker; while ((marker = reader.readMarker()) != MetaMarker.SEPARATOR) { @@ -211,7 +211,9 @@ public static boolean validateTsFileDataCorrectness(TsFileResource resource) { while (dataSize > 0) { valueDecoder.reset(); PageHeader pageHeader = reader.readPageHeader(header.getDataType(), isHasStatistic); - ByteBuffer pageData = reader.readPage(pageHeader, header.getCompressionType()); + ByteBuffer pageData = + reader.readPage( + pageHeader, header.getCompressionType(), pageIndex, header.getChunkOrdinal()); if ((header.getChunkType() & TsFileConstant.TIME_COLUMN_MASK) == TsFileConstant.TIME_COLUMN_MASK) { // Time Chunk diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java index d72fd02e54b88..78b19c36848f0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java @@ -44,6 +44,7 @@ import org.apache.iotdb.db.storageengine.load.splitter.ChunkPayloadRef; import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; +import org.apache.iotdb.db.utils.EncryptDBUtils; import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.common.constant.TsFileConstant; @@ -176,6 +177,11 @@ long getLastCommandTimeMillis() { List writePiece(final List tsFileDataList) throws IOException, PageException { checkNotClosed(); + if (EncryptDBUtils.requiresFileRewrite()) { + throw new IOException( + StorageEngineMessages + .EXCEPTION_PRECALCULATED_CHUNK_STAGING_IS_NOT_SUPPORTED_BY_THE_FILE_ENCRYPTION_POLICY_13DA8C32); + } for (final TsFileData tsFileData : tsFileDataList) { switch (tsFileData.getType()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java index 7354a9d21b07b..1bbfc0bc99b59 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/AlignedChunkData.java @@ -152,7 +152,7 @@ public AlignedChunkData( final TTimePartitionSlot timePartitionSlot, final EncryptParameter encryptParameter) { this(device, timePartitionSlot, encryptParameter); - chunkHeaderList.add(chunkHeader); + chunkHeaderList.add(ChunkData.transportHeader(chunkHeader)); pageNumbers.add(0); currentChunkIndex = 0; addAttrDataSize(); @@ -235,7 +235,7 @@ public boolean isEntireChunk() { } public void addValueChunk(final ChunkHeader chunkHeader) { - this.chunkHeaderList.add(chunkHeader); + this.chunkHeaderList.add(ChunkData.transportHeader(chunkHeader)); this.pageNumbers.add(0); currentChunkIndex++; dataSize += chunkHeader.getSerializedSize(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java index 7deb07b9bd794..62e457932bb84 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/ChunkData.java @@ -42,6 +42,18 @@ /** Represents chunk-level TsFile data during load splitting. */ public interface ChunkData extends TsFileData { + static ChunkHeader transportHeader(ChunkHeader header) { + if (header.getChunkOrdinal() < 0) { + return header; + } + return new ChunkHeader( + header.getChunkType(), + header.getMeasurementID(), + header.getDataSize(), + header.getDataType(), + header.getCompressionType(), + header.getEncodingType()); + } ByteBuffer EMPTY_BUFFER = ByteBuffer.allocate(0); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java index d5342c2918c87..01fe54423139d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/NonAlignedChunkData.java @@ -115,7 +115,7 @@ public NonAlignedChunkData( final EncryptParameter encryptParameter) { this.dataSize = 0; this.device = device; - this.chunkHeader = chunkHeader; + this.chunkHeader = ChunkData.transportHeader(chunkHeader); this.timePartitionSlot = timePartitionSlot; this.encryptParameter = encryptParameter; this.needDecodeChunk = true; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java index 7a976838d7052..0571ecd68b36e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/splitter/TsFileSplitter.java @@ -77,6 +77,7 @@ public class TsFileSplitter { private Map> pageIndex2ChunkData = new LinkedHashMap<>(); private Map pageIndex2Times = new HashMap<>(); private boolean isTimeChunkNeedDecode = true; + private boolean forceDecodeChunks; private IDeviceID curDevice = null; private boolean isAligned; private int timeChunkIndexOfCurrentValueColumn = 0; @@ -120,7 +121,8 @@ public void splitTsFileByDataPartition() } reader.readFileMetadata(); - reader.position((long) TSFileConfig.MAGIC_STRING.getBytes().length + 1); + forceDecodeChunks = reader.hasFileEncryptionHeader(); + reader.position(reader.getDataStartOffset()); getChunkMetadata(reader, offset2ChunkMetadata); byte marker; // It should be noted that time chunk and its corresponding value chunk are not necessarily @@ -269,7 +271,9 @@ private void decodeAndWriteTimeChunkOrNonAlignedChunk( } chunkData.writeEntirePage(pageHeader, reader.readCompressedPage(pageHeader)); } else { // split page - ByteBuffer pageData = reader.readPage(pageHeader, header.getCompressionType()); + ByteBuffer pageData = + reader.readPage( + pageHeader, header.getCompressionType(), pageIndex, header.getChunkOrdinal()); Pair tvArray = decodePage(isAligned, pageData, pageHeader, defaultTimeDecoder, valueDecoder, header); long[] times = tvArray.left; @@ -365,14 +369,15 @@ private void processValueChunk(TsFileSequenceReader reader, byte marker) allChunkData.add(alignedChunkData); } } - if (alignedChunkDataList.size() == 1) { // write entire page + if (alignedChunkDataList.size() == 1 && !forceDecodeChunks) { // write entire page // write the entire page if it's not an empty page. alignedChunkDataList .get(0) .writeEntirePage(pageHeader, reader.readCompressedPage(pageHeader)); } else { // decode page long[] times = pageIndex2Times.get(pageIndex); - TsPrimitiveType[] values = decodeValuePage(reader, header, pageHeader, times, valueDecoder); + TsPrimitiveType[] values = + decodeValuePage(reader, header, pageHeader, times, valueDecoder, pageIndex); for (AlignedChunkData alignedChunkData : alignedChunkDataList) { alignedChunkData.writeDecodeValuePage(times, values, header.getDataType()); } @@ -538,11 +543,15 @@ private void consumeChunkData(String measurement, long offset, ChunkData chunkDa } private boolean needDecodeChunk(IChunkMetadata chunkMetadata) { - return !TimePartitionUtils.getTimePartitionSlot(chunkMetadata.getStartTime()) - .equals(TimePartitionUtils.getTimePartitionSlot(chunkMetadata.getEndTime())); + return forceDecodeChunks + || !TimePartitionUtils.getTimePartitionSlot(chunkMetadata.getStartTime()) + .equals(TimePartitionUtils.getTimePartitionSlot(chunkMetadata.getEndTime())); } private boolean needDecodePage(PageHeader pageHeader, IChunkMetadata chunkMetadata) { + if (forceDecodeChunks) { + return true; + } if (pageHeader.getStatistics() == null) { return !TimePartitionUtils.getTimePartitionSlot(chunkMetadata.getStartTime()) .equals(TimePartitionUtils.getTimePartitionSlot(chunkMetadata.getEndTime())); @@ -609,7 +618,8 @@ private TsPrimitiveType[] decodeValuePage( ChunkHeader chunkHeader, PageHeader pageHeader, long[] times, - Decoder valueDecoder) + Decoder valueDecoder, + int pageIndex) throws IOException { // An empty page stands for an all-null slice of the column. It has to expand to one null per // time of the page, otherwise the decoded values would not line up with the decoded times. @@ -618,7 +628,9 @@ private TsPrimitiveType[] decodeValuePage( } valueDecoder.reset(); - ByteBuffer pageData = reader.readPage(pageHeader, chunkHeader.getCompressionType()); + ByteBuffer pageData = + reader.readPage( + pageHeader, chunkHeader.getCompressionType(), pageIndex, chunkHeader.getChunkOrdinal()); ValuePageReader valuePageReader = new ValuePageReader(pageHeader, pageData, chunkHeader.getDataType(), valueDecoder); return valuePageReader.nextValueBatch( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/TsFileSplitByPartitionTool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/TsFileSplitByPartitionTool.java index 9dae1181295c8..6e9331c7b399d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/TsFileSplitByPartitionTool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/TsFileSplitByPartitionTool.java @@ -148,6 +148,7 @@ public void parseAndRewriteFile(List rewrittenResources) DataNodeMiscMessages .MISC_EXCEPTION_THE_VERSION_OF_THIS_TSFILE_IS_TOO_LOW_PLEASE_UPGRADE_IT_19CC276C); } + reader.position(reader.getDataStartOffset()); // start to scan chunks and chunkGroups byte marker; IDeviceID deviceId = null; @@ -178,16 +179,24 @@ public void parseAndRewriteFile(List rewrittenResources) List dataInChunk = new ArrayList<>(); List needToDecodeInfo = new ArrayList<>(); int dataSize = header.getDataSize(); + int pageIndex = 0; while (dataSize > 0) { // a new Page PageHeader pageHeader = reader.readPageHeader(dataType, header.getChunkType() == MetaMarker.CHUNK_HEADER); - boolean needToDecode = checkIfNeedToDecode(measurementSchema, deviceId, pageHeader); + boolean needToDecode = + reader.hasFileEncryptionHeader() + || checkIfNeedToDecode(measurementSchema, deviceId, pageHeader); needToDecodeInfo.add(needToDecode); ByteBuffer pageData = !needToDecode ? reader.readCompressedPage(pageHeader) - : reader.readPage(pageHeader, header.getCompressionType()); + : reader.readPage( + pageHeader, + header.getCompressionType(), + pageIndex, + header.getChunkOrdinal()); + pageIndex++; pageHeadersInChunk.add(pageHeader); dataInChunk.add(pageData); dataSize -= pageHeader.getSerializedPageSize(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/TsFileSplitTool.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/TsFileSplitTool.java index e5fafca708bb4..e932ac993a23e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/TsFileSplitTool.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/TsFileSplitTool.java @@ -184,6 +184,7 @@ public void run() throws IOException { int dataSize = chunkHeader.getDataSize(); Decoder valueDecoder = Decoder.getDecoderByType(chunkHeader.getEncodingType(), dataType); + int pageIndex = 0; Decoder defaultTimeDecoder = Decoder.getDecoderByType( TSEncoding.valueOf(TSFileDescriptor.getInstance().getConfig().getTimeEncoder()), @@ -197,7 +198,12 @@ public void run() throws IOException { reader.readPageHeader( dataType, ((byte) (chunkHeader.getChunkType() & 0x3F)) == MetaMarker.CHUNK_HEADER); - ByteBuffer pageData = reader.readPage(pageHeader, chunkHeader.getCompressionType()); + ByteBuffer pageData = + reader.readPage( + pageHeader, + chunkHeader.getCompressionType(), + pageIndex++, + chunkHeader.getChunkOrdinal()); PageReader pageReader = new PageReader(pageData, dataType, valueDecoder, defaultTimeDecoder); BatchData batchData = pageReader.getAllSatisfiedPageData(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/utils/TsFileSequenceScan.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/utils/TsFileSequenceScan.java index a260a6e533987..003783d99bcd0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/utils/TsFileSequenceScan.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/tools/utils/TsFileSequenceScan.java @@ -20,7 +20,6 @@ import org.apache.iotdb.db.utils.EncryptDBUtils; -import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.file.MetaMarker; import org.apache.tsfile.file.header.ChunkGroupHeader; import org.apache.tsfile.file.header.ChunkHeader; @@ -57,7 +56,7 @@ public TsFileSequenceScan() {} */ protected boolean onFileOpen(File file) throws IOException { // skip head magic string - reader.position((long) TSFileConfig.MAGIC_STRING.getBytes().length + 1); + reader.position(reader.getDataStartOffset()); this.file = file; return true; } @@ -78,12 +77,18 @@ protected void onChunk(PageVisitor pageVisitor) throws IOException { } protected void onChunkData(PageVisitor pageVisitor, int dataSize) throws IOException { + int pageIndex = 0; while (dataSize > 0) { PageHeader pageHeader = reader.readPageHeader( currChunkHeader.getDataType(), (currChunkHeader.getChunkType() & 0x3F) == MetaMarker.CHUNK_HEADER); - ByteBuffer pageData = reader.readPage(pageHeader, currChunkHeader.getCompressionType()); + ByteBuffer pageData = + reader.readPage( + pageHeader, + currChunkHeader.getCompressionType(), + pageIndex++, + currChunkHeader.getChunkOrdinal()); pageVisitor.onPage(pageHeader, pageData, currChunkHeader); dataSize -= pageHeader.getSerializedPageSize(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/EncryptDBUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/EncryptDBUtils.java index 52396d149ebeb..1c55154d40977 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/EncryptDBUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/EncryptDBUtils.java @@ -20,15 +20,22 @@ import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.i18n.StorageEngineMessages; import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.common.conf.TSFileDescriptor; import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.encrypt.EncryptUtils; +import org.apache.tsfile.write.writer.RestorableTsFileIOWriter; +import org.apache.tsfile.write.writer.TsFileIOWriter; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.io.File; +import java.io.IOException; +import java.util.Iterator; import java.util.Map; +import java.util.ServiceLoader; public class EncryptDBUtils { @@ -38,6 +45,83 @@ public class EncryptDBUtils { private static final Logger LOGGER = LoggerFactory.getLogger(EncryptDBUtils.class); private static volatile EncryptParameter defaultFirstEncryptParam; + private static final FileEncryptionParameterProvider LEGACY_PROVIDER = + (file, database, legacy) -> legacy; + + private static class ProviderHolder { + private static final FileEncryptionParameterProvider PROVIDER = loadProvider(); + } + + private static FileEncryptionParameterProvider loadProvider() { + Iterator providers = + ServiceLoader.load(FileEncryptionParameterProvider.class).iterator(); + if (!providers.hasNext()) { + return LEGACY_PROVIDER; + } + FileEncryptionParameterProvider provider = providers.next(); + if (providers.hasNext()) { + throw new IllegalStateException( + StorageEngineMessages + .EXCEPTION_MULTIPLE_FILE_ENCRYPTION_PARAMETER_PROVIDERS_ARE_CONFIGURED_5E7DDA5B); + } + return provider; + } + + public static EncryptParameter createEncryptParameterForNewFile( + File file, String database, EncryptParameter legacyParameter) { + EncryptParameter parameter = + ProviderHolder.PROVIDER.createForNewFile(file, database, legacyParameter); + if (parameter == null && ProviderHolder.PROVIDER != LEGACY_PROVIDER) { + throw new IllegalStateException( + StorageEngineMessages + .EXCEPTION_FILE_ENCRYPTION_PARAMETER_PROVIDER_RETURNED_NULL_15C41723); + } + return parameter; + } + + public static boolean requiresFileRewrite() { + return ProviderHolder.PROVIDER.requiresFileRewrite(); + } + + public static boolean canCopyChunk(EncryptParameter source, EncryptParameter target) { + return !requiresFileRewrite() + && (source == null || !source.isTdePageAead()) + && (target == null || !target.isTdePageAead()); + } + + public static RestorableTsFileIOWriter createRestorableWriterForNewFile( + File file, String database) throws IOException { + return createRestorableWriterForNewFile( + file, database, getFirstEncryptParamFromDatabase(database)); + } + + public static RestorableTsFileIOWriter createRestorableWriterForNewFile( + File file, String database, EncryptParameter legacyParameter) throws IOException { + EncryptParameter parameter = createEncryptParameterForNewFile(file, database, legacyParameter); + try { + return new RestorableTsFileIOWriter(file, parameter); + } finally { + closeFileEncryptionParameter(parameter); + } + } + + public static TsFileIOWriter createWriterForNewFile(File file, String database) + throws IOException { + EncryptParameter parameter = + createEncryptParameterForNewFile( + file, database, getFirstEncryptParamFromDatabase(database)); + try { + return new TsFileIOWriter(file, parameter); + } finally { + closeFileEncryptionParameter(parameter); + } + } + + public static void closeFileEncryptionParameter(EncryptParameter parameter) { + if (parameter != null && parameter.isTdePageAead()) { + parameter.close(); + } + } public static EncryptParameter getFirstEncryptParamFromDatabase(String database) { if (database == null || database.isEmpty()) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/FileEncryptionParameterProvider.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/FileEncryptionParameterProvider.java new file mode 100644 index 0000000000000..ee56f4e646730 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/FileEncryptionParameterProvider.java @@ -0,0 +1,39 @@ +/* + * 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.utils; + +import org.apache.tsfile.encrypt.EncryptParameter; + +import java.io.File; + +/** Optional file-level policy, discovered once through {@link java.util.ServiceLoader}. */ +@FunctionalInterface +public interface FileEncryptionParameterProvider { + /** + * Returns non-null parameters for exactly one new file. The legacy parameter is borrowed and must + * not be modified or closed. A returned page-AEAD descriptor is owned by the caller; generation + * and key resolution belong to the provider, not to the storage engine. + */ + EncryptParameter createForNewFile(File file, String database, EncryptParameter legacyParameter); + + /** Whether new target files require rewriting rather than copying source chunks verbatim. */ + default boolean requiresFileRewrite() { + return false; + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryChunkReadersTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryChunkReadersTest.java new file mode 100644 index 0000000000000..f515f293fb128 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/fragment/QueryChunkReadersTest.java @@ -0,0 +1,88 @@ +/* + * 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.queryengine.execution.fragment; + +import org.apache.tsfile.read.reader.IChunkReader; +import org.apache.tsfile.read.reader.IPageReader; +import org.junit.Test; + +import java.util.Arrays; +import java.util.Collections; + +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; + +public class QueryChunkReadersTest { + @Test + public void testLastPageReleasesOwnerOnce() throws Exception { + IChunkReader owner = mock(IChunkReader.class); + IPageReader first = mock(IPageReader.class); + IPageReader second = mock(IPageReader.class); + try (QueryChunkReaders readers = new QueryChunkReaders()) { + readers.retain(owner, Arrays.asList(first, second)); + readers.release(first); + verify(owner, never()).close(); + readers.release(second); + readers.release(second); + verify(owner, times(1)).close(); + } + verify(owner, times(1)).close(); + } + + @Test + public void testEarlyCloseReleasesRemainingAndEmptyOwners() throws Exception { + IChunkReader firstOwner = mock(IChunkReader.class); + IChunkReader emptyOwner = mock(IChunkReader.class); + QueryChunkReaders readers = new QueryChunkReaders(); + readers.retain(firstOwner, Arrays.asList(mock(IPageReader.class), mock(IPageReader.class))); + readers.retain(emptyOwner, Collections.emptyList()); + readers.close(); + readers.close(); + verify(firstOwner, times(1)).close(); + verify(emptyOwner, times(1)).close(); + } + + @Test + public void testRetainAfterCloseCannotLeakOwner() throws Exception { + IChunkReader owner = mock(IChunkReader.class); + QueryChunkReaders readers = new QueryChunkReaders(); + readers.close(); + org.junit.Assert.assertThrows( + java.io.IOException.class, + () -> readers.retain(owner, Collections.singletonList(mock(IPageReader.class)))); + verify(owner, times(1)).close(); + } + + @Test + public void testCloseAttemptsEveryOwnerAfterIOException() throws Exception { + IChunkReader failed = mock(IChunkReader.class); + IChunkReader healthy = mock(IChunkReader.class); + doThrow(new java.io.IOException()).when(failed).close(); + QueryChunkReaders readers = new QueryChunkReaders(); + readers.retain(failed, Collections.singletonList(mock(IPageReader.class))); + readers.retain(healthy, Collections.singletonList(mock(IPageReader.class))); + org.junit.Assert.assertThrows(java.io.IOException.class, readers::close); + readers.close(); + verify(failed, times(1)).close(); + verify(healthy, times(1)).close(); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TabletRewriteLoadStrategyTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TabletRewriteLoadStrategyTest.java new file mode 100644 index 0000000000000..6f33a1c7a1a10 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TabletRewriteLoadStrategyTest.java @@ -0,0 +1,177 @@ +/* + * 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.queryengine.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LoadTsFile; +import org.apache.iotdb.db.queryengine.plan.statement.crud.LoadTsFileStatement; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.load.converter.LoadTsFileDataTypeConverter; +import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.mockito.ArgumentCaptor; + +import java.io.File; +import java.util.Collections; +import java.util.Optional; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class TabletRewriteLoadStrategyTest { + + @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder(); + + private LoadTsFileDataTypeConverter converter; + private LoadTsFileCostMetricsSet costMetrics; + private LoadSingleTsFileNode node; + private File sourceFile; + + @Before + public void setUp() throws Exception { + converter = mock(LoadTsFileDataTypeConverter.class); + when(converter.isSuccessful(any())).thenCallRealMethod(); + costMetrics = mock(LoadTsFileCostMetricsSet.class); + sourceFile = temporaryFolder.newFile("source.tsfile"); + node = mock(LoadSingleTsFileNode.class); + when(node.getTsFileResource()).thenReturn(new TsFileResource(sourceFile)); + when(node.isDeleteAfterLoad()).thenReturn(true); + } + + @Test + public void testTreePipeRewritePreservesDatabaseAndLeavesDeletionToScheduler() throws Exception { + when(node.getDatabase()).thenReturn("root.test.sg"); + when(converter.convertForTreeModel(any())) + .thenReturn(Optional.of(status(TSStatusCode.SUCCESS_STATUS))); + + assertTrue(strategy(true).execute(node)); + + final ArgumentCaptor statement = + ArgumentCaptor.forClass(LoadTsFileStatement.class); + verify(converter).convertForTreeModel(statement.capture()); + assertEquals("root.test.sg", statement.getValue().getDatabase()); + assertEquals(2, statement.getValue().getDatabaseLevel()); + assertTrue(statement.getValue().isGeneratedByPipe()); + assertFalse(statement.getValue().isDeleteAfterLoad()); + assertEquals(Collections.singletonList(sourceFile), statement.getValue().getTsFiles()); + assertTrue(sourceFile.exists()); + verify(node, never()).clean(); + verifyMetric(); + } + + @Test + public void testFailedConversionRetainsSourceAndRecordsMetric() throws Exception { + when(converter.convertForTreeModel(any())) + .thenReturn(Optional.of(status(TSStatusCode.LOAD_FILE_ERROR))); + + assertFalse(strategy(false).execute(node)); + + final ArgumentCaptor statement = + ArgumentCaptor.forClass(LoadTsFileStatement.class); + verify(converter).convertForTreeModel(statement.capture()); + assertFalse(statement.getValue().isDeleteAfterLoad()); + assertFalse(statement.getValue().isGeneratedByPipe()); + verify(node, never()).clean(); + assertTrue(sourceFile.exists()); + verifyMetric(); + } + + @Test + public void testMissingConversionStatusFailsWithoutDeletingSource() throws Exception { + when(converter.convertForTreeModel(any())).thenReturn(Optional.empty()); + + assertFalse(strategy(false).execute(node)); + assertTrue(sourceFile.exists()); + verify(node, never()).clean(); + verifyMetric(); + } + + @Test + public void testConversionExceptionIsWrappedAndRetainsSource() { + final IllegalStateException cause = new IllegalStateException("conversion failed"); + when(converter.convertForTreeModel(any())).thenThrow(cause); + + final LoadFileException failure = + assertThrows(LoadFileException.class, () -> strategy(false).execute(node)); + + assertSame(cause, failure.getCause()); + assertTrue(sourceFile.exists()); + verify(node, never()).clean(); + verifyMetric(); + } + + @Test + public void testFallbackConversionStillHonorsItsDeleteAfterLoadArgument() throws Exception { + when(converter.convertForTreeModel(any())) + .thenReturn(Optional.of(status(TSStatusCode.SUCCESS_STATUS))); + + LoadFallbackHandler.convertTsFileToTablets(converter, node, false, true); + + final ArgumentCaptor statement = + ArgumentCaptor.forClass(LoadTsFileStatement.class); + verify(converter).convertForTreeModel(statement.capture()); + assertTrue(statement.getValue().isDeleteAfterLoad()); + } + + private TabletRewriteLoadStrategy strategy(final boolean isGeneratedByPipe) { + return new TabletRewriteLoadStrategy(converter, isGeneratedByPipe, costMetrics); + } + + @Test + public void testTablePipeRewritePreservesDatabaseAndDefersDeletion() throws Exception { + when(node.isTableModel()).thenReturn(true); + when(node.getDatabase()).thenReturn("table_db"); + when(converter.convertForTableModel(any())) + .thenReturn(Optional.of(status(TSStatusCode.SUCCESS_STATUS))); + assertTrue(strategy(true).execute(node)); + ArgumentCaptor statement = ArgumentCaptor.forClass(LoadTsFile.class); + verify(converter).convertForTableModel(statement.capture()); + assertEquals("table_db", statement.getValue().getDatabase()); + assertTrue(statement.getValue().isGeneratedByPipe()); + assertFalse(statement.getValue().isDeleteAfterLoad()); + assertTrue(sourceFile.exists()); + } + + private void verifyMetric() { + verify(costMetrics) + .recordPhaseTimeCost(eq(LoadTsFileCostMetricsSet.SCHEDULER_CAST_TABLETS), anyLong()); + } + + private static TSStatus status(final TSStatusCode code) { + return new TSStatus(code.getStatusCode()); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriterTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriterTest.java index 673ab709df51d..a2ca64fc59a3d 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriterTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriterTest.java @@ -49,6 +49,7 @@ import java.io.IOException; import java.lang.reflect.Method; import java.nio.ByteBuffer; +import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -72,6 +73,28 @@ public void testBinarySizeCheckpointTriggersChunkSizeCheckBeforePointCheckpoint( Assert.assertEquals(1, chunkWriter.chunkSizeCheckCount); } + @Test + public void testChunkWriterFactoryCanSwitchFileEncryptionParameter() { + TestCompactionWriter compactionWriter = new TestCompactionWriter(); + EncryptParameter first = new EncryptParameter("UNENCRYPTED", null); + EncryptParameter second = new EncryptParameter("UNENCRYPTED", null); + List createdWith = new ArrayList<>(); + compactionWriter.setEncryptParameter(first); + + compactionWriter.startMeasurement( + "s1", + encryptParameter -> { + createdWith.add(encryptParameter); + return new CountingChunkWriter(); + }, + SUB_TASK_ID); + compactionWriter.switchEncryptionParameter(second); + + Assert.assertEquals(2, createdWith.size()); + Assert.assertSame(first, createdWith.get(0)); + Assert.assertSame(second, createdWith.get(1)); + } + @Test public void testEstimateWrittenPointTotalSize() throws Exception { TestCompactionWriter compactionWriter = new TestCompactionWriter(); @@ -197,6 +220,8 @@ public boolean checkIsChunkSizeOverThreshold( private static class TestCompactionWriter extends AbstractCompactionWriter { + private EncryptParameter encryptParameter; + private void writePoint(TsPrimitiveType value, IChunkWriter chunkWriter) { writeDataPoint(1, value, chunkWriter, SUB_TASK_ID); } @@ -213,6 +238,14 @@ private int getCompressedSizeToReachSizeCheckpoint() { return (int) Math.max(targetChunkSize / 10, 1L); } + private void setEncryptParameter(EncryptParameter encryptParameter) { + this.encryptParameter = encryptParameter; + } + + private void switchEncryptionParameter(EncryptParameter encryptParameter) { + switchChunkWriterEncryptionParameter(encryptParameter, SUB_TASK_ID); + } + @Override public void startChunkGroup(IDeviceID deviceId, boolean isAlign) {} @@ -241,7 +274,7 @@ public void checkAndMayFlushChunkMetadata() {} @Override public EncryptParameter getEncryptParameter() { - return null; + return encryptParameter; } @Override diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/FileEncryptionContextTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/FileEncryptionContextTest.java new file mode 100644 index 0000000000000..3668b92ec0a22 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/FileEncryptionContextTest.java @@ -0,0 +1,408 @@ +/* + * 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.utils; + +import org.apache.iotdb.db.queryengine.execution.fragment.QueryContext; +import org.apache.iotdb.db.storageengine.buffer.ChunkCache; +import org.apache.iotdb.db.storageengine.dataregion.compaction.io.CompactionTsFileWriter; +import org.apache.iotdb.db.storageengine.dataregion.compaction.schedule.constant.CompactionType; +import org.apache.iotdb.db.storageengine.dataregion.memtable.PrimitiveMemTable; +import org.apache.iotdb.db.storageengine.dataregion.read.control.FileReaderManager; +import org.apache.iotdb.db.storageengine.dataregion.read.reader.chunk.DiskAlignedChunkLoader; +import org.apache.iotdb.db.storageengine.dataregion.read.reader.chunk.DiskChunkLoader; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileSplitter; + +import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptUtils; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.AbstractAlignedChunkMetadata; +import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.apache.tsfile.read.common.BatchData; +import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.read.reader.IChunkReader; +import org.apache.tsfile.read.reader.IPageReader; +import org.apache.tsfile.read.reader.chunk.ChunkReader; +import org.apache.tsfile.write.chunk.AlignedChunkWriterImpl; +import org.apache.tsfile.write.chunk.ChunkWriterImpl; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.apache.tsfile.write.writer.TsFileIOWriter; +import org.awaitility.Awaitility; +import org.junit.After; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.io.File; +import java.lang.reflect.Field; +import java.nio.file.Files; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +public class FileEncryptionContextTest { + @Rule public final TemporaryFolder folder = new TemporaryFolder(); + private static final IDeviceID DEVICE = + IDeviceID.Factory.DEFAULT_FACTORY.create("root.context.d1"); + private static final MeasurementSchema SCHEMA = + new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.PLAIN, CompressionType.UNCOMPRESSED); + + @After + public void tearDown() throws Exception { + FileReaderManager.getInstance().closeAndRemoveAllOpenedReaders(); + ChunkCache.getInstance().clear(); + TestFileEncryptionParameterProvider.disable(); + } + + @Test + public void testDefaultWriterPreservesPlainFileFormat() throws Exception { + File file = file("1-1-0-0.tsfile"); + try (TsFileIOWriter writer = EncryptDBUtils.createWriterForNewFile(file, "root.context")) { + write(writer); + } + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + assertFalse(reader.hasFileEncryptionHeader()); + assertEquals(TSFileConfig.MAGIC_STRING.length() + 1, reader.getDataStartOffset()); + } + assertFalse(EncryptDBUtils.requiresFileRewrite()); + } + + @Test + public void testEachWriterGetsIndependentFileContext() throws Exception { + TestFileEncryptionParameterProvider.enable(); + byte[][] ids = new byte[2][]; + for (int i = 0; i < 2; i++) { + File file = file((i + 1) + "-1-0-0.tsfile"); + try (TsFileIOWriter writer = EncryptDBUtils.createWriterForNewFile(file, "root.context")) { + ids[i] = writer.getEncryptParameter().getFileCryptoId(); + write(writer); + } + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + assertTrue(reader.hasFileEncryptionHeader()); + assertTrue(reader.getDataStartOffset() > TSFileConfig.MAGIC_STRING.length() + 1); + assertArrayEquals(ids[i], reader.getEncryptParam().getFileCryptoId()); + } + } + assertFalse(Arrays.equals(ids[0], ids[1])); + assertTrue(EncryptDBUtils.requiresFileRewrite()); + assertEquals( + TestFileEncryptionParameterProvider.CREATED.get(), + TestFileEncryptionParameterProvider.CLOSED.get()); + } + + @Test + public void testLazyChunkOutlivesFileReaderAndClosesAfterLastPage() throws Exception { + TestFileEncryptionParameterProvider.enable(); + File file = file("1-1-0-0.tsfile"); + try (TsFileIOWriter writer = EncryptDBUtils.createWriterForNewFile(file, "root.context")) { + write(writer); + } + QueryContext context = externalContext(); + ChunkMetadata metadata = metadata(file); + IChunkReader reader = + new DiskChunkLoader(context, new TsFileResource(file)).getChunkReader(metadata, null); + List pages = reader.loadPageReaderList(); + assertEquals(2, pages.size()); + assertEquals(2, trackedPages(context)); + FileReaderManager.getInstance().closeAndRemoveAllOpenedReaders(); + int row = 0; + for (IPageReader page : pages) { + BatchData data = page.getAllSatisfiedPageData(); + while (data.hasCurrent()) { + assertEquals(row, data.currentTime()); + assertEquals(row * 10L, data.getLong()); + row++; + data.next(); + } + context.getQueryChunkReaders().release(page); + } + assertEquals(20, row); + assertEquals(0, trackedPages(context)); + context.getQueryChunkReaders().close(); + assertEquals( + TestFileEncryptionParameterProvider.CREATED.get(), + TestFileEncryptionParameterProvider.CLOSED.get()); + } + + @Test + public void testEarlyQueryCloseReleasesUnconsumedContext() throws Exception { + TestFileEncryptionParameterProvider.enable(); + File file = file("1-1-0-0.tsfile"); + try (TsFileIOWriter writer = EncryptDBUtils.createWriterForNewFile(file, "root.context")) { + write(writer); + } + QueryContext context = externalContext(); + new DiskChunkLoader(context, new TsFileResource(file)).getChunkReader(metadata(file), null); + assertEquals(2, trackedPages(context)); + FileReaderManager.getInstance().closeAndRemoveAllOpenedReaders(); + context.getQueryChunkReaders().close(); + assertEquals( + TestFileEncryptionParameterProvider.CREATED.get(), + TestFileEncryptionParameterProvider.CLOSED.get()); + } + + @Test + public void testPlainChunksAreNotRetainedByQuery() throws Exception { + File file = file("1-1-0-0.tsfile"); + try (TsFileIOWriter writer = EncryptDBUtils.createWriterForNewFile(file, "root.context")) { + write(writer); + } + QueryContext context = externalContext(); + try (IChunkReader reader = + new DiskChunkLoader(context, new TsFileResource(file)) + .getChunkReader(metadata(file), null)) { + assertEquals(2, reader.loadPageReaderList().size()); + assertEquals(0, trackedPages(context)); + } + } + + @Test + public void testMemTableKeepsExplicitContextAcrossDatabaseAssignmentAndCopy() { + try (EncryptParameter parameter = TestFileEncryptionParameterProvider.newParameter()) { + PrimitiveMemTable table = new PrimitiveMemTable(null, null); + table.setEncryptParameter(EncryptUtils.getEncryptParameter(parameter)); + table.setDatabaseAndDataRegionId("root.context", "0"); + assertTrue(table.getEncryptParameter().sharesPageAeadFileContext(parameter)); + assertSame(table.getEncryptParameter(), table.copy().getEncryptParameter()); + } + } + + @Test + public void testDirectChunkCopyRejectsFileBoundContexts() { + EncryptParameter plain = EncryptDBUtils.getDefaultFirstEncryptParam(); + try (EncryptParameter file = TestFileEncryptionParameterProvider.newParameter()) { + assertTrue(EncryptDBUtils.canCopyChunk(plain, plain)); + assertFalse(EncryptDBUtils.canCopyChunk(file, plain)); + assertFalse(EncryptDBUtils.canCopyChunk(plain, file)); + assertFalse(EncryptDBUtils.canCopyChunk(file, file)); + } + } + + private static QueryContext externalContext() { + + return new QueryContext(false, false) { + @Override + public boolean isExternalTsFileScan() { + return true; + } + }; + } + + @Test + public void testNullProviderCannotSilentlyCreatePlainFile() throws Exception { + TestFileEncryptionParameterProvider.RETURN_NULL.set(true); + File file = file("1-1-0-0.tsfile"); + assertThrows( + IllegalStateException.class, + () -> EncryptDBUtils.createWriterForNewFile(file, "root.context")); + assertFalse(file.exists()); + } + + @Test + public void testCacheEvictionDoesNotInvalidateRetainedLazyPages() throws Exception { + TestFileEncryptionParameterProvider.enable(); + File file = file("1-1-0-0.tsfile"); + try (TsFileIOWriter writer = EncryptDBUtils.createWriterForNewFile(file, "root.context")) { + write(writer); + } + QueryContext context = new QueryContext(false, false); + IChunkReader chunkReader = + new DiskChunkLoader(context, new TsFileResource(file)).getChunkReader(metadata(file), null); + ChunkCache.getInstance().clear(); + FileReaderManager.getInstance().closeAndRemoveAllOpenedReaders(); + int rows = 0; + for (IPageReader page : chunkReader.loadPageReaderList()) { + BatchData data = page.getAllSatisfiedPageData(); + rows += data.length(); + context.getQueryChunkReaders().release(page); + } + assertEquals(20, rows); + Awaitility.await() + .atMost(Duration.ofSeconds(5)) + .untilAsserted( + () -> + assertEquals( + TestFileEncryptionParameterProvider.CREATED.get(), + TestFileEncryptionParameterProvider.CLOSED.get())); + } + + @Test + public void testAlignedReaderOwnsTimeAndValueContextsUntilQueryClose() throws Exception { + TestFileEncryptionParameterProvider.enable(); + File file = file("1-1-0-0.tsfile"); + try (TsFileIOWriter writer = EncryptDBUtils.createWriterForNewFile(file, "root.context")) { + writer.startChunkGroup(DEVICE); + AlignedChunkWriterImpl chunk = + new AlignedChunkWriterImpl( + Collections.singletonList(SCHEMA), + EncryptUtils.getEncryptParameter(writer.getEncryptParameter())); + for (int i = 0; i < 20; i++) { + chunk.getTimeChunkWriter().write(i); + chunk.getValueChunkWriterByIndex(0).write(i, i * 10L, false); + if (i == 9) { + chunk.sealCurrentPage(); + } + } + chunk.writeToFileWriter(writer); + writer.endChunkGroup(); + writer.endFile(); + } + AbstractAlignedChunkMetadata metadata; + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + metadata = reader.getAlignedChunkMetadata(DEVICE, false).get(0); + } + QueryContext context = externalContext(); + IChunkReader chunkReader = + new DiskAlignedChunkLoader(context, new TsFileResource(file), false) + .getChunkReader(metadata, null); + assertEquals(2, trackedPages(context)); + FileReaderManager.getInstance().closeAndRemoveAllOpenedReaders(); + int rows = 0; + for (IPageReader page : chunkReader.loadPageReaderList()) { + rows += page.getAllSatisfiedPageData().length(); + } + assertEquals(20, rows); + context.getQueryChunkReaders().close(); + assertEquals( + TestFileEncryptionParameterProvider.CREATED.get(), + TestFileEncryptionParameterProvider.CLOSED.get()); + } + + @Test + public void testCompactionWriterUsesAndReleasesItsOwnFileContext() throws Exception { + TestFileEncryptionParameterProvider.enable(); + File file = file("1-1-0-0.tsfile"); + try (CompactionTsFileWriter writer = + CompactionTsFileWriter.create( + file, + 1024 * 1024, + CompactionType.INNER_SEQ_COMPACTION, + EncryptDBUtils.getDefaultFirstEncryptParam())) { + assertTrue(writer.getEncryptParameter().isTdePageAead()); + write(writer); + } + assertEquals( + TestFileEncryptionParameterProvider.CREATED.get(), + TestFileEncryptionParameterProvider.CLOSED.get()); + } + + @Test + public void testMissingProviderFailsClosedWhenOpeningFile() throws Exception { + TestFileEncryptionParameterProvider.enable(); + File file = file("1-1-0-0.tsfile"); + try (TsFileIOWriter writer = EncryptDBUtils.createWriterForNewFile(file, "root.context")) { + write(writer); + } + TestFileEncryptionParameterProvider.disable(); + assertThrows( + Exception.class, + () -> { + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + reader.readMemChunk(metadata(file)); + } + }); + } + + @Test + public void testEncryptedSplitterAuthenticatesAndRewritesTransportChunks() throws Exception { + TestFileEncryptionParameterProvider.enable(); + File file = file("1-1-0-0.tsfile"); + try (TsFileIOWriter writer = EncryptDBUtils.createWriterForNewFile(file, "root.context")) { + write(writer); + } + List chunks = new ArrayList<>(); + new TsFileSplitter( + file, + data -> { + if (data instanceof ChunkData) { + chunks.addAll(((ChunkData) data).getChunks()); + } + return true; + }) + .splitTsFileByDataPartition(); + assertFalse(chunks.isEmpty()); + int row = 0; + for (Chunk chunk : chunks) { + assertTrue(chunk.getHeader().getChunkOrdinal() < 0); + try (ChunkReader reader = new ChunkReader(chunk)) { + while (reader.hasNextSatisfiedPage()) { + BatchData data = reader.nextPageData(); + while (data.hasCurrent()) { + assertEquals(row, data.currentTime()); + assertEquals(row * 10L, data.getLong()); + row++; + data.next(); + } + } + } + } + assertEquals(20, row); + } + + private static int trackedPages(QueryContext context) throws Exception { + Field field = context.getQueryChunkReaders().getClass().getDeclaredField("pages"); + field.setAccessible(true); + return ((Map) field.get(context.getQueryChunkReaders())).size(); + } + + private static ChunkMetadata metadata(File file) throws Exception { + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + return (ChunkMetadata) + reader.getAllTimeseriesMetadata(true).get(DEVICE).get(0).getChunkMetadataList().get(0); + } + } + + private File file(String name) throws Exception { + File file = new File(folder.getRoot(), "sequence/root.context/0/0/" + name); + Files.createDirectories(file.toPath().getParent()); + return file; + } + + private static void write(TsFileIOWriter writer) throws Exception { + writer.startChunkGroup(DEVICE); + ChunkWriterImpl chunk = + new ChunkWriterImpl(SCHEMA, EncryptUtils.getEncryptParameter(writer.getEncryptParameter())); + for (int i = 0; i < 20; i++) { + chunk.write(i, i * 10L); + if (i == 9) { + chunk.sealCurrentPage(); + } + } + chunk.writeToFileWriter(writer); + writer.endChunkGroup(); + writer.endFile(); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/TestFileEncryptionParameterProvider.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/TestFileEncryptionParameterProvider.java new file mode 100644 index 0000000000000..b15384c4f20b6 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/TestFileEncryptionParameterProvider.java @@ -0,0 +1,181 @@ +/* + * 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.utils; + +import org.apache.tsfile.encrypt.ChunkCryptoContext; +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.encrypt.EncryptionProviderRegistry; +import org.apache.tsfile.encrypt.IDecryptor; +import org.apache.tsfile.encrypt.IEncrypt; +import org.apache.tsfile.encrypt.IEncryptProvider; +import org.apache.tsfile.encrypt.IEncryptor; +import org.apache.tsfile.encrypt.PageCryptoContext; +import org.apache.tsfile.exception.encrypt.EncryptException; + +import java.io.File; +import java.nio.ByteBuffer; +import java.security.MessageDigest; +import java.util.Arrays; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Test-only framing and context checks. This is deliberately not a cryptographic implementation. + */ +public class TestFileEncryptionParameterProvider implements FileEncryptionParameterProvider { + public static final String ID = "iotdb-test-file-context"; + private static final ThreadLocal ENABLED = ThreadLocal.withInitial(() -> false); + public static final ThreadLocal RETURN_NULL = ThreadLocal.withInitial(() -> false); + private static final AtomicLong IDS = new AtomicLong(); + public static final AtomicInteger CREATED = new AtomicInteger(); + public static final AtomicInteger CLOSED = new AtomicInteger(); + + public static void enable() { + ENABLED.set(true); + CREATED.set(0); + CLOSED.set(0); + EncryptionProviderRegistry.registerProvider(new FramingProvider()); + } + + public static void disable() { + ENABLED.remove(); + RETURN_NULL.remove(); + EncryptionProviderRegistry.unregisterProvider(ID); + } + + @Override + public boolean requiresFileRewrite() { + return ENABLED.get(); + } + + @Override + public EncryptParameter createForNewFile(File file, String database, EncryptParameter legacy) { + if (RETURN_NULL.get()) { + return null; + } + return ENABLED.get() ? newParameter() : legacy; + } + + public static EncryptParameter newParameter() { + return EncryptParameter.pageAeadBuilder() + .providerId(ID) + .profileId("TEST_FRAMING") + .keyId("test") + .keyVersion("1") + .wrappedDataKey(new byte[] {1}) + .fileCryptoId(ByteBuffer.allocate(16).putLong(IDS.incrementAndGet()).array()) + .build(); + } + + private static byte[] digest(byte[] context, byte[] data) throws Exception { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + digest.update(context); + return digest.digest(data); + } + + public static class FramingProvider implements IEncryptProvider { + @Override + public String getProviderId() { + return ID; + } + + @Override + public IEncrypt create(EncryptParameter parameter) { + CREATED.incrementAndGet(); + IEncryptor encryptor = mock(IEncryptor.class); + IDecryptor decryptor = mock(IDecryptor.class); + when(encryptor.encryptPage(any(), anyInt(), anyInt(), any())) + .thenAnswer( + call -> { + byte[] data = call.getArgument(0); + int offset = call.getArgument(1); + int size = call.getArgument(2); + PageCryptoContext context = call.getArgument(3); + byte[] plaintext = Arrays.copyOfRange(data, offset, offset + size); + return ByteBuffer.allocate(size + 32) + .put(plaintext) + .put(digest(context.getAssociatedData(), plaintext)) + .array(); + }); + when(decryptor.decryptPage(any(), anyInt(), anyInt(), any())) + .thenAnswer( + call -> { + byte[] data = call.getArgument(0); + int offset = call.getArgument(1); + int size = call.getArgument(2); + PageCryptoContext context = call.getArgument(3); + byte[] plaintext = Arrays.copyOfRange(data, offset, offset + size - 32); + if (!Arrays.equals( + digest(context.getAssociatedData(), plaintext), + Arrays.copyOfRange(data, offset + size - 32, offset + size))) { + throw new EncryptException("test page context mismatch"); + } + return plaintext; + }); + return new IEncrypt() { + @Override + public IEncryptor getEncryptor() { + return encryptor; + } + + @Override + public IDecryptor getDecryptor() { + return decryptor; + } + + @Override + public int getPageBodyOverhead() { + return 32; + } + + @Override + public int getChunkAuthenticationSize() { + return 32; + } + + @Override + public byte[] authenticateChunk(ChunkCryptoContext context) { + try { + return digest(context.getAssociatedData(), new byte[0]); + } catch (Exception e) { + throw new AssertionError(e); + } + } + + @Override + public void verifyChunkAuthentication(byte[] proof, ChunkCryptoContext context) { + if (!Arrays.equals(proof, authenticateChunk(context))) { + throw new EncryptException("test chunk context mismatch"); + } + } + + @Override + public void close() { + CLOSED.incrementAndGet(); + } + }; + } + } +} diff --git a/iotdb-core/datanode/src/test/resources/META-INF/services/org.apache.iotdb.db.utils.FileEncryptionParameterProvider b/iotdb-core/datanode/src/test/resources/META-INF/services/org.apache.iotdb.db.utils.FileEncryptionParameterProvider new file mode 100644 index 0000000000000..48c7b765cb45e --- /dev/null +++ b/iotdb-core/datanode/src/test/resources/META-INF/services/org.apache.iotdb.db.utils.FileEncryptionParameterProvider @@ -0,0 +1,20 @@ +# +# 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. +# + +org.apache.iotdb.db.utils.TestFileEncryptionParameterProvider diff --git a/pom.xml b/pom.xml index c6bbd90cdda4b..3bcfc40929bec 100644 --- a/pom.xml +++ b/pom.xml @@ -145,7 +145,7 @@ 0.25.0 1.9 1.5.6-3 - 2.4.1-261009-SNAPSHOT + 2.5.1-261023-SNAPSHOT en From 45925d7f178a27b5eca826eb94526e877599651d Mon Sep 17 00:00:00 2001 From: HTHou Date: Fri, 23 Oct 2026 17:07:35 +0800 Subject: [PATCH 2/6] Limit compaction context adaptation and preserve repair initialization --- .../impl/FastCompactionPerformer.java | 8 ----- .../impl/ReadChunkCompactionPerformer.java | 22 +------------ ...dChunkAlignedSeriesCompactionExecutor.java | 6 ---- ...dChunkAlignedSeriesCompactionExecutor.java | 12 +------ .../SingleSeriesCompactionExecutor.java | 7 ---- .../writer/AbstractCompactionWriter.java | 17 ++++------ .../writer/FastInnerCompactionWriter.java | 23 ------------- .../apache/iotdb/db/utils/EncryptDBUtils.java | 6 ---- .../FileEncryptionParameterProvider.java | 9 ++++-- .../writer/AbstractCompactionWriterTest.java | 32 +++++++++++++++++++ .../db/utils/FileEncryptionContextTest.java | 11 ------- 11 files changed, 48 insertions(+), 105 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/FastCompactionPerformer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/FastCompactionPerformer.java index 7a86bf60701f6..4af319684885a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/FastCompactionPerformer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/FastCompactionPerformer.java @@ -49,7 +49,6 @@ import org.apache.iotdb.db.storageengine.dataregion.compaction.selector.estimator.FastCrossSpaceCompactionEstimator; import org.apache.iotdb.db.storageengine.dataregion.modification.ModEntry; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; -import org.apache.iotdb.db.utils.EncryptDBUtils; import org.apache.iotdb.db.utils.datastructure.PatternTreeMapFactory; import org.apache.tsfile.common.conf.TSFileDescriptor; @@ -158,13 +157,6 @@ public FastCompactionPerformer(boolean isCrossCompaction, EncryptParameter encry @Override public void perform() throws Exception { - if (EncryptDBUtils.requiresFileRewrite()) { - ReadPointCompactionPerformer rewriting = - new ReadPointCompactionPerformer(seqFiles, unseqFiles, targetFiles, encryptParameter); - rewriting.setSummary(subTaskSummary); - rewriting.perform(); - return; - } this.subTaskSummary.setTemporalFileNum(targetFiles.size()); try (MultiTsFileDeviceIterator deviceIterator = new MultiTsFileDeviceIterator(seqFiles, unseqFiles, readerCacheMap); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadChunkCompactionPerformer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadChunkCompactionPerformer.java index 63cb9916838e8..82bdbd3e6496d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadChunkCompactionPerformer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/performer/impl/ReadChunkCompactionPerformer.java @@ -137,26 +137,6 @@ public void perform() InterruptedException, StorageEngineException, PageException { - if (EncryptDBUtils.requiresFileRewrite()) { - ReadPointCompactionPerformer rewriting = - new ReadPointCompactionPerformer( - seqFiles, Collections.emptyList(), targetResources, firstEncryptParameter); - rewriting.setSummary(summary); - try { - rewriting.perform(); - } catch (IOException - | MetadataException - | InterruptedException - | StorageEngineException - | PageException e) { - throw e; - } catch (RuntimeException e) { - throw e; - } catch (Exception e) { - throw new IOException(e); - } - return; - } try (MultiTsFileDeviceIterator deviceIterator = new MultiTsFileDeviceIterator(seqFiles)) { schema = CompactionTableSchemaCollector.collectSchema( @@ -233,7 +213,7 @@ private void rollCompactionFileWriter() throws IOException { private void useNewWriter() throws IOException { currentWriter = - CompactionTsFileWriter.create( + new CompactionTsFileWriter( targetResources.get(currentTargetFileIndex).getTsFile(), memoryBudgetForFileWriter, CompactionType.INNER_SEQ_COMPACTION, diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/batch/BatchedReadChunkAlignedSeriesCompactionExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/batch/BatchedReadChunkAlignedSeriesCompactionExecutor.java index f50b5d168b513..c2bfe18e12ca1 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/batch/BatchedReadChunkAlignedSeriesCompactionExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/batch/BatchedReadChunkAlignedSeriesCompactionExecutor.java @@ -430,9 +430,6 @@ public FollowingBatchReadChunkAlignedSeriesCompactionFlushController() { @Override public boolean canCompactCurrentChunkByDirectlyFlush( ChunkLoader timeChunk, List valueChunks) throws IOException { - if (requiresPageRewrite()) { - return false; - } CompactChunkPlan compactChunkPlan = batchCompactionPlan.getCompactChunkPlan(currentCompactChunk); boolean isCurrentChunkCompactedByDirectlyFlush = @@ -455,9 +452,6 @@ protected boolean canFlushCurrentChunkWriter() { @Override protected boolean canCompactCurrentPageByDirectlyFlush( PageLoader timePage, List valuePages) { - if (requiresPageRewrite()) { - return false; - } int currentPage = ((FollowingBatchCompactionAlignedChunkWriter) chunkWriter).getCurrentPage(); for (int i = 0; i < valuePages.size(); i++) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/ReadChunkAlignedSeriesCompactionExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/ReadChunkAlignedSeriesCompactionExecutor.java index 9b8fe593842dc..90181e111a2ca 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/ReadChunkAlignedSeriesCompactionExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/ReadChunkAlignedSeriesCompactionExecutor.java @@ -38,7 +38,6 @@ import org.apache.iotdb.db.storageengine.dataregion.compaction.io.CompactionTsFileReader; import org.apache.iotdb.db.storageengine.dataregion.compaction.io.CompactionTsFileWriter; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; -import org.apache.iotdb.db.utils.EncryptDBUtils; import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.common.conf.TSFileDescriptor; @@ -572,9 +571,7 @@ public ReadChunkAlignedSeriesCompactionFlushController(int compactionFileLevel) public boolean canCompactCurrentChunkByDirectlyFlush( ChunkLoader timeChunk, List valueChunks) throws IOException { - return !requiresPageRewrite() - && canFlushCurrentChunkWriter() - && canFlushChunk(timeChunk, valueChunks); + return canFlushCurrentChunkWriter() && canFlushChunk(timeChunk, valueChunks); } protected boolean canFlushCurrentChunkWriter() { @@ -615,9 +612,6 @@ private boolean canFlushChunk(ChunkLoader timeChunk, List valueChun protected boolean canCompactCurrentPageByDirectlyFlush( PageLoader timePage, List valuePages) { - if (requiresPageRewrite()) { - return false; - } boolean isHighLevelCompaction = compactionTargetFileLevel > largeFileLevelSeparator; if (isHighLevelCompaction) { return canFlushPage(timePage, valuePages); @@ -630,10 +624,6 @@ private boolean canSealCurrentPageWriter() { return chunkWriter.checkIsUnsealedPageOverThreshold(targetPageSize, targetPagePointNum, true); } - protected boolean requiresPageRewrite() { - return !EncryptDBUtils.canCopyChunk(null, writer.getEncryptParameter()); - } - private boolean canFlushPage(PageLoader timePage, List valuePages) { if (timePage.getModifiedStatus() == ModifiedStatus.PARTIAL_DELETED) { return false; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/SingleSeriesCompactionExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/SingleSeriesCompactionExecutor.java index 62fa4e00d31d9..44f8427a6f905 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/SingleSeriesCompactionExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/executor/readchunk/SingleSeriesCompactionExecutor.java @@ -25,7 +25,6 @@ import org.apache.iotdb.db.storageengine.dataregion.compaction.execute.task.CompactionTaskSummary; import org.apache.iotdb.db.storageengine.dataregion.compaction.io.CompactionTsFileWriter; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; -import org.apache.iotdb.db.utils.EncryptDBUtils; import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.file.MetaMarker; @@ -156,12 +155,6 @@ public void execute() throws IOException { constructChunkWriterFromReadChunk(currentChunk); } - if (!EncryptDBUtils.canCopyChunk( - currentChunk.getEncryptParam(), fileWriter.getEncryptParameter())) { - processModifiedChunk(currentChunk); - continue; - } - // if this chunk is modified, deserialize it into points if (chunkMetadata.getDeleteIntervalList() != null) { processModifiedChunk(currentChunk); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriter.java index 02ad8a6f0de9b..6d075f9d93814 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriter.java @@ -143,23 +143,20 @@ public ModEntry getTTLLowerBoundForCurrentDevice() { public void startMeasurement(String measurement, IChunkWriter chunkWriter, int subTaskId) { chunkWriterFactories[subTaskId] = null; - initializeChunkWriter(measurement, chunkWriter, subTaskId); + resetChunkWriterStatistics(subTaskId); + lastTimeSet[subTaskId] = false; + chunkWriters[subTaskId] = chunkWriter; + measurementId[subTaskId] = measurement; + hasVariableLengthTypeArray[subTaskId] = containsVariableLengthType(chunkWriter); } public void startMeasurement( String measurement, Function chunkWriterFactory, int subTaskId) { + // Preserve measurement initialization performed by specialized writers, including repair. + startMeasurement(measurement, chunkWriterFactory.apply(getEncryptParameter()), subTaskId); chunkWriterFactories[subTaskId] = chunkWriterFactory; - initializeChunkWriter(measurement, chunkWriterFactory.apply(getEncryptParameter()), subTaskId); - } - - private void initializeChunkWriter(String measurement, IChunkWriter chunkWriter, int subTaskId) { - resetChunkWriterStatistics(subTaskId); - lastTimeSet[subTaskId] = false; - chunkWriters[subTaskId] = chunkWriter; - measurementId[subTaskId] = measurement; - hasVariableLengthTypeArray[subTaskId] = containsVariableLengthType(chunkWriter); } protected void switchChunkWriterEncryptionParameter( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/FastInnerCompactionWriter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/FastInnerCompactionWriter.java index 2615cfea74273..ccac7b4c63516 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/FastInnerCompactionWriter.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/FastInnerCompactionWriter.java @@ -78,9 +78,6 @@ public void write(TsBlock tsBlock, int subTaskId) throws IOException { @Override public boolean flushNonAlignedChunk(Chunk chunk, ChunkMetadata chunkMetadata, int subTaskId) throws IOException { - if (requiresPageRewrite()) { - return false; - } checkPreviousTimestamp(chunkMetadata.getStartTime(), subTaskId); if (chunkPointNumArray[subTaskId] != 0 && chunkWriters[subTaskId].checkIsChunkSizeOverThreshold( @@ -110,9 +107,6 @@ public boolean flushNonAlignedChunk(Chunk chunk, ChunkMetadata chunkMetadata, in @Override public boolean flushAlignedChunk(ChunkMetadataElement chunkMetadataElement, int subTaskId) throws IOException { - if (requiresPageRewrite()) { - return false; - } AbstractAlignedChunkMetadata alignedChunkMetadata = (AbstractAlignedChunkMetadata) chunkMetadataElement.chunkMetadata; IChunkMetadata timeChunkMetadata = alignedChunkMetadata.getTimeChunkMetadata(); @@ -146,9 +140,6 @@ public boolean flushBatchedValueChunk( int subTaskId, AbstractCompactionFlushController flushController) throws IOException { - if (requiresPageRewrite()) { - return false; - } AbstractAlignedChunkMetadata alignedChunkMetadata = (AbstractAlignedChunkMetadata) chunkMetadataElement.chunkMetadata; IChunkMetadata timeChunkMetadata = alignedChunkMetadata.getTimeChunkMetadata(); @@ -185,9 +176,6 @@ public boolean flushBatchedValueChunk( */ public boolean flushAlignedPage(AlignedPageElement alignedPageElement, int subTaskId) throws IOException, PageException { - if (requiresPageRewrite()) { - return false; - } PageHeader timePageHeader = alignedPageElement.getTimePageHeader(); List valuePageHeaders = alignedPageElement.getValuePageHeaders(); ByteBuffer compressedTimePageData = alignedPageElement.getTimePageData(); @@ -224,9 +212,6 @@ public boolean flushBatchedValuePage( int subTaskId, AbstractCompactionFlushController flushController) throws PageException, IOException { - if (requiresPageRewrite()) { - return false; - } PageHeader timePageHeader = alignedPageElement.getTimePageHeader(); List valuePageHeaders = alignedPageElement.getValuePageHeaders(); ByteBuffer compressedTimePageData = alignedPageElement.getTimePageData(); @@ -270,9 +255,6 @@ public boolean flushBatchedValuePage( public boolean flushNonAlignedPage( ByteBuffer compressedPageData, PageHeader pageHeader, int subTaskId) throws IOException, PageException { - if (requiresPageRewrite()) { - return false; - } checkPreviousTimestamp(pageHeader.getStartTime(), subTaskId); boolean isUnsealedPageOverThreshold = chunkWriters[subTaskId].checkIsUnsealedPageOverThreshold( @@ -332,9 +314,4 @@ private boolean checkIsPageLargeEnough(PageHeader pageHeader) { return pageHeader.getStatistics().getCount() >= pagePointNumLowerBoundInCompaction || pageHeader.getSerializedPageSize() >= pageSizeLowerBoundInCompaction; } - - private boolean requiresPageRewrite() { - return fileWriter != null - && !EncryptDBUtils.canCopyChunk(null, fileWriter.getEncryptParameter()); - } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/EncryptDBUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/EncryptDBUtils.java index 1c55154d40977..a8e24b90e4fc0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/EncryptDBUtils.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/EncryptDBUtils.java @@ -83,12 +83,6 @@ public static boolean requiresFileRewrite() { return ProviderHolder.PROVIDER.requiresFileRewrite(); } - public static boolean canCopyChunk(EncryptParameter source, EncryptParameter target) { - return !requiresFileRewrite() - && (source == null || !source.isTdePageAead()) - && (target == null || !target.isTdePageAead()); - } - public static RestorableTsFileIOWriter createRestorableWriterForNewFile( File file, String database) throws IOException { return createRestorableWriterForNewFile( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/FileEncryptionParameterProvider.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/FileEncryptionParameterProvider.java index ee56f4e646730..a8a20bdfe810a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/FileEncryptionParameterProvider.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/FileEncryptionParameterProvider.java @@ -22,7 +22,12 @@ import java.io.File; -/** Optional file-level policy, discovered once through {@link java.util.ServiceLoader}. */ +/** + * Optional file-level policy, discovered once through {@link java.util.ServiceLoader}. + * + *

File-scoped compaction contexts are supported only by ReadPoint compaction. An extension using + * them must select ReadPoint; this SPI does not adapt or redirect Fast/ReadChunk compaction. + */ @FunctionalInterface public interface FileEncryptionParameterProvider { /** @@ -32,7 +37,7 @@ public interface FileEncryptionParameterProvider { */ EncryptParameter createForNewFile(File file, String database, EncryptParameter legacyParameter); - /** Whether new target files require rewriting rather than copying source chunks verbatim. */ + /** Whether Load must rewrite input files through the regular tablet write pipeline. */ default boolean requiresFileRewrite() { return false; } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriterTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriterTest.java index a2ca64fc59a3d..7e178a0cc2feb 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriterTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/compaction/execute/utils/writer/AbstractCompactionWriterTest.java @@ -95,6 +95,38 @@ public void testChunkWriterFactoryCanSwitchFileEncryptionParameter() { Assert.assertSame(second, createdWith.get(1)); } + @Test + public void testChunkWriterFactoryPreservesSubclassMeasurementInitialization() { + List initializedWriters = new ArrayList<>(); + TestCompactionWriter compactionWriter = + new TestCompactionWriter() { + @Override + public void startMeasurement( + String measurement, IChunkWriter chunkWriter, int subTaskId) { + initializedWriters.add(chunkWriter); + super.startMeasurement(measurement, chunkWriter, subTaskId); + } + }; + CountingChunkWriter first = new CountingChunkWriter(); + List createdWith = new ArrayList<>(); + compactionWriter.startMeasurement( + "s1", + parameter -> { + createdWith.add(parameter); + return createdWith.size() == 1 ? first : new CountingChunkWriter(); + }, + SUB_TASK_ID); + + Assert.assertEquals(1, initializedWriters.size()); + Assert.assertSame(first, initializedWriters.get(0)); + + EncryptParameter next = new EncryptParameter("UNENCRYPTED", null); + compactionWriter.switchEncryptionParameter(next); + Assert.assertEquals(2, createdWith.size()); + Assert.assertSame(next, createdWith.get(1)); + Assert.assertEquals(1, initializedWriters.size()); + } + @Test public void testEstimateWrittenPointTotalSize() throws Exception { TestCompactionWriter compactionWriter = new TestCompactionWriter(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/FileEncryptionContextTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/FileEncryptionContextTest.java index 3668b92ec0a22..d750bf6d80862 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/FileEncryptionContextTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/utils/FileEncryptionContextTest.java @@ -200,17 +200,6 @@ public void testMemTableKeepsExplicitContextAcrossDatabaseAssignmentAndCopy() { } } - @Test - public void testDirectChunkCopyRejectsFileBoundContexts() { - EncryptParameter plain = EncryptDBUtils.getDefaultFirstEncryptParam(); - try (EncryptParameter file = TestFileEncryptionParameterProvider.newParameter()) { - assertTrue(EncryptDBUtils.canCopyChunk(plain, plain)); - assertFalse(EncryptDBUtils.canCopyChunk(file, plain)); - assertFalse(EncryptDBUtils.canCopyChunk(plain, file)); - assertFalse(EncryptDBUtils.canCopyChunk(file, file)); - } - } - private static QueryContext externalContext() { return new QueryContext(false, false) { From ddb352570ec3af745ce9a24d3d2a568a661b738c Mon Sep 17 00:00:00 2001 From: HTHou Date: Mon, 26 Oct 2026 16:04:40 +0800 Subject: [PATCH 3/6] [Storage] Fix Pipe retries and WAL recovery file contexts Retain independent aligned chunks and buffer views for each reader construction attempt so failed attempts do not invalidate pending groups. Release pending value chunk references after successful construction. Propagate the recovered writer context to initial and snapshot recovery memtables. Cover constructor failure retries, partial retain cleanup, and aligned/non-aligned WAL insert and snapshot recovery. --- .../scan/TsFileInsertionEventScanParser.java | 50 +++- .../file/UnsealedTsFileRecoverPerformer.java | 8 + ...leInsertionEventScanParserContextTest.java | 223 ++++++++++++++++ ...tedUnsealedTsFileRecoverPerformerTest.java | 247 ++++++++++++++++++ 4 files changed, 521 insertions(+), 7 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParserContextTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/EncryptedUnsealedTsFileRecoverPerformerTest.java 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 f3748abc9ad37..b3c5401eeb42f 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 @@ -913,17 +913,20 @@ private boolean useNextPendingAlignedChunk(final byte marker) throws IOException : pageEstimatedMemoryUsageInBytesList.get(0); resizePageDataMemoryIfNeeded(maxPageEstimatedMemoryUsageInBytes); final IChunkReader nextChunkReader = - nextIsMultiPage - ? new MemoryControlledChunkReader( - new AlignedChunkReader( - retainChunk(timeChunk), pendingAlignedChunkGroup.valueChunkList, filter), - pageEstimatedMemoryUsageInBytesList) - : new AlignedSinglePageWholeChunkReader( - retainChunk(timeChunk), pendingAlignedChunkGroup.valueChunkList, null); + createAlignedChunkReader( + timeChunk, + pendingAlignedChunkGroup.valueChunkList, + nextIsMultiPage, + pageEstimatedMemoryUsageInBytesList); // Publish the state transition only after all memory reservations and reader construction // succeed. Otherwise a retry would lose the already-read aligned chunk group. pendingAlignedChunkGroups.remove(0); + for (Chunk valueChunk : pendingAlignedChunkGroup.valueChunkList) { + if (valueChunk != null) { + valueChunk.close(); + } + } pendingAlignedChunkSize = Math.max(0, pendingAlignedChunkSize - pendingAlignedChunkGroup.chunkSize); currentMeasurements.clear(); @@ -945,6 +948,39 @@ private boolean useNextPendingAlignedChunk(final byte marker) throws IOException return false; } + IChunkReader createAlignedChunkReader( + Chunk timeChunk, + List valueChunks, + boolean multiPage, + List pageEstimatedMemoryUsageInBytesList) + throws IOException { + // Pending chunks must survive a failed constructor so the same group can be retried. + Chunk retainedTimeChunk = null; + List retainedValueChunks = new ArrayList<>(valueChunks.size()); + try { + retainedTimeChunk = retainChunk(timeChunk); + for (Chunk valueChunk : valueChunks) { + retainedValueChunks.add(valueChunk == null ? null : retainChunk(valueChunk)); + } + return multiPage + ? new MemoryControlledChunkReader( + new AlignedChunkReader(retainedTimeChunk, retainedValueChunks, filter), + pageEstimatedMemoryUsageInBytesList) + : new AlignedSinglePageWholeChunkReader(retainedTimeChunk, retainedValueChunks, null); + } catch (IOException | RuntimeException | Error failure) { + // Chunk.close() is idempotent, including when constructor cleanup has already run. + if (retainedTimeChunk != null) { + retainedTimeChunk.close(); + } + for (Chunk valueChunk : retainedValueChunks) { + if (valueChunk != null) { + valueChunk.close(); + } + } + throw failure; + } + } + private boolean shouldReturnPendingAlignedChunkBeforeCaching( final CachedAlignedValueChunk valueChunk) throws IOException { validateAlignedValueChunkTimeIndex(valueChunk.timeChunkIndex); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/UnsealedTsFileRecoverPerformer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/UnsealedTsFileRecoverPerformer.java index d8369e2c406dd..cdfdc13dba669 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/UnsealedTsFileRecoverPerformer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/UnsealedTsFileRecoverPerformer.java @@ -47,6 +47,7 @@ import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALRecoverException; import org.apache.iotdb.db.storageengine.dataregion.wal.utils.listener.WALRecoverListener; +import org.apache.tsfile.encrypt.EncryptUtils; import org.apache.tsfile.file.metadata.ChunkMetadata; import org.apache.tsfile.file.metadata.IDeviceID; import org.apache.tsfile.read.common.TimeRange; @@ -112,6 +113,9 @@ public void startRecovery() throws DataRegionException, IOException { super.recoverWithWriter(); if (hasCrashed()) { + walRedoer + .getRecoveryMemTable() + .setEncryptParameter(EncryptUtils.getEncryptParameter(writer.getEncryptParameter())); // tsfile has crashed due to failure, // the last ChunkGroup may contain the same data as the WALs, // so the time map must be updated first to avoid duplicated insertion @@ -213,6 +217,10 @@ public void redoLog(WALEntry walEntry) { } // update memtable's database and dataRegionId memTable.setDatabaseAndDataRegionId(databaseName, dataRegionId); + if (writer != null) { + memTable.setEncryptParameter( + EncryptUtils.getEncryptParameter(writer.getEncryptParameter())); + } break; case INSERT_ROW_NODE: case INSERT_TABLET_NODE: diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParserContextTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParserContextTest.java new file mode 100644 index 0000000000000..7481329196c00 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/event/common/tsfile/parser/scan/TsFileInsertionEventScanParserContextTest.java @@ -0,0 +1,223 @@ +/* + * 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.scan; + +import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException; +import org.apache.iotdb.db.utils.TestFileEncryptionParameterProvider; + +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.header.ChunkHeader; +import org.apache.tsfile.file.metadata.AbstractAlignedChunkMetadata; +import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.apache.tsfile.read.common.BatchData; +import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.read.reader.IChunkReader; +import org.apache.tsfile.write.chunk.AlignedChunkWriterImpl; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.apache.tsfile.write.writer.TsFileIOWriter; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.mockito.stubbing.Answer; + +import java.io.File; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; + +public class TsFileInsertionEventScanParserContextTest { + @Rule public final TemporaryFolder folder = new TemporaryFolder(); + private static final IDeviceID DEVICE = + IDeviceID.Factory.DEFAULT_FACTORY.create("root.context.d1"); + private static final MeasurementSchema SCHEMA = + new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.PLAIN, CompressionType.UNCOMPRESSED); + + @Before + public void setUp() { + TestFileEncryptionParameterProvider.enable(); + } + + @After + public void tearDown() { + try { + assertEquals( + TestFileEncryptionParameterProvider.CREATED.get(), + TestFileEncryptionParameterProvider.CLOSED.get()); + } finally { + TestFileEncryptionParameterProvider.disable(); + } + } + + @Test + public void testSinglePageConstructorFailureCanRetry() throws Exception { + assertConstructorFailureCanRetry(false); + } + + @Test + public void testMultiPageConstructorFailureCanRetry() throws Exception { + assertConstructorFailureCanRetry(true); + } + + private void assertConstructorFailureCanRetry(boolean multiPage) throws Exception { + TsFileInsertionEventScanParser parser = + mock(TsFileInsertionEventScanParser.class, CALLS_REAL_METHODS); + List chunks = readChunks(multiPage); + try { + Chunk time = chunks.get(0); + Chunk originalValue = chunks.get(1); + ChunkHeader valueHeader = spy(originalValue.getHeader()); + Chunk value = + new Chunk( + valueHeader, + originalValue.getData().duplicate(), + originalValue.getDeleteIntervalList(), + originalValue.getChunkStatistic(), + originalValue.getEncryptParam()); + chunks.set(1, value); + originalValue.close(); + AtomicBoolean fail = new AtomicBoolean(true); + PipeRuntimeOutOfMemoryCriticalException expected = + new PipeRuntimeOutOfMemoryCriticalException("test constructor memory pressure"); + Answer failOnce = + call -> { + if (fail.getAndSet(false)) { + throw expected; + } + return call.callRealMethod(); + }; + if (multiPage) { + doAnswer(failOnce).when(valueHeader).calculateDecoderForNonTimeChunk(); + } else { + doAnswer(failOnce).when(valueHeader).getEncodingType(); + } + assertSame( + expected, + assertThrows( + PipeRuntimeOutOfMemoryCriticalException.class, + () -> { + try (IChunkReader ignored = + parser.createAlignedChunkReader( + time, + Collections.singletonList(value), + multiPage, + Collections.emptyList())) { + // Unexpected success must not leak a reader while reporting the assertion. + } + })); + assertFalse(time.getEncryptParam().isDestroyed()); + assertFalse(value.getEncryptParam().isDestroyed()); + assertEquals(0, time.getData().position()); + assertEquals(0, value.getData().position()); + + try (IChunkReader reader = + parser.createAlignedChunkReader( + time, Collections.singletonList(value), multiPage, Collections.emptyList())) { + // Once construction succeeds, the pending group can release its original value chunks. + value.close(); + int row = 0; + while (reader.hasNextSatisfiedPage()) { + BatchData data = reader.nextPageData(); + while (data.hasCurrent()) { + assertEquals(row, data.currentTime()); + assertEquals(row * 10L, data.getVector()[0].getLong()); + row++; + data.next(); + } + } + assertEquals(20, row); + } + assertFalse(time.getEncryptParam().isDestroyed()); + } finally { + chunks.forEach(Chunk::close); + } + } + + @Test + public void testPartialRetainFailureReleasesAttemptOnly() throws Exception { + TsFileInsertionEventScanParser parser = + mock(TsFileInsertionEventScanParser.class, CALLS_REAL_METHODS); + List chunks = readChunks(false); + try { + Chunk broken = spy(chunks.get(1)); + IllegalStateException expected = new IllegalStateException("test retaining failure"); + doThrow(expected).when(broken).getData(); + assertSame( + expected, + assertThrows( + IllegalStateException.class, + () -> + parser.createAlignedChunkReader( + chunks.get(0), + Arrays.asList(chunks.get(1), broken), + false, + Collections.emptyList()))); + for (Chunk chunk : chunks) { + assertFalse(chunk.getEncryptParam().isDestroyed()); + } + } finally { + chunks.forEach(Chunk::close); + } + } + + private List readChunks(boolean multiPage) throws Exception { + File file = new File(folder.getRoot(), "aligned.tsfile"); + try (EncryptParameter parameter = TestFileEncryptionParameterProvider.newParameter(); + TsFileIOWriter writer = new TsFileIOWriter(file, parameter)) { + writer.startChunkGroup(DEVICE); + AlignedChunkWriterImpl chunk = + new AlignedChunkWriterImpl( + Collections.singletonList(SCHEMA), writer.getEncryptParameter()); + for (int i = 0; i < 20; i++) { + chunk.getTimeChunkWriter().write(i); + chunk.getValueChunkWriterByIndex(0).write(i, i * 10L, false); + if (multiPage && i == 9) { + chunk.sealCurrentPage(); + } + } + chunk.writeToFileWriter(writer); + writer.endChunkGroup(); + writer.endFile(); + } + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + AbstractAlignedChunkMetadata metadata = reader.getAlignedChunkMetadata(DEVICE, false).get(0); + return Arrays.asList( + reader.readMemChunk((ChunkMetadata) metadata.getTimeChunkMetadata()), + reader.readMemChunk((ChunkMetadata) metadata.getValueChunkMetadataList().get(0))); + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/EncryptedUnsealedTsFileRecoverPerformerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/EncryptedUnsealedTsFileRecoverPerformerTest.java new file mode 100644 index 0000000000000..6f0f2fd14564b --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/EncryptedUnsealedTsFileRecoverPerformerTest.java @@ -0,0 +1,247 @@ +/* + * 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.dataregion.wal.recover.file; + +import org.apache.iotdb.commons.path.PartialPath; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; +import org.apache.iotdb.db.storageengine.dataregion.memtable.PrimitiveMemTable; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntry; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALInfoEntry; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.TsFileUtilsForRecoverTest; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALByteBufferForTest; +import org.apache.iotdb.db.utils.EnvironmentUtils; +import org.apache.iotdb.db.utils.TestFileEncryptionParameterProvider; + +import org.apache.tsfile.encrypt.EncryptParameter; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.metadata.AbstractAlignedChunkMetadata; +import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.enums.CompressionType; +import org.apache.tsfile.file.metadata.enums.TSEncoding; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.apache.tsfile.read.common.BatchData; +import org.apache.tsfile.read.common.Chunk; +import org.apache.tsfile.read.common.Path; +import org.apache.tsfile.read.reader.IChunkReader; +import org.apache.tsfile.read.reader.chunk.AlignedChunkReader; +import org.apache.tsfile.read.reader.chunk.ChunkReader; +import org.apache.tsfile.write.chunk.AlignedChunkWriterImpl; +import org.apache.tsfile.write.chunk.ChunkWriterImpl; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.apache.tsfile.write.writer.TsFileIOWriter; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import java.io.ByteArrayInputStream; +import java.io.DataInputStream; +import java.io.File; +import java.io.FileOutputStream; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.Files; +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class EncryptedUnsealedTsFileRecoverPerformerTest { + private static final String DATABASE = "root.context_recovery"; + private static final IDeviceID DEVICE = + IDeviceID.Factory.DEFAULT_FACTORY.create(DATABASE + ".d1"); + private static final MeasurementSchema SCHEMA = + new MeasurementSchema("s1", TSDataType.INT64, TSEncoding.PLAIN, CompressionType.UNCOMPRESSED); + private static final String FILE_NAME = + TsFileUtilsForRecoverTest.getTestTsFilePath(DATABASE, 0, 0, 1); + private TsFileResource resource; + + @Before + public void setUp() throws Exception { + EnvironmentUtils.cleanDir(new File(FILE_NAME).getParent()); + EnvironmentUtils.envSetUp(); + TestFileEncryptionParameterProvider.enable(); + } + + @After + public void tearDown() throws Exception { + try { + if (resource != null) { + resource.close(); + } + EnvironmentUtils.cleanDir(new File(FILE_NAME).getParent()); + EnvironmentUtils.cleanEnv(); + } finally { + TestFileEncryptionParameterProvider.disable(); + } + } + + @Test + public void testNonAlignedInsertRecovery() throws Exception { + assertRecovery(false, false); + } + + @Test + public void testAlignedInsertRecovery() throws Exception { + assertRecovery(true, false); + } + + @Test + public void testNonAlignedSnapshotRecovery() throws Exception { + assertRecovery(false, true); + } + + @Test + public void testAlignedSnapshotRecovery() throws Exception { + assertRecovery(true, true); + } + + private void assertRecovery(boolean aligned, boolean snapshot) throws Exception { + File file = new File(FILE_NAME); + Files.createDirectories(file.getParentFile().toPath()); + byte[] originalFileId; + long dataEnd; + try (EncryptParameter parameter = TestFileEncryptionParameterProvider.newParameter(); + TsFileIOWriter writer = new TsFileIOWriter(file, parameter)) { + originalFileId = writer.getEncryptParameter().getFileCryptoId(); + writer.startChunkGroup(DEVICE); + if (aligned) { + AlignedChunkWriterImpl chunk = + new AlignedChunkWriterImpl( + Collections.singletonList(SCHEMA), writer.getEncryptParameter()); + chunk.getTimeChunkWriter().write(1); + chunk.getValueChunkWriterByIndex(0).write(1, 10L, false); + chunk.writeToFileWriter(writer); + } else { + ChunkWriterImpl chunk = new ChunkWriterImpl(SCHEMA, writer.getEncryptParameter()); + chunk.write(1, 10L); + chunk.writeToFileWriter(writer); + } + writer.endChunkGroup(); + writer.writePlanIndices(); + dataEnd = writer.getPos(); + writer.endFile(); + } + try (FileChannel channel = new FileOutputStream(file, true).getChannel()) { + channel.truncate(dataEnd); + } + + InsertRowNode row = + new InsertRowNode( + new PlanNodeId(""), + new PartialPath(DEVICE), + aligned, + new String[] {"s1"}, + new TSDataType[] {TSDataType.INT64}, + 2, + new Object[] {20L}, + false); + row.setMeasurementSchemas(new MeasurementSchema[] {SCHEMA}); + WALEntry entry; + if (snapshot) { + PrimitiveMemTable table = new PrimitiveMemTable(DATABASE, "0"); + if (aligned) { + table.insertAlignedRow(row); + } else { + table.insert(row); + } + entry = roundTrip(new WALInfoEntry(1, table)); + assertFalse(((PrimitiveMemTable) entry.getValue()).getEncryptParameter().isTdePageAead()); + } else { + entry = roundTrip(new WALInfoEntry(1, row)); + } + + resource = new TsFileResource(file); + try (UnsealedTsFileRecoverPerformer performer = + new UnsealedTsFileRecoverPerformer(resource, true, ignored -> {})) { + performer.startRecovery(); + assertTrue(performer.hasCrashed()); + assertArrayEquals( + originalFileId, performer.getWriter().getEncryptParameter().getFileCryptoId()); + performer.redoLog(entry); + if (snapshot) { + assertTrue( + ((PrimitiveMemTable) entry.getValue()) + .getEncryptParameter() + .sharesPageAeadFileContext(performer.getWriter().getEncryptParameter())); + } + performer.endRecovery(); + } + + Set ordinals = new HashSet<>(); + int rows = 0; + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + assertArrayEquals(originalFileId, reader.getEncryptParam().getFileCryptoId()); + if (aligned) { + for (AbstractAlignedChunkMetadata metadata : + reader.getAlignedChunkMetadata(DEVICE, false)) { + try (Chunk time = reader.readMemChunk((ChunkMetadata) metadata.getTimeChunkMetadata()); + Chunk value = + reader.readMemChunk((ChunkMetadata) metadata.getValueChunkMetadataList().get(0)); + IChunkReader chunkReader = + new AlignedChunkReader(time, Collections.singletonList(value), null)) { + assertTrue(ordinals.add(time.getHeader().getChunkOrdinal())); + assertTrue(ordinals.add(value.getHeader().getChunkOrdinal())); + rows = assertRows(chunkReader, rows, true); + } + } + } else { + for (ChunkMetadata metadata : reader.getChunkMetadataList(new Path(DEVICE, "s1", true))) { + try (Chunk chunk = reader.readMemChunk(metadata); + IChunkReader chunkReader = new ChunkReader(chunk)) { + assertTrue(ordinals.add(chunk.getHeader().getChunkOrdinal())); + rows = assertRows(chunkReader, rows, false); + } + } + } + } + assertEquals(2, rows); + assertEquals(aligned ? 4 : 2, ordinals.size()); + assertEquals(2L, (long) resource.getEndTime(DEVICE).get()); + } + + private static WALEntry roundTrip(WALInfoEntry entry) throws Exception { + ByteBuffer buffer = ByteBuffer.allocate(entry.serializedSize()); + entry.serialize(new WALByteBufferForTest(buffer)); + try (DataInputStream input = + new DataInputStream(new ByteArrayInputStream(buffer.array(), 0, buffer.position()))) { + return WALEntry.deserialize(input); + } + } + + private static int assertRows(IChunkReader reader, int rows, boolean aligned) throws Exception { + while (reader.hasNextSatisfiedPage()) { + BatchData data = reader.nextPageData(); + while (data.hasCurrent()) { + rows++; + assertEquals(rows, data.currentTime()); + assertEquals(rows * 10L, aligned ? data.getVector()[0].getLong() : data.getLong()); + data.next(); + } + } + return rows; + } +} From e6e0ff79bf2a40446972909fad6b5a5c0b6e1bdb Mon Sep 17 00:00:00 2001 From: HTHou Date: Mon, 26 Oct 2026 17:14:42 +0800 Subject: [PATCH 4/6] [Storage] Preserve encryption policy when recovering header-only TsFiles --- .../file/AbstractTsFileRecoverPerformer.java | 55 ++++- ...tedUnsealedTsFileRecoverPerformerTest.java | 192 ++++++++++++++++++ 2 files changed, 244 insertions(+), 3 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/AbstractTsFileRecoverPerformer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/AbstractTsFileRecoverPerformer.java index 5b470ee1c2bb9..804b3377ddba6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/AbstractTsFileRecoverPerformer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/AbstractTsFileRecoverPerformer.java @@ -19,6 +19,7 @@ package org.apache.iotdb.db.storageengine.dataregion.wal.recover.file; +import org.apache.iotdb.commons.utils.IOUtils; import org.apache.iotdb.db.exception.DataRegionException; import org.apache.iotdb.db.i18n.StorageEngineMessages; import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; @@ -27,6 +28,7 @@ import org.apache.iotdb.db.utils.EncryptDBUtils; import org.apache.tsfile.common.conf.TSFileConfig; +import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.exception.NotCompatibleTsFileException; import org.apache.tsfile.external.commons.io.FileUtils; import org.apache.tsfile.read.TsFileSequenceReader; @@ -38,6 +40,11 @@ import java.io.Closeable; import java.io.File; import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.StandardOpenOption; +import java.util.Arrays; /** This class is used to help recover TsFile. */ public abstract class AbstractTsFileRecoverPerformer implements Closeable { @@ -63,6 +70,8 @@ protected AbstractTsFileRecoverPerformer(TsFileResource tsFileResource) { */ protected void recoverWithWriter() throws DataRegionException, IOException { File tsFile = tsFileResource.getTsFile(); + boolean reinitialize = + tsFile.exists() && !tsFileResource.resourceFileExists() && hasNoPersistedData(tsFile); File chunkMetadataTempFile = new File(tsFile.getAbsolutePath() + TsFileIOWriter.CHUNK_METADATA_TEMP_FILE_SUFFIX); if (chunkMetadataTempFile.exists()) { @@ -82,8 +91,11 @@ protected void recoverWithWriter() throws DataRegionException, IOException { } writer = - new RestorableTsFileIOWriter( - tsFile, EncryptDBUtils.getFirstEncryptParamFromTSFilePath(tsFile.getAbsolutePath())); + reinitialize + ? reinitializeWriter(tsFile) + : new RestorableTsFileIOWriter( + tsFile, + EncryptDBUtils.getFirstEncryptParamFromTSFilePath(tsFile.getAbsolutePath())); if (writer.hasCrashed()) { byte versionNumber; try (TsFileSequenceReader sequenceReader = @@ -100,8 +112,9 @@ protected void recoverWithWriter() throws DataRegionException, IOException { writer.close(); org.apache.iotdb.commons.utils.FileUtils.deleteFileIfExist(tsFile); writer = - new RestorableTsFileIOWriter( + EncryptDBUtils.createRestorableWriterForNewFile( tsFile, + tsFileResource.getDatabaseName(), EncryptDBUtils.getFirstEncryptParamFromTSFilePath(tsFile.getAbsolutePath())); } } else { @@ -109,6 +122,42 @@ protected void recoverWithWriter() throws DataRegionException, IOException { } } + private static boolean hasNoPersistedData(File file) throws IOException { + byte[] magic = TSFileConfig.MAGIC_STRING.getBytes(StandardCharsets.UTF_8); + try (FileChannel input = FileChannel.open(file.toPath(), StandardOpenOption.READ)) { + if (input.size() == 0) { + return true; + } + if (input.size() != magic.length + Byte.BYTES) { + return false; + } + byte[] expected = Arrays.copyOf(magic, magic.length + Byte.BYTES); + expected[magic.length] = TSFileConfig.VERSION_NUMBER; + ByteBuffer header = ByteBuffer.allocate(expected.length); + IOUtils.readFully(input, header); + return Arrays.equals(expected, header.array()); + } + } + + private RestorableTsFileIOWriter reinitializeWriter(File file) throws IOException { + // Resolve the new-file policy before modifying even a header-only file. + EncryptParameter parameter = + EncryptDBUtils.createEncryptParameterForNewFile( + file, + tsFileResource.getDatabaseName(), + EncryptDBUtils.getFirstEncryptParamFromTSFilePath(file.getAbsolutePath())); + try { + try (FileChannel output = FileChannel.open(file.toPath(), StandardOpenOption.WRITE)) { + output.truncate(0); + } + return new RestorableTsFileIOWriter(file, parameter); + } finally { + if (parameter != null && parameter.isTdePageAead()) { + parameter.close(); + } + } + } + private void loadResourceFile() throws IOException { try { tsFileResource.deserialize(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/EncryptedUnsealedTsFileRecoverPerformerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/EncryptedUnsealedTsFileRecoverPerformerTest.java index 6f0f2fd14564b..8ba805f1bb995 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/EncryptedUnsealedTsFileRecoverPerformerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/file/EncryptedUnsealedTsFileRecoverPerformerTest.java @@ -31,8 +31,10 @@ import org.apache.iotdb.db.utils.EnvironmentUtils; import org.apache.iotdb.db.utils.TestFileEncryptionParameterProvider; +import org.apache.tsfile.common.conf.TSFileConfig; import org.apache.tsfile.encrypt.EncryptParameter; import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.file.MetaMarker; import org.apache.tsfile.file.metadata.AbstractAlignedChunkMetadata; import org.apache.tsfile.file.metadata.ChunkMetadata; import org.apache.tsfile.file.metadata.IDeviceID; @@ -59,7 +61,9 @@ import java.io.FileOutputStream; import java.nio.ByteBuffer; import java.nio.channels.FileChannel; +import java.nio.charset.StandardCharsets; import java.nio.file.Files; +import java.util.Arrays; import java.util.Collections; import java.util.HashSet; import java.util.Set; @@ -67,6 +71,8 @@ import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; public class EncryptedUnsealedTsFileRecoverPerformerTest { @@ -78,6 +84,7 @@ public class EncryptedUnsealedTsFileRecoverPerformerTest { private static final String FILE_NAME = TsFileUtilsForRecoverTest.getTestTsFilePath(DATABASE, 0, 0, 1); private TsFileResource resource; + private int nextFileId = 100; @Before public void setUp() throws Exception { @@ -119,6 +126,191 @@ public void testAlignedSnapshotRecovery() throws Exception { assertRecovery(true, true); } + @Test + public void testEmptyFileInsertRecoveryUsesProvider() throws Exception { + assertNoDataRecovery(false, false, false); + assertNoDataRecovery(false, true, false); + } + + @Test + public void testBaseHeaderInsertRecoveryUsesProvider() throws Exception { + assertNoDataRecovery(true, false, false); + assertNoDataRecovery(true, true, false); + } + + @Test + public void testEmptyFileSnapshotRecoveryUsesProvider() throws Exception { + assertNoDataRecovery(false, false, true); + assertNoDataRecovery(false, true, true); + } + + @Test + public void testBaseHeaderSnapshotRecoveryUsesProvider() throws Exception { + assertNoDataRecovery(true, false, true); + assertNoDataRecovery(true, true, true); + } + + private void assertNoDataRecovery(boolean baseHeaderOnly, boolean aligned, boolean snapshot) + throws Exception { + File file = newRecoveryFile(); + byte[] oldFileId; + try (EncryptParameter parameter = newFileParameter(); + TsFileIOWriter writer = new TsFileIOWriter(file, parameter)) { + oldFileId = writer.getEncryptParameter().getFileCryptoId(); + } + try (FileChannel channel = new FileOutputStream(file, true).getChannel()) { + channel.truncate(baseHeaderOnly ? baseHeader().length : 0); + } + resource = new TsFileResource(file); + byte[] recoveredFileId; + try (UnsealedTsFileRecoverPerformer performer = + new UnsealedTsFileRecoverPerformer(resource, true, ignored -> {})) { + performer.startRecovery(); + assertTrue(performer.hasCrashed()); + EncryptParameter parameter = performer.getWriter().getEncryptParameter(); + assertTrue(parameter.isTdePageAead()); + recoveredFileId = parameter.getFileCryptoId(); + assertFalse(Arrays.equals(oldFileId, recoveredFileId)); + performer.redoLog(firstEntry(aligned, snapshot)); + performer.endRecovery(); + } + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + assertTrue(reader.hasFileEncryptionHeader()); + assertArrayEquals(recoveredFileId, reader.getEncryptParam().getFileCryptoId()); + assertSingleRow(reader, aligned); + } + resource.close(); + } + + @Test + public void testProviderFailureDoesNotModifyUninitializedFiles() throws Exception { + TestFileEncryptionParameterProvider.RETURN_NULL.set(true); + for (boolean baseHeaderOnly : new boolean[] {false, true}) { + File file = newRecoveryFile(); + byte[] original = baseHeaderOnly ? baseHeader() : new byte[0]; + Files.write(file.toPath(), original); + resource = new TsFileResource(file); + try (UnsealedTsFileRecoverPerformer performer = + new UnsealedTsFileRecoverPerformer(resource, true, ignored -> {})) { + assertThrows(RuntimeException.class, performer::startRecovery); + assertNull(performer.getWriter()); + } + assertArrayEquals(original, Files.readAllBytes(file.toPath())); + assertFalse(resource.resourceFileExists()); + resource.close(); + } + } + + @Test + public void testPartialHeadersCannotBeReinitialized() throws Exception { + byte[] base = baseHeader(); + byte[] partialEncryptionHeader = Arrays.copyOf(base, base.length + 1); + partialEncryptionHeader[base.length] = MetaMarker.ENCRYPTION_HEADER; + for (byte[] original : + new byte[][] {Arrays.copyOf(base, base.length - 1), partialEncryptionHeader}) { + File file = newRecoveryFile(); + Files.write(file.toPath(), original); + resource = new TsFileResource(file); + try (UnsealedTsFileRecoverPerformer performer = + new UnsealedTsFileRecoverPerformer(resource, true, ignored -> {})) { + assertThrows(Exception.class, performer::startRecovery); + } + assertArrayEquals(original, Files.readAllBytes(file.toPath())); + resource.close(); + } + } + + @Test + public void testValidEncryptionHeaderWithoutDataKeepsOriginalContext() throws Exception { + File file = newRecoveryFile(); + byte[] oldFileId; + long headerEnd; + try (EncryptParameter parameter = newFileParameter(); + TsFileIOWriter writer = new TsFileIOWriter(file, parameter)) { + oldFileId = writer.getEncryptParameter().getFileCryptoId(); + headerEnd = writer.getPos(); + writer.endFile(); + } + try (FileChannel channel = new FileOutputStream(file, true).getChannel()) { + channel.truncate(headerEnd); + } + resource = new TsFileResource(file); + try (UnsealedTsFileRecoverPerformer performer = + new UnsealedTsFileRecoverPerformer(resource, true, ignored -> {})) { + performer.startRecovery(); + assertArrayEquals(oldFileId, performer.getWriter().getEncryptParameter().getFileCryptoId()); + performer.redoLog(firstEntry(false, false)); + performer.endRecovery(); + } + try (TsFileSequenceReader reader = new TsFileSequenceReader(file.getPath())) { + assertArrayEquals(oldFileId, reader.getEncryptParam().getFileCryptoId()); + assertSingleRow(reader, false); + } + } + + private File newRecoveryFile() throws Exception { + File file = new File(new File(FILE_NAME).getParentFile(), nextFileId++ + "-1-0-0.tsfile"); + Files.createDirectories(file.getParentFile().toPath()); + Files.createFile(file.toPath()); + return file; + } + + private static byte[] baseHeader() { + byte[] magic = TSFileConfig.MAGIC_STRING.getBytes(StandardCharsets.UTF_8); + return ByteBuffer.allocate(magic.length + Byte.BYTES) + .put(magic) + .put(TSFileConfig.VERSION_NUMBER) + .array(); + } + + private static EncryptParameter newFileParameter() { + return TestFileEncryptionParameterProvider.newParameter(); + } + + private static WALEntry firstEntry(boolean aligned, boolean snapshot) throws Exception { + InsertRowNode row = + new InsertRowNode( + new PlanNodeId(""), + new PartialPath(DEVICE), + aligned, + new String[] {"s1"}, + new TSDataType[] {TSDataType.INT64}, + 1, + new Object[] {10L}, + false); + row.setMeasurementSchemas(new MeasurementSchema[] {SCHEMA}); + if (!snapshot) { + return roundTrip(new WALInfoEntry(1, row)); + } + PrimitiveMemTable table = new PrimitiveMemTable(DATABASE, "0"); + if (aligned) { + table.insertAlignedRow(row); + } else { + table.insert(row); + } + return roundTrip(new WALInfoEntry(1, table)); + } + + private static void assertSingleRow(TsFileSequenceReader reader, boolean aligned) + throws Exception { + if (aligned) { + AbstractAlignedChunkMetadata metadata = reader.getAlignedChunkMetadata(DEVICE, false).get(0); + try (Chunk time = reader.readMemChunk((ChunkMetadata) metadata.getTimeChunkMetadata()); + Chunk value = + reader.readMemChunk((ChunkMetadata) metadata.getValueChunkMetadataList().get(0)); + IChunkReader chunks = + new AlignedChunkReader(time, Collections.singletonList(value), null)) { + assertEquals(1, assertRows(chunks, 0, true)); + } + } else { + ChunkMetadata metadata = reader.getChunkMetadataList(new Path(DEVICE, "s1", true)).get(0); + try (Chunk chunk = reader.readMemChunk(metadata); + IChunkReader chunks = new ChunkReader(chunk)) { + assertEquals(1, assertRows(chunks, 0, false)); + } + } + } + private void assertRecovery(boolean aligned, boolean snapshot) throws Exception { File file = new File(FILE_NAME); Files.createDirectories(file.getParentFile().toPath()); From 8bcd93be8dde7af4cd743f06e2ecbeb302c52d6b Mon Sep 17 00:00:00 2001 From: HTHou Date: Mon, 26 Oct 2026 18:04:06 +0800 Subject: [PATCH 5/6] [Storage] Retain WAL files when TsFile recovery fails --- .../iotdb/db/i18n/StorageEngineMessages.java | 2 + .../iotdb/db/i18n/StorageEngineMessages.java | 2 + .../wal/recover/WALNodeRecoverTask.java | 32 ++- .../wal/recover/WALNodeRecoverTaskTest.java | 229 ++++++++++++++++++ 4 files changed, 260 insertions(+), 5 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTaskTest.java diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java index 9563ec23f1920..20c33ffa73598 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -20,6 +20,8 @@ package org.apache.iotdb.db.i18n; public final class StorageEngineMessages { + public static final String LOG_WAL_NODE_RECOVERY_IN_ARG_WAS_INCOMPLETE_RETAINING_WAL_AND_CHECKPOINT_FILES_FOR_RETRY_48C1F7BF = + "WAL node recovery in {} was incomplete; retaining WAL and checkpoint files for retry."; public static final String EXCEPTION_FILE_ENCRYPTION_PARAMETER_PROVIDER_RETURNED_NULL_15C41723 = "File encryption parameter provider returned null"; public static final String EXCEPTION_PRECALCULATED_CHUNK_STAGING_IS_NOT_SUPPORTED_BY_THE_FILE_ENCRYPTION_POLICY_13DA8C32 = diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java index 6a4b3aaf90de6..6efd9d0898dec 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -20,6 +20,8 @@ package org.apache.iotdb.db.i18n; public final class StorageEngineMessages { + public static final String LOG_WAL_NODE_RECOVERY_IN_ARG_WAS_INCOMPLETE_RETAINING_WAL_AND_CHECKPOINT_FILES_FOR_RETRY_48C1F7BF = + "目录 {} 中的 WAL 节点恢复未完成；保留 WAL 和 checkpoint 文件以便重试。"; public static final String EXCEPTION_FILE_ENCRYPTION_PARAMETER_PROVIDER_RETURNED_NULL_15C41723 = "文件加密参数 Provider 返回 null"; public static final String EXCEPTION_PRECALCULATED_CHUNK_STAGING_IS_NOT_SUPPORTED_BY_THE_FILE_ENCRYPTION_POLICY_13DA8C32 = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTask.java index 73d429e3ccb07..ba064879d7cef 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTask.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTask.java @@ -54,6 +54,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.Map; +import java.util.Set; import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicLong; @@ -72,7 +73,9 @@ public class WALNodeRecoverTask implements Runnable { // version id of first valid .wal file private long firstValidVersionId = Long.MAX_VALUE; private Map memTableId2Info; - private Map memTableId2RecoverPerformer; + private Map memTableId2RecoverPerformer = new HashMap<>(); + private final Set failedMemTables = new HashSet<>(); + private boolean recoveryFailed; public WALNodeRecoverTask(File logDirectory, CountDownLatch allNodesRecoveredLatch) { this.logDirectory = logDirectory; @@ -92,13 +95,14 @@ public void run() { recoverInfoFromCheckpoints(); recoverTsFiles(); } catch (Exception e) { + recoveryFailed = true; for (UnsealedTsFileRecoverPerformer recoverPerformer : memTableId2RecoverPerformer.values()) { recoverPerformer.getRecoverListener().fail(e); } } finally { for (UnsealedTsFileRecoverPerformer recoverPerformer : memTableId2RecoverPerformer.values()) { try { - if (!recoverPerformer.canWrite()) { + if (recoveryFailed || !recoverPerformer.canWrite()) { recoverPerformer.close(); } } catch (Exception e) { @@ -108,6 +112,13 @@ public void run() { } try { + if (recoveryFailed) { + logger.warn( + StorageEngineMessages + .LOG_WAL_NODE_RECOVERY_IN_ARG_WAS_INCOMPLETE_RETAINING_WAL_AND_CHECKPOINT_FILES_FOR_RETRY_48C1F7BF, + logDirectory); + return; + } if (config.getDataRegionConsensusProtocolClass().equals(ConsensusFactory.IOT_CONSENSUS)) { // delete checkpoint info to avoid repeated recover File[] checkpointFiles = CheckpointFileUtils.listAllCheckpointFiles(logDirectory); @@ -279,10 +290,14 @@ private void recoverTsFiles() { return; } // make preparation for recovery - for (UnsealedTsFileRecoverPerformer recoverPerformer : memTableId2RecoverPerformer.values()) { + for (Map.Entry entry : + memTableId2RecoverPerformer.entrySet()) { + UnsealedTsFileRecoverPerformer recoverPerformer = entry.getValue(); try { recoverPerformer.startRecovery(); } catch (Exception e) { + recoveryFailed = true; + failedMemTables.add(entry.getKey()); recoverPerformer.getRecoverListener().fail(e); } } @@ -310,7 +325,7 @@ private void recoverTsFiles() { // see WALInfoEntry#serialize, entry type buffer.position(Byte.BYTES); long memTableId = buffer.getLong(); - if (!memTableId2Info.containsKey(memTableId)) { + if (!memTableId2Info.containsKey(memTableId) || failedMemTables.contains(memTableId)) { continue; } buffer.clear(); @@ -349,11 +364,18 @@ private void recoverTsFiles() { private void endRecovery() { // end recovering all recover performers - for (UnsealedTsFileRecoverPerformer recoverPerformer : memTableId2RecoverPerformer.values()) { + for (Map.Entry entry : + memTableId2RecoverPerformer.entrySet()) { + if (failedMemTables.contains(entry.getKey())) { + continue; + } + UnsealedTsFileRecoverPerformer recoverPerformer = entry.getValue(); try { recoverPerformer.endRecovery(); recoverPerformer.getRecoverListener().succeed(); } catch (Exception e) { + recoveryFailed = true; + failedMemTables.add(entry.getKey()); recoverPerformer.getRecoverListener().fail(e); } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTaskTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTaskTest.java new file mode 100644 index 0000000000000..cb9ea842bfafc --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTaskTest.java @@ -0,0 +1,229 @@ +/* + * 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.dataregion.wal.recover; + +import org.apache.iotdb.commons.path.PartialPath; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.db.conf.IoTDBConfig; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; +import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntry; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALInfoEntry; +import org.apache.iotdb.db.storageengine.dataregion.wal.checkpoint.Checkpoint; +import org.apache.iotdb.db.storageengine.dataregion.wal.checkpoint.CheckpointType; +import org.apache.iotdb.db.storageengine.dataregion.wal.checkpoint.MemTableInfo; +import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALRecoverException; +import org.apache.iotdb.db.storageengine.dataregion.wal.io.CheckpointWriter; +import org.apache.iotdb.db.storageengine.dataregion.wal.io.WALMetaData; +import org.apache.iotdb.db.storageengine.dataregion.wal.io.WALWriter; +import org.apache.iotdb.db.storageengine.dataregion.wal.recover.file.UnsealedTsFileRecoverPerformer; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALByteBufferForTest; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.listener.WALRecoverListener; +import org.apache.iotdb.db.utils.EnvironmentUtils; + +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.mockito.ArgumentCaptor; + +import java.io.File; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class WALNodeRecoverTaskTest { + @Rule public final TemporaryFolder folder = new TemporaryFolder(); + private final WALRecoverManager recovery = WALRecoverManager.getInstance(); + private final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig(); + private String previousConsensus; + + @Before + public void setUp() throws Exception { + EnvironmentUtils.envSetUp(); + recovery.clear(); + previousConsensus = config.getDataRegionConsensusProtocolClass(); + config.setDataRegionConsensusProtocolClass(ConsensusFactory.SIMPLE_CONSENSUS); + } + + @After + public void tearDown() throws Exception { + recovery.clear(); + config.setDataRegionConsensusProtocolClass(previousConsensus); + EnvironmentUtils.cleanEnv(); + } + + @Test + public void testInitializationIoFailureKeepsLogsAndAllowsRetry() throws Exception { + assertFailedRecovery(new IOException("injected initialization failure"), false, true); + } + + @Test + public void testInitializationRuntimeFailureKeepsLogsAndAllowsRetry() throws Exception { + assertFailedRecovery(new IllegalStateException("injected provider failure"), false, true); + } + + @Test + public void testFinalizationFailureKeepsLogsAndAllowsRetry() throws Exception { + assertFailedRecovery( + new WALRecoverException( + "injected flush failure", new IOException("injected flush failure")), + true, + true); + } + + @Test + public void testIoTConsensusFailureKeepsCheckpoint() throws Exception { + config.setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS); + assertFailedRecovery(new IOException("injected initialization failure"), false, false); + } + + @Test + public void testIoTConsensusV2FailureKeepsLogs() throws Exception { + config.setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS_V2); + assertFailedRecovery(new IOException("injected initialization failure"), false, false); + } + + private void assertFailedRecovery(Exception failure, boolean finalization, boolean retry) + throws Exception { + File node = folder.newFolder("node"); + List files = new ArrayList<>(); + List performers = new ArrayList<>(); + List tables = new ArrayList<>(); + List entries = new ArrayList<>(); + for (long id = 1; id <= 2; id++) { + File file = new File(folder.getRoot(), "sequence/root.recovery/0/0/" + id + "-1-0-0.tsfile"); + files.add(file); + performers.add(register(file)); + IMemTable table = mock(IMemTable.class); + when(table.getMemTableId()).thenReturn(id); + when(table.getDataRegionId()).thenReturn("0"); + tables.add(new MemTableInfo(table, file.getAbsolutePath(), 0)); + InsertRowNode row = + new InsertRowNode( + new PlanNodeId(""), + new PartialPath("root.recovery.d1"), + false, + new String[] {"s1"}, + new TSDataType[] {TSDataType.INT64}, + id, + new Object[] {id * 10}, + false); + row.setMeasurementSchemas( + new MeasurementSchema[] {new MeasurementSchema("s1", TSDataType.INT64)}); + entries.add(new WALInfoEntry(id, row)); + } + Checkpoint checkpoint = new Checkpoint(CheckpointType.GLOBAL_MEMORY_TABLE_INFO, tables); + ByteBuffer checkpointBytes = ByteBuffer.allocate(Long.BYTES + checkpoint.serializedSize()); + checkpointBytes.putLong(2); + checkpoint.serialize(checkpointBytes); + File checkpointFile = new File(node, "_0.checkpoint"); + try (CheckpointWriter writer = new CheckpointWriter(checkpointFile)) { + writer.write(checkpointBytes, false); + } + File wal = new File(node, "_0-0-0.wal"); + try (WALWriter writer = new WALWriter(wal)) { + for (WALEntry entry : entries) { + ByteBuffer data = ByteBuffer.allocate(entry.serializedSize()); + entry.serialize(new WALByteBufferForTest(data)); + WALMetaData metadata = new WALMetaData(); + metadata.add(data.position(), -1, entry.getMemTableId()); + writer.write(data, metadata); + } + } + byte[] walBefore = Files.readAllBytes(wal.toPath()); + byte[] checkpointBefore = Files.readAllBytes(checkpointFile.toPath()); + UnsealedTsFileRecoverPerformer failed = performers.get(0); + when(failed.canWrite()).thenReturn(true); + if (finalization) { + doThrow(failure).when(failed).endRecovery(); + } else { + doThrow(failure).when(failed).startRecovery(); + } + CountDownLatch latch = new CountDownLatch(1); + new WALNodeRecoverTask(node, latch).run(); + assertEquals(0, latch.getCount()); + assertEquals(WALRecoverListener.Status.FAILURE, failed.getRecoverListener().waitForResult()); + assertSame(failure, failed.getRecoverListener().getCause()); + verify(failed).close(); + if (!finalization) { + verify(failed, never()).redoLog(any()); + verify(failed, never()).endRecovery(); + } + UnsealedTsFileRecoverPerformer healthy = performers.get(1); + assertEquals(WALRecoverListener.Status.SUCCESS, healthy.getRecoverListener().waitForResult()); + ArgumentCaptor replay = ArgumentCaptor.forClass(WALEntry.class); + verify(healthy).redoLog(replay.capture()); + assertEquals(2, replay.getValue().getMemTableId()); + verify(healthy).endRecovery(); + assertArrayEquals(walBefore, Files.readAllBytes(wal.toPath())); + assertArrayEquals(checkpointBefore, Files.readAllBytes(checkpointFile.toPath())); + if (retry) { + recovery.clear(); + List retried = new ArrayList<>(); + for (File file : files) { + retried.add(register(file)); + } + CountDownLatch retryLatch = new CountDownLatch(1); + new WALNodeRecoverTask(node, retryLatch).run(); + assertEquals(0, retryLatch.getCount()); + for (UnsealedTsFileRecoverPerformer performer : retried) { + assertEquals( + WALRecoverListener.Status.SUCCESS, performer.getRecoverListener().waitForResult()); + verify(performer).redoLog(any()); + verify(performer).endRecovery(); + } + assertFalse(node.exists()); + assertEquals(WALRecoverListener.Status.FAILURE, failed.getRecoverListener().waitForResult()); + } else { + assertTrue(node.isDirectory()); + } + } + + private UnsealedTsFileRecoverPerformer register(File file) { + UnsealedTsFileRecoverPerformer performer = mock(UnsealedTsFileRecoverPerformer.class); + WALRecoverListener listener = new WALRecoverListener(file.getAbsolutePath()); + when(performer.getTsFileResource()).thenReturn(new TsFileResource(file)); + when(performer.getTsFileAbsolutePath()).thenReturn(file.getAbsolutePath()); + when(performer.getRecoverListener()).thenReturn(listener); + recovery.addRecoverPerformer(performer); + return performer; + } +} From 23f23022dd09d5b665bff2c9d6006811fa30074d Mon Sep 17 00:00:00 2001 From: HTHou Date: Mon, 26 Oct 2026 18:54:29 +0800 Subject: [PATCH 6/6] [Storage] Block reuse of WAL nodes after failed recovery --- .../iotdb/db/i18n/StorageEngineMessages.java | 2 + .../iotdb/db/i18n/StorageEngineMessages.java | 2 + .../dataregion/wal/WALManager.java | 25 ++++++ .../wal/recover/WALNodeRecoverTask.java | 3 + .../wal/recover/WALNodeRecoverTaskTest.java | 78 +++++++++++++++++-- 5 files changed, 103 insertions(+), 7 deletions(-) diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java index 20c33ffa73598..ea81ebb83b05c 100644 --- a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -20,6 +20,8 @@ package org.apache.iotdb.db.i18n; public final class StorageEngineMessages { + public static final String EXCEPTION_CANNOT_ENABLE_WAL_NODE_ARG_BECAUSE_ITS_PREVIOUS_RECOVERY_FAILED_RESTART_AFTER_RESOLVING_THE_RECOVERY_FAILURE_56DF9750 = + "Cannot enable WAL node %s because its previous recovery failed; restart after resolving the recovery failure."; public static final String LOG_WAL_NODE_RECOVERY_IN_ARG_WAS_INCOMPLETE_RETAINING_WAL_AND_CHECKPOINT_FILES_FOR_RETRY_48C1F7BF = "WAL node recovery in {} was incomplete; retaining WAL and checkpoint files for retry."; public static final String EXCEPTION_FILE_ENCRYPTION_PARAMETER_PROVIDER_RETURNED_NULL_15C41723 = diff --git a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java index 6efd9d0898dec..fa62bf6ea9095 100644 --- a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java +++ b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java @@ -20,6 +20,8 @@ package org.apache.iotdb.db.i18n; public final class StorageEngineMessages { + public static final String EXCEPTION_CANNOT_ENABLE_WAL_NODE_ARG_BECAUSE_ITS_PREVIOUS_RECOVERY_FAILED_RESTART_AFTER_RESOLVING_THE_RECOVERY_FAILURE_56DF9750 = + "WAL 节点 %s 上次恢复失败，无法启用；请解决恢复失败原因后重启。"; public static final String LOG_WAL_NODE_RECOVERY_IN_ARG_WAS_INCOMPLETE_RETAINING_WAL_AND_CHECKPOINT_FILES_FOR_RETRY_48C1F7BF = "目录 {} 中的 WAL 节点恢复未完成；保留 WAL 和 checkpoint 文件以便重试。"; public static final String EXCEPTION_FILE_ENCRYPTION_PARAMETER_PROVIDER_RETURNED_NULL_15C41723 = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/WALManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/WALManager.java index 8fafee304b2b6..e30055f123325 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/WALManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/WALManager.java @@ -47,6 +47,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Future; @@ -63,6 +64,7 @@ public class WALManager implements IService { // manage all wal nodes and decide how to allocate them private final NodeAllocationStrategy walNodesManager; + private final Map failedRecoveryNodes = new ConcurrentHashMap<>(); // single thread to delete old .wal files private ScheduledExecutorService walDeleteThread; // total disk usage of wal files @@ -94,6 +96,7 @@ public static String getApplicantUniqueId(String dataRegionName, boolean sequenc /** Apply for a wal node. */ public IWALNode applyForWALNode(String applicantUniqueId) { + checkRecoveryFailure(applicantUniqueId); if (config.getWalMode() == WALMode.DISABLE) { return WALFakeNode.getSuccessInstance(); } @@ -104,6 +107,7 @@ public IWALNode applyForWALNode(String applicantUniqueId) { /** WAL node will be registered only when using iot series consensus protocol. */ public void registerWALNode( String applicantUniqueId, String logDirectory, long startFileVersion, long startSearchIndex) { + checkRecoveryFailure(applicantUniqueId); if (config.getWalMode() == WALMode.DISABLE || (!config.getDataRegionConsensusProtocolClass().equals(ConsensusFactory.IOT_CONSENSUS) && !config @@ -117,6 +121,26 @@ public void registerWALNode( WritingMetrics.getInstance().createWALNodeInfoMetrics(applicantUniqueId); } + /** Called during recovery, before failed listeners can permit subsequent node allocation. */ + public void markRecoveryFailed(String applicantUniqueId, Exception failure) { + if (config.getDataRegionConsensusProtocolClass().equals(ConsensusFactory.IOT_CONSENSUS) + || config.getDataRegionConsensusProtocolClass().equals(ConsensusFactory.IOT_CONSENSUS_V2)) { + failedRecoveryNodes.putIfAbsent(applicantUniqueId, failure); + } + } + + private void checkRecoveryFailure(String applicantUniqueId) { + Exception failure = failedRecoveryNodes.get(applicantUniqueId); + if (failure != null) { + throw new IllegalStateException( + String.format( + StorageEngineMessages + .EXCEPTION_CANNOT_ENABLE_WAL_NODE_ARG_BECAUSE_ITS_PREVIOUS_RECOVERY_FAILED_RESTART_AFTER_RESOLVING_THE_RECOVERY_FAILURE_56DF9750, + applicantUniqueId), + failure); + } + } + /** WAL node will be deleted only when using iot series consensus protocol. */ public void deleteWALNode(String applicantUniqueId) { if (config.getWalMode() == WALMode.DISABLE @@ -324,6 +348,7 @@ public void syncDeleteOutdatedFilesInWALNodes() { public void clear() { totalDiskUsage.set(0); walNodesManager.clear(); + failedRecoveryNodes.clear(); } @Override diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTask.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTask.java index ba064879d7cef..fbb5724be98dc 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTask.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTask.java @@ -96,6 +96,7 @@ public void run() { recoverTsFiles(); } catch (Exception e) { recoveryFailed = true; + WALManager.getInstance().markRecoveryFailed(logDirectory.getName(), e); for (UnsealedTsFileRecoverPerformer recoverPerformer : memTableId2RecoverPerformer.values()) { recoverPerformer.getRecoverListener().fail(e); } @@ -297,6 +298,7 @@ private void recoverTsFiles() { recoverPerformer.startRecovery(); } catch (Exception e) { recoveryFailed = true; + WALManager.getInstance().markRecoveryFailed(logDirectory.getName(), e); failedMemTables.add(entry.getKey()); recoverPerformer.getRecoverListener().fail(e); } @@ -375,6 +377,7 @@ private void endRecovery() { recoverPerformer.getRecoverListener().succeed(); } catch (Exception e) { recoveryFailed = true; + WALManager.getInstance().markRecoveryFailed(logDirectory.getName(), e); failedMemTables.add(entry.getKey()); recoverPerformer.getRecoverListener().fail(e); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTaskTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTaskTest.java index cb9ea842bfafc..f7780a7fa9059 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTaskTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/dataregion/wal/recover/WALNodeRecoverTaskTest.java @@ -26,6 +26,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; import org.apache.iotdb.db.storageengine.dataregion.memtable.IMemTable; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.dataregion.wal.WALManager; import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntry; import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALInfoEntry; import org.apache.iotdb.db.storageengine.dataregion.wal.checkpoint.Checkpoint; @@ -35,8 +36,10 @@ import org.apache.iotdb.db.storageengine.dataregion.wal.io.CheckpointWriter; import org.apache.iotdb.db.storageengine.dataregion.wal.io.WALMetaData; import org.apache.iotdb.db.storageengine.dataregion.wal.io.WALWriter; +import org.apache.iotdb.db.storageengine.dataregion.wal.node.WALNode; import org.apache.iotdb.db.storageengine.dataregion.wal.recover.file.UnsealedTsFileRecoverPerformer; import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALByteBufferForTest; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.WALMode; import org.apache.iotdb.db.storageengine.dataregion.wal.utils.listener.WALRecoverListener; import org.apache.iotdb.db.utils.EnvironmentUtils; @@ -54,6 +57,8 @@ import java.nio.ByteBuffer; import java.nio.file.Files; import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; import java.util.List; import java.util.concurrent.CountDownLatch; @@ -61,6 +66,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doThrow; @@ -77,9 +83,10 @@ public class WALNodeRecoverTaskTest { @Before public void setUp() throws Exception { + previousConsensus = config.getDataRegionConsensusProtocolClass(); + config.setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS); EnvironmentUtils.envSetUp(); recovery.clear(); - previousConsensus = config.getDataRegionConsensusProtocolClass(); config.setDataRegionConsensusProtocolClass(ConsensusFactory.SIMPLE_CONSENSUS); } @@ -112,18 +119,34 @@ public void testFinalizationFailureKeepsLogsAndAllowsRetry() throws Exception { @Test public void testIoTConsensusFailureKeepsCheckpoint() throws Exception { config.setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS); - assertFailedRecovery(new IOException("injected initialization failure"), false, false); + assertFailedRecovery(new IOException("injected initialization failure"), false, true); } @Test public void testIoTConsensusV2FailureKeepsLogs() throws Exception { config.setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS_V2); - assertFailedRecovery(new IOException("injected initialization failure"), false, false); + assertFailedRecovery(new IOException("injected initialization failure"), false, true); + } + + @Test + public void testIoTConsensusFinalizationFailureBlocksNodeReuse() throws Exception { + config.setDataRegionConsensusProtocolClass(ConsensusFactory.IOT_CONSENSUS); + assertFailedRecovery( + new WALRecoverException( + "injected flush failure", new IOException("injected flush failure")), + true, + false); } private void assertFailedRecovery(Exception failure, boolean finalization, boolean retry) throws Exception { File node = folder.newFolder("node"); + boolean stableNode = + config.getDataRegionConsensusProtocolClass().equals(ConsensusFactory.IOT_CONSENSUS) + || config + .getDataRegionConsensusProtocolClass() + .equals(ConsensusFactory.IOT_CONSENSUS_V2); + long firstVersion = stableNode ? 7 : 0; List files = new ArrayList<>(); List performers = new ArrayList<>(); List tables = new ArrayList<>(); @@ -135,7 +158,7 @@ private void assertFailedRecovery(Exception failure, boolean finalization, boole IMemTable table = mock(IMemTable.class); when(table.getMemTableId()).thenReturn(id); when(table.getDataRegionId()).thenReturn("0"); - tables.add(new MemTableInfo(table, file.getAbsolutePath(), 0)); + tables.add(new MemTableInfo(table, file.getAbsolutePath(), firstVersion)); InsertRowNode row = new InsertRowNode( new PlanNodeId(""), @@ -148,6 +171,9 @@ private void assertFailedRecovery(Exception failure, boolean finalization, boole false); row.setMeasurementSchemas( new MeasurementSchema[] {new MeasurementSchema("s1", TSDataType.INT64)}); + if (stableNode) { + row.setSearchIndex(10 + id); + } entries.add(new WALInfoEntry(id, row)); } Checkpoint checkpoint = new Checkpoint(CheckpointType.GLOBAL_MEMORY_TABLE_INFO, tables); @@ -158,13 +184,14 @@ private void assertFailedRecovery(Exception failure, boolean finalization, boole try (CheckpointWriter writer = new CheckpointWriter(checkpointFile)) { writer.write(checkpointBytes, false); } - File wal = new File(node, "_0-0-0.wal"); + File wal = new File(node, stableNode ? "_7-11-1.wal" : "_0-0-0.wal"); try (WALWriter writer = new WALWriter(wal)) { for (WALEntry entry : entries) { ByteBuffer data = ByteBuffer.allocate(entry.serializedSize()); entry.serialize(new WALByteBufferForTest(data)); WALMetaData metadata = new WALMetaData(); - metadata.add(data.position(), -1, entry.getMemTableId()); + metadata.add( + data.position(), stableNode ? 10 + entry.getMemTableId() : -1, entry.getMemTableId()); writer.write(data, metadata); } } @@ -195,7 +222,38 @@ private void assertFailedRecovery(Exception failure, boolean finalization, boole verify(healthy).endRecovery(); assertArrayEquals(walBefore, Files.readAllBytes(wal.toPath())); assertArrayEquals(checkpointBefore, Files.readAllBytes(checkpointFile.toPath())); + if (stableNode) { + WALManager manager = WALManager.getInstance(); + for (int i = 0; i < 3; i++) { + IllegalStateException allocation = + assertThrows( + IllegalStateException.class, () -> manager.applyForWALNode(node.getName())); + assertSame(failure, allocation.getCause()); + IllegalStateException registration = + assertThrows( + IllegalStateException.class, + () -> manager.registerWALNode(node.getName(), node.getAbsolutePath(), 8, 12)); + assertSame(failure, registration.getCause()); + } + WALMode previousMode = config.getWalMode(); + try { + config.setWalMode(WALMode.DISABLE); + assertThrows(IllegalStateException.class, () -> manager.applyForWALNode(node.getName())); + } finally { + config.setWalMode(previousMode); + } + assertArrayEquals(walBefore, Files.readAllBytes(wal.toPath())); + assertArrayEquals(checkpointBefore, Files.readAllBytes(checkpointFile.toPath())); + assertEquals( + new HashSet<>(Arrays.asList(1L, 2L)), + CheckpointRecoverUtils.recoverMemTableInfo(node).getMemTableId2Info().keySet()); + assertTrue(manager.applyForWALNode("unrelated-healthy-node") instanceof WALNode); + } if (retry) { + if (stableNode) { + // A fresh WAL service must recover the retained state before accepting this identifier. + WALManager.getInstance().clear(); + } recovery.clear(); List retried = new ArrayList<>(); for (File file : files) { @@ -210,7 +268,13 @@ private void assertFailedRecovery(Exception failure, boolean finalization, boole verify(performer).redoLog(any()); verify(performer).endRecovery(); } - assertFalse(node.exists()); + if (stableNode) { + WALNode restored = (WALNode) WALManager.getInstance().applyForWALNode(node.getName()); + assertEquals(8, restored.getCurrentWALFileVersion()); + assertEquals(12, restored.getCurrentSearchIndex()); + } else { + assertFalse(node.exists()); + } assertEquals(WALRecoverListener.Status.FAILURE, failed.getRecoverListener().waitForResult()); } else { assertTrue(node.isDirectory());