From ad273954325f67a6938ec9e72f0e4ac3e62a45d7 Mon Sep 17 00:00:00 2001 From: luoluoyuyu Date: Fri, 21 Aug 2026 11:19:03 +0800 Subject: [PATCH] feat(load): route decoded tsfile pieces through consensus --- .../iotdb/db/i18n/DataNodeQueryMessages.java | 6 + .../iotdb/db/i18n/StorageEngineMessages.java | 58 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 6 + .../iotdb/db/i18n/StorageEngineMessages.java | 56 + .../org/apache/iotdb/db/conf/IoTDBConfig.java | 2 +- .../dataregion/DataExecutionVisitor.java | 14 + .../impl/DataNodeInternalRPCServiceImpl.java | 46 +- .../executor/RegionWriteExecutor.java | 12 + .../node/DataNodePlanNodeDeserializer.java | 7 + .../plan/planner/plan/node/PlanVisitor.java | 6 + .../node/load/LoadTsFileConsensusNode.java | 782 ++++++++++ .../plan/node/load/LoadTsFileConsensusOp.java | 50 + .../load/DataPartitionBatchFetcher.java | 114 ++ .../scheduler/load/DataPartitionRouter.java | 94 ++ .../load/LoadConsensusSubmitter.java | 202 +++ .../scheduler/load/LoadFallbackHandler.java | 212 +++ .../load/LoadTsFileDispatcherImpl.java | 4 + .../scheduler/load/LoadTsFileScheduler.java | 1053 +++---------- .../scheduler/load/LocalLoadStrategy.java | 203 +++ .../scheduler/load/MemoryBoundedBuffer.java | 78 + .../plan/scheduler/load/PieceDispatcher.java | 234 +++ .../load/RegionConsensusContext.java | 89 ++ .../scheduler/load/TsFileLoadStrategy.java | 46 + .../scheduler/load/TsFileSplitConsumer.java | 137 ++ .../load/TwoPhaseConsensusLoadStrategy.java | 378 +++++ .../iotdb/db/storageengine/StorageEngine.java | 4 + .../dataregion/snapshot/SnapshotLoader.java | 20 + .../dataregion/snapshot/SnapshotTaker.java | 25 +- .../dataregion/wal/buffer/WALEntry.java | 6 + .../dataregion/wal/buffer/WALEntryType.java | 5 +- .../dataregion/wal/buffer/WALInfoEntry.java | 6 + .../dataregion/wal/node/IWALNode.java | 4 + .../dataregion/wal/node/WALFakeNode.java | 6 + .../dataregion/wal/node/WALNode.java | 7 + .../storageengine/load/DataPartitionInfo.java | 76 + .../load/LoadCleanupScheduler.java | 129 ++ .../load/LoadSnapshotManager.java | 332 ++++ .../storageengine/load/LoadTaskRegistry.java | 117 ++ .../load/LoadTsFileChecksumUtils.java | 64 + .../storageengine/load/LoadTsFileManager.java | 1370 +++++++++-------- .../storageengine/load/PartitionContext.java | 650 ++++++++ .../load/TsFileWriterManager.java | 1146 ++++++++++++++ .../load/LoadTsFileConsensusNodeTest.java | 133 ++ .../load/LoadTsFileSchedulerTest.java | 47 +- .../load/LoadTsFileSnapshotMetaTest.java | 88 ++ .../load/TsFileWriterManagerTest.java | 142 ++ .../plan/planner/plan/node/PlanNodeType.java | 1 + 47 files changed, 6765 insertions(+), 1502 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusOp.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionBatchFetcher.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadConsensusSubmitter.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadFallbackHandler.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LocalLoadStrategy.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/MemoryBoundedBuffer.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileLoadStrategy.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/DataPartitionInfo.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadCleanupScheduler.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadSnapshotManager.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTaskRegistry.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileChecksumUtils.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/PartitionContext.java create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileSnapshotMetaTest.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManagerTest.java 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 0e9e8158ac7b0..a246a1f80d03b 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 @@ -1160,6 +1160,9 @@ public final class DataNodeQueryMessages { "Start load TsFile {} locally."; public static final String LOAD_ALL_FAILED_TSFILES_ARE_CONVERTED_TO_TABLETS = "Load: all failed TsFiles are converted to tablets and inserted."; + public static final String + LOG_LOAD_FAILED_TO_LOAD_SOME_TSFILES_BY_CONVERTING_THEM_INTO_TABLETS_FAILED_TSFILES_ARG_7D9DB9C3 = + "Load: failed to load some TsFiles by converting them into tablets. Failed TsFiles: %s"; // --- Plan / Statement --- @@ -2692,6 +2695,8 @@ public final class DataNodeQueryMessages { "Parse or send TsFile %s error."; public static final String DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG = "Dispatch one piece to ReplicaSet {} error. Result status code {}. "; + public static final String LOG_LOAD_CONSENSUS_SUBMIT_TRANSIENT_FAILURE_RETRY_D7E1D9A6 = + "Transient failure while submitting LOAD consensus {} (load {}) to {}, will retry ({}/{}): {}"; public static final String RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG = "Result status message {}. Dispatch piece node error:%n{}"; public static final String SUB_STATUS_CODE_ARG_SUB_STATUS_MESSAGE_ARG = @@ -3793,6 +3798,7 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_THE_SECOND_ARGUMENT_OF_PERCENTILE_FUNCTION_PERCENTAGE_MUST_BE_A_DOUBLE_LITERAL_D9464B46 = "The second argument of 'percentile' function percentage must be a double literal"; public static final String EXCEPTION_DATA_TYPE_MISMATCH_FOR_MEASUREMENT_ARGARGARG_TYPE_IN_TSFILE_ARG_TYPE_IN_IOTDB_ARG_C5BA7DBD = "Data type mismatch for measurement %s%s%s, type in TsFile: %s, type in IoTDB: %s"; public static final String MESSAGE_FAILED_TO_RELEASE_EXTERNAL_TSFILE_QUERY_RESOURCE_712EE978 = "Failed to release external TsFile query resource"; + public static final String EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 = "Unknown LoadTsFileConsensusOp ordinal: "; public static final String EXCEPTION_OUTER_QUERY_TIMEOUT_EXCEEDED_BEFORE_IOTDBLOCAL_QUERY_STARTS_800BFA63 = "Outer query timeout exceeded before IoTDBLocal query starts"; public static final String MESSAGE_FAILED_TO_CLOSE_UDF_RESULT_SET_AT_INDEX_ARG_A293B7EC = "Failed to close UDF result set at index {}"; public static final String EXCEPTION_INTERNAL_QUERY_EXECUTION_NOT_FOUND_62642542 = "Internal query execution not found"; 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 a104bf3576cc6..46d4688bf12b0 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 @@ -29,6 +29,62 @@ private StorageEngineMessages() {} // ======================== StorageEngine ======================== public static final String FAIL_TO_RECOVER_WAL = "Fail to recover wal."; + public static final String LOG_LOAD_CONSENSUS_WRITE_TO_REGION_ARG_VIA_PROTOCOL_ARG_EBB55042 = + "Write LOAD consensus node to region {} via protocol {}"; + public static final String LOG_LOAD_CONSENSUS_WRITE_TO_REGION_ARG_VIA_PEER_ARG_FAILED_TRYING_NEXT_REPLICA_ARG_39217580 = + "LOAD consensus write to region {} via peer {} failed, trying next replica: {}"; + public static final String LOG_LOAD_CONSENSUS_REFRESH_REPLICA_SET_FAILED_7C244C63 = + "Failed to refresh LOAD consensus replica set for region {}, using cached set: {}"; + public static final String MESSAGE_LOAD_CONSENSUS_PIECE_CHECKSUM_MISMATCH_CF261675 = + "LOAD consensus piece checksum mismatch, loadId: %s, pieceIndex: %d"; + public static final String MESSAGE_LOAD_CONSENSUS_WAL_FLUSH_FAILED_8BE1375A = + "Failed to flush LOAD consensus WAL entry"; + public static final String MESSAGE_LOAD_CONSENSUS_RATIS_NOT_SUPPORTED_D371E344 = + "LOAD consensus is not supported on Ratis"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_EOF_8743387D = + "Unexpected end of file when reading staged piece %s at offset %s."; + public static final String MESSAGE_LOAD_CONSENSUS_PIECE_NOT_CONTINUOUS_AFTER_FAILOVER_D6FFAC6C = + "LOAD piece %d of load %s cannot be applied because the previous pieces have not all been " + + "applied (the staged state may have been lost by a leader failover)."; + public static final String MESSAGE_LOAD_CONSENSUS_PIECE_DATA_MISSING_AFTER_PULL_8269CB0B = + "LOAD piece %d data of load %s is still missing after pulling from the write node."; + public static final String LOG_LOAD_CONSENSUS_FORWARD_PIECE_FAILED_34F9EBE7 = + "Failed to forward LOAD piece {} of load {} to follower {}: {}"; + public static final String LOG_LOAD_CONSENSUS_PULL_PIECE_FAILED_AFB003D5 = + "Failed to pull LOAD piece {} of load {} from write node {}: {}"; + public static final String MESSAGE_LOAD_CONSENSUS_PULL_WITHOUT_SOURCE_ENDPOINT_3B20D9E9 = + "LOAD pull request has no source endpoint."; + public static final String MESSAGE_LOAD_CONSENSUS_PULL_WITHOUT_RETAINED_PIECE_AD3C9D4F = + "LOAD piece %d of load %s is not retained on the write node."; + public static final String MESSAGE_LOAD_CONSENSUS_PULL_PUSH_BACK_FAILED_1A90C2B9 = + "Failed to push LOAD piece %d of load %s back to %s: %s"; + public static final String LOG_LOAD_CONSENSUS_ABORT_MARKER_FAILED_6A218023 = + "Failed to log the LOAD ABORT marker of load {}: {}"; + public static final String EXCEPTION_LOAD_CONSENSUS_PIECE_DATA_MISSING_OR_CHECKSUM_MISMATCH_AFTER_PULL_35F4972E = + "LOAD task %s piece %d data is missing or its checksum mismatches after pull."; + public static final String LOG_LOAD_CONSENSUS_RETAINED_PIECE_READ_FAILED_0659D19B = + "Failed to read retained LOAD piece {} of load {} from {}: {}"; + public static final String LOG_LOAD_CONSENSUS_RETAINED_PIECE_WRITE_FAILED_99697608 = + "Failed to write retained LOAD piece {} of load {} to {}: {}"; + public static final String LOG_LOAD_CONSENSUS_APPLIED_PIECE_RESTORE_FAILED_5BC74BBA = + "Failed to restore the applied LOAD piece entry {} of load {}."; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19 = + "Staged file %s of load %s is not continuous: expected offset %d but current file length is %d."; + public static final String MESSAGE_LOAD_CONSENSUS_PREPARE_WITHOUT_STAGED_DATA_FE8ADC37 = + "Cannot prepare load %s because no staged data exists on this node."; + public static final String MESSAGE_LOAD_CONSENSUS_PREPARE_VERIFICATION_FAILED_B3865A82 = + "LOAD prepare verification failed for load %s: expected %d pieces with checksum %d, found %d pieces with checksum %d"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_INCOMPLETE_1CDE954B = + "Staged file %s of load %s is incomplete and cannot be committed."; + public static final String LOG_LOAD_CONSENSUS_SNAPSHOT_TAKEN_09A7DD4C = + "Snapshotted %d in-progress LOAD task(s) with %d staged file(s) for region %s into %s."; + public static final String LOG_LOAD_CONSENSUS_SNAPSHOT_RESTORED_90ABC1BF = + "Restored %d in-progress LOAD task(s) with %d staged file(s) from snapshot %s."; + public static final String EXCEPTION_LOAD_CONSENSUS_SNAPSHOT_RESTORE_FAILED_F8C29C64 = + "Failed to restore LOAD snapshot from %s: %s"; + public static final String EXCEPTION_LOAD_TSFILE_ALIGNED_VALUE_CHUNK_TIME_CHUNK_EEB00760 = + "Cannot attach value chunk of measurement %s in file %s: expected exactly one buffered " + + "aligned time chunk, found %d."; public static final String STORAGE_ENGINE_FAILED_TO_SET_UP = "Storage engine failed to set up."; public static final String SEQ_MEMTABLE_FLUSH_CHECK_THREAD_STARTED = "start sequence memtable timed flush check thread successfully."; public static final String UNSEQ_MEMTABLE_FLUSH_CHECK_THREAD_STARTED = "start unsequence memtable timed flush check thread successfully."; @@ -477,6 +533,8 @@ private StorageEngineMessages() {} public static final String CANNOT_CREATE_TSFILE_FOR_WRITING = "Can not create TsFile {} for writing."; public static final String CLOSE_TSFILE_IO_WRITER_ERROR = "Close TsFileIOWriter {} error."; public static final String CLOSE_MODIFICATION_FILE_ERROR = "Close ModificationFile {} error."; + public static final String EXCEPTION_TABLE_ARG_ARG_DOES_NOT_EXIST_WHEN_APPLYING_LOAD_CHUNK_DATA_IT_MAY_HAVE_BEEN_DROPPED_AFTER_THE_LOAD_WAS_ANALYZED_DDB35F93 = + "Table '%s.%s' does not exist when applying LOAD chunk data. It may have been dropped after the LOAD was analyzed."; public static final String TASK_DIR_NOT_EMPTY_SKIP_DELETE = "Task dir {} is not empty, skip deleting."; public static final String LOAD_CLEANUP_TASK_CANCELED = "Load cleanup task {} is canceled."; public static final String LOAD_CLEANUP_TASK_STARTS = "Load cleanup task {} starts."; 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 1ab922fa121c8..e13f6705ff972 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 @@ -1141,6 +1141,9 @@ public final class DataNodeQueryMessages { "开始本地加载 TsFile {}。"; public static final String LOAD_ALL_FAILED_TSFILES_ARE_CONVERTED_TO_TABLETS = "加载：所有失败的 TsFile 已转换为 Tablet 并插入。"; + public static final String + LOG_LOAD_FAILED_TO_LOAD_SOME_TSFILES_BY_CONVERTING_THEM_INTO_TABLETS_FAILED_TSFILES_ARG_7D9DB9C3 = + "加载：部分 TsFile 通过转换为 Tablet 仍加载失败。失败的 TsFile：%s"; // --- Plan / Statement --- @@ -3174,6 +3177,8 @@ public final class DataNodeQueryMessages { public static final String DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG = "分发 TsFile 片段到 ReplicaSet {} 出错。结果状态码 {}。 "; + public static final String LOG_LOAD_CONSENSUS_SUBMIT_TRANSIENT_FAILURE_RETRY_D7E1D9A6 = + "提交 LOAD 共识 {}（load {}）到 {} 时遇到瞬时失败，将重试（{}/{}）：{}"; public static final String RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG = "结果状态消息 {}。分发片段节点出错：%n{}"; @@ -4549,6 +4554,7 @@ private DataNodeQueryMessages() {} public static final String EXCEPTION_THE_SECOND_ARGUMENT_OF_PERCENTILE_FUNCTION_PERCENTAGE_MUST_BE_A_DOUBLE_LITERAL_D9464B46 = "'percentile' 函数的第二个参数 percentage 必须是 double 字面量"; public static final String EXCEPTION_DATA_TYPE_MISMATCH_FOR_MEASUREMENT_ARGARGARG_TYPE_IN_TSFILE_ARG_TYPE_IN_IOTDB_ARG_C5BA7DBD = "测点 %s%s%s 的数据类型不匹配，TsFile 中类型：%s，IoTDB 中类型：%s"; public static final String MESSAGE_FAILED_TO_RELEASE_EXTERNAL_TSFILE_QUERY_RESOURCE_712EE978 = "释放外部 TsFile 查询资源失败"; + public static final String EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 = "未知的 LoadTsFileConsensusOp 序号："; public static final String EXCEPTION_OUTER_QUERY_TIMEOUT_EXCEEDED_BEFORE_IOTDBLOCAL_QUERY_STARTS_800BFA63 = "在 IoTDBLocal 查询开始前，外层查询已超时"; public static final String MESSAGE_FAILED_TO_CLOSE_UDF_RESULT_SET_AT_INDEX_ARG_A293B7EC = "关闭索引 {} 处的 UDF 结果集失败"; public static final String EXCEPTION_INTERNAL_QUERY_EXECUTION_NOT_FOUND_62642542 = "未找到内部查询执行"; 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 ba857ecfa91e5..668ea2c841cb0 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 @@ -29,6 +29,60 @@ private StorageEngineMessages() {} // ======================== StorageEngine ======================== public static final String FAIL_TO_RECOVER_WAL = "WAL 恢复失败。"; + public static final String LOG_LOAD_CONSENSUS_WRITE_TO_REGION_ARG_VIA_PROTOCOL_ARG_EBB55042 = + "通过协议 {} 向 Region {} 写入 LOAD 共识节点"; + public static final String LOG_LOAD_CONSENSUS_WRITE_TO_REGION_ARG_VIA_PEER_ARG_FAILED_TRYING_NEXT_REPLICA_ARG_39217580 = + "向 Region {} 经节点 {} 写入 LOAD 共识失败，尝试下一个副本：{}"; + public static final String LOG_LOAD_CONSENSUS_REFRESH_REPLICA_SET_FAILED_7C244C63 = + "刷新 Region {} 的 LOAD 共识副本集失败，使用缓存的副本集：{}"; + public static final String MESSAGE_LOAD_CONSENSUS_PIECE_CHECKSUM_MISMATCH_CF261675 = + "LOAD 共识分片校验和不一致，loadId: %s，pieceIndex: %d"; + public static final String MESSAGE_LOAD_CONSENSUS_WAL_FLUSH_FAILED_8BE1375A = + "LOAD 共识 WAL 记录落盘失败"; + public static final String MESSAGE_LOAD_CONSENSUS_RATIS_NOT_SUPPORTED_D371E344 = + "Ratis 暂不支持 LOAD 共识"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_EOF_8743387D = + "读取已暂存的分片 %s 时意外到达文件末尾，offset: %s。"; + public static final String MESSAGE_LOAD_CONSENSUS_PIECE_NOT_CONTINUOUS_AFTER_FAILOVER_D6FFAC6C = + "LOAD 分片 %d（load %s）无法应用：之前的分片尚未全部应用（暂存状态可能在主备切换时丢失）。"; + public static final String MESSAGE_LOAD_CONSENSUS_PIECE_DATA_MISSING_AFTER_PULL_8269CB0B = + "LOAD 分片 %d（load %s）的数据在向写节点回补后仍然缺失。"; + public static final String LOG_LOAD_CONSENSUS_FORWARD_PIECE_FAILED_34F9EBE7 = + "向副本 {} 转发 LOAD 分片 {}（load {}）失败：{}"; + public static final String LOG_LOAD_CONSENSUS_PULL_PIECE_FAILED_AFB003D5 = + "向写节点 {} 回补 LOAD 分片 {}（load {}）失败：{}"; + public static final String MESSAGE_LOAD_CONSENSUS_PULL_WITHOUT_SOURCE_ENDPOINT_3B20D9E9 = + "LOAD 回补请求缺少源节点地址。"; + public static final String MESSAGE_LOAD_CONSENSUS_PULL_WITHOUT_RETAINED_PIECE_AD3C9D4F = + "LOAD 分片 %d（load %s）未在写节点保留。"; + public static final String MESSAGE_LOAD_CONSENSUS_PULL_PUSH_BACK_FAILED_1A90C2B9 = + "将 LOAD 分片 %d（load %s）推回给 %s 失败：%s"; + public static final String LOG_LOAD_CONSENSUS_ABORT_MARKER_FAILED_6A218023 = + "写入 LOAD 中止（ABORT）标记（load {}）失败：{}"; + public static final String EXCEPTION_LOAD_CONSENSUS_PIECE_DATA_MISSING_OR_CHECKSUM_MISMATCH_AFTER_PULL_35F4972E = + "回补后 LOAD 任务 %s 的分片 %d 数据仍缺失或校验和不一致。"; + public static final String LOG_LOAD_CONSENSUS_RETAINED_PIECE_READ_FAILED_0659D19B = + "读取保留的 LOAD 分片 {}（load {}）失败，文件：{}，原因：{}"; + public static final String LOG_LOAD_CONSENSUS_RETAINED_PIECE_WRITE_FAILED_99697608 = + "写入保留的 LOAD 分片 {}（load {}）失败，文件：{}，原因：{}"; + public static final String LOG_LOAD_CONSENSUS_APPLIED_PIECE_RESTORE_FAILED_5BC74BBA = + "恢复 load {} 的已应用 LOAD 分片条目 {} 失败。"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19 = + "load %s 的暂存文件 %s 不连续：期望偏移 %d，但当前文件长度为 %d。"; + public static final String MESSAGE_LOAD_CONSENSUS_PREPARE_WITHOUT_STAGED_DATA_FE8ADC37 = + "无法准备（PREPARE）load %s，因为该节点上不存在暂存数据。"; + public static final String MESSAGE_LOAD_CONSENSUS_PREPARE_VERIFICATION_FAILED_B3865A82 = + "LOAD PREPARE 校验失败，load %s：预期 %d 片、checksum %d，实际 %d 片、checksum %d"; + public static final String EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_INCOMPLETE_1CDE954B = + "load %s 的暂存文件 %s 不完整，无法提交（COMMIT）。"; + public static final String LOG_LOAD_CONSENSUS_SNAPSHOT_TAKEN_09A7DD4C = + "已将 region %s 的 %d 个进行中的 LOAD 任务（共 %d 个暂存文件）纳入快照 %s。"; + public static final String LOG_LOAD_CONSENSUS_SNAPSHOT_RESTORED_90ABC1BF = + "已从快照 %s 恢复 %d 个进行中的 LOAD 任务（共 %d 个暂存文件）。"; + public static final String EXCEPTION_LOAD_CONSENSUS_SNAPSHOT_RESTORE_FAILED_F8C29C64 = + "从 %s 恢复 LOAD 快照失败：%s"; + public static final String EXCEPTION_LOAD_TSFILE_ALIGNED_VALUE_CHUNK_TIME_CHUNK_EEB00760 = + "无法将测量 %s 的值 Chunk 挂载到文件 %s：预期恰好一个已缓冲的 Aligned 时间 Chunk，实际发现 %d 个。"; public static final String STORAGE_ENGINE_FAILED_TO_SET_UP = "存储引擎启动失败。"; public static final String SEQ_MEMTABLE_FLUSH_CHECK_THREAD_STARTED = "顺序 memtable 定时 flush 检查线程启动成功。"; public static final String UNSEQ_MEMTABLE_FLUSH_CHECK_THREAD_STARTED = "乱序 memtable 定时 flush 检查线程启动成功。"; @@ -477,6 +531,8 @@ private StorageEngineMessages() {} public static final String CANNOT_CREATE_TSFILE_FOR_WRITING = "无法创建 TsFile {} 用于写入。"; public static final String CLOSE_TSFILE_IO_WRITER_ERROR = "关闭 TsFileIOWriter {} 出错。"; public static final String CLOSE_MODIFICATION_FILE_ERROR = "关闭修改文件 {} 出错。"; + public static final String EXCEPTION_TABLE_ARG_ARG_DOES_NOT_EXIST_WHEN_APPLYING_LOAD_CHUNK_DATA_IT_MAY_HAVE_BEEN_DROPPED_AFTER_THE_LOAD_WAS_ANALYZED_DDB35F93 = + "应用 LOAD chunk 数据时表 '%s.%s' 不存在，可能在 LOAD 分析之后被删除了。"; public static final String TASK_DIR_NOT_EMPTY_SKIP_DELETE = "任务目录 {} 非空，跳过删除。"; public static final String LOAD_CLEANUP_TASK_CANCELED = "加载清理任务 {} 已取消。"; public static final String LOAD_CLEANUP_TASK_STARTS = "加载清理任务 {} 开始。"; diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java index 1d726366ddd2b..7b7312b8debea 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java @@ -4351,7 +4351,7 @@ public int getLoadTsFileSpiltPartitionMaxSize() { } public void setLoadTsFileSpiltPartitionMaxSize(int loadTsFileSpiltPartitionMaxSize) { - if (loadTsFileSpiltPartitionMaxSize <= 0) { + if (loadTsFileSpiltPartitionMaxSize < 0) { throw new IllegalArgumentException( DataNodeMiscMessages .MISC_EXCEPTION_LOADTSFILESPILTPARTITIONMAXSIZE_SHOULD_BE_GREATER_THAN_OR_95B4DB23); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java index 96bbb6ccc16ba..7892b32604052 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/consensus/statemachine/dataregion/DataExecutionVisitor.java @@ -34,6 +34,7 @@ import org.apache.iotdb.db.i18n.DataNodeMiscMessages; import org.apache.iotdb.db.i18n.DataNodePipeMessages; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedDeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.pipe.PipeEnrichedInsertNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; @@ -47,6 +48,7 @@ import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertRowNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertRowsNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.RelationalInsertTabletNode; +import org.apache.iotdb.db.storageengine.StorageEngine; import org.apache.iotdb.db.storageengine.dataregion.DataRegion; import org.apache.iotdb.rpc.RpcUtils; import org.apache.iotdb.rpc.TSStatusCode; @@ -316,4 +318,16 @@ public TSStatus visitPipeEnrichedDeleteDataNode( public TSStatus visitWriteObjectFile(ObjectNode node, DataRegion dataRegion) { throw new UnsupportedOperationException(); } + + @Override + public TSStatus visitLoadTsFileConsensus(LoadTsFileConsensusNode node, DataRegion dataRegion) { + try { + return StorageEngine.getInstance() + .getLoadTsFileManager() + .applyConsensusRequest(dataRegion, node); + } catch (Exception e) { + LOGGER.error(DataNodeMiscMessages.ERROR_EXECUTING_PLAN_NODE, node, e); + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()).setMessage(e.getMessage()); + } + } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java index a06d10413fd42..77b0f6c59a434 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java @@ -119,6 +119,7 @@ import org.apache.iotdb.db.exception.StorageEngineException; import org.apache.iotdb.db.i18n.DataNodeMiscMessages; import org.apache.iotdb.db.i18n.DataNodeSchemaMessages; +import org.apache.iotdb.db.i18n.StorageEngineMessages; import org.apache.iotdb.db.partition.DataPartitionTableGenerator; import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; import org.apache.iotdb.db.protocol.client.ConfigNodeInfo; @@ -162,6 +163,7 @@ import org.apache.iotdb.db.queryengine.plan.expression.leaf.TimestampOperand; import org.apache.iotdb.db.queryengine.plan.parser.StatementGenerator; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.write.AlterEncodingCompressorNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.write.AlterTimeSeriesNode; @@ -625,15 +627,53 @@ public TLoadResp sendTsFilePieceNode(final TTsFilePieceReq req) { final ConsensusGroupId groupId = ConsensusGroupId.Factory.createFromTConsensusGroupId(req.consensusGroupId); - final LoadTsFilePieceNode pieceNode = (LoadTsFilePieceNode) PlanNodeType.deserialize(req.body); - if (pieceNode == null) { + final PlanNode planNode = PlanNodeType.deserialize(req.body); + if (planNode == null) { return createTLoadResp( new TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode())); } + if (planNode instanceof LoadTsFileConsensusNode) { + // LOAD consensus piece delivery outside the consensus log: a PULL makes the write node push + // one retained piece back to the requester (a follower that applied the WAL marker without + // having the chunk bytes); the pushed PIECE is cached until the marker's apply order allows + // it to be written. + final LoadTsFileConsensusNode loadNode = (LoadTsFileConsensusNode) planNode; + final DataRegion dataRegion = + StorageEngine.getInstance().getDataRegion((DataRegionId) groupId); + if (dataRegion == null) { + return createTLoadResp( + new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + StorageEngineMessages + .STORAGE_LOG_DATAREGION_NOT_FOUND_ON_THIS_DATANODE_WHEN_WRITING_PIECE_E5B5A888 + + groupId)); + } + try { + switch (loadNode.getOp()) { + case PULL: + return createTLoadResp( + StorageEngine.getInstance() + .getLoadTsFileManager() + .handlePullPiece(dataRegion, loadNode)); + case PIECE: + return createTLoadResp( + StorageEngine.getInstance() + .getLoadTsFileManager() + .cacheConsensusPiece(dataRegion, loadNode)); + default: + return createTLoadResp( + new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage(String.valueOf(loadNode.getOp()))); + } + } catch (Exception e) { + return createTLoadResp( + new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()).setMessage(e.getMessage())); + } + } + final LoadTsFilePieceNode pieceNode = (LoadTsFilePieceNode) planNode; final TSStatus resultStatus = StorageEngine.getInstance() .writeLoadTsFileNode((DataRegionId) groupId, pieceNode, req.uuid); - return createTLoadResp(resultStatus); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/executor/RegionWriteExecutor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/executor/RegionWriteExecutor.java index 6fc521a7f724c..57a06679b6a5c 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/executor/RegionWriteExecutor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/executor/RegionWriteExecutor.java @@ -46,6 +46,7 @@ import org.apache.iotdb.db.protocol.thrift.impl.DataNodeRegionManager; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.WritePlanNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.write.ActivateTemplateNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.write.AlterTimeSeriesNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.write.BatchActivateTemplateNode; @@ -387,6 +388,17 @@ public RegionExecutionResult visitWriteObjectFile( } } + @Override + public RegionExecutionResult visitLoadTsFileConsensus( + final LoadTsFileConsensusNode node, final WritePlanNodeExecutionContext context) { + context.getRegionWriteValidationRWLock().writeLock().lock(); + try { + return PlanVisitor.super.visitLoadTsFileConsensus(node, context); + } finally { + context.getRegionWriteValidationRWLock().writeLock().unlock(); + } + } + @Override public RegionExecutionResult visitDeleteTimeseries( final DeleteTimeSeriesNode node, final WritePlanNodeExecutionContext context) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java index 3e5e46a3c0c87..713d910183eab 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/DataNodePlanNodeDeserializer.java @@ -23,6 +23,7 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.plan.analyze.TypeProvider; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.read.CountSchemaMergeNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.metadata.read.DeviceSchemaFetchScanNode; @@ -172,6 +173,8 @@ public PlanNode deserializeFromWAL(DataInputStream stream) throws IOException { return RelationalDeleteDataNode.deserializeFromWAL(stream); case 2004: return ObjectNode.deserializeFromWAL(stream); + case 2010: + return LoadTsFileConsensusNode.deserializeFromWAL(stream); default: throw new IllegalArgumentException(DataNodeQueryMessages.INVALID_NODE_TYPE + nodeType); } @@ -201,6 +204,8 @@ public PlanNode deserializeFromWAL(ByteBuffer buffer) { return RelationalDeleteDataNode.deserializeFromWAL(buffer); case 2004: return ObjectNode.deserialize(buffer); + case 2010: + return LoadTsFileConsensusNode.deserialize(buffer); default: throw new IllegalArgumentException(DataNodeQueryMessages.INVALID_NODE_TYPE + nodeType); } @@ -481,6 +486,8 @@ public PlanNode deserialize(ByteBuffer buffer, short nodeType) { return RelationalDeleteDataNode.deserialize(buffer); case 2004: return ObjectNode.deserialize(buffer); + case 2010: + return LoadTsFileConsensusNode.deserialize(buffer); default: return super.deserialize(buffer, nodeType); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java index a5b886ff4b058..d3d26897e93eb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/PlanVisitor.java @@ -621,6 +621,12 @@ default R visitWriteObjectFile(ObjectNode node, C context) { return visitPlan(node, context); } + default R visitLoadTsFileConsensus( + org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode node, + C context) { + return visitPlan(node, context); + } + ///////////////////////////////////////////////////////////////////////////////////////////////// // Pipe Related Node ///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java new file mode 100644 index 0000000000000..4d7f7cf95b61f --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNode.java @@ -0,0 +1,782 @@ +/* + * 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.planner.plan.node.load; + +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.consensus.index.ProgressIndex; +import org.apache.iotdb.commons.exception.IllegalPathException; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.IPlanVisitor; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.plan.analyze.IAnalysis; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanVisitor; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.WritePlanNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.SearchNode; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.IWALByteBufferView; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntryType; +import org.apache.iotdb.db.storageengine.dataregion.wal.buffer.WALEntryValue; +import org.apache.iotdb.db.storageengine.load.LoadTsFileManager; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; + +import org.apache.tsfile.exception.NotImplementedException; +import org.apache.tsfile.exception.write.PageException; +import org.apache.tsfile.utils.PublicBAOS; +import org.apache.tsfile.utils.ReadWriteIOUtils; + +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.File; +import java.io.IOException; +import java.io.InputStream; +import java.io.RandomAccessFile; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** + * Consensus-backed LOAD request carrying Begin / Piece / Seal / Commit phases. Submitted once to + * the DataRegion write peer (or Ratis leader); replicas apply via the consensus state machine. + */ +public class LoadTsFileConsensusNode extends SearchNode implements WALEntryValue { + + private LoadTsFileConsensusOp op; + private String loadId; + private String tsFileId; + private boolean isTableModel; + private String database; + private int expectedPieceCount = -1; + private long pieceIndex = -1; + private long pieceOffset = -1; + private List pieceRefs = new ArrayList<>(); + private long dataSize; + private long checksum; + private int pieceCount = -1; + private long totalBytes; + private boolean isGeneratedByPipe; + private boolean deleteAfterLoad; + private List tsFileDataList = new ArrayList<>(); + private Map timePartition2ProgressIndex = new HashMap<>(); + private TRegionReplicaSet regionReplicaSet; + private ProgressIndex progressIndex; + private boolean isGeneratedByRemoteConsensusLeader; + + /** Endpoint ("ip:port") of the follower that issued a PULL, used to push the piece back. */ + private String pullSourceEndPoint; + + public LoadTsFileConsensusNode(PlanNodeId id) { + super(id); + } + + public static LoadTsFileConsensusNode begin( + PlanNodeId id, + String loadId, + String tsFileId, + boolean isTableModel, + String database, + int expectedPieceCount) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.BEGIN; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.isTableModel = isTableModel; + node.database = database == null ? "" : database; + node.expectedPieceCount = expectedPieceCount; + return node; + } + + public static LoadTsFileConsensusNode piece( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + long pieceOffset, + List tsFileDataList, + long checksum) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.PIECE; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.pieceIndex = pieceIndex; + node.pieceOffset = pieceOffset; + node.tsFileDataList = + tsFileDataList == null ? new ArrayList<>() : new ArrayList<>(tsFileDataList); + node.dataSize = node.tsFileDataList.stream().mapToLong(TsFileData::getDataSize).sum(); + // The checksum is part of the consensus contract and must be preserved exactly as supplied by + // the caller. Recomputing it here silently changes the value for callers that use a staged + // piece checksum (and made marker/piece validation disagree). + node.checksum = checksum; + return node; + } + + public static LoadTsFileConsensusNode pieceRef( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + String relativePath, + long offset, + long size, + long checksum) { + return pieceRefs( + id, + loadId, + tsFileId, + pieceIndex, + Collections.singletonList(new PieceRef(relativePath, offset, size)), + checksum, + size); + } + + public static LoadTsFileConsensusNode pieceRefs( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + List refs, + long checksum, + long dataSize) { + return pieceRefs(id, loadId, tsFileId, pieceIndex, refs, checksum, dataSize, null); + } + + public static LoadTsFileConsensusNode pieceRefs( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + List refs, + long checksum, + long dataSize, + List deletions) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.PIECE; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.pieceIndex = pieceIndex; + node.pieceRefs = refs == null ? new ArrayList<>() : new ArrayList<>(refs); + node.dataSize = dataSize; + node.checksum = checksum; + if (deletions != null) { + node.tsFileDataList = new ArrayList<>(deletions); + } + return node; + } + + /** + * Marker-only PIECE: carries the piece metadata (index, checksum, byte size) but neither chunk + * data nor piece refs. This is what the write node logs to the WAL after applying a chunk-data + * piece. The write node keeps the actual chunk bytes in its retained-piece store; a follower + * applies them only when this marker arrives (pulling the retained bytes back on demand), so the + * WAL stays at dozens of bytes per piece instead of the full LOAD bytes. + */ + public static LoadTsFileConsensusNode pieceMarker( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + long checksum, + long dataSize) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.PIECE; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.pieceIndex = pieceIndex; + node.checksum = checksum; + node.dataSize = dataSize; + return node; + } + + /** PULL request: asks the current write node to re-deliver one already-applied piece. */ + public static LoadTsFileConsensusNode pull( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + long checksum, + String pullSourceEndPoint) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.PULL; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.pieceIndex = pieceIndex; + node.checksum = checksum; + node.pullSourceEndPoint = pullSourceEndPoint; + return node; + } + + public static LoadTsFileConsensusNode pieceRef( + PlanNodeId id, + String loadId, + String tsFileId, + long pieceIndex, + String relativePath, + long offset, + long size, + long checksum, + long dataSize) { + final LoadTsFileConsensusNode node = + pieceRef(id, loadId, tsFileId, pieceIndex, relativePath, offset, size, checksum); + node.dataSize = dataSize; + return node; + } + + public static LoadTsFileConsensusNode prepare( + PlanNodeId id, + String loadId, + String tsFileId, + int pieceCount, + long totalBytes, + long checksum, + Map timePartition2ProgressIndex) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.PREPARE; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.pieceCount = pieceCount; + node.totalBytes = totalBytes; + node.checksum = checksum; + node.timePartition2ProgressIndex = + timePartition2ProgressIndex == null + ? new HashMap<>() + : new HashMap<>(timePartition2ProgressIndex); + return node; + } + + public static LoadTsFileConsensusNode abort( + PlanNodeId id, String loadId, String tsFileId, boolean isGeneratedByPipe) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.ABORT; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.isGeneratedByPipe = isGeneratedByPipe; + return node; + } + + public static LoadTsFileConsensusNode commit( + PlanNodeId id, + String loadId, + String tsFileId, + boolean isGeneratedByPipe, + boolean deleteAfterLoad, + Map timePartition2ProgressIndex) { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(id); + node.op = LoadTsFileConsensusOp.COMMIT; + node.loadId = loadId; + node.tsFileId = tsFileId; + node.isGeneratedByPipe = isGeneratedByPipe; + node.deleteAfterLoad = deleteAfterLoad; + node.timePartition2ProgressIndex = + timePartition2ProgressIndex == null + ? new HashMap<>() + : new HashMap<>(timePartition2ProgressIndex); + return node; + } + + public LoadTsFileConsensusOp getOp() { + return op; + } + + public String getLoadId() { + return loadId; + } + + public String getTsFileId() { + return tsFileId; + } + + public boolean isTableModel() { + return isTableModel; + } + + public String getDatabase() { + return database; + } + + public int getExpectedPieceCount() { + return expectedPieceCount; + } + + public long getPieceIndex() { + return pieceIndex; + } + + public long getPieceOffset() { + return pieceOffset; + } + + public String getRelativePath() { + return pieceRefs.isEmpty() ? null : pieceRefs.get(0).relativePath; + } + + public long getSize() { + return pieceRefs.stream().mapToLong(ref -> ref.size).sum(); + } + + public List getPieceRefs() { + return pieceRefs; + } + + public long getDataSize() { + return dataSize; + } + + public long getChecksum() { + return checksum; + } + + public boolean isGeneratedByRemoteConsensusLeader() { + return isGeneratedByRemoteConsensusLeader; + } + + public void markAsGeneratedByRemoteConsensusLeader() { + this.isGeneratedByRemoteConsensusLeader = true; + } + + public int getPieceCount() { + return pieceCount; + } + + public long getTotalBytes() { + return totalBytes; + } + + public boolean isGeneratedByPipe() { + return isGeneratedByPipe; + } + + public boolean isDeleteAfterLoad() { + return deleteAfterLoad; + } + + public List getTsFileDataList() { + return tsFileDataList; + } + + /** Whether this PIECE carries actual chunk/deletion data (write-node path) or is a marker. */ + public boolean hasChunkData() { + return !tsFileDataList.isEmpty(); + } + + public String getPullSourceEndPoint() { + return pullSourceEndPoint; + } + + public void setPullSourceEndPoint(String pullSourceEndPoint) { + this.pullSourceEndPoint = pullSourceEndPoint; + } + + public Map getTimePartition2ProgressIndex() { + return timePartition2ProgressIndex; + } + + public void setRegionReplicaSet(TRegionReplicaSet regionReplicaSet) { + this.regionReplicaSet = regionReplicaSet; + } + + @Override + public TRegionReplicaSet getRegionReplicaSet() { + return regionReplicaSet; + } + + @Override + public ProgressIndex getProgressIndex() { + return progressIndex; + } + + @Override + public void setProgressIndex(ProgressIndex progressIndex) { + this.progressIndex = progressIndex; + } + + @Override + public SearchNode merge(List searchNodes) { + if (searchNodes.size() == 1) { + return searchNodes.get(0); + } + throw new UnsupportedOperationException(DataNodeQueryMessages.MERGE_IS_NOT_SUPPORTED); + } + + @Override + public List splitByPartition(IAnalysis analysis) { + return Collections.singletonList(this); + } + + @Override + public List getChildren() { + return Collections.emptyList(); + } + + @Override + public void addChild(PlanNode child) { + // no children + } + + @Override + public PlanNodeType getType() { + return PlanNodeType.LOAD_TSFILE_CONSENSUS; + } + + @Override + public PlanNode clone() { + throw new NotImplementedException( + DataNodeQueryMessages.CLONE_OF_LOAD_PIECE_TSFILE_IS_NOT_IMPLEMENTED); + } + + @Override + public int allowedChildCount() { + return NO_CHILD_ALLOWED; + } + + @Override + public List getOutputColumnNames() { + return Collections.emptyList(); + } + + @Override + public R accept(IPlanVisitor visitor, C context) { + return ((PlanVisitor) visitor).visitLoadTsFileConsensus(this, context); + } + + @Override + protected void serializeAttributes(ByteBuffer byteBuffer) { + PlanNodeType.LOAD_TSFILE_CONSENSUS.serialize(byteBuffer); + try { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final DataOutputStream stream = new DataOutputStream(baos); + serializeBody(stream, true); + byteBuffer.put(baos.toByteArray()); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + @Override + protected void serializeAttributes(DataOutputStream stream) throws IOException { + PlanNodeType.LOAD_TSFILE_CONSENSUS.serialize(stream); + serializeBody(stream, true); + } + + private void serializeBody(DataOutputStream stream) throws IOException { + serializeBody(stream, false); + } + + private void serializeBody(DataOutputStream stream, boolean includeContent) throws IOException { + ReadWriteIOUtils.write(op.ordinal(), stream); + ReadWriteIOUtils.write(loadId, stream); + ReadWriteIOUtils.write(tsFileId, stream); + ReadWriteIOUtils.write(isTableModel, stream); + ReadWriteIOUtils.write(database == null ? "" : database, stream); + ReadWriteIOUtils.write(expectedPieceCount, stream); + ReadWriteIOUtils.write(pieceIndex, stream); + ReadWriteIOUtils.write(pieceOffset, stream); + ReadWriteIOUtils.write(dataSize, stream); + ReadWriteIOUtils.write(checksum, stream); + ReadWriteIOUtils.write(pieceCount, stream); + ReadWriteIOUtils.write(totalBytes, stream); + ReadWriteIOUtils.write(isGeneratedByPipe, stream); + ReadWriteIOUtils.write(deleteAfterLoad, stream); + ReadWriteIOUtils.write(pieceRefs.size(), stream); + for (PieceRef ref : pieceRefs) { + ReadWriteIOUtils.write(ref.relativePath, stream); + ReadWriteIOUtils.write(ref.offset, stream); + ReadWriteIOUtils.write(ref.size, stream); + if (includeContent) { + final byte[] content = readPieceContent(ref); + ReadWriteIOUtils.write(content.length, stream); + stream.write(content); + } + } + ReadWriteIOUtils.write(tsFileDataList.size(), stream); + for (TsFileData data : tsFileDataList) { + data.serialize(stream); + } + ReadWriteIOUtils.write(timePartition2ProgressIndex.size(), stream); + for (Map.Entry entry : timePartition2ProgressIndex.entrySet()) { + ReadWriteIOUtils.write(entry.getKey().getStartTime(), stream); + ReadWriteIOUtils.write(entry.getValue().length, stream); + stream.write(entry.getValue()); + } + ReadWriteIOUtils.write(pullSourceEndPoint == null ? "" : pullSourceEndPoint, stream); + } + + public static LoadTsFileConsensusNode deserialize(ByteBuffer buffer) { + try { + final LoadTsFilePieceNode.ByteBufferInputStream stream = + new LoadTsFilePieceNode.ByteBufferInputStream(buffer); + final LoadTsFileConsensusNode node = deserializeBody(stream, true); + final PlanNodeId planNodeId = PlanNodeId.deserialize(buffer); + node.setPlanNodeId(planNodeId); + ReadWriteIOUtils.readInt(buffer); + return node; + } catch (IOException | PageException | IllegalPathException e) { + throw new IllegalStateException(e); + } + } + + private static LoadTsFileConsensusNode deserializeBody(InputStream stream) + throws IOException, PageException, IllegalPathException { + return deserializeBody(stream, false); + } + + private static LoadTsFileConsensusNode deserializeBody(InputStream stream, boolean readContent) + throws IOException, PageException, IllegalPathException { + final LoadTsFileConsensusNode node = new LoadTsFileConsensusNode(new PlanNodeId("")); + node.op = LoadTsFileConsensusOp.fromOrdinal(ReadWriteIOUtils.readInt(stream)); + node.loadId = ReadWriteIOUtils.readString(stream); + node.tsFileId = ReadWriteIOUtils.readString(stream); + node.isTableModel = ReadWriteIOUtils.readBool(stream); + node.database = ReadWriteIOUtils.readString(stream); + node.expectedPieceCount = ReadWriteIOUtils.readInt(stream); + node.pieceIndex = ReadWriteIOUtils.readLong(stream); + node.pieceOffset = ReadWriteIOUtils.readLong(stream); + node.dataSize = ReadWriteIOUtils.readLong(stream); + node.checksum = ReadWriteIOUtils.readLong(stream); + node.pieceCount = ReadWriteIOUtils.readInt(stream); + node.totalBytes = ReadWriteIOUtils.readLong(stream); + node.isGeneratedByPipe = ReadWriteIOUtils.readBool(stream); + node.deleteAfterLoad = ReadWriteIOUtils.readBool(stream); + final int refCount = ReadWriteIOUtils.readInt(stream); + node.pieceRefs = new ArrayList<>(refCount); + for (int i = 0; i < refCount; i++) { + final String refPath = ReadWriteIOUtils.readString(stream); + final long refOffset = ReadWriteIOUtils.readLong(stream); + final long refSize = ReadWriteIOUtils.readLong(stream); + byte[] content = null; + if (readContent) { + final int contentLength = ReadWriteIOUtils.readInt(stream); + if (contentLength > 0) { + content = new byte[contentLength]; + int offset = 0; + while (offset < contentLength) { + final int read = stream.read(content, offset, contentLength - offset); + if (read < 0) { + throw new IOException( + DataNodeQueryMessages.EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 + + "pieceContent"); + } + offset += read; + } + } + } + node.pieceRefs.add(new PieceRef(refPath, refOffset, refSize, content)); + } + final int dataCount = ReadWriteIOUtils.readInt(stream); + node.tsFileDataList = new ArrayList<>(dataCount); + for (int i = 0; i < dataCount; i++) { + node.tsFileDataList.add(TsFileData.deserialize(stream)); + } + final int progressCount = ReadWriteIOUtils.readInt(stream); + node.timePartition2ProgressIndex = new HashMap<>(progressCount); + for (int i = 0; i < progressCount; i++) { + final long startTime = ReadWriteIOUtils.readLong(stream); + final int len = ReadWriteIOUtils.readInt(stream); + final byte[] bytes = new byte[len]; + int offset = 0; + while (offset < len) { + final int read = stream.read(bytes, offset, len - offset); + if (read < 0) { + throw new IOException( + DataNodeQueryMessages.EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 + + "progressIndex"); + } + offset += read; + } + node.timePartition2ProgressIndex.put(new TTimePartitionSlot(startTime), bytes); + } + node.pullSourceEndPoint = ReadWriteIOUtils.readString(stream); + return node; + } + + @Override + public void serializeToWAL(IWALByteBufferView buffer) { + serializeToWAL(buffer, getEncodedSearchIndex()); + } + + public void serializeToWAL(IWALByteBufferView buffer, long encodedSearchIndex) { + try { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final DataOutputStream stream = new DataOutputStream(baos); + stream.writeShort(getType().getNodeType()); + stream.writeLong(encodedSearchIndex); + serializeBody(stream); + buffer.put(baos.toByteArray()); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + /** Serialize this node into an IoTConsensusRequest-compatible buffer with content expanded. */ + public ByteBuffer serialize() { + try (PublicBAOS byteArrayOutputStream = new PublicBAOS(); + DataOutputStream stream = new DataOutputStream(byteArrayOutputStream)) { + ReadWriteIOUtils.write(WALEntryType.LOAD_TSFILE_CONSENSUS_NODE.getCode(), stream); + ReadWriteIOUtils.write(-1L, stream); + ReadWriteIOUtils.write(getType().getNodeType(), stream); + serializeBody(stream, true); + getPlanNodeId().serialize(stream); + ReadWriteIOUtils.write(0, stream); + return ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + private byte[] readPieceContent(PieceRef ref) { + if (ref.content != null) { + return ref.content; + } + final byte[] content = new byte[(int) Math.max(0, ref.size)]; + if (content.length == 0) { + return content; + } + final File file = + LoadTsFileManager.findLoadTsFile(ref.relativePath) + .orElseThrow( + () -> + new IllegalStateException( + DataNodeQueryMessages + .EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 + + "piece file not found: " + + ref.relativePath)); + try (RandomAccessFile raf = new RandomAccessFile(file, "r")) { + raf.seek(ref.offset); + raf.readFully(content); + return content; + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + @Override + public int serializedSize() { + try { + final ByteArrayOutputStream baos = new ByteArrayOutputStream(); + final DataOutputStream stream = new DataOutputStream(baos); + serializeBody(stream); + return Short.BYTES + Long.BYTES + baos.size(); + } catch (IOException e) { + throw new IllegalStateException(e); + } + } + + public static LoadTsFileConsensusNode deserializeFromWAL(DataInputStream stream) + throws IOException { + final long searchIndex = stream.readLong(); + try { + final LoadTsFileConsensusNode node = deserializeBody(stream, false); + node.setSearchIndexFromWAL(searchIndex); + return node; + } catch (PageException | IllegalPathException e) { + throw new IOException(e); + } + } + + public static LoadTsFileConsensusNode deserializeFromWAL(ByteBuffer buffer) { + final long searchIndex = buffer.getLong(); + try { + final LoadTsFileConsensusNode node = + deserializeBody(new LoadTsFilePieceNode.ByteBufferInputStream(buffer), false); + node.setSearchIndexFromWAL(searchIndex); + return node; + } catch (IOException | PageException | IllegalPathException e) { + throw new IllegalStateException(e); + } + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof LoadTsFileConsensusNode)) { + return false; + } + final LoadTsFileConsensusNode that = (LoadTsFileConsensusNode) o; + return Objects.equals(loadId, that.loadId) + && Objects.equals(tsFileId, that.tsFileId) + && op == that.op + && pieceIndex == that.pieceIndex + && checksum == that.checksum; + } + + @Override + public int hashCode() { + return Objects.hash(op, loadId, tsFileId, pieceIndex, checksum); + } + + @Override + public String toString() { + return "LoadTsFileConsensusNode{op=" + + op + + ", loadId=" + + loadId + + ", tsFileId=" + + tsFileId + + ", pieceIndex=" + + pieceIndex + + '}'; + } + + /** A staging-file reference recorded by the consensus/WAL entry instead of the raw content. */ + public static class PieceRef { + private final String relativePath; + private final long offset; + private final long size; + private final byte[] content; + + public PieceRef(String relativePath, long size) { + this(relativePath, 0L, size, null); + } + + public PieceRef(String relativePath, long offset, long size) { + this(relativePath, offset, size, null); + } + + public PieceRef(String relativePath, long offset, long size, byte[] content) { + this.relativePath = relativePath; + this.offset = offset; + this.size = size; + this.content = content; + } + + public String getRelativePath() { + return relativePath; + } + + public long getOffset() { + return offset; + } + + public long getSize() { + return size; + } + + public byte[] getContent() { + return content; + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusOp.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusOp.java new file mode 100644 index 0000000000000..39ebe0319568e --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusOp.java @@ -0,0 +1,50 @@ +/* + * 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.planner.plan.node.load; + +/** + * Consensus LOAD request phases, sequenced by the scheduler/client. + * + *

{@link #BEGIN} opens a load, {@link #PIECE} stages one data piece, {@link #PREPARE} seals the + * data of a load, {@link #COMMIT} imports the staged data, {@link #ABORT} discards it, and {@link + * #PULL} asks the current write node to re-deliver the serialized bytes of one already-applied + * piece so a follower that did not receive the data yet can catch up. Phase enforcement and + * idempotency are the client's responsibility; the server applies each request against the staged + * data directory identified by the load id without keeping per-load transaction state in memory. + */ +public enum LoadTsFileConsensusOp { + BEGIN, + PIECE, + PREPARE, + COMMIT, + ABORT, + PULL; + + public static LoadTsFileConsensusOp fromOrdinal(int ordinal) { + final LoadTsFileConsensusOp[] values = values(); + if (ordinal < 0 || ordinal >= values.length) { + throw new IllegalArgumentException( + org.apache.iotdb.db.i18n.DataNodeQueryMessages + .EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 + + ordinal); + } + return values[ordinal]; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionBatchFetcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionBatchFetcher.java new file mode 100644 index 0000000000000..d6e3da78998dc --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionBatchFetcher.java @@ -0,0 +1,114 @@ +/* + * 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.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.partition.DataPartition; +import org.apache.iotdb.commons.partition.DataPartitionQueryParam; +import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.utils.Pair; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Batch partition fetcher for LOAD. It wraps the query-engine {@link IPartitionFetcher} and hides + * two details from the rest of the pipeline: + * + *

    + *
  • Transmit limit. The requested (device, time-partition) pairs are split into batches + * of at most {@code TTimePartitionSlotTransmitLimit} entries; each batch is resolved with one + * {@code getOrCreateDataPartition} call. + *
  • Database hint. {@link #setDatabase(String)} enables the explicit database lookup + * used by table-model loads and pipe-generated tree-model loads. + *
+ * + * {@link #queryDataPartition(List, String)} returns one {@link TRegionReplicaSet} per input pair, + * in the same order, which {@link DataPartitionRouter} maps back onto chunks. + */ +class DataPartitionBatchFetcher { + + private static final int TRANSMIT_LIMIT = + CommonDescriptor.getInstance().getConfig().getTTimePartitionSlotTransmitLimit(); + + private final IPartitionFetcher fetcher; + private String database; + + DataPartitionBatchFetcher(IPartitionFetcher fetcher) { + this.fetcher = fetcher; + } + + void setDatabase(String database) { + this.database = database; + } + + List queryDataPartition( + List> slotList, String userName) { + List replicaSets = new ArrayList<>(slotList.size()); + int size = slotList.size(); + + for (int i = 0; i < size; i += TRANSMIT_LIMIT) { + List> subSlotList = + slotList.subList(i, Math.min(size, i + TRANSMIT_LIMIT)); + DataPartition dataPartition = + fetcher.getOrCreateDataPartition(toQueryParam(subSlotList), userName); + for (final Pair pair : subSlotList) { + // database is an explicit database hint for table-model loads and + // pipe-generated tree-model loads. + replicaSets.add( + database != null + ? dataPartition.getDataRegionReplicaSetForWriting(pair.left, pair.right, database) + : dataPartition.getDataRegionReplicaSetForWriting(pair.left, pair.right)); + } + } + return replicaSets; + } + + private List toQueryParam( + List> slots) { + final Map> device2TimePartitionSlots = new HashMap<>(); + for (final Pair slot : slots) { + device2TimePartitionSlots.computeIfAbsent(slot.left, key -> new HashSet<>()).add(slot.right); + } + + final List queryParams = + new ArrayList<>(device2TimePartitionSlots.size()); + for (final Map.Entry> entry : + device2TimePartitionSlots.entrySet()) { + final DataPartitionQueryParam queryParam = + new DataPartitionQueryParam(entry.getKey(), new ArrayList<>(entry.getValue())); + // database is an explicit database hint for table-model loads and + // pipe-generated tree-model loads. + if (database != null) { + queryParam.setDatabaseName(database); + } + queryParams.add(queryParam); + } + return queryParams; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.java new file mode 100644 index 0000000000000..16c53c7816e03 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/DataPartitionRouter.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.plan.scheduler.load; + +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.utils.Pair; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * LOAD chunk router: maps chunks to their target regions. {@link #route(List)} takes the + * directionless chunk buffer and: + * + *
    + *
  1. deduplicates the (device, time-partition) pairs, so many chunks over the same slot issue + * only one partition query per distinct slot; + *
  2. resolves every distinct pair through {@link DataPartitionBatchFetcher}; + *
  3. returns the replica set of every input chunk, preserving the input order. + *
+ * + * The caller ({@code TsFileSplitConsumer.routeChunkData()}) feeds the result to {@link + * PieceDispatcher}, which performs the replica-set-change detection and appends the chunk to the + * per-region piece. + */ +class DataPartitionRouter { + + private final DataPartitionBatchFetcher partitionFetcher; + private final String userName; + + DataPartitionRouter(DataPartitionBatchFetcher partitionFetcher, String userName) { + this.partitionFetcher = partitionFetcher; + this.userName = userName; + } + + /** + * Returns, for every chunk in the input list (same order), the region replica set it must be + * written to. + */ + List route(List chunkDataList) { + if (chunkDataList.isEmpty()) { + return new ArrayList<>(); + } + + final List> partitionSlotList = new ArrayList<>(); + final int[] chunkPartitionIndexes = new int[chunkDataList.size()]; + final Map> partitionSlotIndexes = new HashMap<>(); + for (int i = 0, size = chunkDataList.size(); i < size; i++) { + final ChunkData chunkData = chunkDataList.get(i); + final IDeviceID device = chunkData.getDevice(); + final TTimePartitionSlot timePartitionSlot = chunkData.getTimePartitionSlot(); + final Map slotIndexes = + partitionSlotIndexes.computeIfAbsent(device, key -> new HashMap<>()); + Integer partitionSlotIndex = slotIndexes.get(timePartitionSlot); + if (partitionSlotIndex == null) { + partitionSlotIndex = partitionSlotList.size(); + slotIndexes.put(timePartitionSlot, partitionSlotIndex); + partitionSlotList.add(new Pair<>(device, timePartitionSlot)); + } + chunkPartitionIndexes[i] = partitionSlotIndex; + } + + final List replicaSets = + partitionFetcher.queryDataPartition(partitionSlotList, userName); + final List routedReplicaSets = new ArrayList<>(chunkDataList.size()); + for (int i = 0, size = chunkDataList.size(); i < size; i++) { + routedReplicaSets.add(replicaSets.get(chunkPartitionIndexes[i])); + } + return routedReplicaSets; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadConsensusSubmitter.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadConsensusSubmitter.java new file mode 100644 index 0000000000000..f428236700eff --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadConsensusSubmitter.java @@ -0,0 +1,202 @@ +/* + * 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.TDataNodeLocation; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient; +import org.apache.iotdb.commons.consensus.ConsensusGroupId; +import org.apache.iotdb.consensus.ConsensusFactory; +import org.apache.iotdb.consensus.common.Peer; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.consensus.DataRegionConsensusImpl; +import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.execution.executor.RegionExecutionResult; +import org.apache.iotdb.db.queryengine.execution.executor.RegionWriteExecutor; +import org.apache.iotdb.db.queryengine.plan.analyze.ClusterPartitionFetcher; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; +import org.apache.iotdb.mpp.rpc.thrift.TPlanNode; +import org.apache.iotdb.mpp.rpc.thrift.TSendBatchPlanNodeReq; +import org.apache.iotdb.mpp.rpc.thrift.TSendSinglePlanNodeReq; +import org.apache.iotdb.mpp.rpc.thrift.TSendSinglePlanNodeResp; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Collections; +import java.util.List; + +/** + * Transport for LOAD consensus commands (BEGIN / PIECE / PREPARE / COMMIT / ABORT). One instance is + * shared by all files of a scheduler; it is stateless besides the local endpoint and the client + * manager. + * + *

{@link #submit(TRegionReplicaSet, LoadTsFileConsensusNode)}: + * + *

    + *
  1. stamps the target {@code regionReplicaSet} onto the node for correlation only - no follower + * endpoints are carried, replicas receive the command through consensus log replication + * exactly like ordinary writes; + *
  2. resolves the single write peer of the partition: the current Ratis leader when the protocol + * is Ratis, otherwise the first replica-set location (the IoTConsensus write node, the same + * target the normal write path dispatches to); + *
  3. writes the command through {@link RegionWriteExecutor} (local) or the internal RPC ({@code + * sendBatchPlanNode}) on that peer, which applies it via {@code + * DataRegionConsensusImpl.write} like any other write plan. + *
+ * + *

IoTConsensus replicates the WAL entries (marker-only for LOAD pieces) to the followers, whose + * own {@code TsFileWriterManager} rebuilds the staged files; the chunk bytes are pulled back from + * the write node on demand. Ratis replicates the full command through its own log, so every replica + * applies the chunk data directly and keeps its own writer. + */ +public class LoadConsensusSubmitter { + + private static final Logger LOGGER = LoggerFactory.getLogger(LoadConsensusSubmitter.class); + + private final String localhostIp; + private final int localhostPort; + private final IClientManager clientManager; + + public LoadConsensusSubmitter( + IClientManager clientManager) { + this.clientManager = clientManager; + this.localhostIp = IoTDBDescriptor.getInstance().getConfig().getInternalAddress(); + this.localhostPort = IoTDBDescriptor.getInstance().getConfig().getInternalPort(); + } + + public TSStatus submit(TRegionReplicaSet replicaSet, LoadTsFileConsensusNode node) { + final ConsensusGroupId regionId = + ConsensusGroupId.Factory.createFromTConsensusGroupId(replicaSet.getRegionId()); + // Follow the freshest partition route: after a write-node switch (leader change / region + // migration) the replica set captured at split time may be stale, so re-resolve it from the + // local partition table (a cache miss fetches the latest route map from the ConfigNode) and + // submit to the current write node, exactly like normal writes. + final TRegionReplicaSet currentReplicaSet = refreshReplicaSet(replicaSet, regionId); + node.setRegionReplicaSet(currentReplicaSet); + + final String protocol = + IoTDBDescriptor.getInstance().getConfig().getDataRegionConsensusProtocolClass(); + LOGGER.info( + StorageEngineMessages.LOG_LOAD_CONSENSUS_WRITE_TO_REGION_ARG_VIA_PROTOCOL_ARG_EBB55042, + regionId, + protocol); + + final TDataNodeLocation writePeer = resolveWritePeer(currentReplicaSet, regionId, protocol); + if (writePeer == null) { + return new TSStatus(TSStatusCode.DISPATCH_ERROR.getStatusCode()) + .setMessage(String.valueOf(replicaSet)); + } + return isLocal(writePeer.getInternalEndPoint()) + ? writeLocal(regionId, node) + : writeRemote(writePeer.getInternalEndPoint(), regionId, node); + } + + /** + * Re-resolves the partition replica set from the local partition table (falling back to the + * passed set when the lookup fails). IoTConsensus V1 has no leader election, so after a write + * node switch the coordinator must route by the refreshed route map instead of the stale set. + */ + private TRegionReplicaSet refreshReplicaSet( + TRegionReplicaSet replicaSet, ConsensusGroupId regionId) { + try { + final List replicaSets = + ClusterPartitionFetcher.getInstance() + .getRegionReplicaSet( + Collections.singletonList(regionId.convertToTConsensusGroupId())); + if (!replicaSets.isEmpty()) { + return replicaSets.get(0); + } + } catch (Exception e) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_REFRESH_REPLICA_SET_FAILED_7C244C63, + regionId, + e.getMessage()); + } + return replicaSet; + } + + /** + * Resolves the single write peer of the partition. Ratis writes must land on the current leader, + * so the leader endpoint is matched against the replica-set locations first (falling back to the + * first location while the leader is not known yet); IoTConsensus routes every write to the first + * replica-set location, which is the partition's write node - the same target normal writes use. + */ + private TDataNodeLocation resolveWritePeer( + TRegionReplicaSet replicaSet, ConsensusGroupId regionId, String protocol) { + final List locations = replicaSet.getDataNodeLocations(); + if (locations == null || locations.isEmpty()) { + return null; + } + if (ConsensusFactory.RATIS_CONSENSUS.equals(protocol)) { + final Peer leader = DataRegionConsensusImpl.getInstance().getLeader(regionId); + if (leader != null) { + for (TDataNodeLocation location : locations) { + final TEndPoint endPoint = location.getInternalEndPoint(); + if (endPoint != null && endPoint.getIp().equals(leader.getEndpoint().getIp())) { + return location; + } + } + } + } + return locations.get(0); + } + + private TSStatus writeLocal(ConsensusGroupId regionId, LoadTsFileConsensusNode node) { + final RegionWriteExecutor executor = new RegionWriteExecutor(); + final RegionExecutionResult result = executor.execute(regionId, node); + return result.getStatus(); + } + + private TSStatus writeRemote( + TEndPoint endPoint, ConsensusGroupId regionId, LoadTsFileConsensusNode node) { + try (SyncDataNodeInternalServiceClient client = clientManager.borrowClient(endPoint)) { + final TSendSinglePlanNodeReq singleReq = + new TSendSinglePlanNodeReq( + new TPlanNode(node.serializeToByteBuffer()), regionId.convertToTConsensusGroupId()); + final TSendBatchPlanNodeReq batchReq = + new TSendBatchPlanNodeReq(Collections.singletonList(singleReq)); + final List responses = + client.sendBatchPlanNode(batchReq).getResponses(); + if (responses == null || responses.isEmpty()) { + return new TSStatus(TSStatusCode.DISPATCH_ERROR.getStatusCode()); + } + final TSendSinglePlanNodeResp resp = responses.get(0); + if (resp.isAccepted()) { + return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()); + } + return resp.getStatus() != null + ? resp.getStatus() + : new TSStatus(TSStatusCode.DISPATCH_ERROR.getStatusCode()); + } catch (Exception e) { + return new TSStatus(TSStatusCode.DISPATCH_ERROR.getStatusCode()).setMessage(e.getMessage()); + } + } + + private boolean isLocal(TEndPoint endPoint) { + return endPoint != null + && localhostIp.equals(endPoint.getIp()) + && localhostPort == endPoint.getPort(); + } +} 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 new file mode 100644 index 0000000000000..c73fb444ff69b --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadFallbackHandler.java @@ -0,0 +1,212 @@ +/* + * 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.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.execution.QueryStateMachine; +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.load.converter.LoadTsFileDataTypeConverter; +import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.FileNotFoundException; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.ListIterator; +import java.util.stream.Collectors; + +/** + * Failure fallback of the LOAD scheduler: when at least one TsFile fails, this handler converts the + * failed files into tablets and retries the insertion, then resolves the final result. Extracted + * from the scheduler so the failure path is a single, testable component. + * + *

Flow ({@link #convertFailedTsFilesToTablets()}): + * + *

    + *
  1. builds the comma-separated list of failed file paths for logging; + *
  2. for every failed file, converts it with {@link LoadTsFileDataTypeConverter} - table-model + * files through {@code convertForTableModel}, tree-model files through {@code + * convertForTreeModel} with a retry statement built by {@code buildRetryTreeLoadStatement}; + *
  3. removes successfully converted files from the failure list; + *
  4. if no failure remains the load is considered successful (FINISHED); otherwise the state + * machine transitions to FAILED with the remaining file list. + *
+ * + * The whole conversion is measured as the {@code SCHEDULER_CAST_TABLETS} phase metric. + */ +public class LoadFallbackHandler { + + private static final Logger LOGGER = LoggerFactory.getLogger(LoadFallbackHandler.class); + + private static final LoadTsFileCostMetricsSet LOAD_TSFILE_COST_METRICS_SET = + LoadTsFileCostMetricsSet.getInstance(); + + private final MPPQueryContext queryContext; + private final boolean isGeneratedByPipe; + private final List tsFileNodeList; + private final List failedTsFileNodeIndexes; + private final QueryStateMachine stateMachine; + + public LoadFallbackHandler( + MPPQueryContext queryContext, + boolean isGeneratedByPipe, + List tsFileNodeList, + List failedTsFileNodeIndexes, + QueryStateMachine stateMachine) { + this.queryContext = queryContext; + this.isGeneratedByPipe = isGeneratedByPipe; + this.tsFileNodeList = tsFileNodeList; + this.failedTsFileNodeIndexes = failedTsFileNodeIndexes; + this.stateMachine = stateMachine; + } + + public void convertFailedTsFilesToTablets() { + final StringBuilder failedTsFiles = + new StringBuilder( + !tsFileNodeList.isEmpty() + ? tsFileNodeList + .get(failedTsFileNodeIndexes.get(0)) + .getTsFileResource() + .getTsFilePath() + : ""); + final ListIterator iterator = failedTsFileNodeIndexes.listIterator(1); + while (iterator.hasNext()) { + failedTsFiles + .append(", ") + .append(tsFileNodeList.get(iterator.next()).getTsFileResource().getTsFilePath()); + } + + final long startTime = System.nanoTime(); + try { + // if failed to load some TsFiles, then try to convert the TsFiles to Tablets + LOGGER.info( + DataNodeQueryMessages + .LOAD_TSFILE_S_FAILED_WILL_TRY_TO_CONVERT_TO_TABLETS_AND_INSERT_FAILED_TSFILES_ARG, + failedTsFiles); + convertFailedTsFilesToTabletsAndRetry(); + } finally { + LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( + LoadTsFileCostMetricsSet.SCHEDULER_CAST_TABLETS, System.nanoTime() - startTime); + } + } + + private void convertFailedTsFilesToTabletsAndRetry() { + final LoadTsFileDataTypeConverter loadTsFileDataTypeConverter = + new LoadTsFileDataTypeConverter(queryContext, isGeneratedByPipe); + + final Iterator iterator = failedTsFileNodeIndexes.listIterator(); + while (iterator.hasNext()) { + final int failedLoadTsFileIndex = iterator.next(); + final LoadSingleTsFileNode failedNode = tsFileNodeList.get(failedLoadTsFileIndex); + final String filePath = failedNode.getTsFileResource().getTsFilePath(); + + try { + final TSStatus status = + failedNode.isTableModel() + ? loadTsFileDataTypeConverter + .convertForTableModel( + (isGeneratedByPipe + ? LoadTsFile.createForPipe(null, filePath, Collections.emptyMap()) + : LoadTsFile.createUnchecked( + null, filePath, Collections.emptyMap())) + .setDatabase(failedNode.getDatabase()) + .setDeleteAfterLoad(failedNode.isDeleteAfterLoad()) + .setConvertOnTypeMismatch(true)) + .orElse(null) + : loadTsFileDataTypeConverter + .convertForTreeModel( + buildRetryTreeLoadStatement( + filePath, + failedNode.isDeleteAfterLoad(), + LoadTsFileScheduler.getPartitionQueryDatabase( + failedNode, isGeneratedByPipe))) + .orElse(null); + + if (loadTsFileDataTypeConverter.isSuccessful(status)) { + iterator.remove(); + LOGGER.info( + DataNodeQueryMessages + .LOAD_SUCCESSFULLY_CONVERTED_TSFILE_ARG_INTO_TABLETS_AND_INSERTED, + failedNode.getTsFileResource().getTsFilePath()); + } else { + LOGGER.warn( + DataNodeQueryMessages.LOAD_FAILED_TO_CONVERT_TO_TABLETS_FROM_TSFILE_ARG_STATUS_ARG, + failedNode.getTsFileResource().getTsFilePath(), + status); + } + } catch (final Exception e) { + LOGGER.warn( + DataNodeQueryMessages.LOAD_FAILED_TO_CONVERT_TO_TABLETS_FROM_TSFILE_ARG_EXCEPTION_ARG, + failedNode.getTsFileResource().getTsFilePath(), + e.getMessage(), + e); + } + } + + // If all failed TsFiles are converted into tablets and inserted, + // we can consider the load process as successful. + if (failedTsFileNodeIndexes.isEmpty()) { + LOGGER.info(DataNodeQueryMessages.LOAD_ALL_FAILED_TSFILES_ARE_CONVERTED_TO_TABLETS); + stateMachine.transitionToFinished(); + } else { + final String failedFiles = + failedTsFileNodeIndexes.stream() + .map(i -> tsFileNodeList.get(i).getTsFileResource().getTsFilePath()) + .collect(Collectors.joining(", ")); + LOGGER.warn( + DataNodeQueryMessages + .LOG_LOAD_FAILED_TO_LOAD_SOME_TSFILES_BY_CONVERTING_THEM_INTO_TABLETS_FAILED_TSFILES_ARG_7D9DB9C3, + failedFiles); + stateMachine.transitionToFailed( + new LoadFileException( + String.format( + DataNodeQueryMessages + .LOG_LOAD_FAILED_TO_LOAD_SOME_TSFILES_BY_CONVERTING_THEM_INTO_TABLETS_FAILED_TSFILES_ARG_7D9DB9C3, + failedFiles))); + } + } + + private LoadTsFileStatement buildRetryTreeLoadStatement( + final String filePath, final boolean deleteAfterLoad, final String database) + throws FileNotFoundException { + final LoadTsFileStatement statement = + (isGeneratedByPipe + ? LoadTsFileStatement.createForPipe(filePath) + : LoadTsFileStatement.createUnchecked(filePath)) + .setDeleteAfterLoad(deleteAfterLoad) + .setConvertOnTypeMismatch(true); + if (database != null) { + statement.setDatabase(database); + statement.updateDatabaseLevelByTreeDatabase(); + } + if (isGeneratedByPipe) { + statement.markIsGeneratedByPipe(); + } + return statement; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java index 3f5020cccf31e..de3963fee03e0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java @@ -74,6 +74,10 @@ import static com.google.common.util.concurrent.Futures.immediateFuture; +/** + * LOAD dispatcher: legacy local dispatcher of the LOAD local-load path (no decode needed) and the + * per-file uuid holder used for executor naming and log correlation. + */ public class LoadTsFileDispatcherImpl implements IFragInstanceDispatcher, AutoCloseable { private static final Logger LOGGER = LoggerFactory.getLogger(LoadTsFileDispatcherImpl.class); 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 da143c6b8cb40..584c6f61db32c 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 @@ -19,32 +19,11 @@ package org.apache.iotdb.db.queryengine.plan.scheduler.load; -import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; import org.apache.iotdb.common.rpc.thrift.TEndPoint; -import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; -import org.apache.iotdb.common.rpc.thrift.TSStatus; -import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; import org.apache.iotdb.commons.client.IClientManager; import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient; -import org.apache.iotdb.commons.conf.CommonDescriptor; -import org.apache.iotdb.commons.consensus.ConsensusGroupId; -import org.apache.iotdb.commons.consensus.DataRegionId; -import org.apache.iotdb.commons.consensus.index.ProgressIndex; -import org.apache.iotdb.commons.exception.IoTDBException; -import org.apache.iotdb.commons.partition.DataPartition; -import org.apache.iotdb.commons.partition.DataPartitionQueryParam; -import org.apache.iotdb.commons.partition.StorageExecutor; -import org.apache.iotdb.commons.service.metric.MetricService; -import org.apache.iotdb.commons.service.metric.enums.Metric; -import org.apache.iotdb.commons.service.metric.enums.Tag; -import org.apache.iotdb.db.conf.IoTDBConfig; -import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.load.LoadFileException; -import org.apache.iotdb.db.exception.load.LoadReadOnlyException; -import org.apache.iotdb.db.exception.load.RegionReplicaSetChangedException; -import org.apache.iotdb.db.exception.mpp.FragmentInstanceDispatchException; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; -import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; import org.apache.iotdb.db.queryengine.common.MPPQueryContext; import org.apache.iotdb.db.queryengine.common.PlanFragmentId; import org.apache.iotdb.db.queryengine.execution.QueryStateMachine; @@ -52,86 +31,171 @@ import org.apache.iotdb.db.queryengine.plan.analyze.IPartitionFetcher; import org.apache.iotdb.db.queryengine.plan.planner.plan.DistributedQueryPlan; import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; -import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; -import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; -import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LoadTsFile; -import org.apache.iotdb.db.queryengine.plan.scheduler.FragInstanceDispatchResult; import org.apache.iotdb.db.queryengine.plan.scheduler.IScheduler; -import org.apache.iotdb.db.queryengine.plan.statement.crud.LoadTsFileStatement; import org.apache.iotdb.db.service.RegionMigrateService; -import org.apache.iotdb.db.storageengine.StorageEngine; -import org.apache.iotdb.db.storageengine.dataregion.flush.MemTableFlushTask; -import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; -import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.ArrayDeviceTimeIndex; -import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.PlainDeviceTimeIndex; -import org.apache.iotdb.db.storageengine.load.converter.LoadTsFileDataTypeConverter; import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileMemoryManager; -import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; -import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; -import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; -import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; -import org.apache.iotdb.db.storageengine.load.splitter.TsFileSplitter; -import org.apache.iotdb.metrics.utils.MetricLevel; -import org.apache.iotdb.mpp.rpc.thrift.TLoadCommandReq; -import org.apache.iotdb.rpc.TSStatusCode; import io.airlift.units.Duration; -import org.apache.tsfile.file.metadata.IDeviceID; -import org.apache.tsfile.file.metadata.StringArrayDeviceID; -import org.apache.tsfile.utils.Pair; -import org.apache.tsfile.utils.PublicBAOS; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.DataOutputStream; -import java.io.File; -import java.io.FileNotFoundException; -import java.io.IOException; -import java.nio.ByteBuffer; import java.util.ArrayList; -import java.util.Collections; -import java.util.Comparator; -import java.util.HashMap; import java.util.HashSet; -import java.util.Iterator; import java.util.List; -import java.util.ListIterator; -import java.util.Map; -import java.util.Objects; -import java.util.Optional; import java.util.Set; -import java.util.UUID; -import java.util.concurrent.CancellationException; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.ExecutionException; -import java.util.concurrent.Future; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.TimeoutException; -import java.util.stream.Collectors; /** - * {@link LoadTsFileScheduler} is used for scheduling {@link LoadSingleTsFileNode} and {@link - * LoadTsFilePieceNode}. because these two nodes need two phases to finish transfer. + * LOAD scheduler: {@link LoadTsFileScheduler} is the coordinator of a batch of {@link + * LoadSingleTsFileNode} loads (one node per source TsFile). It owns only the lifecycle - concurrent + * file lock, per-node guard clauses, state-machine transitions and the failure fallback - and + * routes every file to a {@link TsFileLoadStrategy}. All per-region bookkeeping, memory management + * and consensus submission live in the strategy components. * - *

for more details please check: ...; + *

Overall structure

+ * + *

All components below are part of the LOAD pipeline (LOAD TSFILE): the scheduler, the load + * strategies and every routing/buffering/dispatching helper are LOAD-only classes under {@code + * org.apache.iotdb.db.queryengine.plan.scheduler.load}. + * + * 

{@code
+ * LoadTsFileScheduler.start()
+ *     |
+ *     +--> per LoadSingleTsFileNode
+ *     |       lock file -> empty? -> strategy -> migration check
+ *     |
+ *     +--> needDecodeTsFile?
+ *     |       |-- false -> LocalLoadStrategy
+ *     |       |            `-- FragmentInstance -> local region (no decode)
+ *     |       `-- true  -> TwoPhaseConsensusLoadStrategy
+ *     |                     phase1: TsFileSplitConsumer
+ *     |                       DataPartitionRouter -> MemoryBoundedBuffer
+ *     |                       -> PieceDispatcher
+ *     |                     phase2: BEGIN -> PIECE* -> PREPARE -> COMMIT
+ *     |                       or ABORT, via RegionConsensusContext +
+ *     |                       LoadConsensusSubmitter
+ *     |
+ *     +--> success -> node.clean() + log
+ *     |       failure -> record failed index
+ *     |
+ *     `--> all success -> FINISHED
+ *             else -> LoadFallbackHandler (convert to tablets, retry)
+ *                      -> FINISHED / FAILED
+ * }
+ * + *

Consensus pipeline (phase 1)

+ * + * 
{@code
+ * TsFileSplitter
+ *      |
+ *      | TsFileData (CHUNK / DELETION)
+ *      v
+ * TsFileSplitConsumer
+ *      |
+ *      +-- CHUNK:  buffer -> DataPartitionRouter -> per-region piece
+ *      |           over budget? -> PieceDispatcher: dispatch largest first
+ *      +-- end of file: flush remaining pieces
+ *      |
+ *      v
+ * PieceDispatcher
+ *      | dispatch callback
+ *      v
+ * TwoPhaseConsensusLoadStrategy.dispatchConsensusPiece
+ *      |
+ *      +-- first piece of a region: BEGIN(loadId) then PIECE(0)
+ *      +-- later pieces:            PIECE(1), PIECE(2), ...
+ *      |
+ *      v
+ * RegionConsensusContext.accumulate(bytes, checksum)
+ *      |
+ *      v
+ * LoadConsensusSubmitter (submit to the partition write node; bounded retry)
+ * }
+ * + *

Two-phase protocol timeline

+ * + * 
{@code
+ * coordinator                          region write peer
+ *      |                                      |
+ *      |---- BEGIN(loadId) ------------------>| create staged writer
+ *      |---- PIECE(0, chunks) --------------->| append chunks
+ *      |---- PIECE(1, chunks) --------------->| append chunks
+ *      |---- ...                              |
+ *      |---- PREPARE(count, bytes, checksum)->| seal staged TsFile
+ *      |---- COMMIT ------------------------->| load staged TsFile
+ *      |                                      |
+ *   on failure:
+ *      |---- ABORT -------------------------->| drop staged data
+ * }
+ * + *

Result handling

+ * + * Successful files are cleaned up and logged (debug for pipe-generated loads, info otherwise). + * Failed indexes are collected; when all files are done the scheduler either transitions to + * FINISHED or hands the failures to {@link LoadFallbackHandler}, which converts the failed TsFiles + * into tablets, retries the insertion and finally transitions to FINISHED or FAILED. + * + *

Component responsibilities

+ * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + * + *
Components of the LOAD pipeline
ComponentResponsibility
{@link DataPartitionBatchFetcher}LOAD partition fetcher: batches (device, time-partition) queries with the transmit limit + * and applies the table-model/pipe database hint
{@link DataPartitionRouter}LOAD chunk router: deduplicates (device, slot) pairs and maps every chunk to its target + * {@code TRegionReplicaSet}
{@link MemoryBoundedBuffer}LOAD memory budget: pure memory-pool accounting; emits the "over budget" signal that + * triggers eviction
{@link PieceDispatcher}LOAD piece dispatcher: buffered per-region pieces, largest-first eviction heap and + * flushing
{@link TsFileSplitConsumer}LOAD split consumer: the {@code TsFileDataConsumer} composing router, buffer and + * dispatcher into the route -> buffer -> dispatch pipeline
{@link RegionConsensusContext}LOAD per-region two-phase state: one context per region with load id, piece count, total + * bytes, XOR checksum and BEGIN state
{@link LoadConsensusSubmitter}LOAD consensus transport for BEGIN/PIECE/PREPARE/COMMIT/ABORT: resolves the partition's + * write node and submits via local consensus write or internal RPC, like the normal write + * path
{@link LoadTsFileDispatcherImpl}legacy LOAD local dispatcher (local-load path) and the per-file uuid holder for + * correlation
{@link LoadFallbackHandler}LOAD failure fallback: converts failed TsFiles into tablets and retries, then resolves + * the final state-machine result
*/ public class LoadTsFileScheduler implements IScheduler { private static final Logger LOGGER = LoggerFactory.getLogger(LoadTsFileScheduler.class); - private static final IoTDBConfig CONFIG = IoTDBDescriptor.getInstance().getConfig(); - - private static final LoadTsFileCostMetricsSet LOAD_TSFILE_COST_METRICS_SET = - LoadTsFileCostMetricsSet.getInstance(); - - private static final long SINGLE_SCHEDULER_MAX_MEMORY_SIZE = - IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize() >> 2; - private static final int TRANSMIT_LIMIT = - CommonDescriptor.getInstance().getConfig().getTTimePartitionSlotTransmitLimit(); - private static final Set LOADING_FILE_SET = new HashSet<>(); private final MPPQueryContext queryContext; @@ -141,10 +205,9 @@ public class LoadTsFileScheduler implements IScheduler { private final List tsFileNodeList; private final List failedTsFileNodeIndexes; private final PlanFragmentId fragmentId; - private final Set allReplicaSets; private final boolean isGeneratedByPipe; - private final Map timePartitionSlotToProgressIndex; private final LoadTsFileDataCacheMemoryBlock block; + private final LoadConsensusSubmitter consensusSubmitter; public LoadTsFileScheduler( DistributedQueryPlan distributedQueryPlan, @@ -160,10 +223,9 @@ public LoadTsFileScheduler( this.fragmentId = distributedQueryPlan.getRootSubPlan().getPlanFragment().getId(); this.dispatcher = new LoadTsFileDispatcherImpl(internalServiceClientManager, isGeneratedByPipe); this.partitionFetcher = new DataPartitionBatchFetcher(partitionFetcher); - this.allReplicaSets = new HashSet<>(); this.isGeneratedByPipe = isGeneratedByPipe; - this.timePartitionSlotToProgressIndex = new HashMap<>(); this.block = LoadTsFileMemoryManager.getInstance().allocateDataCacheMemoryBlock(); + this.consensusSubmitter = new LoadConsensusSubmitter(internalServiceClientManager); for (FragmentInstance fragmentInstance : distributedQueryPlan.getInstances()) { tsFileNodeList.add((LoadSingleTsFileNode) fragmentInstance.getFragment().getPlanNodeTree()); @@ -174,151 +236,47 @@ public LoadTsFileScheduler( public void start() { try { stateMachine.transitionToRunning(); - int tsFileNodeListSize = tsFileNodeList.size(); boolean isLoadSuccess = true; - for (int i = 0; i < tsFileNodeListSize; ++i) { + for (int i = 0; i < tsFileNodeList.size(); ++i) { final LoadSingleTsFileNode node = tsFileNodeList.get(i); final String filePath = node.getTsFileResource().getTsFilePath(); + final String userName = queryContext.getSession().getUserName(); partitionFetcher.setDatabase(getPartitionQueryDatabase(node, isGeneratedByPipe)); - boolean isLoadSingleTsFileSuccess = true; - boolean shouldRemoveFileFromLoadingSet = false; - try { - synchronized (LOADING_FILE_SET) { - if (LOADING_FILE_SET.contains(filePath)) { - throw new LoadFileException( - String.format( - DataNodeQueryMessages - .QUERY_EXCEPTION_TSFILE_S_IS_LOADING_BY_ANOTHER_SCHEDULER_55077B82, - filePath)); - } - LOADING_FILE_SET.add(filePath); - } - shouldRemoveFileFromLoadingSet = true; - - final long startTimeMs = System.currentTimeMillis(); - - if (node.isTsFileEmpty()) { - LOGGER.info(DataNodeQueryMessages.LOAD_SKIP_TSFILE_BECAUSE_IT_HAS_NO_DATA, filePath); - } else if (!node.needDecodeTsFile( - slotList -> - partitionFetcher.queryDataPartition( - slotList, queryContext.getSession().getUserName()))) { - // do not decode, load locally - final long startTime = System.nanoTime(); - try { - isLoadSingleTsFileSuccess = loadLocally(node); - } finally { - LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( - LoadTsFileCostMetricsSet.LOAD_LOCALLY, System.nanoTime() - startTime); - } - } else { - // need decode, load locally or remotely, use two phases method - String uuid = UUID.randomUUID().toString(); - dispatcher.setUuid(uuid); - allReplicaSets.clear(); - - long startTime = System.nanoTime(); - final boolean isFirstPhaseSuccess; - try { - isFirstPhaseSuccess = firstPhase(node); - } finally { - LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( - LoadTsFileCostMetricsSet.FIRST_PHASE, System.nanoTime() - startTime); - } - - startTime = System.nanoTime(); - final boolean isSecondPhaseSuccess; - try { - isSecondPhaseSuccess = - secondPhase(isFirstPhaseSuccess, uuid, node.getTsFileResource()); - } finally { - LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( - LoadTsFileCostMetricsSet.SECOND_PHASE, System.nanoTime() - startTime); - } - - if (!isFirstPhaseSuccess || !isSecondPhaseSuccess) { - isLoadSingleTsFileSuccess = false; - } - } - - if (RegionMigrateService.getInstance().getLastNotifyMigratingTime() > startTimeMs - || RegionMigrateService.getInstance().mayHaveMigratingRegions()) { - LOGGER.warn( - DataNodeQueryMessages - .LOADTSFILESCHEDULER_REGION_MIGRATION_WAS_DETECTED_DURING_LOADING_TSFILE_ARG_WILL_CONVERT, - filePath); - isLoadSingleTsFileSuccess = false; - } - - if (isLoadSingleTsFileSuccess) { - node.clean(); - if (isGeneratedByPipe) { - LOGGER.debug( - DataNodeQueryMessages.LOAD_TSFILE_ARG_SUCCESSFULLY_LOAD_PROCESS_ARG_ARG, - filePath, - i + 1, - tsFileNodeListSize); - } else { - LOGGER.info( - DataNodeQueryMessages.LOAD_TSFILE_ARG_SUCCESSFULLY_LOAD_PROCESS_ARG_ARG, - filePath, - i + 1, - tsFileNodeListSize); - } - } else { - isLoadSuccess = false; - failedTsFileNodeIndexes.add(i); - LOGGER.warn( - DataNodeQueryMessages.CAN_NOT_LOAD_TSFILE_ARG_LOAD_PROCESS_ARG_ARG, - filePath, - i + 1, - tsFileNodeListSize); - } - } catch (Exception e) { + if (!processSingleNode(node, i, tsFileNodeList.size(), userName)) { isLoadSuccess = false; failedTsFileNodeIndexes.add(i); - LOGGER.warn(DataNodeQueryMessages.LOADTSFILESCHEDULER_LOADS_TSFILE_ERROR, filePath, e); - } finally { - if (shouldRemoveFileFromLoadingSet) { - synchronized (LOADING_FILE_SET) { - LOADING_FILE_SET.remove(filePath); - } - } + continue; + } + + node.clean(); + if (isGeneratedByPipe) { + LOGGER.debug( + DataNodeQueryMessages.LOAD_TSFILE_ARG_SUCCESSFULLY_LOAD_PROCESS_ARG_ARG, + filePath, + i + 1, + tsFileNodeList.size()); + } else { + LOGGER.info( + DataNodeQueryMessages.LOAD_TSFILE_ARG_SUCCESSFULLY_LOAD_PROCESS_ARG_ARG, + filePath, + i + 1, + tsFileNodeList.size()); } } if (isLoadSuccess) { stateMachine.transitionToFinished(); } else { - final StringBuilder failedTsFiles = - new StringBuilder( - !tsFileNodeList.isEmpty() - ? tsFileNodeList - .get(failedTsFileNodeIndexes.get(0)) - .getTsFileResource() - .getTsFilePath() - : ""); - final ListIterator iterator = failedTsFileNodeIndexes.listIterator(1); - while (iterator.hasNext()) { - failedTsFiles - .append(", ") - .append(tsFileNodeList.get(iterator.next()).getTsFileResource().getTsFilePath()); - } - final long startTime = System.nanoTime(); - try { - // if failed to load some TsFiles, then try to convert the TsFiles to Tablets - LOGGER.info( - DataNodeQueryMessages - .LOAD_TSFILE_S_FAILED_WILL_TRY_TO_CONVERT_TO_TABLETS_AND_INSERT_FAILED_TSFILES_ARG, - failedTsFiles); - convertFailedTsFilesToTabletsAndRetry(); - } finally { - LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( - LoadTsFileCostMetricsSet.SCHEDULER_CAST_TABLETS, System.nanoTime() - startTime); - } + new LoadFallbackHandler( + queryContext, + isGeneratedByPipe, + tsFileNodeList, + failedTsFileNodeIndexes, + stateMachine) + .convertFailedTsFilesToTablets(); } } finally { dispatcher.close(); @@ -326,356 +284,83 @@ public void start() { } } - private boolean firstPhase(LoadSingleTsFileNode node) { - final TsFileDataManager tsFileDataManager = new TsFileDataManager(this, node, block); + /** + * Loads one TsFile: guard clauses first (concurrent-load lock, empty file), then route to the + * strategy and finally detect whether a region migration raced with the load. + */ + private boolean processSingleNode( + LoadSingleTsFileNode node, int index, int listSize, String userName) { + final String filePath = node.getTsFileResource().getTsFilePath(); + final long startTimeMs = System.currentTimeMillis(); + boolean shouldRemoveFileFromLoadingSet = false; try { - new TsFileSplitter( - node.getTsFileResource().getTsFile(), tsFileDataManager::addOrSendTsFileData) - .splitTsFileByDataPartition(); - if (!tsFileDataManager.sendAllTsFileData()) { - return false; - } - } catch (IllegalStateException e) { - LOGGER.warn( - String.format( - DataNodeQueryMessages.DISPATCH_TSFILEDATA_ERROR_WHEN_PARSING_TSFILE_S, - node.getTsFileResource().getTsFile()), - e); - return false; - } catch (Exception e) { - LOGGER.warn( - String.format( - DataNodeQueryMessages.PARSE_OR_SEND_TSFILE_S_ERROR, - node.getTsFileResource().getTsFile()), - e); - return false; - } finally { - tsFileDataManager.clear(); - } - return true; - } - - private boolean dispatchOnePieceNode( - LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet) { - allReplicaSets.add(replicaSet); - FragmentInstance instance = - new FragmentInstance( - new PlanFragment(fragmentId, pieceNode), - fragmentId.genFragmentInstanceId(), - null, - queryContext.getQueryType(), - queryContext.getTimeOut() - (System.currentTimeMillis() - queryContext.getStartTime()), - queryContext.getSession(), - queryContext.isDebug(), - queryContext.isVerbose()); - instance.setExecutorAndHost(new StorageExecutor(replicaSet)); - Future dispatchResultFuture = - dispatcher.dispatch(null, Collections.singletonList(instance)); - - try { - FragInstanceDispatchResult result = - dispatchResultFuture.get( - CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds(), TimeUnit.SECONDS); - if (!result.isSuccessful()) { - LOGGER.warn( - DataNodeQueryMessages.DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG - + DataNodeQueryMessages - .RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG, - replicaSet, - TSStatusCode.representOf(result.getFailureStatus().getCode()).name(), - result.getFailureStatus().getMessage(), - pieceNode); - if (result.getFailureStatus().getSubStatus() != null) { - for (TSStatus status : result.getFailureStatus().getSubStatus()) { - LOGGER.warn( - DataNodeQueryMessages.SUB_STATUS_CODE_ARG_SUB_STATUS_MESSAGE_ARG, - TSStatusCode.representOf(status.getCode()).toString(), - status.getMessage()); - } + synchronized (LOADING_FILE_SET) { + if (LOADING_FILE_SET.contains(filePath)) { + throw new LoadFileException( + String.format( + DataNodeQueryMessages + .QUERY_EXCEPTION_TSFILE_S_IS_LOADING_BY_ANOTHER_SCHEDULER_55077B82, + filePath)); } - TSStatus status = result.getFailureStatus(); - status.setMessage( - String.format( - DataNodeQueryMessages.MESSAGE_LOAD_ARG_PIECE_ERROR_1ST_PHASE_BECAUSE_F3D9672C, - pieceNode.getTsFile()) - + status.getMessage()); - return false; - } - } catch (InterruptedException | ExecutionException | CancellationException e) { - if (e instanceof InterruptedException) { - Thread.currentThread().interrupt(); + LOADING_FILE_SET.add(filePath); } - LOGGER.warn(DataNodeQueryMessages.INTERRUPT_OR_EXECUTION_ERROR, e); - return false; - } catch (TimeoutException e) { - dispatchResultFuture.cancel(true); - LOGGER.warn( - String.format( - DataNodeQueryMessages.WAIT_FOR_LOADING_S_TIME_OUT, - LoadTsFilePieceNode.class.getName()), - e); - return false; - } - return true; - } + shouldRemoveFileFromLoadingSet = true; - private boolean secondPhase( - boolean isFirstPhaseSuccess, String uuid, TsFileResource tsFileResource) { - if (isGeneratedByPipe) { - LOGGER.debug(DataNodeQueryMessages.START_DISPATCHING_LOAD_COMMAND_FOR_UUID, uuid); - } else { - LOGGER.info(DataNodeQueryMessages.START_DISPATCHING_LOAD_COMMAND_FOR_UUID, uuid); - } - final File tsFile = tsFileResource.getTsFile(); - final TLoadCommandReq loadCommandReq = - new TLoadCommandReq( - (isFirstPhaseSuccess ? LoadCommand.EXECUTE : LoadCommand.ROLLBACK).ordinal(), uuid); + if (node.isTsFileEmpty()) { + LOGGER.info(DataNodeQueryMessages.LOAD_SKIP_TSFILE_BECAUSE_IT_HAS_NO_DATA, filePath); + return true; + } - try { - loadCommandReq.setIsGeneratedByPipe(isGeneratedByPipe); - loadCommandReq.setTimePartition2ProgressIndex( - timePartitionSlotToProgressIndex.entrySet().stream() - .collect( - Collectors.toMap( - Map.Entry::getKey, - entry -> { - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); - final DataOutputStream dataOutputStream = - new DataOutputStream(byteArrayOutputStream)) { - entry.getValue().serialize(dataOutputStream); - return ByteBuffer.wrap( - byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()); - } catch (final IOException e) { - throw new RuntimeException( - String.format( - DataNodeQueryMessages - .QUERY_EXCEPTION_SERIALIZE_PROGRESS_INDEX_ERROR_ISFIRSTPHASESUCCESS_S_UUID_690F0419, - isFirstPhaseSuccess, - uuid, - tsFile.getAbsolutePath()), - e); - } - }))); - Future dispatchResultFuture = - dispatcher.dispatchCommand(loadCommandReq, allReplicaSets); + final TsFileLoadStrategy strategy; + if (!node.needDecodeTsFile( + slotList -> partitionFetcher.queryDataPartition(slotList, userName))) { + // do not decode, load locally + strategy = new LocalLoadStrategy(queryContext, fragmentId, dispatcher); + } else { + // need decode, use the consensus two-phase pipeline + strategy = + new TwoPhaseConsensusLoadStrategy( + dispatcher, + partitionFetcher, + block, + consensusSubmitter, + userName, + isGeneratedByPipe); + } + final boolean isLoadSingleTsFileSuccess = strategy.execute(node); - FragInstanceDispatchResult result = dispatchResultFuture.get(); - if (!result.isSuccessful()) { + if (RegionMigrateService.getInstance().getLastNotifyMigratingTime() > startTimeMs + || RegionMigrateService.getInstance().mayHaveMigratingRegions()) { LOGGER.warn( DataNodeQueryMessages - .DISPATCH_LOAD_COMMAND_ARG_OF_TSFILE_ARG_ERROR_TO_REPLICASETS_ARG_ERROR - + DataNodeQueryMessages.RESULT_STATUS_CODE_ARG_RESULT_STATUS_MESSAGE_ARG, - loadCommandReq, - tsFile, - allReplicaSets, - TSStatusCode.representOf(result.getFailureStatus().getCode()).name(), - result.getFailureStatus().getMessage()); - TSStatus status = result.getFailureStatus(); - status.setMessage( - String.format( - DataNodeQueryMessages - .MESSAGE_LOAD_ARG_ERROR_SECOND_PHASE_BECAUSE_ARG_FIRST_PHASE_ARG_CBA980FC, - tsFile, - status.getMessage(), - isFirstPhaseSuccess - ? DataNodeQueryMessages.MESSAGE_SUCCESS_260CA9DD - : DataNodeQueryMessages.MESSAGE_FAILED_26934EB3)); + .LOADTSFILESCHEDULER_REGION_MIGRATION_WAS_DETECTED_DURING_LOADING_TSFILE_ARG_WILL_CONVERT, + filePath); + logCannotLoad(node, index, listSize); return false; } - } catch (InterruptedException | ExecutionException e) { - if (e instanceof InterruptedException) { - Thread.currentThread().interrupt(); + if (!isLoadSingleTsFileSuccess) { + logCannotLoad(node, index, listSize); + return false; } - LOGGER.warn(DataNodeQueryMessages.INTERRUPT_OR_EXECUTION_ERROR, e); - return false; + return true; } catch (Exception e) { - LOGGER.warn( - DataNodeQueryMessages.EXCEPTION_OCCURRED_DURING_SECOND_PHASE_OF_LOADING_TSFILE, - tsFile, - e); + LOGGER.warn(DataNodeQueryMessages.LOADTSFILESCHEDULER_LOADS_TSFILE_ERROR, filePath, e); return false; - } - return true; - } - - private ByteBuffer assignProgressIndex(TsFileResource tsFileResource) throws IOException { - PipeDataNodeAgent.runtime().assignProgressIndexForTsFileLoad(tsFileResource); - - try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); - final DataOutputStream dataOutputStream = new DataOutputStream(byteArrayOutputStream)) { - tsFileResource.getMaxProgressIndex().serialize(dataOutputStream); - return ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0, byteArrayOutputStream.size()); - } - } - - private boolean loadLocally(LoadSingleTsFileNode node) throws IoTDBException { - LOGGER.info( - DataNodeQueryMessages.START_LOAD_TSFILE_LOCALLY, - node.getTsFileResource().getTsFile().getPath()); - - if (CommonDescriptor.getInstance().getConfig().isReadOnly()) { - throw new LoadReadOnlyException(); - } - - // if the time index is PlainDeviceTimeIndex, convert it to ArrayDeviceTimeIndex - if (node.getTsFileResource().getTimeIndex() instanceof PlainDeviceTimeIndex) { - final PlainDeviceTimeIndex timeIndex = - (PlainDeviceTimeIndex) node.getTsFileResource().getTimeIndex(); - final Map convertedDeviceToIndex = new ConcurrentHashMap<>(); - for (final Map.Entry entry : timeIndex.getDeviceToIndex().entrySet()) { - convertedDeviceToIndex.put( - entry.getKey() instanceof StringArrayDeviceID - ? entry.getKey() - : new StringArrayDeviceID(entry.getKey().toString()), - entry.getValue()); - } - node.getTsFileResource() - .setTimeIndex( - new ArrayDeviceTimeIndex( - convertedDeviceToIndex, timeIndex.getStartTimes(), timeIndex.getEndTimes())); - } - - try { - FragmentInstance instance = - new FragmentInstance( - new PlanFragment(fragmentId, node), - fragmentId.genFragmentInstanceId(), - null, - queryContext.getQueryType(), - queryContext.getTimeOut() - - (System.currentTimeMillis() - queryContext.getStartTime()), - queryContext.getSession(), - queryContext.isDebug(), - queryContext.isVerbose()); - instance.setExecutorAndHost(new StorageExecutor(node.getLocalRegionReplicaSet())); - dispatcher.dispatchLocally(instance); - } catch (FragmentInstanceDispatchException e) { - LOGGER.warn( - String.format( - DataNodeQueryMessages.DISPATCH_TSFILE_S_ERROR_TO_LOCAL_ERROR_RESULT_STATUS_CODE_S - + DataNodeQueryMessages.RESULT_STATUS_MESSAGE_S, - node.getTsFileResource().getTsFile(), - TSStatusCode.representOf(e.getFailureStatus().getCode()).name(), - e.getFailureStatus().getMessage())); - return false; - } - - // add metrics - Optional.ofNullable( - StorageEngine.getInstance() - .getDataRegion( - (DataRegionId) - ConsensusGroupId.Factory.createFromTConsensusGroupId( - node.getLocalRegionReplicaSet().getRegionId()))) - .ifPresent( - dataRegion -> - dataRegion - .getNonSystemDatabaseName() - .ifPresent( - databaseName -> { - // Report load tsFile points to IoTDB flush metrics - MemTableFlushTask.recordFlushPointsMetricInternal( - node.getWritePointCount(), - databaseName, - dataRegion.getDataRegionIdString()); - - MetricService.getInstance() - .count( - node.getWritePointCount(), - Metric.QUANTITY.toString(), - MetricLevel.CORE, - Tag.NAME.toString(), - Metric.POINTS_IN.toString(), - Tag.DATABASE.toString(), - databaseName, - Tag.REGION.toString(), - dataRegion.getDataRegionIdString(), - Tag.TYPE.toString(), - Metric.LOAD_POINT_COUNT.toString()); - MetricService.getInstance() - .count( - node.getWritePointCount(), - Metric.LEADER_QUANTITY.toString(), - MetricLevel.CORE, - Tag.NAME.toString(), - Metric.POINTS_IN.toString(), - Tag.DATABASE.toString(), - databaseName, - Tag.REGION.toString(), - dataRegion.getDataRegionIdString(), - Tag.TYPE.toString(), - Metric.LOAD_POINT_COUNT.toString()); - })); - - return true; - } - - private void convertFailedTsFilesToTabletsAndRetry() { - final LoadTsFileDataTypeConverter loadTsFileDataTypeConverter = - new LoadTsFileDataTypeConverter(queryContext, isGeneratedByPipe); - - final Iterator iterator = failedTsFileNodeIndexes.listIterator(); - while (iterator.hasNext()) { - final int failedLoadTsFileIndex = iterator.next(); - final LoadSingleTsFileNode failedNode = tsFileNodeList.get(failedLoadTsFileIndex); - final String filePath = failedNode.getTsFileResource().getTsFilePath(); - - try { - final TSStatus status = - failedNode.isTableModel() - ? loadTsFileDataTypeConverter - .convertForTableModel( - (isGeneratedByPipe - ? LoadTsFile.createForPipe(null, filePath, Collections.emptyMap()) - : LoadTsFile.createUnchecked( - null, filePath, Collections.emptyMap())) - .setDatabase(failedNode.getDatabase()) - .setDeleteAfterLoad(failedNode.isDeleteAfterLoad()) - .setConvertOnTypeMismatch(true)) - .orElse(null) - : loadTsFileDataTypeConverter - .convertForTreeModel( - buildRetryTreeLoadStatement( - filePath, - failedNode.isDeleteAfterLoad(), - getPartitionQueryDatabase(failedNode, isGeneratedByPipe))) - .orElse(null); - - if (loadTsFileDataTypeConverter.isSuccessful(status)) { - iterator.remove(); - LOGGER.info( - DataNodeQueryMessages - .LOAD_SUCCESSFULLY_CONVERTED_TSFILE_ARG_INTO_TABLETS_AND_INSERTED, - failedNode.getTsFileResource().getTsFilePath()); - } else { - LOGGER.warn( - DataNodeQueryMessages.LOAD_FAILED_TO_CONVERT_TO_TABLETS_FROM_TSFILE_ARG_STATUS_ARG, - failedNode.getTsFileResource().getTsFilePath(), - status); + } finally { + if (shouldRemoveFileFromLoadingSet) { + synchronized (LOADING_FILE_SET) { + LOADING_FILE_SET.remove(filePath); } - } catch (final Exception e) { - LOGGER.warn( - DataNodeQueryMessages.LOAD_FAILED_TO_CONVERT_TO_TABLETS_FROM_TSFILE_ARG_EXCEPTION_ARG, - failedNode.getTsFileResource().getTsFilePath(), - e.getMessage(), - e); } } + } - // If all failed TsFiles are converted into tablets and inserted, - // we can consider the load process as successful. - if (failedTsFileNodeIndexes.isEmpty()) { - LOGGER.info(DataNodeQueryMessages.LOAD_ALL_FAILED_TSFILES_ARE_CONVERTED_TO_TABLETS); - stateMachine.transitionToFinished(); - } else { - final String errorMsg = - "Load: failed to load some TsFiles by converting them into tablets. Failed TsFiles: " - + failedTsFileNodeIndexes.stream() - .map(i -> tsFileNodeList.get(i).getTsFileResource().getTsFilePath()) - .collect(Collectors.joining(", ")); - LOGGER.warn(errorMsg); - stateMachine.transitionToFailed(new LoadFileException(errorMsg)); - } + private void logCannotLoad(LoadSingleTsFileNode node, int index, int listSize) { + LOGGER.warn( + DataNodeQueryMessages.CAN_NOT_LOAD_TSFILE_ARG_LOAD_PROCESS_ARG_ARG, + node.getTsFileResource().getTsFilePath(), + index + 1, + listSize); } static String getPartitionQueryDatabase( @@ -683,25 +368,6 @@ static String getPartitionQueryDatabase( return node.isTableModel() || isGeneratedByPipe ? node.getDatabase() : null; } - private LoadTsFileStatement buildRetryTreeLoadStatement( - final String filePath, final boolean deleteAfterLoad, final String database) - throws FileNotFoundException { - final LoadTsFileStatement statement = - (isGeneratedByPipe - ? LoadTsFileStatement.createForPipe(filePath) - : LoadTsFileStatement.createUnchecked(filePath)) - .setDeleteAfterLoad(deleteAfterLoad) - .setConvertOnTypeMismatch(true); - if (database != null) { - statement.setDatabase(database); - statement.updateDatabaseLevelByTreeDatabase(); - } - if (isGeneratedByPipe) { - statement.markIsGeneratedByPipe(); - } - return statement; - } - @Override public void stop(Throwable t) { dispatcher.abort(); @@ -717,265 +383,8 @@ public FragmentInfo getFragmentInfo() { return null; } - private void computeTimePartitionSlotToProgressIndexIfAbsent( - final TTimePartitionSlot timePartitionSlot) { - timePartitionSlotToProgressIndex.putIfAbsent( - timePartitionSlot, PipeDataNodeAgent.runtime().getNextProgressIndexForTsFileLoad()); - } - public enum LoadCommand { EXECUTE, ROLLBACK } - - private static class TsFileDataManager { - private final LoadTsFileScheduler scheduler; - private final LoadSingleTsFileNode singleTsFileNode; - - private long dataSize; - private final Map> - regionId2ReplicaSetAndNode; - private final List nonDirectionalChunkData; - private final LoadTsFileDataCacheMemoryBlock block; - - public TsFileDataManager( - LoadTsFileScheduler scheduler, - LoadSingleTsFileNode singleTsFileNode, - LoadTsFileDataCacheMemoryBlock block) { - this.scheduler = scheduler; - this.singleTsFileNode = singleTsFileNode; - this.dataSize = 0; - this.regionId2ReplicaSetAndNode = new HashMap<>(); - this.nonDirectionalChunkData = new ArrayList<>(); - this.block = block; - } - - private boolean addOrSendTsFileData(TsFileData tsFileData) throws LoadFileException { - switch (tsFileData.getType()) { - case CHUNK: - return addOrSendChunkData((ChunkData) tsFileData); - case DELETION: - return addOrSendDeletionData((DeletionData) tsFileData); - default: - throw new UnsupportedOperationException( - String.format( - DataNodeQueryMessages.QUERY_EXCEPTION_UNSUPPORTED_TSFILEDATATYPE_S_374475FA, - tsFileData.getType())); - } - } - - private boolean isMemoryEnough() { - return dataSize <= SINGLE_SCHEDULER_MAX_MEMORY_SIZE && block.hasEnoughMemory(); - } - - private boolean addOrSendChunkData(ChunkData chunkData) throws LoadFileException { - nonDirectionalChunkData.add(chunkData); - dataSize += chunkData.getDataSize(); - block.addMemoryUsage(chunkData.getDataSize()); - scheduler.computeTimePartitionSlotToProgressIndexIfAbsent(chunkData.getTimePartitionSlot()); - - if (!isMemoryEnough()) { - routeChunkData(); - - // start to dispatch from the biggest TsFilePieceNode - List sortedRegionIds = - regionId2ReplicaSetAndNode.keySet().stream() - .sorted( - Comparator.comparingLong( - o -> regionId2ReplicaSetAndNode.get(o).getRight().getDataSize()) - .reversed()) - .collect(Collectors.toList()); - - for (TConsensusGroupId sortedRegionId : sortedRegionIds) { - final TRegionReplicaSet replicaSet = - regionId2ReplicaSetAndNode.get(sortedRegionId).getLeft(); - final LoadTsFilePieceNode pieceNode = - regionId2ReplicaSetAndNode.get(sortedRegionId).getRight(); - if (pieceNode.getDataSize() == 0) { // total data size has been reduced to 0 - break; - } - final boolean isDispatchSuccess = scheduler.dispatchOnePieceNode(pieceNode, replicaSet); - - regionId2ReplicaSetAndNode.replace( - sortedRegionId, - new Pair<>( - replicaSet, - new LoadTsFilePieceNode( - singleTsFileNode.getPlanNodeId(), - singleTsFileNode - .getTsFileResource() - .getTsFile()))); // can not just remove, because of deletion - releaseMemoryUsage(pieceNode.getDataSize()); - - if (!isDispatchSuccess) { - // Currently there is no retry, so return directly - return false; - } - - if (isMemoryEnough()) { - break; - } - } - } - - return true; - } - - private void routeChunkData() throws LoadFileException { - if (nonDirectionalChunkData.isEmpty()) { - return; - } - - final List> partitionSlotList = new ArrayList<>(); - final int[] chunkPartitionIndexes = new int[nonDirectionalChunkData.size()]; - final Map> partitionSlotIndexes = new HashMap<>(); - for (int i = 0, size = nonDirectionalChunkData.size(); i < size; i++) { - final ChunkData chunkData = nonDirectionalChunkData.get(i); - final IDeviceID device = chunkData.getDevice(); - final TTimePartitionSlot timePartitionSlot = chunkData.getTimePartitionSlot(); - final Map slotIndexes = - partitionSlotIndexes.computeIfAbsent(device, key -> new HashMap<>()); - Integer partitionSlotIndex = slotIndexes.get(timePartitionSlot); - if (partitionSlotIndex == null) { - partitionSlotIndex = partitionSlotList.size(); - slotIndexes.put(timePartitionSlot, partitionSlotIndex); - partitionSlotList.add(new Pair<>(device, timePartitionSlot)); - } - chunkPartitionIndexes[i] = partitionSlotIndex; - } - - List replicaSets = - scheduler.partitionFetcher.queryDataPartition( - partitionSlotList, scheduler.queryContext.getSession().getUserName()); - for (int i = 0, size = nonDirectionalChunkData.size(); i < size; i++) { - final TRegionReplicaSet replicaSet = replicaSets.get(chunkPartitionIndexes[i]); - final TConsensusGroupId regionId = replicaSet.getRegionId(); - if (regionId2ReplicaSetAndNode.containsKey(regionId) - && !Objects.equals(regionId2ReplicaSetAndNode.get(regionId).getLeft(), replicaSet)) { - // Detected region replica set changed (maybe due to region migration), throw an exception - throw new RegionReplicaSetChangedException( - regionId2ReplicaSetAndNode.get(regionId).getLeft(), replicaSet); - } - - regionId2ReplicaSetAndNode - .computeIfAbsent( - replicaSet.getRegionId(), - o -> - new Pair<>( - replicaSet, - new LoadTsFilePieceNode( - singleTsFileNode.getPlanNodeId(), - singleTsFileNode.getTsFileResource().getTsFile()))) - .getRight() - .addTsFileData(nonDirectionalChunkData.get(i)); - } - nonDirectionalChunkData.clear(); - } - - private boolean addOrSendDeletionData(DeletionData deletionData) throws LoadFileException { - routeChunkData(); // ensure chunk data will be added before deletion - - for (Map.Entry> entry : - regionId2ReplicaSetAndNode.entrySet()) { - dataSize += deletionData.getDataSize(); - block.addMemoryUsage(deletionData.getDataSize()); - entry.getValue().getRight().addTsFileData(deletionData); - } - return true; - } - - private boolean sendAllTsFileData() throws LoadFileException { - routeChunkData(); - - boolean isAllSuccess = true; - for (Map.Entry> entry : - regionId2ReplicaSetAndNode.entrySet()) { - releaseMemoryUsage(entry.getValue().getRight().getDataSize()); - if (isAllSuccess - && !scheduler.dispatchOnePieceNode( - entry.getValue().getRight(), entry.getValue().getLeft())) { - LOGGER.warn( - DataNodeQueryMessages.DISPATCH_PIECE_NODE_ARG_OF_TSFILE_ARG_ERROR, - entry.getValue(), - singleTsFileNode.getTsFileResource().getTsFile()); - isAllSuccess = false; - } - } - return isAllSuccess; - } - - private void releaseMemoryUsage(final long memorySize) { - dataSize -= memorySize; - block.reduceMemoryUsage(memorySize); - } - - private void clear() { - if (dataSize > 0) { - block.reduceMemoryUsage(dataSize); - dataSize = 0; - } - nonDirectionalChunkData.clear(); - regionId2ReplicaSetAndNode.clear(); - } - } - - private static class DataPartitionBatchFetcher { - private final IPartitionFetcher fetcher; - private String database; - - public DataPartitionBatchFetcher(IPartitionFetcher fetcher) { - this.fetcher = fetcher; - } - - public void setDatabase(String database) { - this.database = database; - } - - public List queryDataPartition( - List> slotList, String userName) { - List replicaSets = new ArrayList<>(slotList.size()); - int size = slotList.size(); - - for (int i = 0; i < size; i += TRANSMIT_LIMIT) { - List> subSlotList = - slotList.subList(i, Math.min(size, i + TRANSMIT_LIMIT)); - DataPartition dataPartition = - fetcher.getOrCreateDataPartition(toQueryParam(subSlotList), userName); - for (final Pair pair : subSlotList) { - // database is an explicit database hint for table-model loads and - // pipe-generated tree-model loads. - replicaSets.add( - database != null - ? dataPartition.getDataRegionReplicaSetForWriting(pair.left, pair.right, database) - : dataPartition.getDataRegionReplicaSetForWriting(pair.left, pair.right)); - } - } - return replicaSets; - } - - private List toQueryParam( - List> slots) { - final Map> device2TimePartitionSlots = new HashMap<>(); - for (final Pair slot : slots) { - device2TimePartitionSlots - .computeIfAbsent(slot.left, key -> new HashSet<>()) - .add(slot.right); - } - - final List queryParams = - new ArrayList<>(device2TimePartitionSlots.size()); - for (final Map.Entry> entry : - device2TimePartitionSlots.entrySet()) { - final DataPartitionQueryParam queryParam = - new DataPartitionQueryParam(entry.getKey(), new ArrayList<>(entry.getValue())); - // database is an explicit database hint for table-model loads and - // pipe-generated tree-model loads. - if (database != null) { - queryParam.setDatabaseName(database); - } - queryParams.add(queryParam); - } - return queryParams; - } - } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LocalLoadStrategy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LocalLoadStrategy.java new file mode 100644 index 0000000000000..0be512ff696a3 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LocalLoadStrategy.java @@ -0,0 +1,203 @@ +/* + * 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.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.consensus.ConsensusGroupId; +import org.apache.iotdb.commons.consensus.DataRegionId; +import org.apache.iotdb.commons.exception.IoTDBException; +import org.apache.iotdb.commons.partition.StorageExecutor; +import org.apache.iotdb.commons.service.metric.MetricService; +import org.apache.iotdb.commons.service.metric.enums.Metric; +import org.apache.iotdb.commons.service.metric.enums.Tag; +import org.apache.iotdb.db.exception.load.LoadReadOnlyException; +import org.apache.iotdb.db.exception.mpp.FragmentInstanceDispatchException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.common.MPPQueryContext; +import org.apache.iotdb.db.queryengine.common.PlanFragmentId; +import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance; +import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.storageengine.StorageEngine; +import org.apache.iotdb.db.storageengine.dataregion.flush.MemTableFlushTask; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.ArrayDeviceTimeIndex; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.timeindex.PlainDeviceTimeIndex; +import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; +import org.apache.iotdb.metrics.utils.MetricLevel; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.file.metadata.StringArrayDeviceID; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; + +/** + * LOAD local-load strategy: loads a TsFile without decoding it, used when the file's device/time + * ranges map to a single local region. The whole file is wrapped into one {@code FragmentInstance} + * (reusing the scheduler's fragment id) and dispatched to the local data region through {@link + * LoadTsFileDispatcherImpl#dispatchLocally}; nothing crosses the network. + * + *

The strategy also: + * + *

    + *
  • rejects loads while the node is read-only ({@link LoadReadOnlyException}); + *
  • converts a {@code PlainDeviceTimeIndex} into an {@code ArrayDeviceTimeIndex} so the local + * writer can use the time index directly; + *
  • records flush/points metrics on the target data region; + *
  • measures the whole execution as the {@code LOAD_LOCALLY} phase metric. + *
+ */ +public class LocalLoadStrategy implements TsFileLoadStrategy { + + private static final Logger LOGGER = LoggerFactory.getLogger(LocalLoadStrategy.class); + + private static final LoadTsFileCostMetricsSet LOAD_TSFILE_COST_METRICS_SET = + LoadTsFileCostMetricsSet.getInstance(); + + private final MPPQueryContext queryContext; + private final PlanFragmentId fragmentId; + private final LoadTsFileDispatcherImpl dispatcher; + + public LocalLoadStrategy( + MPPQueryContext queryContext, + PlanFragmentId fragmentId, + LoadTsFileDispatcherImpl dispatcher) { + this.queryContext = queryContext; + this.fragmentId = fragmentId; + this.dispatcher = dispatcher; + } + + @Override + public boolean execute(LoadSingleTsFileNode node) throws IoTDBException { + final long startTime = System.nanoTime(); + try { + return loadLocally(node); + } finally { + LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( + LoadTsFileCostMetricsSet.LOAD_LOCALLY, System.nanoTime() - startTime); + } + } + + private boolean loadLocally(LoadSingleTsFileNode node) throws IoTDBException { + LOGGER.info( + DataNodeQueryMessages.START_LOAD_TSFILE_LOCALLY, + node.getTsFileResource().getTsFile().getPath()); + + if (CommonDescriptor.getInstance().getConfig().isReadOnly()) { + throw new LoadReadOnlyException(); + } + + // if the time index is PlainDeviceTimeIndex, convert it to ArrayDeviceTimeIndex + if (node.getTsFileResource().getTimeIndex() instanceof PlainDeviceTimeIndex) { + final PlainDeviceTimeIndex timeIndex = + (PlainDeviceTimeIndex) node.getTsFileResource().getTimeIndex(); + final Map convertedDeviceToIndex = new ConcurrentHashMap<>(); + for (final Map.Entry entry : timeIndex.getDeviceToIndex().entrySet()) { + convertedDeviceToIndex.put( + entry.getKey() instanceof StringArrayDeviceID + ? entry.getKey() + : new StringArrayDeviceID(entry.getKey().toString()), + entry.getValue()); + } + node.getTsFileResource() + .setTimeIndex( + new ArrayDeviceTimeIndex( + convertedDeviceToIndex, timeIndex.getStartTimes(), timeIndex.getEndTimes())); + } + + try { + FragmentInstance instance = + new FragmentInstance( + new PlanFragment(fragmentId, node), + fragmentId.genFragmentInstanceId(), + null, + queryContext.getQueryType(), + queryContext.getTimeOut() + - (System.currentTimeMillis() - queryContext.getStartTime()), + queryContext.getSession(), + queryContext.isDebug(), + queryContext.isVerbose()); + instance.setExecutorAndHost(new StorageExecutor(node.getLocalRegionReplicaSet())); + dispatcher.dispatchLocally(instance); + } catch (FragmentInstanceDispatchException e) { + LOGGER.warn( + String.format( + DataNodeQueryMessages.DISPATCH_TSFILE_S_ERROR_TO_LOCAL_ERROR_RESULT_STATUS_CODE_S + + DataNodeQueryMessages.RESULT_STATUS_MESSAGE_S, + node.getTsFileResource().getTsFile(), + TSStatusCode.representOf(e.getFailureStatus().getCode()).name(), + e.getFailureStatus().getMessage())); + return false; + } + + // add metrics + Optional.ofNullable( + StorageEngine.getInstance() + .getDataRegion( + (DataRegionId) + ConsensusGroupId.Factory.createFromTConsensusGroupId( + node.getLocalRegionReplicaSet().getRegionId()))) + .ifPresent( + dataRegion -> + dataRegion + .getNonSystemDatabaseName() + .ifPresent( + databaseName -> { + // Report load tsFile points to IoTDB flush metrics + MemTableFlushTask.recordFlushPointsMetricInternal( + node.getWritePointCount(), + databaseName, + dataRegion.getDataRegionIdString()); + + MetricService.getInstance() + .count( + node.getWritePointCount(), + Metric.QUANTITY.toString(), + MetricLevel.CORE, + Tag.NAME.toString(), + Metric.POINTS_IN.toString(), + Tag.DATABASE.toString(), + databaseName, + Tag.REGION.toString(), + dataRegion.getDataRegionIdString(), + Tag.TYPE.toString(), + Metric.LOAD_POINT_COUNT.toString()); + MetricService.getInstance() + .count( + node.getWritePointCount(), + Metric.LEADER_QUANTITY.toString(), + MetricLevel.CORE, + Tag.NAME.toString(), + Metric.POINTS_IN.toString(), + Tag.DATABASE.toString(), + databaseName, + Tag.REGION.toString(), + dataRegion.getDataRegionIdString(), + Tag.TYPE.toString(), + Metric.LOAD_POINT_COUNT.toString()); + })); + + return true; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/MemoryBoundedBuffer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/MemoryBoundedBuffer.java new file mode 100644 index 0000000000000..05b0204e58753 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/MemoryBoundedBuffer.java @@ -0,0 +1,78 @@ +/* + * 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.conf.IoTDBDescriptor; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; + +/** + * LOAD memory budget: memory-pool accounting for the pieces of one source TsFile. It keeps two + * views in sync: + * + *
    + *
  • its own {@code dataSize} - the bytes currently buffered into pieces, and + *
  • the shared {@link LoadTsFileDataCacheMemoryBlock} - so the cluster-wide LOAD data cache is + * aware of this file's footprint. + *
+ * + * The budget is {@code thriftMaxFrameSize >> 2}; {@link #isMemoryEnough()} is the signal the + * pipeline polls after every chunk. When it turns false, {@link PieceDispatcher} evicts the largest + * buffered piece first. {@link #add(long)} / {@link #release(long)} keep both views consistent on + * every buffered/dispatched piece, and {@link #clear()} is the last-chance cleanup that returns any + * leftover accounting to the shared block. + */ +class MemoryBoundedBuffer { + + private static final long MAX_MEMORY_SIZE = + IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize() >> 2; + + private final LoadTsFileDataCacheMemoryBlock block; + private long dataSize = 0; + + MemoryBoundedBuffer(LoadTsFileDataCacheMemoryBlock block) { + this.block = block; + } + + boolean isMemoryEnough() { + return dataSize <= MAX_MEMORY_SIZE && block.hasEnoughMemory(); + } + + void add(long memorySize) { + dataSize += memorySize; + block.addMemoryUsage(memorySize); + } + + void release(long memorySize) { + dataSize -= memorySize; + block.reduceMemoryUsage(memorySize); + } + + long getDataSize() { + return dataSize; + } + + /** Returns all buffered accounting to the shared memory block; safe to call multiple times. */ + void clear() { + if (dataSize > 0) { + block.reduceMemoryUsage(dataSize); + dataSize = 0; + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java new file mode 100644 index 0000000000000..8d188e0c8845f --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/PieceDispatcher.java @@ -0,0 +1,234 @@ +/* + * 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.TConsensusGroupId; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.exception.load.RegionReplicaSetChangedException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; +import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; + +import org.apache.tsfile.utils.Pair; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; +import java.util.PriorityQueue; + +/** + * LOAD piece dispatcher: holds the buffered pieces of one source TsFile - one {@link + * LoadTsFilePieceNode} per target region - and decides when they are sent. + * + *

State: + * + *

    + *
  • {@code regionId2ReplicaSetAndNode} - the current piece of every touched region; + *
  • {@code largestPieceRegions} - max-heap of (region, piece size at offer time) used for + * largest-first eviction; entries become stale once a piece grows or is dispatched and are + * skipped lazily on poll. + *
+ * + *

Behaviors: + * + *

    + *
  • {@link #offerChunk(ChunkData, TRegionReplicaSet)} appends a chunk and throws {@link + * RegionReplicaSetChangedException} if the same region suddenly maps to a different replica + * set (region migration); + *
  • {@link #addDeletionToAll(DeletionData)} replicates a deletion into every buffered piece + * (memory is accounted once per region); + *
  • {@link #dispatchLargestUntilMemoryEnough()} evicts largest pieces while over budget; + *
  • {@link #flushAll()} flushes the remainder at end of file. + *
+ * + * Dispatch is delegated back through {@link DispatchCallback}, so this class never talks to + * consensus itself; {@link TwoPhaseConsensusLoadStrategy} implements the callback with its + * BEGIN/PIECE submission logic. + */ +class PieceDispatcher { + + private static final Logger LOGGER = LoggerFactory.getLogger(PieceDispatcher.class); + + @FunctionalInterface + interface DispatchCallback { + boolean dispatch(LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet); + } + + private final LoadSingleTsFileNode singleTsFileNode; + private final MemoryBoundedBuffer memoryBuffer; + private final DispatchCallback dispatchCallback; + + private final Map> + regionId2ReplicaSetAndNode = new HashMap<>(); + + /** + * Max-heap of (regionId, buffered piece size at offer time) used to dispatch the largest pieces + * first when the data cache is over budget. Entries become stale once the piece size changes or + * the piece is dispatched, and are skipped lazily on poll, which avoids re-sorting all buffered + * pieces on every over-budget event. + */ + private final PriorityQueue> largestPieceRegions = + new PriorityQueue<>((a, b) -> Long.compare(b.getValue(), a.getValue())); + + PieceDispatcher( + LoadSingleTsFileNode singleTsFileNode, + MemoryBoundedBuffer memoryBuffer, + DispatchCallback dispatchCallback) { + this.singleTsFileNode = singleTsFileNode; + this.memoryBuffer = memoryBuffer; + this.dispatchCallback = dispatchCallback; + } + + void offerChunk(ChunkData chunkData, TRegionReplicaSet replicaSet) throws LoadFileException { + final TConsensusGroupId regionId = replicaSet.getRegionId(); + if (regionId2ReplicaSetAndNode.containsKey(regionId) + && !Objects.equals(regionId2ReplicaSetAndNode.get(regionId).getLeft(), replicaSet)) { + // Detected region replica set changed (maybe due to region migration), throw an exception + throw new RegionReplicaSetChangedException( + regionId2ReplicaSetAndNode.get(regionId).getLeft(), replicaSet); + } + + regionId2ReplicaSetAndNode + .computeIfAbsent(regionId, o -> new Pair<>(replicaSet, newPieceNode())) + .getRight() + .addTsFileData(chunkData); + offerPieceRegion(regionId); + } + + /** Replicates the deletion into every buffered piece; memory is accounted once per region. */ + void addDeletionToAll(DeletionData deletionData) { + for (Map.Entry> entry : + regionId2ReplicaSetAndNode.entrySet()) { + memoryBuffer.add(deletionData.getDataSize()); + entry.getValue().getRight().addTsFileData(deletionData); + offerPieceRegion(entry.getKey()); + } + } + + /** Dispatches from the biggest buffered piece until the data cache is back under budget. */ + boolean dispatchLargestUntilMemoryEnough() throws LoadFileException { + while (!memoryBuffer.isMemoryEnough()) { + final TConsensusGroupId regionId = pollLargestPieceRegion(); + if (regionId == null) { + // No dispatchable piece remains; the remaining buffered data stays buffered until the + // next flush (end of file, a later deletion, or another over-budget event). + break; + } + final Pair pair = + regionId2ReplicaSetAndNode.get(regionId); + final LoadTsFilePieceNode pieceNode = pair.getRight(); + memoryBuffer.release(pieceNode.getDataSize()); + if (!dispatchOne(pieceNode, pair.getLeft())) { + return false; + } + replacePieceNode(regionId, pair.getLeft()); + } + return true; + } + + /** Dispatches every non-empty buffered piece, e.g. at the end of the source TsFile. */ + boolean flushAll() throws LoadFileException { + for (Map.Entry> entry : + regionId2ReplicaSetAndNode.entrySet()) { + final LoadTsFilePieceNode pieceNode = entry.getValue().getRight(); + if (pieceNode.getDataSize() == 0) { + continue; + } + if (!dispatchPieces(Collections.singleton(entry.getKey()))) { + return false; + } + } + return true; + } + + private boolean dispatchPieces(Collection regionIds) throws LoadFileException { + for (TConsensusGroupId regionId : regionIds) { + final Pair pair = + regionId2ReplicaSetAndNode.get(regionId); + if (pair == null) { + continue; + } + final LoadTsFilePieceNode pieceNode = pair.getRight(); + if (pieceNode.getDataSize() == 0) { + continue; + } + memoryBuffer.release(pieceNode.getDataSize()); + if (!dispatchOne(pieceNode, pair.getLeft())) { + LOGGER.warn( + DataNodeQueryMessages.DISPATCH_PIECE_NODE_ARG_OF_TSFILE_ARG_ERROR, + pieceNode, + singleTsFileNode.getTsFileResource().getTsFile()); + return false; + } + replacePieceNode(regionId, pair.getLeft()); + } + return true; + } + + private boolean dispatchOne(LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet) { + return dispatchCallback.dispatch(pieceNode, replicaSet); + } + + private void replacePieceNode(TConsensusGroupId regionId, TRegionReplicaSet replicaSet) { + regionId2ReplicaSetAndNode.replace(regionId, new Pair<>(replicaSet, newPieceNode())); + } + + private LoadTsFilePieceNode newPieceNode() { + return new LoadTsFilePieceNode( + singleTsFileNode.getPlanNodeId(), singleTsFileNode.getTsFileResource().getTsFile()); + } + + private void offerPieceRegion(final TConsensusGroupId regionId) { + final Pair pair = + regionId2ReplicaSetAndNode.get(regionId); + if (pair != null) { + largestPieceRegions.offer(Map.entry(regionId, pair.getRight().getDataSize())); + } + } + + /** Pops the region with the largest non-empty buffered piece, skipping stale heap entries. */ + private TConsensusGroupId pollLargestPieceRegion() { + while (!largestPieceRegions.isEmpty()) { + final Map.Entry entry = largestPieceRegions.poll(); + final Pair pair = + regionId2ReplicaSetAndNode.get(entry.getKey()); + if (pair == null) { + continue; + } + final long currentSize = pair.getRight().getDataSize(); + if (entry.getValue() == currentSize && currentSize > 0) { + return entry.getKey(); + } + } + return null; + } + + void clear() { + regionId2ReplicaSetAndNode.clear(); + largestPieceRegions.clear(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java new file mode 100644 index 0000000000000..73145b30bec64 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/RegionConsensusContext.java @@ -0,0 +1,89 @@ +/* + * 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 java.util.UUID; + +/** + * One high-cohesion state object per target region of the LOAD consensus two-phase protocol. It + * replaces the five parallel maps ({@code regionPieceCount}, {@code regionPieceTotalBytes}, {@code + * regionPieceChecksum}, {@code begunConsensusRegions}, {@code regionLoadId}) that used to live in + * {@link LoadTsFileScheduler}, so the load id, the BEGIN state and the three accumulated counters + * can never diverge. + * + *

Lifecycle, driven by {@link TwoPhaseConsensusLoadStrategy}: + * + *

    + *
  • Created lazily when the first piece of a region is dispatched ({@code + * consensusContexts.computeIfAbsent}). + *
  • {@link #markBegun()} runs before the BEGIN command is submitted, so a region is begun at + * most once per load. + *
  • {@link #accumulate(long, long)} runs after every successful PIECE: it increments the piece + * count, adds the piece bytes and XORs the piece checksum. + *
  • Phase two reads {@link #getPieceCount()}, {@link #getTotalBytes()} and {@link + * #getChecksum()} to build the PREPARE command. + *
+ * + * The load id is generated per instance, so every region of the same source file is isolated from + * the others on the write nodes. + */ +public class RegionConsensusContext { + + /** Each region gets its own load id so its staged data is isolated from other regions. */ + private final String loadId = UUID.randomUUID().toString(); + + private long pieceCount = 0; + private long totalBytes = 0; + private long checksum = 0; + + /** Whether the BEGIN command has already been sent to this region. */ + private boolean begun = false; + + public String getLoadId() { + return loadId; + } + + public long getPieceCount() { + return pieceCount; + } + + public long getTotalBytes() { + return totalBytes; + } + + public long getChecksum() { + return checksum; + } + + public boolean isBegun() { + return begun; + } + + public void markBegun() { + begun = true; + } + + /** Records one successfully applied piece: increments the count, adds bytes and XORs checksum. */ + public void accumulate(long bytes, long pieceChecksum) { + pieceCount++; + totalBytes += bytes; + checksum ^= pieceChecksum; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileLoadStrategy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileLoadStrategy.java new file mode 100644 index 0000000000000..97b89f8208af1 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileLoadStrategy.java @@ -0,0 +1,46 @@ +/* + * 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.commons.exception.IoTDBException; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; + +/** + * Strategy of the LOAD pipeline for loading one source TsFile. The scheduler ({@link + * LoadTsFileScheduler}) decides which strategy applies via {@code needDecodeTsFile}; each strategy + * owns its complete pipeline, phase bookkeeping and phase time metrics. + * + *

Implementations: + * + *

    + *
  • {@link LocalLoadStrategy} - no decode needed, dispatch the file to the local region; + *
  • {@link TwoPhaseConsensusLoadStrategy} - decode, stream PIECE batches through consensus, + * then PREPARE+COMMIT or ABORT. + *
+ */ +public interface TsFileLoadStrategy { + + /** + * Loads the given TsFile. + * + * @return true if the file was loaded successfully + */ + boolean execute(LoadSingleTsFileNode node) throws IoTDBException; +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java new file mode 100644 index 0000000000000..2fe9e25d0a91e --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TsFileSplitConsumer.java @@ -0,0 +1,137 @@ +/* + * 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.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; +import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileSplitter; + +import java.util.ArrayList; +import java.util.List; +import java.util.function.Consumer; + +/** + * LOAD split consumer: receives every {@link TsFileData} of one source TsFile's split output from + * {@link TsFileSplitter} and routes -> buffers -> dispatches it as per-region consensus + * pieces. + * + *
    + *
  • {@code CHUNK} data is buffered directionless first; before dispatch {@link + * DataPartitionRouter} resolves the target regions and {@link PieceDispatcher} appends every + * chunk to its per-region piece. + *
  • {@code DELETION} data is replicated into every buffered piece (chunks are routed first so a + * deletion never overtakes chunk data). + *
  • {@link MemoryBoundedBuffer} guards the shared data cache; over budget, the largest piece is + * dispatched immediately. + *
  • At end of file {@link #sendAllTsFileData()} flushes the remainder. + *
+ * + * The pipeline also notifies the progress-index callback for every chunk's time partition, so the + * strategy can track pipe progress while splitting. {@link #clear()} releases all buffered + * accounting and piece references. + */ +public class TsFileSplitConsumer implements TsFileSplitter.TsFileDataConsumer { + + private final LoadSingleTsFileNode singleTsFileNode; + private final DataPartitionRouter router; + private final MemoryBoundedBuffer memoryBuffer; + private final PieceDispatcher dispatcher; + private final Consumer progressIndexCallback; + + private final List nonDirectionalChunkData = new ArrayList<>(); + + public TsFileSplitConsumer( + LoadSingleTsFileNode singleTsFileNode, + LoadTsFileDataCacheMemoryBlock block, + DataPartitionBatchFetcher partitionFetcher, + String userName, + Consumer progressIndexCallback, + PieceDispatcher.DispatchCallback dispatchCallback) { + this.singleTsFileNode = singleTsFileNode; + this.router = new DataPartitionRouter(partitionFetcher, userName); + this.memoryBuffer = new MemoryBoundedBuffer(block); + this.dispatcher = new PieceDispatcher(singleTsFileNode, memoryBuffer, dispatchCallback); + this.progressIndexCallback = progressIndexCallback; + } + + @Override + public boolean apply(TsFileData tsFileData) throws LoadFileException { + return switch (tsFileData.getType()) { + case CHUNK -> addOrSendChunkData((ChunkData) tsFileData); + case DELETION -> addOrSendDeletionData((DeletionData) tsFileData); + default -> + throw new UnsupportedOperationException( + String.format( + DataNodeQueryMessages.QUERY_EXCEPTION_UNSUPPORTED_TSFILEDATATYPE_S_374475FA, + tsFileData.getType())); + }; + } + + private boolean addOrSendChunkData(ChunkData chunkData) throws LoadFileException { + nonDirectionalChunkData.add(chunkData); + memoryBuffer.add(chunkData.getDataSize()); + progressIndexCallback.accept(chunkData.getTimePartitionSlot()); + + if (!memoryBuffer.isMemoryEnough()) { + routeChunkData(); + if (!dispatcher.dispatchLargestUntilMemoryEnough()) { + return false; + } + } + return true; + } + + private boolean addOrSendDeletionData(DeletionData deletionData) throws LoadFileException { + routeChunkData(); // ensure chunk data will be added before deletion + dispatcher.addDeletionToAll(deletionData); + return true; + } + + private void routeChunkData() throws LoadFileException { + if (nonDirectionalChunkData.isEmpty()) { + return; + } + + final List replicaSets = router.route(nonDirectionalChunkData); + for (int i = 0, size = nonDirectionalChunkData.size(); i < size; i++) { + dispatcher.offerChunk(nonDirectionalChunkData.get(i), replicaSets.get(i)); + } + nonDirectionalChunkData.clear(); + } + + boolean sendAllTsFileData() throws LoadFileException { + routeChunkData(); + return dispatcher.flushAll(); + } + + /** Last-chance cleanup: returns all buffered accounting and drops every piece reference. */ + void clear() { + memoryBuffer.clear(); + nonDirectionalChunkData.clear(); + dispatcher.clear(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java new file mode 100644 index 0000000000000..5a31b590a2b1a --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/TwoPhaseConsensusLoadStrategy.java @@ -0,0 +1,378 @@ +/* + * 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.TConsensusGroupId; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.consensus.index.ProgressIndex; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; +import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadSingleTsFileNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; +import org.apache.iotdb.db.storageengine.load.memory.LoadTsFileDataCacheMemoryBlock; +import org.apache.iotdb.db.storageengine.load.metrics.LoadTsFileCostMetricsSet; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileSplitter; +import org.apache.iotdb.rpc.TSStatusCode; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Two-phase consensus LOAD strategy for files that need decoding. + * + *

Phase 1 (split & stream). {@link #execute(LoadSingleTsFileNode)} resets the + * per-file state, assigns a fresh uuid to {@link LoadTsFileDispatcherImpl} (executor naming / log + * correlation) and feeds the source TsFile through {@link TsFileSplitter} into {@link + * TsFileSplitConsumer}. Every dispatched piece goes through {@code dispatchConsensusPiece}: the + * first piece of a region sends BEGIN with a fresh per-region load id, then PIECE commands with a + * monotonically increasing {@code pieceIndex}; {@link RegionConsensusContext#accumulate(long, + * long)} records piece count, total bytes and the XOR checksum. Submission goes through {@link + * LoadConsensusSubmitter} with bounded retries for transient failures only. + * + *

Phase 2 (commit or abort). If every region received all its pieces, each touched region + * gets PREPARE (with the accumulated count/bytes/checksum) followed by COMMIT; otherwise every + * touched region gets ABORT so the staged data is dropped. + * + *

Per-file state: {@code allReplicaSets} (the regions to prepare/commit/abort), {@code + * consensusContexts} (per-region two-phase state) and {@code timePartitionSlotToProgressIndex} + * (pipe progress index per time partition, collected while splitting for the upcoming + * progress-index sync). + */ +public class TwoPhaseConsensusLoadStrategy implements TsFileLoadStrategy { + + private static final Logger LOGGER = LoggerFactory.getLogger(TwoPhaseConsensusLoadStrategy.class); + + private static final LoadTsFileCostMetricsSet LOAD_TSFILE_COST_METRICS_SET = + LoadTsFileCostMetricsSet.getInstance(); + + /** + * Bounded retry for transient LOAD consensus submission failures (network errors, region + * migration, transient server errors). The write node deduplicates pieces by (loadId, pieceIndex, + * checksum), so a retried request whose first attempt actually applied is acknowledged as success + * instead of being applied twice. + */ + private static final int LOAD_CONSENSUS_SUBMIT_MAX_RETRIES = 3; + + private static final long LOAD_CONSENSUS_SUBMIT_RETRY_BACKOFF_MS = 100L; + + private final LoadTsFileDispatcherImpl dispatcher; + private final DataPartitionBatchFetcher partitionFetcher; + private final LoadTsFileDataCacheMemoryBlock block; + private final LoadConsensusSubmitter consensusSubmitter; + private final String userName; + private final boolean isGeneratedByPipe; + + /** Regions touched by the current file; used to send ABORT/PREPARE+COMMIT in phase two. */ + private final Set allReplicaSets = new HashSet<>(); + + /** Per-region two-phase state of the current file; replaces the old five parallel maps. */ + private final Map consensusContexts = + new ConcurrentHashMap<>(); + + /** + * Progress index per time partition, assigned while the file is being split. Kept for the + * upcoming progress-index sync with the consensus prepare phase. + */ + private final Map timePartitionSlotToProgressIndex = + new HashMap<>(); + + public TwoPhaseConsensusLoadStrategy( + LoadTsFileDispatcherImpl dispatcher, + DataPartitionBatchFetcher partitionFetcher, + LoadTsFileDataCacheMemoryBlock block, + LoadConsensusSubmitter consensusSubmitter, + String userName, + boolean isGeneratedByPipe) { + this.dispatcher = dispatcher; + this.partitionFetcher = partitionFetcher; + this.block = block; + this.consensusSubmitter = consensusSubmitter; + this.userName = userName; + this.isGeneratedByPipe = isGeneratedByPipe; + } + + @Override + public boolean execute(LoadSingleTsFileNode node) { + dispatcher.setUuid(UUID.randomUUID().toString()); + allReplicaSets.clear(); + consensusContexts.clear(); + timePartitionSlotToProgressIndex.clear(); + + long startTime = System.nanoTime(); + final boolean isFirstPhaseSuccess; + try { + isFirstPhaseSuccess = firstPhase(node); + } finally { + LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( + LoadTsFileCostMetricsSet.FIRST_PHASE, System.nanoTime() - startTime); + } + + startTime = System.nanoTime(); + final boolean isSecondPhaseSuccess; + try { + isSecondPhaseSuccess = secondPhase(isFirstPhaseSuccess); + } finally { + LOAD_TSFILE_COST_METRICS_SET.recordPhaseTimeCost( + LoadTsFileCostMetricsSet.SECOND_PHASE, System.nanoTime() - startTime); + } + + return isFirstPhaseSuccess && isSecondPhaseSuccess; + } + + private boolean firstPhase(LoadSingleTsFileNode node) { + final TsFileSplitConsumer pipeline = + new TsFileSplitConsumer( + node, + block, + partitionFetcher, + userName, + this::computeTimePartitionSlotToProgressIndexIfAbsent, + this::dispatchOnePieceNode); + try { + new TsFileSplitter(node.getTsFileResource().getTsFile(), pipeline) + .splitTsFileByDataPartition(); + return pipeline.sendAllTsFileData(); + } catch (IllegalStateException e) { + LOGGER.warn( + String.format( + DataNodeQueryMessages.DISPATCH_TSFILEDATA_ERROR_WHEN_PARSING_TSFILE_S, + node.getTsFileResource().getTsFile()), + e); + return false; + } catch (Exception e) { + LOGGER.warn( + String.format( + DataNodeQueryMessages.PARSE_OR_SEND_TSFILE_S_ERROR, + node.getTsFileResource().getTsFile()), + e); + return false; + } finally { + pipeline.clear(); + } + } + + private boolean dispatchOnePieceNode( + LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet) { + allReplicaSets.add(replicaSet); + return dispatchConsensusPiece(pieceNode, replicaSet); + } + + /** + * Submits a LOAD consensus request with a bounded number of attempts. Only transient failures are + * retried; permanent rejections (checksum mismatch, missing staged writer) are returned to the + * caller immediately so the scheduler can abort. + */ + private TSStatus submitConsensusWithRetry( + TRegionReplicaSet replicaSet, LoadTsFileConsensusNode node) { + TSStatus status = null; + for (int attempt = 1; attempt <= LOAD_CONSENSUS_SUBMIT_MAX_RETRIES; attempt++) { + status = consensusSubmitter.submit(replicaSet, node); + if (status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode() + || !isTransientConsensusFailure(status) + || attempt == LOAD_CONSENSUS_SUBMIT_MAX_RETRIES) { + break; + } + LOGGER.warn( + DataNodeQueryMessages.LOG_LOAD_CONSENSUS_SUBMIT_TRANSIENT_FAILURE_RETRY_D7E1D9A6, + node.getOp(), + node.getLoadId(), + replicaSet, + attempt, + LOAD_CONSENSUS_SUBMIT_MAX_RETRIES, + status.getMessage()); + try { + Thread.sleep(LOAD_CONSENSUS_SUBMIT_RETRY_BACKOFF_MS * attempt); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + } + return status; + } + + private boolean isTransientConsensusFailure(TSStatus status) { + switch (TSStatusCode.representOf(status.getCode())) { + case DISPATCH_ERROR: + case INTERNAL_SERVER_ERROR: + case NO_AVAILABLE_REGION_GROUP: + case EXECUTE_STATEMENT_ERROR: + return true; + default: + return false; + } + } + + private boolean dispatchConsensusPiece( + LoadTsFilePieceNode pieceNode, TRegionReplicaSet replicaSet) { + final TConsensusGroupId regionId = replicaSet.getRegionId(); + final RegionConsensusContext context = + consensusContexts.computeIfAbsent(regionId, o -> new RegionConsensusContext()); + final String loadId = context.getLoadId(); + + if (!context.isBegun()) { + context.markBegun(); + final LoadTsFileConsensusNode begin = + LoadTsFileConsensusNode.begin( + new PlanNodeId("load-begin-" + loadId), + loadId, + pieceNode.getTsFile() == null ? null : pieceNode.getTsFile().getName(), + false, + "", + -1); + final TSStatus beginStatus = submitConsensusWithRetry(replicaSet, begin); + if (beginStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.warn( + DataNodeQueryMessages.DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG + + DataNodeQueryMessages + .RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG, + replicaSet, + TSStatusCode.representOf(beginStatus.getCode()).name(), + beginStatus.getMessage(), + pieceNode); + return false; + } + } + + final long pieceIndex = context.getPieceCount(); + final LoadTsFileConsensusNode piece = + LoadTsFileConsensusNode.piece( + new PlanNodeId("load-piece-" + loadId + "-" + pieceIndex), + loadId, + pieceNode.getTsFile() == null ? null : pieceNode.getTsFile().getName(), + pieceIndex, + 0L, + pieceNode.getAllTsFileData(), + 0L); + final TSStatus pieceStatus = submitConsensusWithRetry(replicaSet, piece); + if (pieceStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.warn( + DataNodeQueryMessages.DISPATCH_ONE_PIECE_TO_REPLICASET_ARG_ERROR_RESULT_STATUS_CODE_ARG + + DataNodeQueryMessages + .RESULT_STATUS_MESSAGE_ARG_DISPATCH_PIECE_NODE_ERROR_PERCENT_NARG, + replicaSet, + TSStatusCode.representOf(pieceStatus.getCode()).name(), + pieceStatus.getMessage(), + pieceNode); + return false; + } + context.accumulate(piece.getDataSize(), piece.getChecksum()); + return true; + } + + private boolean secondPhase(boolean isFirstPhaseSuccess) { + if (!isFirstPhaseSuccess) { + return abortAllRegions(); + } + return prepareAndCommitAllRegions(); + } + + private boolean abortAllRegions() { + for (TRegionReplicaSet replicaSet : allReplicaSets) { + final String loadId = consensusContexts.get(replicaSet.getRegionId()).getLoadId(); + final LoadTsFileConsensusNode abort = + LoadTsFileConsensusNode.abort( + new PlanNodeId("load-abort-" + loadId), loadId, null, isGeneratedByPipe); + final TSStatus status = consensusSubmitter.submit(replicaSet, abort); + if (status.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.warn( + DataNodeQueryMessages + .DISPATCH_LOAD_COMMAND_ARG_OF_TSFILE_ARG_ERROR_TO_REPLICASETS_ARG_ERROR + + DataNodeQueryMessages.RESULT_STATUS_CODE_ARG_RESULT_STATUS_MESSAGE_ARG, + abort, + loadId, + allReplicaSets, + TSStatusCode.representOf(status.getCode()).name(), + status.getMessage()); + return false; + } + } + return true; + } + + private boolean prepareAndCommitAllRegions() { + for (TRegionReplicaSet replicaSet : allReplicaSets) { + final RegionConsensusContext context = consensusContexts.get(replicaSet.getRegionId()); + final String loadId = context.getLoadId(); + final LoadTsFileConsensusNode prepare = + LoadTsFileConsensusNode.prepare( + new PlanNodeId("load-prepare-" + loadId), + loadId, + null, + (int) context.getPieceCount(), + context.getTotalBytes(), + context.getChecksum(), + Collections.emptyMap()); + final TSStatus prepareStatus = consensusSubmitter.submit(replicaSet, prepare); + if (prepareStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.warn( + DataNodeQueryMessages + .DISPATCH_LOAD_COMMAND_ARG_OF_TSFILE_ARG_ERROR_TO_REPLICASETS_ARG_ERROR + + DataNodeQueryMessages.RESULT_STATUS_CODE_ARG_RESULT_STATUS_MESSAGE_ARG, + prepare, + loadId, + allReplicaSets, + TSStatusCode.representOf(prepareStatus.getCode()).name(), + prepareStatus.getMessage()); + return false; + } + + final LoadTsFileConsensusNode commit = + LoadTsFileConsensusNode.commit( + new PlanNodeId("load-commit-" + loadId), + loadId, + null, + isGeneratedByPipe, + false, + Collections.emptyMap()); + final TSStatus commitStatus = consensusSubmitter.submit(replicaSet, commit); + if (commitStatus.getCode() != TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + LOGGER.warn( + DataNodeQueryMessages + .DISPATCH_LOAD_COMMAND_ARG_OF_TSFILE_ARG_ERROR_TO_REPLICASETS_ARG_ERROR + + DataNodeQueryMessages.RESULT_STATUS_CODE_ARG_RESULT_STATUS_MESSAGE_ARG, + commit, + loadId, + allReplicaSets, + TSStatusCode.representOf(commitStatus.getCode()).name(), + commitStatus.getMessage()); + return false; + } + } + return true; + } + + private void computeTimePartitionSlotToProgressIndexIfAbsent( + final TTimePartitionSlot timePartitionSlot) { + timePartitionSlotToProgressIndex.putIfAbsent( + timePartitionSlot, PipeDataNodeAgent.runtime().getNextProgressIndexForTsFileLoad()); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java index 4a679790dd44c..d276f87ec9d1a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java @@ -171,6 +171,10 @@ public class StorageEngine implements IService { private final LoadTsFileManager loadTsFileManager = new LoadTsFileManager(); + public LoadTsFileManager getLoadTsFileManager() { + return loadTsFileManager; + } + public final AtomicLong objectFileId = new AtomicLong(0); private StorageEngine() {} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLoader.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLoader.java index 3c9d6ed6a8fa2..9ebf332784a0e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLoader.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotLoader.java @@ -29,6 +29,7 @@ import org.apache.iotdb.db.storageengine.dataregion.DataRegion; import org.apache.iotdb.db.storageengine.dataregion.flush.CompressionRatio; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.load.LoadTsFileManager; import org.apache.tsfile.common.constant.TsFileConstant; import org.apache.tsfile.external.commons.io.FileUtils; @@ -156,6 +157,7 @@ private DataRegion loadSnapshotFromMultipleDirs() { // IoTConsensus fragments arrive under different recv folders; do not map each // fragment back to the same disk as its recv path, rely on fileTarget instead. createLinksFromSnapshotDirToDataDirWithoutLog(snapshotDir, fileTarget, false); + restoreLoadTasksFromSnapshotDir(snapshotDir); loadCompressionRatio(snapshotDir); } return loadSnapshot(); @@ -178,6 +180,7 @@ private DataRegion loadSnapshotWithoutLog() { LOGGER.info(StorageEngineMessages.MOVING_SNAPSHOT_FILE_TO_DATA_DIRS); File snapshotDir = new File(snapshotPath); createLinksFromSnapshotDirToDataDirWithoutLog(snapshotDir, new HashMap<>(), true); + restoreLoadTasksFromSnapshotDir(snapshotDir); loadCompressionRatio(snapshotDir); return loadSnapshot(); } catch (IOException | DiskSpaceInsufficientException e) { @@ -236,6 +239,7 @@ private DataRegion loadSnapshotWithLog(File logFile) { deleteAllFilesInDataDirs(); LOGGER.info(StorageEngineMessages.REMOVE_ALL_DATA_FILES_IN_ORIGINAL_DIR); createLinksFromSnapshotDirToDataDirWithLog(); + restoreLoadTasksFromSnapshotDir(new File(snapshotPath)); loadCompressionRatio(new File(snapshotPath)); return loadSnapshot(); } catch (IOException e) { @@ -247,6 +251,22 @@ private DataRegion loadSnapshotWithLog(File logFile) { } } + /** + * Restores the in-progress LOAD staging files carried by this snapshot's {@value + * LoadTsFileManager#LOAD_SNAPSHOT_DIR_NAME} directory, if any. A failure here must fail the whole + * snapshot load: silently dropping the staged files would let a later COMMIT replay hit a + * high-availability hole instead of continuing the load. + */ + private void restoreLoadTasksFromSnapshotDir(File snapshotDir) throws IOException { + final File loadSnapshotDir = new File(snapshotDir, LoadTsFileManager.LOAD_SNAPSHOT_DIR_NAME); + if (!loadSnapshotDir.isDirectory()) { + return; + } + StorageEngine.getInstance() + .getLoadTsFileManager() + .restoreLoadTasksFromSnapshot(loadSnapshotDir); + } + private void deleteAllFilesInDataDirs() throws IOException { String[] dataDirPaths = IoTDBDescriptor.getInstance().getConfig().getLocalDataDirs(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotTaker.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotTaker.java index 6707c3f6cc99b..10b64dab21c2e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotTaker.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/snapshot/SnapshotTaker.java @@ -24,6 +24,7 @@ import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.exception.DirectoryNotLegalException; import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.storageengine.StorageEngine; import org.apache.iotdb.db.storageengine.dataregion.DataRegion; import org.apache.iotdb.db.storageengine.dataregion.flush.CompressionRatio; import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; @@ -104,10 +105,14 @@ public boolean takeFullSnapshot( } success = createSnapshot(seqFiles, tempSnapshotId); success = success && createSnapshot(unseqFiles, tempSnapshotId); - success = success && snapshotCompressionRatio(snapshotDirPath); } finally { readUnlockTheFile(); } + // The LOAD staging files are independent of the TsFileManager, and snapshotting them while + // holding the TsFileManager read lock would invert the lock order against a COMMIT apply + // (loadLock read -> TsFileManager write), so they are copied after the lock is released. + success = success && snapshotLoadTasks(snapshotDir); + success = success && snapshotCompressionRatio(snapshotDirPath); if (!success) { LOGGER.warn( @@ -141,6 +146,24 @@ public boolean takeFullSnapshot( } } + /** + * Includes the in-progress LOAD staging files of this DataRegion into the snapshot so that a + * replica restored from it can continue (or at least explicitly fail) the pending load instead of + * replaying COMMIT as a silent no-op. Only the already-synced byte prefix of every staged file is + * copied; the remaining bytes are delivered by the following PIECE refs. + */ + private boolean snapshotLoadTasks(File snapshotDir) { + try { + StorageEngine.getInstance() + .getLoadTsFileManager() + .snapshotLoadTasksForRegion(dataRegion, snapshotDir); + return true; + } catch (Exception e) { + LOGGER.error(StorageEngineMessages.CATCH_IO_EXCEPTION_CREATING_SNAPSHOT, e); + return false; + } + } + private boolean snapshotCompressionRatio(String snapshotDir) { File compressionRatioFile = CompressionRatio.getInstance().getCompressionRatioFile(dataRegion.getDataRegionIdString()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntry.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntry.java index 8c0a5e353291e..bb3a18fad6199 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntry.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntry.java @@ -22,6 +22,7 @@ import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ContinuousSameSearchIndexSeparatorNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; @@ -81,6 +82,8 @@ protected WALEntry(long memTableId, WALEntryValue value, boolean wait) { this.type = WALEntryType.RELATIONAL_DELETE_DATA_NODE; } else if (value instanceof ObjectNode) { this.type = WALEntryType.OBJECT_FILE_NODE; + } else if (value instanceof LoadTsFileConsensusNode) { + this.type = WALEntryType.LOAD_TSFILE_CONSENSUS_NODE; } else { throw new RuntimeException(StorageEngineMessages.UNKNOWN_WAL_ENTRY_TYPE); } @@ -141,6 +144,9 @@ public static WALEntry deserialize(DataInputStream stream) throws IOException { case OBJECT_FILE_NODE: value = (ObjectNode) PlanNodeType.deserializeFromWAL(stream); break; + case LOAD_TSFILE_CONSENSUS_NODE: + value = (LoadTsFileConsensusNode) PlanNodeType.deserializeFromWAL(stream); + break; default: throw new RuntimeException(StorageEngineMessages.UNKNOWN_WAL_ENTRY_TYPE_WITH_VALUE + type); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntryType.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntryType.java index 2ad72f4a2324d..bbce521baf6a0 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntryType.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALEntryType.java @@ -48,6 +48,8 @@ public enum WALEntryType { MEMORY_TABLE_SNAPSHOT((byte) 10), RELATIONAL_DELETE_DATA_NODE((byte) 11), OBJECT_FILE_NODE((byte) 12), + /** {@link org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode} */ + LOAD_TSFILE_CONSENSUS_NODE((byte) 13), // endregion // region signal entry type // signal wal buffer has been closed @@ -75,7 +77,8 @@ public boolean needSearch() { || this == INSERT_ROWS_NODE || this == DELETE_DATA_NODE || this == RELATIONAL_DELETE_DATA_NODE - || this == OBJECT_FILE_NODE; + || this == OBJECT_FILE_NODE + || this == LOAD_TSFILE_CONSENSUS_NODE; } public static WALEntryType valueOf(byte code) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALInfoEntry.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALInfoEntry.java index 957cf03d89237..38a445bc8fcef 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALInfoEntry.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/buffer/WALInfoEntry.java @@ -22,6 +22,7 @@ import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; @@ -116,6 +117,9 @@ public void serialize(IWALByteBufferView buffer) { case OBJECT_FILE_NODE: ((ObjectNode) value).serializeToWAL(buffer, encodedSearchIndex); break; + case LOAD_TSFILE_CONSENSUS_NODE: + ((LoadTsFileConsensusNode) value).serializeToWAL(buffer, encodedSearchIndex); + break; case MEMORY_TABLE_SNAPSHOT: case CONTINUOUS_SAME_SEARCH_INDEX_SEPARATOR_NODE: value.serializeToWAL(buffer); @@ -210,6 +214,8 @@ public long getMemorySize() { return RamUsageEstimator.sizeOfObject(value); case OBJECT_FILE_NODE: return ((ObjectNode) value).serializedSize(); + case LOAD_TSFILE_CONSENSUS_NODE: + return ((LoadTsFileConsensusNode) value).serializedSize(); default: throw new RuntimeException(StorageEngineMessages.UNSUPPORTED_WAL_ENTRY_TYPE + type); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/IWALNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/IWALNode.java index a8fbbee0dc4f5..b1f22fdf55f0e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/IWALNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/IWALNode.java @@ -21,6 +21,7 @@ import org.apache.iotdb.consensus.common.DataSet; import org.apache.iotdb.consensus.iot.log.ConsensusReqReader; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ContinuousSameSearchIndexSeparatorNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; @@ -57,6 +58,9 @@ public interface IWALNode extends FlushListener, AutoCloseable, ConsensusReqRead WALFlushListener log(long memTableId, ObjectNode objectNode); + /** Log consensus-backed LOAD request. */ + WALFlushListener log(long memTableId, LoadTsFileConsensusNode loadTsFileConsensusNode); + /** Callback when memTable created. */ void onMemTableCreated(IMemTable memTable, String targetTsFile); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALFakeNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALFakeNode.java index d24d1a1d82348..411ddca9d5de2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALFakeNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALFakeNode.java @@ -20,6 +20,7 @@ package org.apache.iotdb.db.storageengine.dataregion.wal.node; import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ContinuousSameSearchIndexSeparatorNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; @@ -89,6 +90,11 @@ public WALFlushListener log(long memTableId, ObjectNode objectNode) { return getResult(); } + @Override + public WALFlushListener log(long memTableId, LoadTsFileConsensusNode loadTsFileConsensusNode) { + return getResult(); + } + private WALFlushListener getResult() { switch (status) { case SUCCESS: diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALNode.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALNode.java index a59ebeaa21db4..b77138efd4201 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALNode.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/dataregion/wal/node/WALNode.java @@ -29,6 +29,7 @@ import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.ContinuousSameSearchIndexSeparatorNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.DeleteDataNode; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.write.InsertRowNode; @@ -208,6 +209,12 @@ public WALFlushListener log(long memTableId, ObjectNode objectNode) { return log(walEntry); } + @Override + public WALFlushListener log(long memTableId, LoadTsFileConsensusNode loadTsFileConsensusNode) { + WALEntry walEntry = new WALInfoEntry(memTableId, loadTsFileConsensusNode); + return log(walEntry); + } + private WALFlushListener log(WALEntry walEntry) { buffer.write(walEntry); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/DataPartitionInfo.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/DataPartitionInfo.java new file mode 100644 index 0000000000000..b1423911c9e11 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/DataPartitionInfo.java @@ -0,0 +1,76 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.storageengine.load; + +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.conf.IoTDBConstant; +import org.apache.iotdb.db.storageengine.dataregion.DataRegion; + +import java.util.Objects; + +/** + * LOAD partition key: immutable identifier of one staged partition file: (DataRegion, time + * partition slot). + */ +final class DataPartitionInfo { + + private final DataRegion dataRegion; + private final TTimePartitionSlot timePartitionSlot; + + DataPartitionInfo(DataRegion dataRegion, TTimePartitionSlot timePartitionSlot) { + this.dataRegion = dataRegion; + this.timePartitionSlot = timePartitionSlot; + } + + DataRegion getDataRegion() { + return dataRegion; + } + + TTimePartitionSlot getTimePartitionSlot() { + return timePartitionSlot; + } + + @Override + public String toString() { + return String.join( + IoTDBConstant.FILE_NAME_SEPARATOR, + dataRegion.getDatabaseName(), + dataRegion.getDataRegionIdString(), + Long.toString(timePartitionSlot.getStartTime())); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + DataPartitionInfo that = (DataPartitionInfo) o; + return Objects.equals(dataRegion, that.dataRegion) + && timePartitionSlot.getStartTime() == that.timePartitionSlot.getStartTime(); + } + + @Override + public int hashCode() { + return Objects.hash(dataRegion, timePartitionSlot.getStartTime()); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadCleanupScheduler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadCleanupScheduler.java new file mode 100644 index 0000000000000..3f75405ef263a --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadCleanupScheduler.java @@ -0,0 +1,129 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.storageengine.load; + +import org.apache.iotdb.db.conf.IoTDBConfig; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Consumer; + +/** + * Background cleanup daemon for abandoned LOAD tasks. A plain hash map with lazy expiry replaces + * the old {@code PriorityBlockingQueue}: {@link #registerOrRefresh} is an O(1) {@code compute}, + * eviction removes exactly the expired entry and never scans the queue, and no monitor is held + * while the eviction action (which closes writers and may block) runs. + */ +final class LoadCleanupScheduler { + + private static final Logger LOGGER = LoggerFactory.getLogger(LoadCleanupScheduler.class); + private static final IoTDBConfig CONFIG = IoTDBDescriptor.getInstance().getConfig(); + + private static final class TaskState { + private volatile long expireTime; + private volatile boolean isRunning; + } + + private final ConcurrentHashMap tasks = new ConcurrentHashMap<>(); + private final Consumer evictionAction; + private final long delayInMs; + + LoadCleanupScheduler(long timeoutSeconds, Consumer evictionAction) { + this.evictionAction = evictionAction; + this.delayInMs = timeoutSeconds * 1000L; + } + + /** Registers the task or refreshes its expiry, without touching any other entry. */ + void registerOrRefresh(String uuid) { + tasks.compute( + uuid, + (id, state) -> { + if (state == null) { + state = new TaskState(); + } + state.expireTime = System.currentTimeMillis() + delayInMs; + return state; + }); + } + + void markRunning(String uuid) { + final TaskState state = tasks.get(uuid); + if (state != null) { + state.isRunning = true; + state.expireTime = System.currentTimeMillis() + delayInMs; + } + } + + void markIdle(String uuid) { + final TaskState state = tasks.get(uuid); + if (state != null) { + state.isRunning = false; + state.expireTime = System.currentTimeMillis() + delayInMs; + } + } + + void remove(String uuid) { + tasks.remove(uuid); + } + + void start() { + PipeDataNodeAgent.runtime() + .registerPeriodicalJob( + "LoadTsFileManager#cleanupTasks", + this::sweep, + CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds() >> 2); + } + + void shutdown() { + tasks.clear(); + } + + private void sweep() { + for (Map.Entry entry : tasks.entrySet()) { + final String uuid = entry.getKey(); + final TaskState state = entry.getValue(); + if (state.isRunning) { + // A live LOAD must never be evicted; defer it like the old queue re-schedule. + state.expireTime = System.currentTimeMillis() + delayInMs; + continue; + } + if (state.expireTime > System.currentTimeMillis()) { + continue; + } + final TaskState removed = tasks.remove(uuid); + if (removed == null) { + continue; + } + LOGGER.info(StorageEngineMessages.LOAD_CLEANUP_TASK_STARTS, uuid); + try { + // Run outside the map iteration and any monitor: the eviction closes writers and may block. + evictionAction.accept(uuid); + } catch (Exception e) { + LOGGER.warn(StorageEngineMessages.LOAD_CLEANUP_TASK_ERROR, uuid, e); + } + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadSnapshotManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadSnapshotManager.java new file mode 100644 index 0000000000000..3608bacd64239 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadSnapshotManager.java @@ -0,0 +1,332 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.storageengine.load; + +import org.apache.iotdb.commons.exception.DiskSpaceInsufficientException; +import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.storageengine.dataregion.DataRegion; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.StandardCopyOption; +import java.util.ArrayList; +import java.util.List; + +/** + * Owns everything about carrying in-progress LOAD staging files through a DataRegion snapshot: + * snapshot inclusion, restore registration and the {@code snapshot.meta} format. Live data writing + * never touches this class, which is what lets it be factored out of the LOAD facade. + */ +final class LoadSnapshotManager { + + private static final Logger LOGGER = LoggerFactory.getLogger(LoadSnapshotManager.class); + + private static final String LOAD_SNAPSHOT_META_NAME = "snapshot.meta"; + private static final String APPLIED_PIECES_PREFIX = "#applied "; + + private final LoadTaskRegistry registry; + private final LoadCleanupScheduler cleanupScheduler; + private final TaskDirAllocator taskDirAllocator; + + LoadSnapshotManager( + LoadTaskRegistry registry, + LoadCleanupScheduler cleanupScheduler, + TaskDirAllocator taskDirAllocator) { + this.registry = registry; + this.cleanupScheduler = cleanupScheduler; + this.taskDirAllocator = taskDirAllocator; + } + + /** + * Includes the in-progress LOAD staging files owned by the given DataRegion into a snapshot. + * + *

Only the already-synced byte prefix {@code [0, syncedOffset)} of every staged partition file + * is copied: bytes written after the last chunk-group boundary are still owned by the write path + * and will be captured by the next PIECE ref, so a replica restored from this snapshot can keep + * appending from exactly the snapshot length without a hole or an overlap. The registry write + * lock plus the per-task lock serializes this against concurrent LOAD applies so the synced + * cursor is stable. + */ + void snapshotLoadTasksForRegion(DataRegion dataRegion, File snapshotDir) throws IOException { + final int[] taskCount = new int[1]; + final int[] stagedFileCount = new int[1]; + registry.snapshot( + writerManager -> { + if (!writerManager.belongsTo(dataRegion)) { + return; + } + final File taskSnapshotDir = + new File( + snapshotDir, + LoadTsFileManager.LOAD_SNAPSHOT_DIR_NAME + + File.separator + + writerManager.getTaskName()); + if (!taskSnapshotDir.exists() && !taskSnapshotDir.mkdirs()) { + throw new IOException( + String.format( + StorageEngineMessages.FAILED_TO_CREATE_DIR, taskSnapshotDir.getAbsolutePath())); + } + final TaskSnapshot taskSnapshot = writerManager.snapshotTask(taskSnapshotDir); + if (taskSnapshot.stagedFiles.isEmpty()) { + return; + } + writeSnapshotMeta(new File(taskSnapshotDir, LOAD_SNAPSHOT_META_NAME), taskSnapshot); + taskCount[0]++; + stagedFileCount[0] += taskSnapshot.stagedFiles.size(); + }); + if (taskCount[0] > 0) { + LOGGER.info( + StorageEngineMessages.LOG_LOAD_CONSENSUS_SNAPSHOT_TAKEN_09A7DD4C, + taskCount[0], + stagedFileCount[0], + dataRegion.getDataRegionIdString(), + snapshotDir); + } + } + + /** + * Restores the in-progress LOAD staging files carried by a snapshot's {@value + * LoadTsFileManager#LOAD_SNAPSHOT_DIR_NAME} directory. The restored task dirs are registered so + * that the coordinator can continue the load (subsequent PIECE refs append to the restored files + * and COMMIT binds them to the DataRegion). + */ + void restoreLoadTasksFromSnapshot(File loadSnapshotDir) throws IOException { + final File[] taskDirs = loadSnapshotDir.listFiles(File::isDirectory); + if (taskDirs == null) { + return; + } + int taskCount = 0; + int stagedFileCount = 0; + for (File taskSnapshotDir : taskDirs) { + final File metaFile = new File(taskSnapshotDir, LOAD_SNAPSHOT_META_NAME); + if (!metaFile.isFile()) { + continue; + } + final String uuid = taskSnapshotDir.getName(); + final TaskSnapshot taskSnapshot = parseSnapshotMeta(metaFile); + TsFileWriterManager writerManager = registry.get(uuid).orElse(null); + if (writerManager == null) { + try { + writerManager = + registry.getOrCreate( + uuid, + id -> { + final File targetTaskDir = taskDirAllocator.allocate(id); + copyLoadSnapshotTaskFiles(taskSnapshotDir, targetTaskDir); + return new TsFileWriterManager(targetTaskDir, false); + }); + } catch (IOException e) { + if (e.getCause() instanceof DiskSpaceInsufficientException) { + throw new IOException( + String.format( + StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_SNAPSHOT_RESTORE_FAILED_F8C29C64, + loadSnapshotDir, + e.getCause().getMessage()), + e); + } + throw e; + } + cleanupScheduler.registerOrRefresh(uuid); + taskCount++; + } else { + // A snapshot may be spread across several receive folders: merge the remaining files + // into the task dir created by an earlier fragment of the same load. + copyLoadSnapshotTaskFiles(taskSnapshotDir, writerManager.getTaskDir()); + } + writerManager.registerRestoredPartitions(taskSnapshot.stagedFiles); + writerManager.restoreAppliedPieces(taskSnapshot.appliedPieces); + cleanupScheduler.markRunning(uuid); + stagedFileCount += taskSnapshot.stagedFiles.size(); + } + if (taskCount > 0) { + LOGGER.info( + StorageEngineMessages.LOG_LOAD_CONSENSUS_SNAPSHOT_RESTORED_90ABC1BF, + taskCount, + stagedFileCount, + loadSnapshotDir); + } + } + + private void copyLoadSnapshotTaskFiles(File sourceTaskDir, File targetTaskDir) + throws IOException { + if (!targetTaskDir.exists() && !targetTaskDir.mkdirs()) { + throw new IOException( + String.format( + StorageEngineMessages.FAILED_TO_CREATE_DIR, targetTaskDir.getAbsolutePath())); + } + final File[] files = sourceTaskDir.listFiles(); + if (files == null) { + return; + } + for (File file : files) { + if (file.getName().equals(LOAD_SNAPSHOT_META_NAME)) { + continue; + } + if (file.isDirectory()) { + copyDirectoryRecursively(file, new File(targetTaskDir, file.getName())); + } else if (file.isFile()) { + Files.copy( + file.toPath(), + new File(targetTaskDir, file.getName()).toPath(), + StandardCopyOption.REPLACE_EXISTING); + } + } + } + + private static void copyDirectoryRecursively(File sourceDir, File targetDir) throws IOException { + if (!targetDir.exists() && !targetDir.mkdirs()) { + throw new IOException( + String.format(StorageEngineMessages.FAILED_TO_CREATE_DIR, targetDir.getAbsolutePath())); + } + final File[] files = sourceDir.listFiles(); + if (files == null) { + return; + } + for (final File file : files) { + if (file.isDirectory()) { + copyDirectoryRecursively(file, new File(targetDir, file.getName())); + } else if (file.isFile()) { + Files.copy( + file.toPath(), + new File(targetDir, file.getName()).toPath(), + StandardCopyOption.REPLACE_EXISTING); + } + } + } + + static void writeSnapshotMeta(File metaFile, TaskSnapshot taskSnapshot) throws IOException { + final StringBuilder sb = new StringBuilder(); + sb.append(APPLIED_PIECES_PREFIX).append(taskSnapshot.appliedPieces).append('\n'); + for (StagedFileSnapshot snapshot : taskSnapshot.stagedFiles) { + sb.append(snapshot.fileName) + .append('\t') + .append(snapshot.database) + .append('\t') + .append(snapshot.regionId) + .append('\t') + .append(snapshot.timePartitionStart) + .append('\t') + .append(snapshot.finalized) + .append('\n'); + } + Files.write(metaFile.toPath(), sb.toString().getBytes(StandardCharsets.UTF_8)); + } + + static TaskSnapshot parseSnapshotMeta(File metaFile) throws IOException { + final List stagedFiles = new ArrayList<>(); + final StringBuilder appliedPieces = new StringBuilder(); + for (String line : Files.readAllLines(metaFile.toPath(), StandardCharsets.UTF_8)) { + if (line.isEmpty()) { + continue; + } + if (line.startsWith(APPLIED_PIECES_PREFIX)) { + appliedPieces.append(line.substring(APPLIED_PIECES_PREFIX.length())); + continue; + } + final String[] parts = line.split("\t", -1); + if (parts.length != 5) { + throw new IOException( + String.format( + StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_SNAPSHOT_RESTORE_FAILED_F8C29C64, + metaFile, + line)); + } + stagedFiles.add( + new StagedFileSnapshot( + parts[0], + parts[1], + parts[2], + Long.parseLong(parts[3]), + Boolean.parseBoolean(parts[4]))); + } + return new TaskSnapshot(stagedFiles, appliedPieces.toString()); + } + + /** One task's snapshot payload: the staged-file metadata plus the applied-piece prefix. */ + static final class TaskSnapshot { + private final List stagedFiles; + private final String appliedPieces; + + TaskSnapshot(List stagedFiles, String appliedPieces) { + this.stagedFiles = stagedFiles; + this.appliedPieces = appliedPieces; + } + + List getStagedFiles() { + return stagedFiles; + } + + String getAppliedPieces() { + return appliedPieces; + } + } + + /** Immutable description of one staged partition file captured by a snapshot. */ + static final class StagedFileSnapshot { + private final String fileName; + private final String database; + private final String regionId; + private final long timePartitionStart; + private final boolean finalized; + + StagedFileSnapshot( + String fileName, + String database, + String regionId, + long timePartitionStart, + boolean finalized) { + this.fileName = fileName; + this.database = database; + this.regionId = regionId; + this.timePartitionStart = timePartitionStart; + this.finalized = finalized; + } + + String getFileName() { + return fileName; + } + + String getDatabase() { + return database; + } + + String getRegionId() { + return regionId; + } + + long getTimePartitionStart() { + return timePartitionStart; + } + + boolean isFinalized() { + return finalized; + } + } + + @FunctionalInterface + interface TaskDirAllocator { + File allocate(String uuid) throws Exception; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTaskRegistry.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTaskRegistry.java new file mode 100644 index 0000000000000..2398947f4c91b --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTaskRegistry.java @@ -0,0 +1,117 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.storageengine.load; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +import static org.apache.iotdb.db.i18n.StorageEngineMessages.STORAGE_EXCEPTION_FAILED_TO_CREATE_TSFILEWRITERMANAGER_FOR_UUID_S_BECAUSE_A0D68950; + +/** + * LOAD task registry: owns the uuid -> {@link TsFileWriterManager} lifecycle mapping of one + * in-progress LOAD task. Single-task operations (get/create/remove) only take the cheap read lock, + * so concurrent LOAD applies of different tasks never contend; only full-set operations (snapshot + * inclusion, stop) take the write lock so the task set is stable while every live task is visited. + */ +final class LoadTaskRegistry { + + private final ConcurrentHashMap tasks = new ConcurrentHashMap<>(); + private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock(); + + TsFileWriterManager getOrCreate(String uuid, LoadTaskFactory factory) throws IOException { + lock.readLock().lock(); + try { + final AtomicReference exception = new AtomicReference<>(); + final TsFileWriterManager writerManager = + tasks.computeIfAbsent( + uuid, + id -> { + try { + return factory.create(id); + } catch (Exception e) { + exception.set(e); + return null; + } + }); + if (exception.get() != null || writerManager == null) { + final String message = + String.format( + STORAGE_EXCEPTION_FAILED_TO_CREATE_TSFILEWRITERMANAGER_FOR_UUID_S_BECAUSE_A0D68950, + uuid); + throw new IOException(message, exception.get()); + } + return writerManager; + } finally { + lock.readLock().unlock(); + } + } + + Optional get(String uuid) { + return Optional.ofNullable(tasks.get(uuid)); + } + + boolean contains(String uuid) { + return tasks.containsKey(uuid); + } + + TsFileWriterManager remove(String uuid) { + lock.readLock().lock(); + try { + return tasks.remove(uuid); + } finally { + lock.readLock().unlock(); + } + } + + /** Visits every task under the write lock so creation/removal cannot happen concurrently. */ + void snapshot(LoadTaskVisitor visitor) throws IOException { + lock.writeLock().lock(); + try { + for (TsFileWriterManager writerManager : new ArrayList<>(tasks.values())) { + visitor.visit(writerManager); + } + } finally { + lock.writeLock().unlock(); + } + } + + void clear() { + lock.writeLock().lock(); + try { + tasks.clear(); + } finally { + lock.writeLock().unlock(); + } + } + + @FunctionalInterface + interface LoadTaskFactory { + TsFileWriterManager create(String uuid) throws Exception; + } + + @FunctionalInterface + interface LoadTaskVisitor { + void visit(TsFileWriterManager writerManager) throws IOException; + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileChecksumUtils.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileChecksumUtils.java new file mode 100644 index 0000000000000..caf054d487c02 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileChecksumUtils.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.storageengine.load; + +import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; + +import org.apache.tsfile.utils.ReadWriteIOUtils; + +import java.io.DataOutputStream; +import java.io.IOException; +import java.io.OutputStream; +import java.security.DigestOutputStream; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.List; + +/** Deterministic content checksum helpers for consensus-backed LOAD pieces. */ +public final class LoadTsFileChecksumUtils { + + private LoadTsFileChecksumUtils() {} + + /** + * Computes a stable checksum for a list of {@link TsFileData}. The value only depends on the + * serialized bytes, so any replica that applies the same consensus record observes the same + * result. + */ + public static long checksum(final List dataList) { + try { + final MessageDigest digest = MessageDigest.getInstance("SHA-256"); + final DataOutputStream stream = + new DataOutputStream(new DigestOutputStream(OutputStream.nullOutputStream(), digest)); + for (TsFileData data : dataList) { + ReadWriteIOUtils.write(data.getType().ordinal(), stream); + data.serialize(stream); + } + stream.flush(); + final byte[] bytes = digest.digest(); + long checksum = 0; + for (int i = 0; i < Long.BYTES; i++) { + checksum = (checksum << 8) | (bytes[i] & 0xFFL); + } + return checksum; + } catch (NoSuchAlgorithmException | IOException e) { + throw new IllegalStateException(e); + } + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java index 220be0cd0c691..f7fe7baa00de2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java @@ -19,86 +19,136 @@ package org.apache.iotdb.db.storageengine.load; +import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId; import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType; +import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation; +import org.apache.iotdb.common.rpc.thrift.TEndPoint; +import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet; +import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; -import org.apache.iotdb.commons.conf.IoTDBConstant; +import org.apache.iotdb.commons.client.ClientPoolFactory; +import org.apache.iotdb.commons.client.IClientManager; +import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient; import org.apache.iotdb.commons.consensus.ConsensusGroupId; +import org.apache.iotdb.commons.consensus.DataRegionId; import org.apache.iotdb.commons.consensus.index.ProgressIndex; import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex; import org.apache.iotdb.commons.disk.FolderManager; import org.apache.iotdb.commons.disk.strategy.DirectoryStrategyType; import org.apache.iotdb.commons.exception.DiskSpaceInsufficientException; -import org.apache.iotdb.commons.file.SystemFileFactory; -import org.apache.iotdb.commons.schema.table.TsFileTableSchemaUtil; -import org.apache.iotdb.commons.schema.table.TsTable; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; import org.apache.iotdb.commons.service.metric.MetricService; import org.apache.iotdb.commons.service.metric.enums.Metric; import org.apache.iotdb.commons.service.metric.enums.Tag; -import org.apache.iotdb.commons.utils.FileUtils; -import org.apache.iotdb.commons.utils.PathUtils; -import org.apache.iotdb.commons.utils.RetryUtils; +import org.apache.iotdb.commons.utils.StatusUtils; +import org.apache.iotdb.consensus.common.Peer; import org.apache.iotdb.db.conf.IoTDBConfig; import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.consensus.DataRegionConsensusImpl; import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.i18n.StorageEngineMessages; -import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; +import org.apache.iotdb.db.queryengine.plan.analyze.ClusterPartitionFetcher; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusOp; import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; -import org.apache.iotdb.db.queryengine.plan.scheduler.load.LoadTsFileScheduler.LoadCommand; -import org.apache.iotdb.db.schemaengine.table.DataNodeTableCache; import org.apache.iotdb.db.storageengine.dataregion.DataRegion; import org.apache.iotdb.db.storageengine.dataregion.flush.MemTableFlushTask; +import org.apache.iotdb.db.storageengine.dataregion.memtable.TsFileProcessor; import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; import org.apache.iotdb.db.storageengine.dataregion.modification.v1.ModificationFileV1; import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; -import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResourceStatus; +import org.apache.iotdb.db.storageengine.dataregion.wal.node.IWALNode; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.listener.AbstractResultListener; +import org.apache.iotdb.db.storageengine.dataregion.wal.utils.listener.WALFlushListener; import org.apache.iotdb.db.storageengine.load.active.ActiveLoadAgent; -import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; -import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; -import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; import org.apache.iotdb.metrics.utils.MetricLevel; +import org.apache.iotdb.mpp.rpc.thrift.TLoadResp; +import org.apache.iotdb.mpp.rpc.thrift.TTsFilePieceReq; +import org.apache.iotdb.rpc.TSStatusCode; import com.github.benmanes.caffeine.cache.Cache; import com.github.benmanes.caffeine.cache.Caffeine; -import org.apache.tsfile.common.constant.TsFileConstant; -import org.apache.tsfile.enums.TSDataType; import org.apache.tsfile.exception.write.PageException; -import org.apache.tsfile.file.metadata.ChunkGroupMetadata; -import org.apache.tsfile.file.metadata.ChunkMetadata; -import org.apache.tsfile.file.metadata.IDeviceID; -import org.apache.tsfile.read.TimeValuePair; -import org.apache.tsfile.utils.Pair; -import org.apache.tsfile.utils.RamUsageEstimator; -import org.apache.tsfile.utils.TsPrimitiveType; -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.nio.file.DirectoryNotEmptyException; +import java.nio.ByteBuffer; import java.nio.file.Files; -import java.nio.file.Path; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.Objects; import java.util.Optional; -import java.util.Set; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.PriorityBlockingQueue; -import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; -import java.util.stream.Collectors; import java.util.stream.Stream; /** - * {@link LoadTsFileManager} is used for dealing with {@link LoadTsFilePieceNode} and {@link - * LoadCommand}. This class turn the content of a piece of loading TsFile into a new TsFile. When - * DataNode finish transfer pieces, this class will flush all TsFile and load them into IoTDB, or - * delete all. + * {@link LoadTsFileManager} is the DataNode-side facade over the LOAD staging machinery. Every + * consensus command of a LOAD task (BEGIN / PIECE / PREPARE / COMMIT / ABORT) enters through {@link + * #applyConsensusRequest(DataRegion, LoadTsFileConsensusNode)}; the facade routes the work to the + * per-task {@link TsFileWriterManager} and keeps the rest of the lifecycle in dedicated components. + * The facade itself carries no parallel maps and no global write lock on the apply path. + * + *

All classes in this structure belong to the LOAD subsystem ({@code storageengine.load}): they + * manage the staged files, WAL/ref bookkeeping, snapshots and cleanup of LOAD TSFILE only. + * + * 

{@code
+ *         LOAD consensus commands (BEGIN / PIECE / PREPARE / COMMIT / ABORT)
+ *                                     |
+ *                                     v
+ *                   +-----------------+-----------------+
+ *                   |   LoadTsFileManager (facade)      |
+ *                   +-----------------+-----------------+
+ *                                     |
+ *        +----------------+-----------+-----------+----------------+
+ *        |                |                       |                |
+ *        v                v                       v                v
+ * LoadTaskRegistry  TsFileWriterManager     LoadSnapshotManager LoadCleanupScheduler
+ * (uuid -> manager) | per-task lock         (include LOAD       (abandoned-task
+ * lifecycle)        | applied/cached/        staging in         eviction,
+ *                    | retained pieces       snapshots)          delayed cleanup)
+ *                    v
+ *              PartitionContext (one per data partition)
+ *              [TsFileIOWriter + TsFileResource + mods]
+ *              writeChunk/writeDeletion, syncedOffset, finalized
+ * }
+ * + *

Command flow

+ * + * 
{@code
+ * BEGIN(loadId)             -> register the task in LoadTaskRegistry
+ * PIECE(idx, checksum)      -> apply chunk/deletion to the PartitionContext writers,
+ *                              retain the serialized bytes, write a marker-only WAL
+ *                              entry; followers apply the marker (via consensus log
+ *                              replication) and pull the retained bytes back on demand
+ * PREPARE(count, bytes, cs) -> finalizeAll(): endChunkGroup + endFile (footer) +
+ *                              captureRefs; write a WAL marker so followers seal their
+ *                              own staged files at the same logical point
+ * COMMIT                    -> write a WAL marker, load every staged file into the
+ *                              DataRegion via loadNewTsFile(progress indexes), clean up
+ * ABORT                     -> write a WAL marker, delete the staged files
+ * }
+ * + *

Startup recovery

+ * + * 
{@code
+ * recover(): scan every configured load directory
+ *   -> for each leftover task dir, rebuild a TsFileWriterManager (unsealed-file
+ *      recovery) and register it with LoadCleanupScheduler for later eviction
+ * }
+ * + *

Staged files live under the configured load directories and are resolved by {@link + * #findLoadTsFile(String)}. The WAL keeps marker-only entries (dozens of bytes per piece); the + * actual chunk bytes stay in the write node's retained-piece store and are pulled back by a + * follower through a DataNode-to-DataNode client ({@code SYNC_DATANODE_CLIENT_MANAGER}) when its + * marker arrives. Under Ratis the full command is replicated through the Ratis log, so every + * replica applies the chunk data directly. */ public class LoadTsFileManager { @@ -106,9 +156,8 @@ public class LoadTsFileManager { private static final IoTDBConfig CONFIG = IoTDBDescriptor.getInstance().getConfig(); - private static final String MESSAGE_WRITER_MANAGER_HAS_BEEN_CLOSED = - "%s TsFileWriterManager has been closed."; - private static final String MESSAGE_DELETE_FAIL = "failed to delete {}."; + /** Snapshot sub-directory that carries the in-progress LOAD staging files of a DataRegion. */ + public static final String LOAD_SNAPSHOT_DIR_NAME = "load"; private static final AtomicReference LOAD_BASE_DIRS = new AtomicReference<>(CONFIG.getLoadTsFileDirs()); @@ -120,99 +169,54 @@ public class LoadTsFileManager { .weigher((String k, String v) -> v.length()) .build(); - private final Map uuid2WriterManager = new ConcurrentHashMap<>(); + private final LoadTaskRegistry taskRegistry = new LoadTaskRegistry(); + private final LoadCleanupScheduler cleanupScheduler = + new LoadCleanupScheduler( + CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds(), this::forceCloseWriterManager); + private final LoadSnapshotManager snapshotManager = + new LoadSnapshotManager(taskRegistry, cleanupScheduler, this::allocateTaskDir); + private final ActiveLoadAgent activeLoadAgent = new ActiveLoadAgent(); - private final Map uuid2CleanupTask = new ConcurrentHashMap<>(); - private final PriorityBlockingQueue cleanupTaskQueue = new PriorityBlockingQueue<>(); + /** DataNode-to-DataNode client used for the LOAD piece pull-back. */ + private static final IClientManager + SYNC_DATANODE_CLIENT_MANAGER = + new IClientManager.Factory() + .createClientManager( + new ClientPoolFactory.SyncDataNodeInternalServiceClientPoolFactory()); - private final ActiveLoadAgent activeLoadAgent = new ActiveLoadAgent(); + private static final long PULL_WAIT_INTERVAL_MS = 100L; + private static final int PULL_WAIT_RETRIES = 50; public LoadTsFileManager() { - registerCleanupTaskExecutor(); + cleanupScheduler.start(); recover(); } - public void start() { - activeLoadAgent.start(); - } - - public void stop() { - activeLoadAgent.stop(); - synchronized (uuid2CleanupTask) { - uuid2CleanupTask.values().forEach(CleanupTask::cancel); - uuid2CleanupTask.clear(); - cleanupTaskQueue.clear(); + /** Resolve a staging file by its relative path under any configured load directory. */ + public static Optional findLoadTsFile(String relativePath) { + if (relativePath == null || relativePath.isEmpty()) { + return Optional.empty(); } - new HashSet<>(uuid2WriterManager.keySet()).forEach(this::forceCloseWriterManager); - } - - private long getCleanupTaskDelayInMs() { - return CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds() * 1000L; - } - - private void createCleanupTaskIfAbsent(final String uuid) { - synchronized (uuid2CleanupTask) { - if (uuid2CleanupTask.containsKey(uuid)) { - return; - } - - final CleanupTask cleanupTask = new CleanupTask(uuid, getCleanupTaskDelayInMs()); - uuid2CleanupTask.put(uuid, cleanupTask); - cleanupTaskQueue.add(cleanupTask); - } - } - - private void rescheduleCleanupTask(final CleanupTask cleanupTask) { - synchronized (uuid2CleanupTask) { - if (uuid2CleanupTask.get(cleanupTask.uuid) != cleanupTask) { - return; + for (String baseDir : LOAD_BASE_DIRS.get()) { + final File file = new File(baseDir, relativePath); + if (file.isFile()) { + return Optional.of(file); } - - cleanupTaskQueue.remove(cleanupTask); - cleanupTask.resetScheduledTime(); - cleanupTaskQueue.add(cleanupTask); } + return Optional.empty(); } - private void registerCleanupTaskExecutor() { - PipeDataNodeAgent.runtime() - .registerPeriodicalJob( - "LoadTsFileManager#cleanupTasks", - this::cleanupTasks, - CONFIG.getLoadCleanupTaskExecutionDelayTimeSeconds() >> 2); + public void start() { + activeLoadAgent.start(); } - private void cleanupTasks() { - while (!cleanupTaskQueue.isEmpty()) { - synchronized (uuid2CleanupTask) { - if (cleanupTaskQueue.isEmpty()) { - continue; - } - - final CleanupTask cleanupTask = cleanupTaskQueue.peek(); - if (cleanupTask.scheduledTime <= System.currentTimeMillis()) { - if (cleanupTask.isLoadTaskRunning) { - cleanupTaskQueue.poll(); - cleanupTask.resetScheduledTime(); - cleanupTaskQueue.add(cleanupTask); - continue; - } - - cleanupTask.run(); - - uuid2CleanupTask.remove(cleanupTask.uuid); - cleanupTaskQueue.poll(); - } else { - final long waitTimeInMs = cleanupTask.scheduledTime - System.currentTimeMillis(); - LOGGER.info( - StorageEngineMessages - .STORAGE_LOG_NEXT_LOAD_CLEANUP_TASK_IS_NOT_READY_TO_RUN_WAIT_FOR_AT_LEAST_CBE0023F, - cleanupTask.uuid, - waitTimeInMs, - waitTimeInMs / 1000.0); - return; - } - } + public void stop() { + activeLoadAgent.stop(); + cleanupScheduler.shutdown(); + try { + taskRegistry.snapshot(TsFileWriterManager::close); + } catch (IOException e) { + LOGGER.warn(StorageEngineMessages.LOAD_CLEANUP_TASK_ERROR, "all", e); } } @@ -240,66 +244,592 @@ private void recover() { .parallel() .forEach( taskDir -> { - final String uuid = taskDir.getName(); final TsFileWriterManager writerManager = new TsFileWriterManager(taskDir); - - uuid2WriterManager.put(uuid, writerManager); writerManager.close(); - - createCleanupTaskIfAbsent(uuid); + cleanupScheduler.registerOrRefresh(taskDir.getName()); }); } + private TsFileWriterManager createWriterManager(String uuid) throws Exception { + return getFolderManager() + .getNextWithRetry(folder -> new TsFileWriterManager(new File(folder, uuid))); + } + + private File allocateTaskDir(String uuid) throws Exception { + return getFolderManager().getNextWithRetry(folder -> new File(folder, uuid)); + } + public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNode, String uuid) - throws IOException, PageException { - createCleanupTaskIfAbsent(uuid); + throws IOException, PageException, LoadFileException { + cleanupScheduler.registerOrRefresh(uuid); + cleanupScheduler.markRunning(uuid); + try { + final TsFileWriterManager writerManager = getOrCreateWriterManager(uuid); + writerManager.writePieceNode(dataRegion, pieceNode); + } finally { + cleanupScheduler.markIdle(uuid); + } + } + + private TsFileWriterManager getOrCreateWriterManager(String uuid) throws IOException { + return taskRegistry.getOrCreate(uuid, this::createWriterManager); + } + + /** Whether pieces {@code 0..pieceIndex-1} were already applied contiguously on this node. */ + private boolean isContinuous(final String uuid, final long pieceIndex) { + return taskRegistry + .get(uuid) + .map(writerManager -> writerManager.hasAppliedAllUpTo(pieceIndex - 1)) + .orElse(pieceIndex == 0); + } + + /** Apply a consensus LOAD request on the DataRegion state machine path. */ + public TSStatus applyConsensusRequest(DataRegion dataRegion, LoadTsFileConsensusNode node) + throws IOException, PageException, LoadFileException { + switch (node.getOp()) { + case BEGIN: + return beginConsensus(node); + case PIECE: + return appendConsensusPiece(dataRegion, node); + case PREPARE: + return prepareConsensus(dataRegion, node); + case COMMIT: + return commitConsensus(dataRegion, node); + case ABORT: + return abortConsensus(dataRegion, node); + default: + return new TSStatus(TSStatusCode.ILLEGAL_PARAMETER.getStatusCode()) + .setMessage( + DataNodeQueryMessages.EXCEPTION_UNKNOWN_LOADTSFILECONSENSUSOP_ORDINAL_ARG_62848FC2 + + node.getOp()); + } + } + + private TSStatus beginConsensus(LoadTsFileConsensusNode node) { + cleanupScheduler.registerOrRefresh(node.getLoadId()); + return StatusUtils.OK; + } - final Optional cleanupTask = Optional.ofNullable(uuid2CleanupTask.get(uuid)); - cleanupTask.ifPresent(CleanupTask::markLoadTaskRunning); + private TSStatus appendConsensusPiece(DataRegion dataRegion, LoadTsFileConsensusNode node) + throws IOException, PageException, LoadFileException { + final String uuid = node.getLoadId(); + cleanupScheduler.registerOrRefresh(uuid); + cleanupScheduler.markRunning(uuid); try { - final AtomicReference exception = new AtomicReference<>(); - final TsFileWriterManager writerManager = - uuid2WriterManager.computeIfAbsent( - uuid, - o -> { - try { - return getFolderManager() - .getNextWithRetry(folder -> new TsFileWriterManager(new File(folder, uuid))); - } catch (DiskSpaceInsufficientException e) { - exception.set(e); - return null; - } - }); - - if (exception.get() != null || writerManager == null) { - throw new IOException( + if (!node.getPieceRefs().isEmpty()) { + // Legacy raw-ref PIECE (previous format): the refs are contiguous from offset 0, so a + // replica can rebuild the staged file from the WAL without a local writer. New entries no + // longer use this form, but entries logged by an older leader must stay applicable during a + // rolling upgrade. + final TsFileWriterManager writerManager = getOrCreateWriterManager(uuid); + writerManager.appendRawTsFilePieces(dataRegion, node.getPieceRefs()); + writerManager.applyDeletion(dataRegion, node.getTsFileDataList()); + return StatusUtils.OK; + } + + if (!node.hasChunkData()) { + // Marker-only PIECE replicated through the WAL. The marker is the ordering authority of the + // load: a follower applies the chunk data (pulled back from the write node, or retained + // locally) only when its marker arrives (and only after every previous marker was + // applied), so consensus order and local apply order can never diverge. + return applyPieceMarker(dataRegion, node); + } + + // Chunk-data PIECE submitted by the coordinator to the write node (or to a caught-up new + // leader after failover). Every node maintains its own applied-piece prefix, so the failover + // fence is continuity: pieceIndex is accepted only when 0..pieceIndex-1 were applied locally, + // which a follower-turned-leader satisfies automatically because it built its own writers + // while applying the markers. A node without the prefix must fail instead of silently + // rebuilding the file, which would fork the replicas. + if (!isContinuous(uuid, node.getPieceIndex())) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages + .MESSAGE_LOAD_CONSENSUS_PIECE_NOT_CONTINUOUS_AFTER_FAILOVER_D6FFAC6C, + node.getPieceIndex(), + uuid)); + } + + final TsFileWriterManager writerManager = getOrCreateWriterManager(uuid); + // Idempotent apply guard: a scheduler retry after a lost response may re-deliver the same + // piece. Without deduplication the chunk data would be appended twice and the staged file + // would diverge from the followers. + if (writerManager.isPieceAlreadyApplied(node.getPieceIndex(), node.getChecksum())) { + return StatusUtils.OK; + } + if (writerManager.isPieceConflicting(node.getPieceIndex(), node.getChecksum())) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PIECE_CHECKSUM_MISMATCH_CF261675, + uuid, + node.getPieceIndex())); + } + if (node.getChecksum() != LoadTsFileChecksumUtils.checksum(node.getTsFileDataList())) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PIECE_CHECKSUM_MISMATCH_CF261675, + uuid, + node.getPieceIndex())); + } + + writerManager.appendChunkPieceAndRecord( + dataRegion, node.getTsFileDataList(), node.getPieceIndex(), node.getChecksum()); + // Retain the serialized piece until COMMIT/ABORT as the backfill source for a follower that + // pulls it back on demand. + writerManager.retainPiece(node.getPieceIndex(), serializeNode(node)); + // Only the write node logs the marker: the marker-only WAL entry is what IoTConsensus + // replicates to the followers, which then pull the retained chunk bytes back. A follower + // applying the same command through consensus log replication skips the local WAL write, + // exactly like ordinary writes on a follower. + if (!node.isGeneratedByRemoteConsensusLeader()) { + logPieceMarkerToWal(dataRegion, node, writerManager); + } + return StatusUtils.OK; + } finally { + // An applied PIECE leaves the task idle again so an abandoned load (no COMMIT/ABORT ever + // arrives) is eventually reclaimed by the sweeper after the configured delay. + cleanupScheduler.markIdle(uuid); + } + } + + /** + * Applies a marker-only PIECE replicated through the WAL. The chunk data (pulled back from the + * write node, or retained locally when this node was the write node before a restart) is written + * into this node's own partition writers exactly like the write node does; a still-missing piece + * is pulled back from the current write node before failing. + */ + private TSStatus applyPieceMarker(DataRegion dataRegion, LoadTsFileConsensusNode node) + throws IOException, PageException, LoadFileException { + final String uuid = node.getLoadId(); + final TsFileWriterManager writerManager = getOrCreateWriterManager(uuid); + if (writerManager.isPieceAlreadyApplied(node.getPieceIndex(), node.getChecksum())) { + return StatusUtils.OK; + } + if (writerManager.isPieceConflicting(node.getPieceIndex(), node.getChecksum())) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PIECE_CHECKSUM_MISMATCH_CF261675, + uuid, + node.getPieceIndex())); + } + if (!writerManager.hasAppliedAllUpTo(node.getPieceIndex() - 1)) { + // The WAL delivers markers in order, so a hole here means the marker log is corrupt or the + // task state was reset (e.g. restore without the applied-piece prefix). Failing loudly beats + // skipping a piece and forking the staged file. + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages + .MESSAGE_LOAD_CONSENSUS_PIECE_NOT_CONTINUOUS_AFTER_FAILOVER_D6FFAC6C, + node.getPieceIndex(), + uuid)); + } + if (writerManager.hasCachedPiece(node.getPieceIndex(), node.getChecksum())) { + writerManager.applyCachedPiece(dataRegion, node.getPieceIndex(), node.getChecksum()); + return StatusUtils.OK; + } + // The write node's own retained store may still hold the serialized piece (durable on disk), + // e.g. this node was the write node before a restart and is now replaying its own markers. + // Backfilling locally avoids a self-RPC round trip and works even when no peer is reachable. + final Optional retained = writerManager.getRetainedPiece(node.getPieceIndex()); + if (retained.isPresent() && cacheLocalRetainedPiece(writerManager, node, retained.get())) { + writerManager.applyCachedPiece(dataRegion, node.getPieceIndex(), node.getChecksum()); + return StatusUtils.OK; + } + // The piece was not delivered yet (there is no out-of-band push; the marker itself carries no + // chunk bytes). Pull the piece back from the current write node, which retains the serialized + // bytes until COMMIT/ABORT. + if (pullPieceFromLeader(dataRegion, node, writerManager)) { + writerManager.applyCachedPiece(dataRegion, node.getPieceIndex(), node.getChecksum()); + return StatusUtils.OK; + } + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( String.format( - StorageEngineMessages - .STORAGE_EXCEPTION_FAILED_TO_CREATE_TSFILEWRITERMANAGER_FOR_UUID_S_BECAUSE_A0D68950, - uuid), - exception.get()); + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PIECE_DATA_MISSING_AFTER_PULL_8269CB0B, + node.getPieceIndex(), + uuid)); + } + + /** Caches a locally retained serialized PIECE if it matches the marker being applied. */ + private boolean cacheLocalRetainedPiece( + TsFileWriterManager writerManager, LoadTsFileConsensusNode marker, byte[] serializedPiece) { + final PlanNode planNode = PlanNodeType.deserialize(ByteBuffer.wrap(serializedPiece)); + if (!(planNode instanceof LoadTsFileConsensusNode)) { + return false; + } + final LoadTsFileConsensusNode chunkPiece = (LoadTsFileConsensusNode) planNode; + if (chunkPiece.getOp() != LoadTsFileConsensusOp.PIECE + || !chunkPiece.hasChunkData() + || !chunkPiece.getLoadId().equals(marker.getLoadId()) + || chunkPiece.getPieceIndex() != marker.getPieceIndex() + || chunkPiece.getChecksum() != marker.getChecksum()) { + return false; + } + try { + return writerManager.cachePiece( + chunkPiece.getPieceIndex(), chunkPiece.getChecksum(), chunkPiece.getTsFileDataList()); + } catch (IOException e) { + return false; + } + } + + /** + * Logs a marker-only PIECE (metadata without chunk bytes) to the WAL. The WAL is the ordering + * authority replicated to the followers; the actual chunk data is retained by the write node and + * pulled back by a follower when the marker arrives, so the WAL stays at dozens of bytes per + * piece instead of the full LOAD bytes. + */ + private void logPieceMarkerToWal( + DataRegion dataRegion, LoadTsFileConsensusNode node, TsFileWriterManager writerManager) + throws IOException { + final Optional walNodeOptional = dataRegion.getWALNode(); + if (!walNodeOptional.isPresent()) { + return; + } + // The refs are only local bookkeeping (they advance the synced cursor used by snapshots); they + // are no longer replicated, so drain and discard them. + writerManager.drainPendingPieceRefs(); + final LoadTsFileConsensusNode marker = + LoadTsFileConsensusNode.pieceMarker( + new PlanNodeId("load-wal-marker-" + node.getLoadId() + "-" + node.getPieceIndex()), + node.getLoadId(), + node.getTsFileId(), + node.getPieceIndex(), + node.getChecksum(), + node.getDataSize()); + final WALFlushListener listener = + walNodeOptional.get().log(TsFileProcessor.MEMTABLE_NOT_EXIST, marker); + if (listener.waitForResult() == AbstractResultListener.Status.FAILURE) { + throw new IOException( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_WAL_FLUSH_FAILED_8BE1375A, + listener.getCause()); + } + } + + /** Logs a PREPARE/COMMIT/ABORT op itself as the WAL marker so followers apply the same phase. */ + private void logOpMarkerToWal(DataRegion dataRegion, LoadTsFileConsensusNode node) + throws IOException { + final Optional walNodeOptional = dataRegion.getWALNode(); + if (!walNodeOptional.isPresent()) { + return; + } + final WALFlushListener listener = + walNodeOptional.get().log(TsFileProcessor.MEMTABLE_NOT_EXIST, node); + if (listener.waitForResult() == AbstractResultListener.Status.FAILURE) { + throw new IOException( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_WAL_FLUSH_FAILED_8BE1375A, + listener.getCause()); + } + } + + /** + * Pulls one missing piece from the current write node. The write node receives a PULL request, + * reads the retained serialized piece and pushes it back to {@code pullSourceEndPoint}; this side + * waits (bounded) for the pushed piece to land in the cache. + */ + private boolean pullPieceFromLeader( + DataRegion dataRegion, LoadTsFileConsensusNode marker, TsFileWriterManager writerManager) { + final ConsensusGroupId groupId = + new DataRegionId(Integer.parseInt(dataRegion.getDataRegionIdString())); + final TEndPoint leaderEndPoint = resolveWriteNodeEndPoint(dataRegion, groupId, marker); + if (leaderEndPoint == null) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_PULL_PIECE_FAILED_AFB003D5, + marker.getPieceIndex(), + marker.getLoadId(), + "unknown", + "cannot resolve the current write node from the partition table"); + return false; + } + final String localEndPoint = CONFIG.getInternalAddress() + ":" + CONFIG.getInternalPort(); + final LoadTsFileConsensusNode pull = + LoadTsFileConsensusNode.pull( + new PlanNodeId("load-pull-" + marker.getLoadId() + "-" + marker.getPieceIndex()), + marker.getLoadId(), + marker.getTsFileId(), + marker.getPieceIndex(), + marker.getChecksum(), + localEndPoint); + try (final SyncDataNodeInternalServiceClient client = + SYNC_DATANODE_CLIENT_MANAGER.borrowClient(leaderEndPoint)) { + final TLoadResp resp = + client.sendTsFilePieceNode( + new TTsFilePieceReq( + pull.serializeToByteBuffer(), + marker.getLoadId(), + groupId.convertToTConsensusGroupId())); + if (!resp.isAccepted()) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_PULL_PIECE_FAILED_AFB003D5, + marker.getPieceIndex(), + marker.getLoadId(), + leaderEndPoint, + resp.getMessage()); + return false; + } + } catch (Exception e) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_PULL_PIECE_FAILED_AFB003D5, + marker.getPieceIndex(), + marker.getLoadId(), + leaderEndPoint, + e.getMessage()); + return false; + } + for (int i = 0; i < PULL_WAIT_RETRIES; i++) { + if (writerManager.hasCachedPiece(marker.getPieceIndex(), marker.getChecksum())) { + return true; + } + try { + Thread.sleep(PULL_WAIT_INTERVAL_MS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return false; } + } + return false; + } - for (TsFileData tsFileData : pieceNode.getAllTsFileData()) { - switch (tsFileData.getType()) { - case CHUNK: - ChunkData chunkData = (ChunkData) tsFileData; - writerManager.write( - new DataPartitionInfo(dataRegion, chunkData.getTimePartitionSlot()), chunkData); - break; - case DELETION: - writerManager.writeDeletion(dataRegion, (DeletionData) tsFileData); - break; - default: - throw new IOException( - StorageEngineMessages.UNSUPPORTED_TSFILE_DATA_TYPE + tsFileData.getType()); + /** + * Resolves the internal endpoint of the partition's current write node from the local partition + * table (the same replica-set routing the normal write path uses). This is the node that retains + * the applied piece bytes, so a follower missing a delivery pulls them back from here. + */ + private TEndPoint resolveWriteNodeEndPoint( + DataRegion dataRegion, ConsensusGroupId groupId, LoadTsFileConsensusNode marker) { + try { + final List replicaSets = + ClusterPartitionFetcher.getInstance() + .getRegionReplicaSet(Collections.singletonList(groupId.convertToTConsensusGroupId())); + if (!replicaSets.isEmpty()) { + final List locations = replicaSets.get(0).getDataNodeLocations(); + if (locations != null && !locations.isEmpty()) { + final TEndPoint writeNodeEndPoint = locations.get(0).getInternalEndPoint(); + if (writeNodeEndPoint != null) { + return writeNodeEndPoint; + } } } + } catch (Exception e) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_PULL_PIECE_FAILED_AFB003D5, + marker.getPieceIndex(), + marker.getLoadId(), + "unknown", + "failed to resolve the current write node from the partition table: " + e.getMessage()); + } + // Fall back to the consensus leader lookup (meaningful for Ratis; IoTConsensus reports the + // local node, which is still correct for the local retained-piece case). + final Peer leader = DataRegionConsensusImpl.getInstance().getLeader(groupId); + return leader == null + ? null + : new TEndPoint(leader.getEndpoint().getIp(), leader.getEndpoint().getPort()); + } + + /** The write node's side of a PULL: push the retained piece bytes back to the requester. */ + public TSStatus handlePullPiece(DataRegion dataRegion, LoadTsFileConsensusNode pullNode) { + final String uuid = pullNode.getLoadId(); + final long pieceIndex = pullNode.getPieceIndex(); + final TEndPoint target = parseEndPoint(pullNode.getPullSourceEndPoint()); + if (target == null) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PULL_WITHOUT_SOURCE_ENDPOINT_3B20D9E9); + } + final Optional retained = + taskRegistry.get(uuid).flatMap(writerManager -> writerManager.getRetainedPiece(pieceIndex)); + if (!retained.isPresent()) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PULL_WITHOUT_RETAINED_PIECE_AD3C9D4F, + pieceIndex, + uuid)); + } + final byte[] serialized = retained.get(); + final PlanNode planNode = PlanNodeType.deserialize(ByteBuffer.wrap(serialized)); + if (!(planNode instanceof LoadTsFileConsensusNode) + || ((LoadTsFileConsensusNode) planNode).getOp() != LoadTsFileConsensusOp.PIECE) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PULL_WITHOUT_RETAINED_PIECE_AD3C9D4F, + pieceIndex, + uuid)); + } + final TConsensusGroupId groupId = + new DataRegionId(Integer.parseInt(dataRegion.getDataRegionIdString())) + .convertToTConsensusGroupId(); + try (final SyncDataNodeInternalServiceClient client = + SYNC_DATANODE_CLIENT_MANAGER.borrowClient(target)) { + final TLoadResp resp = + client.sendTsFilePieceNode( + new TTsFilePieceReq(ByteBuffer.wrap(serialized), uuid, groupId)); + return resp.isAccepted() + ? StatusUtils.OK + : new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PULL_PUSH_BACK_FAILED_1A90C2B9, + pieceIndex, + uuid, + target, + resp.getMessage())); + } catch (Exception e) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PULL_PUSH_BACK_FAILED_1A90C2B9, + pieceIndex, + uuid, + target, + e.getMessage())); + } + } + + /** + * Caches a chunk-data PIECE pushed back by the write node in response to a PULL (or delivered by + * a legacy out-of-band push). The data is not applied until the corresponding WAL marker arrives, + * because the marker (not the delivery order) decides the apply order on this node. + */ + public TSStatus cacheConsensusPiece(DataRegion dataRegion, LoadTsFileConsensusNode node) + throws IOException { + final String uuid = node.getLoadId(); + cleanupScheduler.registerOrRefresh(uuid); + cleanupScheduler.markRunning(uuid); + try { + if (!node.hasChunkData()) { + return StatusUtils.OK; + } + final TsFileWriterManager writerManager = getOrCreateWriterManager(uuid); + if (writerManager.isPieceAlreadyApplied(node.getPieceIndex(), node.getChecksum())) { + // The marker already applied this piece; the redundant delivery is dropped. + return StatusUtils.OK; + } + if (!writerManager.cachePiece( + node.getPieceIndex(), node.getChecksum(), node.getTsFileDataList())) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PIECE_CHECKSUM_MISMATCH_CF261675, + uuid, + node.getPieceIndex())); + } + return StatusUtils.OK; + } finally { + cleanupScheduler.markIdle(uuid); + } + } + + private TEndPoint parseEndPoint(final String endPointString) { + if (endPointString == null || endPointString.isEmpty()) { + return null; + } + final int separatorIndex = endPointString.lastIndexOf(':'); + if (separatorIndex <= 0 || separatorIndex == endPointString.length() - 1) { + return null; + } + try { + return new TEndPoint( + endPointString.substring(0, separatorIndex), + Integer.parseInt(endPointString.substring(separatorIndex + 1))); + } catch (NumberFormatException e) { + return null; + } + } + + private static byte[] serializeNode(LoadTsFileConsensusNode node) { + final ByteBuffer buffer = node.serializeToByteBuffer(); + final byte[] bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + return bytes; + } + + private TSStatus prepareConsensus(DataRegion dataRegion, LoadTsFileConsensusNode node) + throws IOException { + if (!taskRegistry.contains(node.getLoadId())) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages.MESSAGE_LOAD_CONSENSUS_PREPARE_WITHOUT_STAGED_DATA_FE8ADC37, + node.getLoadId())); + } + final String uuid = node.getLoadId(); + cleanupScheduler.markRunning(uuid); + try { + final TsFileWriterManager writerManager = getOrCreateWriterManager(uuid); + // Reconcile before sealing: the staged file must contain exactly the pieces the coordinator + // sent. A write-node switch mid-load can leave this node with a hole in its applied prefix + // that the per-piece continuity fence cannot detect (no further PIECE arrives); sealing and + // loading such a file would silently fork the replicas, so fail loudly instead. + if (!writerManager.isLegacyRawRefTask() + && !writerManager.verifyAppliedPieces(node.getPieceCount(), node.getChecksum())) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + String.format( + StorageEngineMessages + .MESSAGE_LOAD_CONSENSUS_PREPARE_VERIFICATION_FAILED_B3865A82, + uuid, + node.getPieceCount(), + node.getChecksum(), + writerManager.getAppliedPieceCount(), + writerManager.getAppliedPiecesChecksum())); + } + writerManager.finalizeAll(); + if (!node.isGeneratedByRemoteConsensusLeader()) { + // Replicate the PREPARE marker so every follower seals its own staged files at the same + // logical point before COMMIT. + logOpMarkerToWal(dataRegion, node); + } + return StatusUtils.OK; } finally { - cleanupTask.ifPresent(CleanupTask::markLoadTaskNotRunning); + cleanupScheduler.markIdle(uuid); } } + private TSStatus commitConsensus(DataRegion dataRegion, LoadTsFileConsensusNode node) + throws IOException, LoadFileException { + final Map progressIndexes = new HashMap<>(); + for (Map.Entry entry : + node.getTimePartition2ProgressIndex().entrySet()) { + final ProgressIndex progressIndex = MinimumProgressIndex.INSTANCE; + progressIndexes.put(entry.getKey(), progressIndex); + } + if (!node.isGeneratedByRemoteConsensusLeader()) { + // Replicate the COMMIT marker before loading so followers import their own staged files too. + logOpMarkerToWal(dataRegion, node); + } + if (!loadAll(node.getLoadId(), dataRegion, node.isGeneratedByPipe(), progressIndexes)) { + return new TSStatus(TSStatusCode.LOAD_FILE_ERROR.getStatusCode()) + .setMessage( + StorageEngineMessages + .MESSAGE_NO_LOAD_TSFILE_UUID_ARG_RECORDED_EXECUTE_LOAD_COMMAND_ARG_66722D80 + + node.getLoadId()); + } + return StatusUtils.OK; + } + + private TSStatus abortConsensus(DataRegion dataRegion, LoadTsFileConsensusNode node) { + if (!node.isGeneratedByRemoteConsensusLeader()) { + try { + // Replicate the ABORT marker so followers discard their staged files as well. + logOpMarkerToWal(dataRegion, node); + } catch (IOException e) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_ABORT_MARKER_FAILED_6A218023, + node.getLoadId(), + e.getMessage()); + } + } + deleteAll(node.getLoadId()); + return StatusUtils.OK; + } + private FolderManager getFolderManager() throws DiskSpaceInsufficientException { if (CONFIG.getLoadTsFileDirs() != LOAD_BASE_DIRS.get()) { synchronized (FOLDER_MANAGER) { @@ -332,18 +862,34 @@ public boolean loadAll( boolean isGeneratedByPipe, Map timePartitionProgressIndexMap) throws IOException, LoadFileException { - if (!uuid2WriterManager.containsKey(uuid)) { + return loadAll(uuid, null, isGeneratedByPipe, timePartitionProgressIndexMap); + } + + /** + * Loads the staged data of the given load into the DataRegion. The consensus COMMIT path passes + * the current {@code dataRegion} so that staged files restored from a snapshot (which have no + * in-memory writer) can be bound and loaded; the legacy direct-load path passes {@code null} and + * only ever contains writer-managed files. + */ + public boolean loadAll( + String uuid, + DataRegion dataRegion, + boolean isGeneratedByPipe, + Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + final Optional writerManagerOptional = taskRegistry.get(uuid); + if (!writerManagerOptional.isPresent()) { return false; } - createCleanupTaskIfAbsent(uuid); - - final Optional cleanupTask = Optional.ofNullable(uuid2CleanupTask.get(uuid)); - cleanupTask.ifPresent(CleanupTask::markLoadTaskRunning); + cleanupScheduler.registerOrRefresh(uuid); + cleanupScheduler.markRunning(uuid); try { - uuid2WriterManager.get(uuid).loadAll(isGeneratedByPipe, timePartitionProgressIndexMap); + writerManagerOptional + .get() + .loadAll(dataRegion, isGeneratedByPipe, timePartitionProgressIndexMap); } finally { - cleanupTask.ifPresent(CleanupTask::markLoadTaskNotRunning); + cleanupScheduler.markIdle(uuid); } clean(uuid); @@ -351,7 +897,7 @@ public boolean loadAll( } public boolean deleteAll(String uuid) { - if (!uuid2WriterManager.containsKey(uuid)) { + if (!taskRegistry.contains(uuid)) { return false; } clean(uuid); @@ -359,24 +905,26 @@ public boolean deleteAll(String uuid) { } private void clean(String uuid) { - synchronized (uuid2CleanupTask) { - final CleanupTask cleanupTask = uuid2CleanupTask.remove(uuid); - if (cleanupTask != null) { - cleanupTask.cancel(); - cleanupTaskQueue.remove(cleanupTask); - } - } - + cleanupScheduler.remove(uuid); forceCloseWriterManager(uuid); } private void forceCloseWriterManager(String uuid) { - final TsFileWriterManager writerManager = uuid2WriterManager.remove(uuid); - if (Objects.nonNull(writerManager)) { + final TsFileWriterManager writerManager = taskRegistry.remove(uuid); + if (writerManager != null) { writerManager.close(); } } + public void snapshotLoadTasksForRegion(DataRegion dataRegion, File snapshotDir) + throws IOException { + snapshotManager.snapshotLoadTasksForRegion(dataRegion, snapshotDir); + } + + public void restoreLoadTasksFromSnapshot(File loadSnapshotDir) throws IOException { + snapshotManager.restoreLoadTasksFromSnapshot(loadSnapshotDir); + } + public static void updateWritePointCountMetrics( final DataRegion dataRegion, final String databaseName, @@ -437,466 +985,4 @@ public static void cleanTsFile(final File tsFile) { LOGGER.warn(StorageEngineMessages.DELETE_AFTER_LOADING_ERROR, tsFile, e); } } - - private static class TsFileWriterManager { - - private final File taskDir; - private Map dataPartition2Writer; - private Map dataPartition2Resource; - private Map dataPartition2LastDevice; - private Map dataPartition2ModificationFile; - private Map> device2Partition; - private boolean isClosed; - - private TsFileWriterManager(File taskDir) { - this.taskDir = taskDir; - this.dataPartition2Writer = new HashMap<>(); - this.dataPartition2Resource = new HashMap<>(); - this.dataPartition2LastDevice = new HashMap<>(); - this.dataPartition2ModificationFile = new HashMap<>(); - device2Partition = new HashMap<>(); - this.isClosed = false; - - clearDir(taskDir); - } - - private void clearDir(File dir) { - if (dir.exists()) { - FileUtils.deleteFileOrDirectoryWithRetry(dir); - } - if (dir.mkdirs()) { - LOGGER.info(StorageEngineMessages.LOAD_TSFILE_DIR_CREATED, dir.getPath()); - } - } - - /** - * It should be noted that all AlignedChunkData of the same partition split from a source file - * should be guaranteed to be written to the same new file. Otherwise, for detached - * BatchedAlignedChunkData, it may result in no data for the time column in the new file. - */ - @SuppressWarnings("squid:S3824") - private void write(DataPartitionInfo partitionInfo, ChunkData chunkData) - throws IOException, PageException { - if (isClosed) { - throw new IOException(String.format(MESSAGE_WRITER_MANAGER_HAS_BEEN_CLOSED, taskDir)); - } - if (!dataPartition2Writer.containsKey(partitionInfo)) { - File newTsFile = - SystemFileFactory.INSTANCE.getFile( - taskDir, partitionInfo.toString() + TsFileConstant.TSFILE_SUFFIX); - if (!newTsFile.createNewFile()) { - LOGGER.error(StorageEngineMessages.CANNOT_CREATE_TSFILE_FOR_WRITING, newTsFile.getPath()); - return; - } - - final long chunkMetadataMaxSizeForEachWriter = - CONFIG.getLoadChunkMetadataMemorySizeInBytes() / (dataPartition2Writer.size() + 1); - final TsFileIOWriter writer = - new TsFileIOWriter(newTsFile, chunkMetadataMaxSizeForEachWriter); - final TsFileResource resource = new TsFileResource(writer.getFile()); - writer.addFlushListener( - // Update time index by chunk groups going to be flushed to temp file - sortedChunkMetadataList -> - sortedChunkMetadataList.forEach( - pair -> { - final IDeviceID deviceId = pair.left.left; - pair.getRight() - .forEach( - chunkMetadata -> { - resource.updateStartTime(deviceId, chunkMetadata.getStartTime()); - resource.updateEndTime(deviceId, chunkMetadata.getEndTime()); - }); - })); - - // When a new writer is added, we need to reduce the metadata size limit of all existing - // writers for memory control - for (final TsFileIOWriter existingWriter : dataPartition2Writer.values()) { - existingWriter.setMaxMetadataSize(chunkMetadataMaxSizeForEachWriter); - } - dataPartition2Writer.put(partitionInfo, writer); - dataPartition2Resource.put(partitionInfo, resource); - } - TsFileIOWriter writer = dataPartition2Writer.get(partitionInfo); - - // Table model needs to register TableSchema - final String tableName = - chunkData.getDevice() != null ? chunkData.getDevice().getTableName() : null; - if (tableName != null - && PathUtils.isTableModelDatabase(partitionInfo.getDataRegion().getDatabaseName())) { - // If the table does not exist, it means that the table is all deleted by mods - final TsTable table = - DataNodeTableCache.getInstance() - .getTable(partitionInfo.getDataRegion().getDatabaseName(), tableName, false); - if (Objects.nonNull(table)) { - writer - .getSchema() - .getTableSchemaMap() - .computeIfAbsent( - tableName, t -> TsFileTableSchemaUtil.toTsFileTableSchemaNoAttribute(table)); - } - } - - IDeviceID device = chunkData.getDevice(); - IDeviceID lastDevice = dataPartition2LastDevice.get(partitionInfo); - - if (!Objects.equals(device, lastDevice)) { - if (lastDevice != null && device2Partition.containsKey(lastDevice)) { - Set partitions = device2Partition.get(lastDevice); - for (DataPartitionInfo partition : partitions) { - TsFileIOWriter w = dataPartition2Writer.get(partition); - if (dataPartition2LastDevice.containsKey(partition) && w != null) { - w.endChunkGroup(); - w.checkMetadataSizeAndMayFlush(); - } - } - device2Partition.remove(lastDevice); - } - if (writer.isWritingChunkGroup()) { - LOGGER.warn( - StorageEngineMessages - .STORAGE_LOG_WRITER_FOR_PARTITION_IS_ALREADY_WRITING_CHUNK_GROUP_FOR_903B1D66, - writer.getFile().getAbsolutePath(), - partitionInfo, - device, - lastDevice); - } - writer.startChunkGroup(device); - dataPartition2LastDevice.put(partitionInfo, device); - device2Partition.computeIfAbsent(device, k -> new HashSet<>()).add(partitionInfo); - } - - chunkData.writeToFileWriter(writer); - } - - private void writeDeletion(DataRegion dataRegion, DeletionData deletionData) - throws IOException { - if (isClosed) { - throw new IOException(String.format(MESSAGE_WRITER_MANAGER_HAS_BEEN_CLOSED, taskDir)); - } - for (Map.Entry entry : dataPartition2Writer.entrySet()) { - final DataPartitionInfo partitionInfo = entry.getKey(); - if (partitionInfo.getDataRegion().equals(dataRegion)) { - final TsFileIOWriter writer = entry.getValue(); - if (!dataPartition2ModificationFile.containsKey(partitionInfo)) { - File newModificationFile = ModificationFile.getExclusiveMods(writer.getFile()); - if (!newModificationFile.createNewFile()) { - LOGGER.error( - StorageEngineMessages - .STORAGE_LOG_CAN_NOT_CREATE_MODIFICATIONFILE_FOR_WRITING_17D14C11, - newModificationFile.getPath()); - return; - } - - dataPartition2ModificationFile.put( - partitionInfo, new ModificationFile(newModificationFile, false)); - } - ModificationFile modificationFile = dataPartition2ModificationFile.get(partitionInfo); - writer.flush(); - deletionData.writeToModificationFile(modificationFile); - } - } - } - - private void loadAll( - boolean isGeneratedByPipe, - Map timePartitionProgressIndexMap) - throws IOException, LoadFileException { - if (isClosed) { - throw new IOException(String.format(MESSAGE_WRITER_MANAGER_HAS_BEEN_CLOSED, taskDir)); - } - for (final Map.Entry entry : - dataPartition2ModificationFile.entrySet()) { - entry.getValue().close(); - } - for (final Map.Entry entry : - dataPartition2Writer.entrySet()) { - final TsFileIOWriter writer = entry.getValue(); - if (writer.isWritingChunkGroup()) { - writer.endChunkGroup(); - } - writer.endFile(); - - final DataRegion dataRegion = entry.getKey().getDataRegion(); - final TsFileResource tsFileResource = dataPartition2Resource.get(entry.getKey()); - tsFileResource.setGeneratedByPipe(isGeneratedByPipe); - endTsFileResource( - writer, - tsFileResource, - timePartitionProgressIndexMap.getOrDefault( - entry.getKey().getTimePartitionSlot(), MinimumProgressIndex.INSTANCE)); - dataRegion.loadNewTsFile( - tsFileResource, - true, - isGeneratedByPipe, - false, - Optional.ofNullable(writer.getTableSizeMap())); - - // Metrics - dataRegion - .getNonSystemDatabaseName() - .ifPresent( - databaseName -> - updateWritePointCountMetrics( - dataRegion, databaseName, getTsFileWritePointCount(writer), false)); - } - } - - private void endTsFileResource( - TsFileIOWriter writer, TsFileResource tsFileResource, ProgressIndex progressIndex) - throws IOException { - // Update time index by chunk groups still in memory - Map> deviceLastValues = null; - if (IoTDBDescriptor.getInstance().getConfig().isCacheLastValuesForLoad()) { - deviceLastValues = new HashMap<>(); - } - AtomicLong lastValuesMemCost = new AtomicLong(0); - - for (final ChunkGroupMetadata chunkGroupMetadata : writer.getChunkGroupMetadataList()) { - final IDeviceID device = chunkGroupMetadata.getDevice(); - for (final ChunkMetadata chunkMetadata : chunkGroupMetadata.getChunkMetadataList()) { - tsFileResource.updateStartTime(device, chunkMetadata.getStartTime()); - tsFileResource.updateEndTime(device, chunkMetadata.getEndTime()); - if (deviceLastValues != null) { - Map deviceMap = - deviceLastValues.computeIfAbsent( - device, - d -> { - Map map = new HashMap<>(); - lastValuesMemCost.addAndGet(RamUsageEstimator.shallowSizeOf(map)); - lastValuesMemCost.addAndGet(device.ramBytesUsed()); - return map; - }); - int prevSize = deviceMap.size(); - deviceMap.compute( - chunkMetadata.getMeasurementUid(), - (m, oldPair) -> { - if (oldPair != null && oldPair.getTimestamp() > chunkMetadata.getEndTime()) { - return oldPair; - } - TsPrimitiveType lastValue = - chunkMetadata.getStatistics() != null - && chunkMetadata.getDataType() != TSDataType.BLOB - ? TsPrimitiveType.getByType( - chunkMetadata.getDataType() == TSDataType.VECTOR - ? TSDataType.INT64 - : chunkMetadata.getDataType(), - chunkMetadata.getStatistics().getLastValue()) - : null; - TimeValuePair timeValuePair = - lastValue != null - ? new TimeValuePair(chunkMetadata.getEndTime(), lastValue) - : null; - if (oldPair != null) { - lastValuesMemCost.addAndGet(-oldPair.getSize()); - } - if (timeValuePair != null) { - lastValuesMemCost.addAndGet(timeValuePair.getSize()); - } - return timeValuePair; - }); - int afterSize = deviceMap.size(); - lastValuesMemCost.addAndGet( - (afterSize - prevSize) * RamUsageEstimator.HASHTABLE_RAM_BYTES_PER_ENTRY); - if (lastValuesMemCost.get() - > IoTDBDescriptor.getInstance() - .getConfig() - .getCacheLastValuesMemoryBudgetInByte()) { - deviceLastValues = null; - } - } - } - } - if (deviceLastValues != null) { - Map>> finalDeviceLastValues; - finalDeviceLastValues = new HashMap<>(deviceLastValues.size()); - for (final Map.Entry> entry : - deviceLastValues.entrySet()) { - final IDeviceID device = entry.getKey(); - Map lastValues = entry.getValue(); - List> pairList = - lastValues.entrySet().stream() - .map(e -> new Pair<>(e.getKey(), e.getValue())) - .collect(Collectors.toList()); - finalDeviceLastValues.put(device, pairList); - } - tsFileResource.setLastValues(finalDeviceLastValues); - } - tsFileResource.setStatus(TsFileResourceStatus.NORMAL); - tsFileResource.setProgressIndex(progressIndex); - tsFileResource.serialize(); - } - - private long getTsFileWritePointCount(TsFileIOWriter writer) { - return writer.getChunkGroupMetadataList().stream() - .flatMap(chunkGroupMetadata -> chunkGroupMetadata.getChunkMetadataList().stream()) - .mapToLong(chunkMetadata -> chunkMetadata.getStatistics().getCount()) - .sum(); - } - - private void close() { - if (isClosed) { - return; - } - if (dataPartition2Writer != null) { - for (Map.Entry entry : dataPartition2Writer.entrySet()) { - try { - final TsFileIOWriter writer = entry.getValue(); - if (writer.canWrite()) { - writer.close(); - } - final Path writerPath = writer.getFile().toPath(); - if (Files.exists(writerPath)) { - RetryUtils.retryOnException( - () -> { - Files.delete(writerPath); - return null; - }); - } - } catch (IOException e) { - LOGGER.warn( - StorageEngineMessages.CLOSE_TSFILE_IO_WRITER_ERROR, - entry.getValue().getFile().getPath(), - e); - } - } - } - if (dataPartition2ModificationFile != null) { - for (Map.Entry entry : - dataPartition2ModificationFile.entrySet()) { - try { - final ModificationFile modificationFile = entry.getValue(); - modificationFile.close(); - final Path modificationFilePath = modificationFile.getFile().toPath(); - if (Files.exists(modificationFilePath)) { - RetryUtils.retryOnException( - () -> { - Files.delete(modificationFilePath); - return null; - }); - } - } catch (IOException e) { - LOGGER.warn( - StorageEngineMessages.CLOSE_MODIFICATION_FILE_ERROR, entry.getValue().getFile(), e); - } - } - } - try { - RetryUtils.retryOnException( - () -> { - Files.delete(taskDir.toPath()); - return null; - }); - } catch (DirectoryNotEmptyException e) { - LOGGER.info(StorageEngineMessages.TASK_DIR_NOT_EMPTY_SKIP_DELETE, taskDir.getPath()); - } catch (IOException e) { - LOGGER.warn(MESSAGE_DELETE_FAIL, taskDir.getPath(), e); - } - dataPartition2Writer = null; - dataPartition2Resource = null; - dataPartition2LastDevice = null; - dataPartition2ModificationFile = null; - device2Partition = null; - isClosed = true; - } - } - - private class CleanupTask implements Runnable, Comparable { - - private final String uuid; - - private final long delayInMs; - private long scheduledTime; - - private volatile boolean isLoadTaskRunning = false; - private volatile boolean isCanceled = false; - - private CleanupTask(String uuid, long delayInMs) { - this.uuid = uuid; - this.delayInMs = delayInMs; - resetScheduledTime(); - } - - public void markLoadTaskRunning() { - isLoadTaskRunning = true; - rescheduleCleanupTask(this); - } - - public void markLoadTaskNotRunning() { - isLoadTaskRunning = false; - rescheduleCleanupTask(this); - } - - public void resetScheduledTime() { - scheduledTime = System.currentTimeMillis() + delayInMs; - } - - public void cancel() { - isCanceled = true; - } - - @Override - public void run() { - if (isCanceled) { - LOGGER.info(StorageEngineMessages.LOAD_CLEANUP_TASK_CANCELED, uuid); - } else { - LOGGER.info(StorageEngineMessages.LOAD_CLEANUP_TASK_STARTS, uuid); - try { - forceCloseWriterManager(uuid); - } catch (Exception e) { - LOGGER.warn(StorageEngineMessages.LOAD_CLEANUP_TASK_ERROR, uuid, e); - } - } - } - - @Override - public int compareTo(CleanupTask that) { - return Long.compare(this.scheduledTime, that.scheduledTime); - } - } - - private static class DataPartitionInfo { - - private final DataRegion dataRegion; - private final TTimePartitionSlot timePartitionSlot; - - private DataPartitionInfo(DataRegion dataRegion, TTimePartitionSlot timePartitionSlot) { - this.dataRegion = dataRegion; - this.timePartitionSlot = timePartitionSlot; - } - - public DataRegion getDataRegion() { - return dataRegion; - } - - public TTimePartitionSlot getTimePartitionSlot() { - return timePartitionSlot; - } - - @Override - public String toString() { - return String.join( - IoTDBConstant.FILE_NAME_SEPARATOR, - dataRegion.getDatabaseName(), - dataRegion.getDataRegionIdString(), - Long.toString(timePartitionSlot.getStartTime())); - } - - @Override - public boolean equals(Object o) { - if (this == o) { - return true; - } - if (o == null || getClass() != o.getClass()) { - return false; - } - DataPartitionInfo that = (DataPartitionInfo) o; - return Objects.equals(dataRegion, that.dataRegion) - && timePartitionSlot.getStartTime() == that.timePartitionSlot.getStartTime(); - } - - @Override - public int hashCode() { - return Objects.hash(dataRegion, timePartitionSlot.getStartTime()); - } - } } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/PartitionContext.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/PartitionContext.java new file mode 100644 index 0000000000000..f67cfcad242c8 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/PartitionContext.java @@ -0,0 +1,650 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.storageengine.load; + +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.consensus.index.ProgressIndex; +import org.apache.iotdb.commons.file.SystemFileFactory; +import org.apache.iotdb.commons.schema.table.TsFileTableSchemaUtil; +import org.apache.iotdb.commons.schema.table.TsTable; +import org.apache.iotdb.commons.utils.PathUtils; +import org.apache.iotdb.commons.utils.RetryUtils; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; +import org.apache.iotdb.db.schemaengine.table.DataNodeTableCache; +import org.apache.iotdb.db.storageengine.dataregion.DataRegion; +import org.apache.iotdb.db.storageengine.dataregion.modification.ModificationFile; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResourceStatus; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; +import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; + +import org.apache.tsfile.common.constant.TsFileConstant; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.exception.write.PageException; +import org.apache.tsfile.file.metadata.ChunkGroupMetadata; +import org.apache.tsfile.file.metadata.ChunkMetadata; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.read.TimeValuePair; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.utils.RamUsageEstimator; +import org.apache.tsfile.utils.TsPrimitiveType; +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.nio.channels.FileChannel; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.nio.file.StandardOpenOption; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicLong; +import java.util.stream.Collectors; + +/** + * LOAD partition writer context: cohesive per-partition state of one in-progress LOAD task. This + * replaces the parallel per-partition maps of the old writer manager: the staged {@link + * TsFileIOWriter}, its {@link TsFileResource}, the modification file, the device currently being + * written, the already-synced byte cursor and the finalized flag now live together, and every + * single-partition behavior (write chunk / deletion, chunk-group end, ref capture, snapshot prefix + * copy, final load and close) is owned by this object. + */ +final class PartitionContext { + + private static final Logger LOGGER = LoggerFactory.getLogger(PartitionContext.class); + + /** Immutable identity of this partition: (DataRegion, time partition slot). */ + private final DataPartitionInfo partitionInfo; + + /** Task directory that holds the staged file of this partition (and of sibling partitions). */ + private final File taskDir; + + /** + * One-shot, write-only staged file builder. Chunk bytes are buffered in memory and only reach the + * file on flush(); its footer/metadata state cannot be rebuilt from the file after close. + */ + private final TsFileIOWriter writer; + + /** Resource describing the staged file; filled with time index and serialized before loading. */ + private final TsFileResource resource; + + /** Lazy-created {@code .mods} file that records deletions applied to this staged file. */ + private ModificationFile modificationFile; + + /** Device of the currently open chunk group; {@code null} when no chunk group is open. */ + private IDeviceID currentWritingDevice; + + /** + * Byte cursor of the durable prefix {@code [0, syncedOffset)} already captured by PieceRefs. Refs + * advance continuously from offset 0, which is what lets followers rebuild the file from the WAL + * and snapshots copy exactly the synced prefix without holes or overlaps. + */ + private long syncedOffset = 0L; + + /** Whether the file footer has been written (PREPARE seals the file once). */ + private boolean finalized = false; + + private PartitionContext( + DataPartitionInfo partitionInfo, + File taskDir, + TsFileIOWriter writer, + TsFileResource resource) { + this.partitionInfo = partitionInfo; + this.taskDir = taskDir; + this.writer = writer; + this.resource = resource; + } + + /** + * Creates the staged partition file and its writer. Returns {@code null} when the target file + * already exists (the original behavior logs the error and skips the chunk instead of failing the + * whole piece); IO failures are propagated. + */ + static PartitionContext create( + DataPartitionInfo partitionInfo, File taskDir, long chunkMetadataMaxSizeForEachWriter) + throws IOException { + // One staged file per (database, region, time partition): the partition's toString is a unique + // file name, e.g. "root.sg_1_0.tsfile". + final File newTsFile = + SystemFileFactory.INSTANCE.getFile( + taskDir, partitionInfo.toString() + TsFileConstant.TSFILE_SUFFIX); + if (!newTsFile.createNewFile()) { + // createNewFile returns false when the file already exists: re-creating it would truncate an + // existing staged file, so the chunk is skipped (mirrors the historical behavior). + LOGGER.error(StorageEngineMessages.CANNOT_CREATE_TSFILE_FOR_WRITING, newTsFile.getPath()); + return null; + } + + // chunkMetadataMaxSizeForEachWriter bounds how much chunk metadata one writer may keep in + // memory; the manager divides the configured budget across all concurrent partition writers. + final TsFileIOWriter writer = new TsFileIOWriter(newTsFile, chunkMetadataMaxSizeForEachWriter); + final TsFileResource resource = new TsFileResource(writer.getFile()); + addResourceFlushListener(writer, resource); + return new PartitionContext(partitionInfo, taskDir, writer, resource); + } + + /** + * Rebuilds a partition context from an unsealed staged file restored from a snapshot. {@link + * RestorableTsFileIOWriter} is the writer used by the DataRegion's unsealed-TsFile recovery: it + * scans the file, truncates it to the last complete chunk-group boundary and positions the writer + * so further chunks can be appended to the same file. The already-durable prefix becomes the new + * synced cursor; a {@code .mods} restored next to the file is re-attached lazily. + */ + static PartitionContext restore( + DataPartitionInfo partitionInfo, + File taskDir, + File restoredFile, + long chunkMetadataMaxSizeForEachWriter) + throws IOException { + if (!restoredFile.isFile() || restoredFile.length() == 0) { + return null; + } + final RestorableTsFileIOWriter writer = new RestorableTsFileIOWriter(restoredFile); + final TsFileResource resource = new TsFileResource(writer.getFile()); + addResourceFlushListener(writer, resource); + final PartitionContext context = new PartitionContext(partitionInfo, taskDir, writer, resource); + context.syncedOffset = writer.getFile().length(); + final File modsFile = ModificationFile.getExclusiveMods(restoredFile); + if (modsFile.isFile()) { + context.modificationFile = new ModificationFile(modsFile, false); + } + return context; + } + + private static void addResourceFlushListener(TsFileIOWriter writer, TsFileResource resource) { + // TsFileIOWriter calls back with the chunk groups it is about to flush to disk. Update the + // resource's per-device start/end time here, so the time index is already correct when the + // file is loaded and we do not have to re-scan it. + writer.addFlushListener( + sortedChunkMetadataList -> + sortedChunkMetadataList.forEach( + pair -> { + // pair is (device, chunk metadata list) of one flushed chunk group. + final IDeviceID deviceId = pair.left.left; + pair.getRight() + .forEach( + chunkMetadata -> { + resource.updateStartTime(deviceId, chunkMetadata.getStartTime()); + resource.updateEndTime(deviceId, chunkMetadata.getEndTime()); + }); + })); + } + + boolean belongsTo(DataRegion dataRegion) { + // DataRegion instances are singletons per region, so identity comparison is sufficient. + return partitionInfo.getDataRegion() == dataRegion; + } + + TsFileIOWriter getWriter() { + return writer; + } + + IDeviceID getCurrentWritingDevice() { + return currentWritingDevice; + } + + long getTimePartitionStart() { + return partitionInfo.getTimePartitionSlot().getStartTime(); + } + + TTimePartitionSlot getTimePartitionSlot() { + return partitionInfo.getTimePartitionSlot(); + } + + /** + * Starts a new chunk group for the given device, warning on the inconsistent state where this + * writer still has an open chunk group (it should have been ended by the device fan-out). + * + *

A chunk group groups all measurements of one device in a time range; the device switch is + * handled by the task manager across every partition of the old device before this is called. + */ + void startChunkGroup(IDeviceID device) throws IOException { + if (writer.isWritingChunkGroup()) { + LOGGER.warn( + StorageEngineMessages + .STORAGE_LOG_WRITER_FOR_PARTITION_IS_ALREADY_WRITING_CHUNK_GROUP_FOR_903B1D66, + writer.getFile().getAbsolutePath(), + partitionInfo, + device, + currentWritingDevice); + } + writer.startChunkGroup(device); + currentWritingDevice = device; + } + + /** + * Seals the open chunk group of this partition and asks the writer to flush if its buffered + * metadata exceeds the size bound. Used by the device fan-out: when a source file switches to a + * new device, every partition that was writing the old device must end its chunk group at the + * same logical point so aligned chunks stay consistent across partitions. + */ + void endChunkGroupAndCheckMetadataSize() throws IOException { + if (writer.isWritingChunkGroup()) { + writer.endChunkGroup(); + } + writer.checkMetadataSizeAndMayFlush(); + } + + /** + * Registers the table schema for table-model databases and writes the chunk to the writer. The + * staged TsFile must carry the table schema in its footer for downstream readers (query, + * compaction, pipe) to interpret the table-model chunks; a table missing from the DataNode cache + * means it was dropped after the LOAD statement was analyzed, which must fail the load explicitly + * instead of writing schema-less chunks or silently dropping the data. + */ + void writeChunk(ChunkData chunkData) throws IOException, PageException, LoadFileException { + final String tableName = + chunkData.getDevice() != null ? chunkData.getDevice().getTableName() : null; + if (tableName != null + && PathUtils.isTableModelDatabase(partitionInfo.getDataRegion().getDatabaseName())) { + final TsTable table = + DataNodeTableCache.getInstance() + .getTable(partitionInfo.getDataRegion().getDatabaseName(), tableName, false); + if (Objects.nonNull(table)) { + writer + .getSchema() + .getTableSchemaMap() + .computeIfAbsent( + tableName, t -> TsFileTableSchemaUtil.toTsFileTableSchemaNoAttribute(table)); + } else { + throw new LoadFileException( + String.format( + StorageEngineMessages + .EXCEPTION_TABLE_ARG_ARG_DOES_NOT_EXIST_WHEN_APPLYING_LOAD_CHUNK_DATA_IT_MAY_HAVE_BEEN_DROPPED_AFTER_THE_LOAD_WAS_ANALYZED_DDB35F93, + partitionInfo.getDataRegion().getDatabaseName(), + tableName)); + } + } + chunkData.writeToFileWriter(writer); + } + + /** + * Applies one deletion to this partition's modification file. The {@code .mods} file is created + * lazily next to the staged file on the first deletion; the DataRegion reads it when the file is + * loaded, so the deletion never has to touch the already-written chunk bytes. + */ + void writeDeletion(DeletionData deletionData) throws IOException { + if (modificationFile == null) { + final File newModificationFile = ModificationFile.getExclusiveMods(writer.getFile()); + if (!newModificationFile.isFile() && !newModificationFile.createNewFile()) { + // The file may already exist because it was restored from a snapshot together with the + // staged file; createNewFile returns false in that case, which is not an error. + if (!newModificationFile.isFile()) { + LOGGER.error( + StorageEngineMessages + .STORAGE_LOG_CAN_NOT_CREATE_MODIFICATIONFILE_FOR_WRITING_17D14C11, + newModificationFile.getPath()); + return; + } + } + modificationFile = new ModificationFile(newModificationFile, false); + } + writer.flush(); + // Flush the chunk file first so the deletion (which references point ranges of the file) is + // recorded after the corresponding bytes are durable. + deletionData.writeToModificationFile(modificationFile); + } + + /** + * Ends the partition file (chunk groups + footer) and captures the final byte range. The footer + * capture matters: nodes rebuilding the file from WAL refs only get a complete, readable file if + * the refs cover the footer too. The {@code finalized} flag makes this idempotent because PREPARE + * may be applied once per task. + */ + void finalizeFile(List pendingPieceRefs) throws IOException { + if (finalized) { + return; + } + if (writer.isWritingChunkGroup()) { + writer.endChunkGroup(); + } + writer.endFile(); + captureRefs(pendingPieceRefs); + finalized = true; + } + + private void captureRefs(List pendingPieceRefs) + throws IOException { + // TsFileIOWriter buffers chunk bytes in memory: endChunkGroup() only appends the + // ChunkGroupFooter to the buffer, and the bytes do not reach the file until flush(). Read the + // length only after a flush so the captured [startOffset, endOffset) always covers a sealed, + // durably written chunk group. + writer.flush(); + // startOffset is the previous synced cursor; every ref starts exactly there and extends to the + // current file length, so the sequence of refs of one staged file is contiguous from offset 0. + final long startOffset = syncedOffset; + final long endOffset = writer.getFile().length(); + if (endOffset <= startOffset) { + return; + } + final int length = (int) (endOffset - startOffset); + pendingPieceRefs.add( + new LoadTsFileConsensusNode.PieceRef( + taskDir.getName() + File.separator + writer.getFile().getName(), startOffset, length)); + syncedOffset = endOffset; + } + + /** + * Copies the already-synced byte prefix of this staged partition file into the snapshot task dir + * and returns its snapshot metadata, or {@code null} when nothing has been synced yet. + * + *

Only {@code [0, syncedOffset)} is copied: bytes after the last chunk-group boundary are + * still owned by the writer buffer and will be covered by the next PIECE ref, so a node restored + * from this snapshot can keep appending from exactly the snapshot length without a hole or an + * overlap. + */ + LoadSnapshotManager.StagedFileSnapshot snapshotTo(File targetDir) throws IOException { + if (syncedOffset <= 0) { + return null; + } + final File stagedFile = writer.getFile(); + copyPrefix(stagedFile, new File(targetDir, stagedFile.getName()), syncedOffset); + copyModsIfPresent(targetDir); + return new LoadSnapshotManager.StagedFileSnapshot( + stagedFile.getName(), + partitionInfo.getDataRegion().getDatabaseName(), + partitionInfo.getDataRegion().getDataRegionIdString(), + getTimePartitionStart(), + finalized); + } + + private void copyPrefix(File source, File target, long length) throws IOException { + if (length <= 0) { + return; + } + // transferTo may copy fewer bytes than requested, so loop until the whole prefix is copied and + // fail loudly on EOF (a truncated staged file would silently corrupt the replica otherwise). + try (final FileChannel in = FileChannel.open(source.toPath(), StandardOpenOption.READ); + final FileChannel out = + FileChannel.open( + target.toPath(), + StandardOpenOption.CREATE, + StandardOpenOption.WRITE, + StandardOpenOption.TRUNCATE_EXISTING)) { + long transferred = 0; + while (transferred < length) { + final long count = in.transferTo(transferred, length - transferred, out); + if (count <= 0) { + throw new IOException( + String.format( + StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_EOF_8743387D, + source, + transferred)); + } + transferred += count; + } + out.force(true); + } + } + + private void copyModsIfPresent(File targetDir) throws IOException { + // Deletions must travel with the snapshot, otherwise the restored staged file would resurrect + // data the load had already marked as deleted. + if (modificationFile != null && modificationFile.getFile().isFile()) { + Files.copy( + modificationFile.getFile().toPath(), + new File(targetDir, modificationFile.getFile().getName()).toPath(), + StandardCopyOption.REPLACE_EXISTING); + } + } + + private void forceFile(final File file) throws IOException { + // fsync the file so the refs logged to the WAL cover bytes that are actually durable. + try (final FileChannel channel = FileChannel.open(file.toPath(), StandardOpenOption.WRITE)) { + channel.force(true); + } + } + + void closeModificationFile() throws IOException { + // Closed before loadNewTsFile so the DataRegion can read the mods of the staged file; the file + // itself is deleted later by close(). + if (modificationFile != null) { + modificationFile.close(); + } + } + + /** + * Loads this staged partition file into its DataRegion: seals the file if PREPARE has not run, + * validates it is a complete TsFile, binds the resource (time index, last values, progress) and + * hands it to the DataRegion, then updates the load point-count metrics. + */ + void loadIntoRegion(boolean isGeneratedByPipe, ProgressIndex progressIndex) + throws IOException, LoadFileException { + if (!finalized) { + if (writer.isWritingChunkGroup()) { + writer.endChunkGroup(); + } + writer.endFile(); + } + validateStagedFileComplete(writer.getFile()); + + final DataRegion partitionRegion = partitionInfo.getDataRegion(); + resource.setGeneratedByPipe(isGeneratedByPipe); + endTsFileResource(writer, resource, progressIndex); + partitionRegion.loadNewTsFile( + resource, true, isGeneratedByPipe, false, Optional.ofNullable(writer.getTableSizeMap())); + + // Metrics + partitionRegion + .getNonSystemDatabaseName() + .ifPresent( + databaseName -> + LoadTsFileManager.updateWritePointCountMetrics( + partitionRegion, databaseName, getTsFileWritePointCount(writer), false)); + } + + private void validateStagedFileComplete(File stagedFile) throws LoadFileException { + // isComplete() verifies the file has a valid magic header and footer, i.e. it is not a + // truncated staged file left behind by a failed transfer. + try (final TsFileSequenceReader reader = + new TsFileSequenceReader(stagedFile.getAbsolutePath(), true)) { + if (!reader.isComplete()) { + throw new LoadFileException( + String.format( + StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_INCOMPLETE_1CDE954B, + stagedFile, + taskDir.getName())); + } + } catch (IOException e) { + throw new LoadFileException( + String.format( + StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_INCOMPLETE_1CDE954B, + stagedFile, + taskDir.getName()), + e); + } + } + + private void endTsFileResource( + TsFileIOWriter writer, TsFileResource tsFileResource, ProgressIndex progressIndex) + throws IOException { + // Build the time index from every chunk group (still in the writer's memory) and optionally + // cache the last value of each measurement for fast "last query" after load. + Map> deviceLastValues = null; + if (IoTDBDescriptor.getInstance().getConfig().isCacheLastValuesForLoad()) { + deviceLastValues = new HashMap<>(); + } + // Tracks the estimated memory of the last-value cache; the cache is disabled as soon as the + // configured budget is exceeded so LOAD cannot blow up the heap. + AtomicLong lastValuesMemCost = new AtomicLong(0); + + for (final ChunkGroupMetadata chunkGroupMetadata : writer.getChunkGroupMetadataList()) { + final IDeviceID device = chunkGroupMetadata.getDevice(); + for (final ChunkMetadata chunkMetadata : chunkGroupMetadata.getChunkMetadataList()) { + // Per-device min start time / max end time across all chunks. + tsFileResource.updateStartTime(device, chunkMetadata.getStartTime()); + tsFileResource.updateEndTime(device, chunkMetadata.getEndTime()); + if (deviceLastValues != null) { + // deviceMap: measurement uid -> (timestamp, value) of the last point in this file. + Map deviceMap = + deviceLastValues.computeIfAbsent( + device, + d -> { + // Account for the per-device map and the device id memory when computing the + // budget, so the estimate tracks the real allocation. + Map map = new HashMap<>(); + lastValuesMemCost.addAndGet(RamUsageEstimator.shallowSizeOf(map)); + lastValuesMemCost.addAndGet(device.ramBytesUsed()); + return map; + }); + int prevSize = deviceMap.size(); + deviceMap.compute( + chunkMetadata.getMeasurementUid(), + (m, oldPair) -> { + // Keep the existing (later) value if it is still newer than this chunk's end. + if (oldPair != null && oldPair.getTimestamp() > chunkMetadata.getEndTime()) { + return oldPair; + } + // Reconstruct the last value from the chunk statistics; VECTOR chunks use the time + // column (INT64) because the vector itself has no scalar statistics. + TsPrimitiveType lastValue = + chunkMetadata.getStatistics() != null + && chunkMetadata.getDataType() != TSDataType.BLOB + ? TsPrimitiveType.getByType( + chunkMetadata.getDataType() == TSDataType.VECTOR + ? TSDataType.INT64 + : chunkMetadata.getDataType(), + chunkMetadata.getStatistics().getLastValue()) + : null; + TimeValuePair timeValuePair = + lastValue != null + ? new TimeValuePair(chunkMetadata.getEndTime(), lastValue) + : null; + // Adjust the budget by the size difference of the replaced entry. + if (oldPair != null) { + lastValuesMemCost.addAndGet(-oldPair.getSize()); + } + if (timeValuePair != null) { + lastValuesMemCost.addAndGet(timeValuePair.getSize()); + } + return timeValuePair; + }); + int afterSize = deviceMap.size(); + lastValuesMemCost.addAndGet( + (afterSize - prevSize) * RamUsageEstimator.HASHTABLE_RAM_BYTES_PER_ENTRY); + // Give up caching once the budget is exceeded; the data is still loaded correctly, only + // the last-value cache is dropped. + if (lastValuesMemCost.get() + > IoTDBDescriptor.getInstance().getConfig().getCacheLastValuesMemoryBudgetInByte()) { + deviceLastValues = null; + } + } + } + } + if (deviceLastValues != null) { + // Flatten device -> {measurement -> last pair} into device -> [(measurement, pair)] for the + // resource's compact last-value representation. + Map>> finalDeviceLastValues; + finalDeviceLastValues = new HashMap<>(deviceLastValues.size()); + for (final Map.Entry> entry : + deviceLastValues.entrySet()) { + final IDeviceID device = entry.getKey(); + Map lastValues = entry.getValue(); + List> pairList = + lastValues.entrySet().stream() + .map(e -> new Pair<>(e.getKey(), e.getValue())) + .collect(Collectors.toList()); + finalDeviceLastValues.put(device, pairList); + } + tsFileResource.setLastValues(finalDeviceLastValues); + } + tsFileResource.setStatus(TsFileResourceStatus.NORMAL); + tsFileResource.setProgressIndex(progressIndex); + // Serialize the .tsfile.resource metadata file next to the staged file. + tsFileResource.serialize(); + } + + private long getTsFileWritePointCount(TsFileIOWriter writer) { + // Sum of the row counts of every chunk, used only for the load point-count metric. + return writer.getChunkGroupMetadataList().stream() + .flatMap(chunkGroupMetadata -> chunkGroupMetadata.getChunkMetadataList().stream()) + .mapToLong(chunkMetadata -> chunkMetadata.getStatistics().getCount()) + .sum(); + } + + /** + * Closes and deletes the writer file and the modification file, tolerating per-file errors so one + * failing file cannot block cleanup of the remaining task directory. + * + *

Closing a stream and deleting its file are independent best-effort steps: a close failure + * (e.g. disk full while writing the footer) must never skip the file deletion, otherwise the + * abandoned staged file would leak as garbage in the data directory. + */ + void close() { + // canWrite() is false once endFile() sealed the writer; only an unsealed writer is closed. + if (writer.canWrite()) { + try { + writer.close(); + } catch (IOException e) { + LOGGER.warn( + StorageEngineMessages.CLOSE_TSFILE_IO_WRITER_ERROR, writer.getFile().getPath(), e); + } + } + try { + final Path writerPath = writer.getFile().toPath(); + if (Files.exists(writerPath)) { + RetryUtils.retryOnException( + () -> { + Files.delete(writerPath); + return null; + }); + } + } catch (Exception e) { + LOGGER.warn( + StorageEngineMessages.FAILED_TO_DELETE_FILE_OR_DIR, writer.getFile().getPath(), e); + } + if (modificationFile != null) { + try { + modificationFile.close(); + } catch (IOException e) { + LOGGER.warn( + StorageEngineMessages.CLOSE_MODIFICATION_FILE_ERROR, + modificationFile.getFile().getPath(), + e); + } + try { + final Path modificationFilePath = modificationFile.getFile().toPath(); + if (Files.exists(modificationFilePath)) { + RetryUtils.retryOnException( + () -> { + Files.delete(modificationFilePath); + return null; + }); + } + } catch (Exception e) { + LOGGER.warn( + StorageEngineMessages.FAILED_TO_DELETE_FILE_OR_DIR, + modificationFile.getFile().getPath(), + e); + } + } + } +} 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 new file mode 100644 index 0000000000000..0840ee4303020 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManager.java @@ -0,0 +1,1146 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.storageengine.load; + +import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot; +import org.apache.iotdb.commons.consensus.index.ProgressIndex; +import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex; +import org.apache.iotdb.commons.file.SystemFileFactory; +import org.apache.iotdb.commons.utils.FileUtils; +import org.apache.iotdb.commons.utils.RetryUtils; +import org.apache.iotdb.db.conf.IoTDBConfig; +import org.apache.iotdb.db.conf.IoTDBDescriptor; +import org.apache.iotdb.db.exception.load.LoadFileException; +import org.apache.iotdb.db.i18n.StorageEngineMessages; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFileConsensusNode; +import org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode; +import org.apache.iotdb.db.storageengine.dataregion.DataRegion; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResource; +import org.apache.iotdb.db.storageengine.dataregion.tsfile.TsFileResourceStatus; +import org.apache.iotdb.db.storageengine.dataregion.utils.TsFileResourceUtils; +import org.apache.iotdb.db.storageengine.load.LoadSnapshotManager.StagedFileSnapshot; +import org.apache.iotdb.db.storageengine.load.splitter.ChunkData; +import org.apache.iotdb.db.storageengine.load.splitter.DeletionData; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileDataType; + +import org.apache.tsfile.exception.write.PageException; +import org.apache.tsfile.file.metadata.IDeviceID; +import org.apache.tsfile.read.TsFileSequenceReader; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.File; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.channels.FileChannel; +import java.nio.file.DirectoryNotEmptyException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.nio.file.StandardOpenOption; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.locks.ReentrantLock; + +/** + * Lifecycle of one LOAD task: owns the task directory, the {@link PartitionContext} set of this + * task and the WAL/ref bookkeeping (raw pieces, restored snapshot partitions, applied piece + * checksums). All mutations are serialized by the per-task {@link #taskLock}, so different LOAD + * tasks never contend on a global manager lock. + */ +final class TsFileWriterManager { + + private static final Logger LOGGER = LoggerFactory.getLogger(TsFileWriterManager.class); + private static final IoTDBConfig CONFIG = IoTDBDescriptor.getInstance().getConfig(); + + private static final String MESSAGE_WRITER_MANAGER_HAS_BEEN_CLOSED = + "%s TsFileWriterManager has been closed."; + private static final String MESSAGE_DELETE_FAIL = "failed to delete {}."; + + /** Sub-directory holding the retained serialized PIECE bytes (the backfill source). */ + private static final String RETAINED_PIECES_DIR_NAME = "pieces"; + + private final File taskDir; + private final ReentrantLock taskLock = new ReentrantLock(); + + /** One cohesive context per staged partition file: writer, resource, mods, cursor, finalized. */ + private final Map partitionContexts = new HashMap<>(); + + /** Devices currently routed to partition files, used to end one chunk group across partitions. */ + private final Map> device2Partition = new HashMap<>(); + + private final List pendingPieceRefs = new ArrayList<>(); + private final List rawTsFiles = new ArrayList<>(); + private final Set rawTsFilePaths = new HashSet<>(); + private final Map restoredPartitions = new HashMap<>(); + + /** + * Piece data delivered to this node (pulled back from the write node) but not yet applied. A + * follower applies the cached data only when the corresponding WAL marker arrives, because the + * marker (not the delivery order) is the ordering authority of the load. + */ + private final Map cachedPieces = new HashMap<>(); + + /** + * Serialized bytes of every chunk-data piece applied by this (write) node, retained until COMMIT + * or ABORT so a follower that missed the delivery can pull the piece back. + */ + private final Map retainedPieces = new HashMap<>(); + + /** + * Tracks the checksum of every chunk-data piece already applied to this task so that a scheduler + * retry (or a duplicated consensus log entry) is acknowledged idempotently instead of appending + * the same bytes twice. Keyed by the coordinator-assigned piece index, which is unique per load. + */ + private final Map appliedPieceIndex2Checksum = new HashMap<>(); + + /** + * Length of the contiguous applied-piece prefix ({@code 0..appliedContiguousCount-1}). Kept in + * sync with {@link #appliedPieceIndex2Checksum} so the failover fence is O(1). + */ + private long appliedContiguousCount; + + private boolean isClosed; + + TsFileWriterManager(File taskDir) { + this(taskDir, true); + } + + TsFileWriterManager(File taskDir, boolean clearExistingDir) { + this.taskDir = taskDir; + if (clearExistingDir) { + clearDir(taskDir); + } else { + // A task dir restored from a snapshot may already carry the retained PIECE bytes that were + // durable before the crash; reload them so marker replay can be backfilled locally instead + // of depending on a live write node. + loadRetainedPiecesFromDisk(); + } + } + + String getTaskName() { + return taskDir.getName(); + } + + File getTaskDir() { + return taskDir; + } + + private void clearDir(File dir) { + if (dir.exists()) { + FileUtils.deleteFileOrDirectoryWithRetry(dir); + } + if (dir.mkdirs()) { + LOGGER.info(StorageEngineMessages.LOAD_TSFILE_DIR_CREATED, dir.getPath()); + } + } + + boolean hasLiveWriter() { + taskLock.lock(); + try { + return !partitionContexts.isEmpty(); + } finally { + taskLock.unlock(); + } + } + + boolean isPieceAlreadyApplied(long pieceIndex, long checksum) { + taskLock.lock(); + try { + return appliedPieceIndex2Checksum.containsKey(pieceIndex) + && appliedPieceIndex2Checksum.get(pieceIndex) == checksum; + } finally { + taskLock.unlock(); + } + } + + boolean isPieceConflicting(long pieceIndex, long checksum) { + taskLock.lock(); + try { + return appliedPieceIndex2Checksum.containsKey(pieceIndex) + && appliedPieceIndex2Checksum.get(pieceIndex) != checksum; + } finally { + taskLock.unlock(); + } + } + + /** Whether every piece {@code 0..pieceIndex} (inclusive) has been applied contiguously. */ + boolean hasAppliedAllUpTo(long pieceIndex) { + taskLock.lock(); + try { + // appliedContiguousCount tracks how many pieces (starting from 0) are already applied, so + // the failover fence is a single comparison instead of an O(n) scan per piece. + return pieceIndex < 0 || pieceIndex < appliedContiguousCount; + } finally { + taskLock.unlock(); + } + } + + /** + * Caches a chunk-data PIECE pushed back by the write node in response to a PULL. Returns {@code + * false} when a piece with the same index is already cached with a different checksum (a + * divergent delivery). + */ + boolean cachePiece(long pieceIndex, long checksum, List dataList) throws IOException { + taskLock.lock(); + try { + checkNotClosed(); + final CachedPiece existing = cachedPieces.get(pieceIndex); + if (existing != null) { + return existing.checksum == checksum; + } + cachedPieces.put(pieceIndex, new CachedPiece(pieceIndex, checksum, dataList)); + return true; + } finally { + taskLock.unlock(); + } + } + + boolean hasCachedPiece(long pieceIndex, long checksum) { + taskLock.lock(); + try { + final CachedPiece cached = cachedPieces.get(pieceIndex); + return cached != null && cached.checksum == checksum; + } finally { + taskLock.unlock(); + } + } + + /** + * Applies the cached data of {@code pieceIndex} (whose WAL marker has arrived) to this task's own + * writers, exactly like the write node applies the chunk-data PIECE, then records the piece as + * applied so retries and the continuity check are idempotent. + */ + void applyCachedPiece(DataRegion dataRegion, long pieceIndex, long checksum) + throws IOException, PageException, LoadFileException { + taskLock.lock(); + try { + checkNotClosed(); + if (isPieceAlreadyApplied(pieceIndex, checksum)) { + cachedPieces.remove(pieceIndex); + return; + } + final CachedPiece cached = cachedPieces.get(pieceIndex); + if (cached == null || cached.checksum != checksum) { + throw new IOException( + String.format( + StorageEngineMessages + .EXCEPTION_LOAD_CONSENSUS_PIECE_DATA_MISSING_OR_CHECKSUM_MISMATCH_AFTER_PULL_35F4972E, + getTaskName(), + pieceIndex)); + } + appendChunkPieceAndRecord(dataRegion, cached.dataList, pieceIndex, checksum); + cachedPieces.remove(pieceIndex); + } finally { + taskLock.unlock(); + } + } + + void retainPiece(long pieceIndex, byte[] serializedPiece) { + taskLock.lock(); + try { + retainedPieces.put(pieceIndex, serializedPiece); + writeRetainedPieceToDisk(pieceIndex, serializedPiece); + } finally { + taskLock.unlock(); + } + } + + Optional getRetainedPiece(long pieceIndex) { + taskLock.lock(); + try { + final byte[] inMemory = retainedPieces.get(pieceIndex); + if (inMemory != null) { + return Optional.of(inMemory); + } + final File pieceFile = retainedPieceFile(pieceIndex); + if (!pieceFile.isFile()) { + return Optional.empty(); + } + try { + final byte[] fromDisk = Files.readAllBytes(pieceFile.toPath()); + retainedPieces.put(pieceIndex, fromDisk); + return Optional.of(fromDisk); + } catch (IOException e) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_RETAINED_PIECE_READ_FAILED_0659D19B, + pieceIndex, + getTaskName(), + pieceFile, + e.getMessage()); + return Optional.empty(); + } + } finally { + taskLock.unlock(); + } + } + + void clearRetainedPieces() { + taskLock.lock(); + try { + retainedPieces.clear(); + cachedPieces.clear(); + deleteRetainedPieceFiles(); + } finally { + taskLock.unlock(); + } + } + + void recordAppliedPiece(long pieceIndex, long checksum) { + taskLock.lock(); + try { + recordAppliedPieceUnlocked(pieceIndex, checksum); + } finally { + taskLock.unlock(); + } + } + + /** Number of chunk pieces already applied to this task, for the PREPARE reconciliation. */ + long getAppliedPieceCount() { + taskLock.lock(); + try { + return appliedPieceIndex2Checksum.size(); + } finally { + taskLock.unlock(); + } + } + + /** XOR of every applied piece checksum, for the PREPARE reconciliation. */ + long getAppliedPiecesChecksum() { + taskLock.lock(); + try { + long checksum = 0; + for (final long pieceChecksum : appliedPieceIndex2Checksum.values()) { + checksum ^= pieceChecksum; + } + return checksum; + } finally { + taskLock.unlock(); + } + } + + /** + * Verifies that the pieces applied on this node match the PREPARE summary accumulated by the + * coordinator: the applied piece count and the XOR of every applied piece checksum must equal the + * expected values. A mismatch means a piece was lost or replaced on this node (e.g. the write + * node switched mid-load and this node never received some markers), so the staged file must not + * be sealed or loaded silently. + */ + boolean verifyAppliedPieces(int expectedCount, long expectedChecksum) { + taskLock.lock(); + try { + if (appliedPieceIndex2Checksum.size() != expectedCount) { + return false; + } + long checksum = 0; + for (final long pieceChecksum : appliedPieceIndex2Checksum.values()) { + checksum ^= pieceChecksum; + } + return checksum == expectedChecksum; + } finally { + taskLock.unlock(); + } + } + + /** Whether this task was rebuilt from raw byte refs (legacy WAL format without piece records). */ + boolean isLegacyRawRefTask() { + taskLock.lock(); + try { + return !rawTsFiles.isEmpty(); + } finally { + taskLock.unlock(); + } + } + + private void recordAppliedPieceUnlocked(long pieceIndex, long checksum) { + appliedPieceIndex2Checksum.put(pieceIndex, checksum); + while (appliedPieceIndex2Checksum.containsKey(appliedContiguousCount)) { + appliedContiguousCount++; + } + } + + private File retainedPiecesDir() { + return new File(taskDir, RETAINED_PIECES_DIR_NAME); + } + + private File retainedPieceFile(long pieceIndex) { + return new File(retainedPiecesDir(), "piece-" + pieceIndex + ".bin"); + } + + private void writeRetainedPieceToDisk(long pieceIndex, byte[] bytes) { + final File piecesDir = retainedPiecesDir(); + if (!piecesDir.isDirectory() && !piecesDir.mkdirs()) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_RETAINED_PIECE_WRITE_FAILED_99697608, + pieceIndex, + getTaskName(), + piecesDir, + "failed to create directory"); + return; + } + final File target = retainedPieceFile(pieceIndex); + final File tmp = new File(piecesDir, target.getName() + ".tmp"); + try { + Files.write( + tmp.toPath(), bytes, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING); + Files.move( + tmp.toPath(), + target.toPath(), + StandardCopyOption.REPLACE_EXISTING, + StandardCopyOption.ATOMIC_MOVE); + } catch (IOException e) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_RETAINED_PIECE_WRITE_FAILED_99697608, + pieceIndex, + getTaskName(), + target, + e.getMessage()); + } + } + + private void loadRetainedPiecesFromDisk() { + final File piecesDir = retainedPiecesDir(); + final File[] files = piecesDir.isDirectory() ? piecesDir.listFiles() : null; + if (files == null) { + return; + } + for (final File file : files) { + if (!file.isFile() + || !file.getName().startsWith("piece-") + || !file.getName().endsWith(".bin")) { + continue; + } + final String name = file.getName(); + try { + final long pieceIndex = + Long.parseLong(name.substring("piece-".length(), name.length() - ".bin".length())); + retainedPieces.put(pieceIndex, Files.readAllBytes(file.toPath())); + } catch (IOException | NumberFormatException e) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_RETAINED_PIECE_READ_FAILED_0659D19B, + name, + getTaskName(), + file, + e.getMessage()); + } + } + } + + private void deleteRetainedPieceFiles() { + final File piecesDir = retainedPiecesDir(); + if (!piecesDir.isDirectory()) { + return; + } + final File[] files = piecesDir.listFiles(); + if (files != null) { + for (final File file : files) { + try { + Files.deleteIfExists(file.toPath()); + } catch (IOException e) { + LOGGER.warn(MESSAGE_DELETE_FAIL, file, e); + } + } + } + try { + Files.deleteIfExists(piecesDir.toPath()); + } catch (IOException e) { + LOGGER.warn(MESSAGE_DELETE_FAIL, piecesDir, e); + } + } + + boolean belongsTo(DataRegion dataRegion) { + taskLock.lock(); + try { + for (PartitionContext context : partitionContexts.values()) { + if (context.belongsTo(dataRegion)) { + return true; + } + } + return false; + } finally { + taskLock.unlock(); + } + } + + public void writeChunk(ChunkData chunkData, DataRegion dataRegion) + throws IOException, PageException, LoadFileException { + taskLock.lock(); + try { + final DataPartitionInfo partitionInfo = + new DataPartitionInfo(dataRegion, chunkData.getTimePartitionSlot()); + write(partitionInfo, chunkData); + } finally { + taskLock.unlock(); + } + } + + public void writeDeletion(DeletionData deletionData, DataRegion dataRegion) throws IOException { + taskLock.lock(); + try { + applyDeletionToContexts(dataRegion, deletionData); + } finally { + taskLock.unlock(); + } + } + + /** Legacy direct-load path: writes every piece datum through the typed write helpers. */ + void writePieceNode(DataRegion dataRegion, LoadTsFilePieceNode pieceNode) + throws IOException, PageException, LoadFileException { + taskLock.lock(); + try { + checkNotClosed(); + for (TsFileData tsFileData : pieceNode.getAllTsFileData()) { + if (tsFileData.getType() == TsFileDataType.CHUNK) { + writeChunk((ChunkData) tsFileData, dataRegion); + } else if (tsFileData.getType() == TsFileDataType.DELETION) { + writeDeletion((DeletionData) tsFileData, dataRegion); + } else { + throw new IOException( + StorageEngineMessages.UNSUPPORTED_TSFILE_DATA_TYPE + tsFileData.getType()); + } + } + } finally { + taskLock.unlock(); + } + } + + /** + * It should be noted that all AlignedChunkData of the same partition split from a source file + * should be guaranteed to be written to the same new file. Otherwise, for detached + * BatchedAlignedChunkData, it may result in no data for the time column in the new file. + */ + private void write(DataPartitionInfo partitionInfo, ChunkData chunkData) + throws IOException, PageException, LoadFileException { + checkNotClosed(); + final PartitionContext context = getOrCreatePartitionContext(partitionInfo); + if (context == null) { + // The staged file already exists; the original behavior logs the error and skips the chunk. + return; + } + + final IDeviceID device = chunkData.getDevice(); + final IDeviceID lastDevice = context.getCurrentWritingDevice(); + if (!Objects.equals(device, lastDevice)) { + if (lastDevice != null && device2Partition.containsKey(lastDevice)) { + final Set partitions = device2Partition.get(lastDevice); + for (DataPartitionInfo partition : new ArrayList<>(partitions)) { + final PartitionContext partitionContext = partitionContexts.get(partition); + if (partitionContext != null && partitionContext.getCurrentWritingDevice() != null) { + partitionContext.endChunkGroupAndCheckMetadataSize(); + } + } + device2Partition.remove(lastDevice); + } + context.startChunkGroup(device); + device2Partition.computeIfAbsent(device, k -> new HashSet<>()).add(partitionInfo); + } + + context.writeChunk(chunkData); + } + + private PartitionContext getOrCreatePartitionContext(DataPartitionInfo partitionInfo) + throws IOException { + PartitionContext context = partitionContexts.get(partitionInfo); + if (context != null) { + return context; + } + + final long chunkMetadataMaxSizeForEachWriter = + CONFIG.getLoadChunkMetadataMemorySizeInBytes() / (partitionContexts.size() + 1); + context = PartitionContext.create(partitionInfo, taskDir, chunkMetadataMaxSizeForEachWriter); + if (context == null) { + return null; + } + + // When a new writer is added, we need to reduce the metadata size limit of all existing + // writers for memory control + for (final PartitionContext existingContext : partitionContexts.values()) { + existingContext.getWriter().setMaxMetadataSize(chunkMetadataMaxSizeForEachWriter); + } + partitionContexts.put(partitionInfo, context); + return context; + } + + private void applyDeletionToContexts(DataRegion dataRegion, DeletionData deletionData) + throws IOException { + checkNotClosed(); + for (final PartitionContext context : partitionContexts.values()) { + if (context.belongsTo(dataRegion)) { + context.writeDeletion(deletionData); + } + } + } + + /** + * Applies a consensus PIECE to this task. Chunk data is written into the single writer of each + * affected partition, then the chunk group is ended and the newly written byte range captured so + * it can be synced to replicas (or rebuilt from the WAL) at the exact offset of the final file; + * deletion data is routed to the modification files of the matching partitions. + */ + void appendChunkPiece(DataRegion dataRegion, List dataList) + throws IOException, PageException, LoadFileException { + taskLock.lock(); + try { + checkNotClosed(); + appendChunkPieceUnlocked(dataRegion, dataList); + } finally { + taskLock.unlock(); + } + } + + /** + * Appends a consensus PIECE and records it as applied under the same task lock. Making the file + * capture (inside {@link #appendChunkPieceUnlocked}) and the applied-piece map atomic with + * respect to snapshotting guarantees a snapshot never observes the staged-file prefix ahead of + * the applied-piece prefix (or vice versa), which would otherwise fork the restored file on + * replay. + */ + void appendChunkPieceAndRecord( + DataRegion dataRegion, List dataList, long pieceIndex, long checksum) + throws IOException, PageException, LoadFileException { + taskLock.lock(); + try { + checkNotClosed(); + appendChunkPieceUnlocked(dataRegion, dataList); + recordAppliedPieceUnlocked(pieceIndex, checksum); + } finally { + taskLock.unlock(); + } + } + + private void appendChunkPieceUnlocked(DataRegion dataRegion, List dataList) + throws IOException, PageException, LoadFileException { + restorePartitionWriters(dataRegion); + for (final TsFileData data : dataList) { + if (data.getType() == TsFileDataType.CHUNK) { + final ChunkData chunkData = (ChunkData) data; + final DataPartitionInfo partitionInfo = + new DataPartitionInfo(dataRegion, chunkData.getTimePartitionSlot()); + write(partitionInfo, chunkData); + } else if (data.getType() == TsFileDataType.DELETION) { + writeDeletion((DeletionData) data, dataRegion); + } else { + throw new IOException(StorageEngineMessages.UNSUPPORTED_TSFILE_DATA_TYPE + data.getType()); + } + } + } + + /** + * Rebuilds an in-memory {@link TsFileIOWriter} for every non-finalized staged partition file that + * was restored from a snapshot, so this node can keep appending subsequent chunk-data pieces to + * the same file. The rebuild uses the unsealed-TsFile recovery mechanism ({@link + * org.apache.tsfile.write.writer.RestorableTsFileIOWriter}), the same writer class the DataRegion + * uses to continue an interrupted flush, and truncates the restored file to its last complete + * chunk-group boundary. Finalized files (footer already written) are left untouched and loaded + * directly at COMMIT. + */ + private void restorePartitionWriters(DataRegion dataRegion) throws IOException { + if (restoredPartitions.isEmpty()) { + return; + } + final List> toRestore = new ArrayList<>(); + for (final Map.Entry entry : + restoredPartitions.entrySet()) { + final RestoredPartitionKey key = entry.getKey(); + final RestoredLoadFile restored = entry.getValue(); + if (restored.finalized + || !dataRegion.getDatabaseName().equals(restored.database) + || !dataRegion.getDataRegionIdString().equals(restored.regionId)) { + continue; + } + if (hasPartitionContext(dataRegion, key.timePartitionStart)) { + // This partition already has a live writer (e.g. a snapshot fragment merged into a task + // that kept writing); leave the live context in charge. + continue; + } + toRestore.add(entry); + } + if (toRestore.isEmpty()) { + return; + } + final long chunkMetadataMaxSizeForEachWriter = + CONFIG.getLoadChunkMetadataMemorySizeInBytes() / Math.max(1, toRestore.size()); + for (final Map.Entry entry : toRestore) { + final RestoredPartitionKey key = entry.getKey(); + final RestoredLoadFile restored = entry.getValue(); + final DataPartitionInfo partitionInfo = + new DataPartitionInfo(dataRegion, new TTimePartitionSlot(key.timePartitionStart)); + final PartitionContext context = + PartitionContext.restore( + partitionInfo, taskDir, restored.file, chunkMetadataMaxSizeForEachWriter); + if (context == null) { + continue; + } + partitionContexts.put(partitionInfo, context); + // The file is now writer-managed; the COMMIT path loads it through the writer contexts. + restoredPartitions.remove(entry.getKey()); + } + } + + private boolean hasPartitionContext(DataRegion dataRegion, long timePartitionStart) { + for (final DataPartitionInfo partitionInfo : partitionContexts.keySet()) { + if (partitionInfo.getDataRegion() == dataRegion + && partitionInfo.getTimePartitionSlot().getStartTime() == timePartitionStart) { + return true; + } + } + return false; + } + + /** Applies only the deletion data of a raw-ref PIECE to the matching partition files. */ + void applyDeletion(DataRegion dataRegion, List dataList) + throws IOException, PageException, LoadFileException { + taskLock.lock(); + try { + checkNotClosed(); + for (final TsFileData data : dataList) { + if (data.getType() == TsFileDataType.DELETION) { + writeDeletion((DeletionData) data, dataRegion); + } + } + } finally { + taskLock.unlock(); + } + } + + /** Ends every partition file (chunk groups + footer) and captures the final byte ranges. */ + void finalizeAll() throws IOException { + taskLock.lock(); + try { + checkNotClosed(); + for (final PartitionContext context : partitionContexts.values()) { + context.finalizeFile(pendingPieceRefs); + } + } finally { + taskLock.unlock(); + } + } + + void appendRawTsFilePieces( + DataRegion dataRegion, List pieceRefs) throws IOException { + taskLock.lock(); + try { + checkNotClosed(); + for (final LoadTsFileConsensusNode.PieceRef ref : pieceRefs) { + final String relativePath = ref.getRelativePath(); + final File targetFile = + SystemFileFactory.INSTANCE.getFile(taskDir, new File(relativePath).getName()); + if (isManagedByLiveWriter(targetFile)) { + // The partition file was already written by this node's single partition writer from the + // chunk-data PIECE; the reference only needs to be applied by nodes that rebuild the file + // from the WAL (no in-memory writer). + continue; + } + byte[] content = ref.getContent(); + if (content == null) { + content = readFileRange(relativePath, ref.getOffset(), (int) ref.getSize()); + } + if (targetFile.getParentFile() != null) { + Files.createDirectories(targetFile.getParentFile().toPath()); + } + try (final FileChannel channel = + FileChannel.open( + targetFile.toPath(), StandardOpenOption.CREATE, StandardOpenOption.WRITE)) { + final long currentLength = channel.size(); + if (currentLength != ref.getOffset()) { + // Last line of defense on the replica side: the leader guarantees the refs advance + // continuously from offset 0, so a hole or overlap here means the staged file cannot + // be repaired by appending and must be aborted instead of corrupted. + throw new IOException( + String.format( + StorageEngineMessages + .EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_NOT_CONTINUOUS_F9408C19, + targetFile, + taskDir.getName(), + ref.getOffset(), + currentLength)); + } + channel.position(ref.getOffset()); + channel.write(ByteBuffer.wrap(content)); + channel.force(true); + } + if (rawTsFilePaths.add(targetFile.getAbsolutePath())) { + rawTsFiles.add(new RawTsFile(targetFile, dataRegion)); + } + } + } finally { + taskLock.unlock(); + } + } + + private boolean isManagedByLiveWriter(File targetFile) { + for (final PartitionContext context : partitionContexts.values()) { + if (context.getWriter().getFile().getAbsolutePath().equals(targetFile.getAbsolutePath())) { + return true; + } + } + return false; + } + + private byte[] readFileRange(String relativePath, long offset, int length) throws IOException { + final byte[] content = new byte[length]; + if (length == 0) { + return content; + } + final String message = + String.format( + StorageEngineMessages + .MESSAGE_NO_LOAD_TSFILE_UUID_ARG_RECORDED_EXECUTE_LOAD_COMMAND_ARG_66722D80, + relativePath); + final File file = + LoadTsFileManager.findLoadTsFile(relativePath).orElseThrow(() -> new IOException(message)); + try (final FileChannel channel = FileChannel.open(file.toPath(), StandardOpenOption.READ)) { + channel.position(offset); + final ByteBuffer buffer = ByteBuffer.wrap(content); + while (buffer.hasRemaining()) { + final int read = channel.read(buffer); + if (read < 0) { + throw new IOException( + String.format( + StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_EOF_8743387D, + relativePath, + offset)); + } + } + } + return content; + } + + List drainPendingPieceRefs() { + taskLock.lock(); + try { + final List refs = new ArrayList<>(pendingPieceRefs); + pendingPieceRefs.clear(); + return refs; + } finally { + taskLock.unlock(); + } + } + + /** + * Copies the already-synced byte prefix of every staged partition file of this task into the + * snapshot task dir and returns the metadata needed to restore it. + */ + LoadSnapshotManager.TaskSnapshot snapshotTask(File targetDir) throws IOException { + taskLock.lock(); + try { + final List snapshots = new ArrayList<>(); + for (final PartitionContext context : partitionContexts.values()) { + final StagedFileSnapshot snapshot = context.snapshotTo(targetDir); + if (snapshot != null) { + snapshots.add(snapshot); + } + } + final StringBuilder appliedPieces = new StringBuilder(); + for (final Map.Entry entry : appliedPieceIndex2Checksum.entrySet()) { + appliedPieces.append(entry.getKey()).append(':').append(entry.getValue()).append(','); + } + return new LoadSnapshotManager.TaskSnapshot(snapshots, appliedPieces.toString()); + } finally { + taskLock.unlock(); + } + } + + /** + * Seeds the applied-piece prefix captured by a snapshot. The staged files restored from the + * snapshot already contain the data of every piece in the prefix, so the continuity fence must + * treat them as applied; otherwise the first replayed marker after a failover would be rejected + * as non-contiguous even though the file is complete up to that point. + */ + void restoreAppliedPieces(String serialized) { + taskLock.lock(); + try { + if (serialized == null || serialized.isEmpty()) { + return; + } + for (final String entry : serialized.split(",")) { + if (entry.isEmpty()) { + continue; + } + final int separator = entry.indexOf(':'); + if (separator <= 0) { + continue; + } + try { + appliedPieceIndex2Checksum.put( + Long.parseLong(entry.substring(0, separator)), + Long.parseLong(entry.substring(separator + 1))); + } catch (NumberFormatException e) { + LOGGER.warn( + StorageEngineMessages.LOG_LOAD_CONSENSUS_APPLIED_PIECE_RESTORE_FAILED_5BC74BBA, + getTaskName(), + entry); + } + } + appliedContiguousCount = 0; + while (appliedPieceIndex2Checksum.containsKey(appliedContiguousCount)) { + appliedContiguousCount++; + } + } finally { + taskLock.unlock(); + } + } + + void registerRestoredPartitions(List stagedFiles) { + taskLock.lock(); + try { + for (StagedFileSnapshot snapshot : stagedFiles) { + final File file = new File(taskDir, snapshot.getFileName()); + if (!file.isFile() || file.length() == 0) { + continue; + } + restoredPartitions.put( + new RestoredPartitionKey( + snapshot.getDatabase(), snapshot.getRegionId(), snapshot.getTimePartitionStart()), + new RestoredLoadFile( + file, snapshot.getDatabase(), snapshot.getRegionId(), snapshot.isFinalized())); + } + } finally { + taskLock.unlock(); + } + } + + void loadAll( + DataRegion dataRegion, + boolean isGeneratedByPipe, + Map timePartitionProgressIndexMap) + throws IOException, LoadFileException { + taskLock.lock(); + try { + checkNotClosed(); + for (final PartitionContext context : partitionContexts.values()) { + context.closeModificationFile(); + } + for (final PartitionContext context : partitionContexts.values()) { + context.loadIntoRegion( + isGeneratedByPipe, + timePartitionProgressIndexMap.getOrDefault( + context.getTimePartitionSlot(), MinimumProgressIndex.INSTANCE)); + } + for (final RawTsFile rawTsFile : rawTsFiles) { + if (isManagedByLiveWriter(rawTsFile.file)) { + continue; + } + final TsFileResource tsFileResource = new TsFileResource(rawTsFile.file); + try (final TsFileSequenceReader reader = + new TsFileSequenceReader(rawTsFile.file.getAbsolutePath(), true)) { + if (!reader.isComplete()) { + throw new LoadFileException( + String.format( + StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_INCOMPLETE_1CDE954B, + rawTsFile.file, + taskDir.getName())); + } + TsFileResourceUtils.updateTsFileResource(reader, tsFileResource); + } + tsFileResource.setGeneratedByPipe(isGeneratedByPipe); + tsFileResource.setStatus(TsFileResourceStatus.NORMAL); + tsFileResource.setProgressIndex(MinimumProgressIndex.INSTANCE); + rawTsFile.dataRegion.loadNewTsFile( + tsFileResource, true, isGeneratedByPipe, false, Optional.empty()); + } + // Cached pieces and leader-retained bytes are no longer needed after COMMIT. + clearRetainedPieces(); + // Staged files restored from a snapshot have no in-memory writer. Files that were re-synced + // by PIECE refs after the restore were already added to rawTsFiles and are skipped here; + // the remaining files (e.g. the load was sealed before the snapshot) are loaded directly. + for (final Map.Entry entry : + restoredPartitions.entrySet()) { + final RestoredLoadFile restored = entry.getValue(); + if (rawTsFilePaths.contains(restored.file.getAbsolutePath())) { + continue; + } + if (dataRegion == null + || !dataRegion.getDatabaseName().equals(restored.database) + || !dataRegion.getDataRegionIdString().equals(restored.regionId)) { + continue; + } + final TsFileResource tsFileResource = new TsFileResource(restored.file); + try (final TsFileSequenceReader reader = + new TsFileSequenceReader(restored.file.getAbsolutePath(), true)) { + if (!reader.isComplete()) { + throw new LoadFileException( + String.format( + StorageEngineMessages.EXCEPTION_LOAD_CONSENSUS_STAGED_FILE_INCOMPLETE_1CDE954B, + restored.file, + taskDir.getName())); + } + TsFileResourceUtils.updateTsFileResource(reader, tsFileResource); + } + tsFileResource.setGeneratedByPipe(isGeneratedByPipe); + tsFileResource.setStatus(TsFileResourceStatus.NORMAL); + tsFileResource.setProgressIndex(MinimumProgressIndex.INSTANCE); + dataRegion.loadNewTsFile(tsFileResource, true, isGeneratedByPipe, false, Optional.empty()); + rawTsFilePaths.add(restored.file.getAbsolutePath()); + } + } finally { + taskLock.unlock(); + } + } + + /** Closes every writer, deletes the staged/raw files and finally the task directory. */ + void close() { + taskLock.lock(); + try { + if (isClosed) { + return; + } + clearRetainedPieces(); + for (final PartitionContext context : partitionContexts.values()) { + context.close(); + } + partitionContexts.clear(); + + for (final RawTsFile rawTsFile : rawTsFiles) { + try { + final Path path = rawTsFile.file.toPath(); + if (Files.exists(path)) { + RetryUtils.retryOnException( + () -> { + Files.delete(path); + return null; + }); + } + } catch (IOException e) { + LOGGER.warn( + StorageEngineMessages.CLOSE_TSFILE_IO_WRITER_ERROR, rawTsFile.file.getPath(), e); + } + } + rawTsFiles.clear(); + rawTsFilePaths.clear(); + + for (final RestoredLoadFile restored : restoredPartitions.values()) { + if (rawTsFilePaths.contains(restored.file.getAbsolutePath())) { + continue; + } + try { + final Path path = restored.file.toPath(); + if (Files.exists(path)) { + RetryUtils.retryOnException( + () -> { + Files.delete(path); + return null; + }); + } + } catch (IOException e) { + LOGGER.warn( + StorageEngineMessages.CLOSE_TSFILE_IO_WRITER_ERROR, restored.file.getPath(), e); + } + } + restoredPartitions.clear(); + + try { + RetryUtils.retryOnException( + () -> { + Files.delete(taskDir.toPath()); + return null; + }); + } catch (DirectoryNotEmptyException e) { + LOGGER.info(StorageEngineMessages.TASK_DIR_NOT_EMPTY_SKIP_DELETE, taskDir.getPath()); + } catch (IOException e) { + LOGGER.warn(MESSAGE_DELETE_FAIL, taskDir.getPath(), e); + } + isClosed = true; + } finally { + taskLock.unlock(); + } + } + + private void checkNotClosed() throws IOException { + if (isClosed) { + throw new IOException(String.format(MESSAGE_WRITER_MANAGER_HAS_BEEN_CLOSED, taskDir)); + } + } + + private static class RawTsFile { + private final File file; + private final DataRegion dataRegion; + + private RawTsFile(File file, DataRegion dataRegion) { + this.file = file; + this.dataRegion = dataRegion; + } + } + + private static class RestoredLoadFile { + private final File file; + private final String database; + private final String regionId; + private final boolean finalized; + + private RestoredLoadFile(File file, String database, String regionId, boolean finalized) { + this.file = file; + this.database = database; + this.regionId = regionId; + this.finalized = finalized; + } + } + + /** One piece cached on a follower, waiting for its WAL marker to arrive. */ + private static class CachedPiece { + private final long pieceIndex; + private final long checksum; + private final List dataList; + + private CachedPiece(long pieceIndex, long checksum, List dataList) { + this.pieceIndex = pieceIndex; + this.checksum = checksum; + this.dataList = dataList; + } + } + + private static class RestoredPartitionKey { + private final String database; + private final String regionId; + private final long timePartitionStart; + + private RestoredPartitionKey(String database, String regionId, long timePartitionStart) { + this.database = database; + this.regionId = regionId; + this.timePartitionStart = timePartitionStart; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + final RestoredPartitionKey that = (RestoredPartitionKey) o; + return timePartitionStart == that.timePartitionStart + && database.equals(that.database) + && regionId.equals(that.regionId); + } + + @Override + public int hashCode() { + return Objects.hash(database, regionId, timePartitionStart); + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java new file mode 100644 index 0000000000000..7173f5cdc5884 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/planner/plan/node/load/LoadTsFileConsensusNodeTest.java @@ -0,0 +1,133 @@ +/* + * 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.planner.plan.node.load; + +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNode; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId; +import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType; + +import org.junit.Assert; +import org.junit.Test; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Collections; + +public class LoadTsFileConsensusNodeTest { + + @Test + public void testBeginSerializeRoundTrip() { + final LoadTsFileConsensusNode begin = + LoadTsFileConsensusNode.begin(new PlanNodeId("n1"), "load-1", "file-1", true, "db", 3); + final ByteBuffer buffer = begin.serializeToByteBuffer(); + final PlanNode deserialized = PlanNodeType.deserialize(buffer); + Assert.assertTrue(deserialized instanceof LoadTsFileConsensusNode); + final LoadTsFileConsensusNode node = (LoadTsFileConsensusNode) deserialized; + Assert.assertEquals(LoadTsFileConsensusOp.BEGIN, node.getOp()); + Assert.assertEquals("load-1", node.getLoadId()); + Assert.assertEquals("file-1", node.getTsFileId()); + Assert.assertTrue(node.isTableModel()); + Assert.assertEquals("db", node.getDatabase()); + Assert.assertEquals(3, node.getExpectedPieceCount()); + } + + @Test + public void testPieceSerializeRoundTrip() { + final LoadTsFileConsensusNode commit = + LoadTsFileConsensusNode.commit( + new PlanNodeId("n2"), "load-2", "file-2", true, false, Collections.emptyMap()); + final ByteBuffer buffer = commit.serializeToByteBuffer(); + final PlanNode deserialized = PlanNodeType.deserialize(buffer); + Assert.assertTrue(deserialized instanceof LoadTsFileConsensusNode); + final LoadTsFileConsensusNode node = (LoadTsFileConsensusNode) deserialized; + Assert.assertEquals(LoadTsFileConsensusOp.COMMIT, node.getOp()); + Assert.assertTrue(node.isGeneratedByPipe()); + } + + @Test + public void testPiecePreservesCallerChecksum() { + final LoadTsFileConsensusNode piece = + LoadTsFileConsensusNode.piece( + new PlanNodeId("checksum"), + "load", + "file", + 0L, + 0L, + Collections.emptyList(), + 987654321L); + + Assert.assertEquals(987654321L, piece.getChecksum()); + } + + @Test + public void testPieceRefSerializeRoundTrip() { + final LoadTsFileConsensusNode.PieceRef ref = + new LoadTsFileConsensusNode.PieceRef( + "task-1/partition.tsfile", 4096L, 3L, new byte[] {1, 2, 3}); + final LoadTsFileConsensusNode piece = + LoadTsFileConsensusNode.pieceRefs( + new PlanNodeId("n3"), "load-3", "file-3", 0L, Collections.singletonList(ref), 0L, 3L); + final ByteBuffer buffer = piece.serializeToByteBuffer(); + final PlanNode deserialized = PlanNodeType.deserialize(buffer); + Assert.assertTrue(deserialized instanceof LoadTsFileConsensusNode); + final LoadTsFileConsensusNode node = (LoadTsFileConsensusNode) deserialized; + Assert.assertEquals(LoadTsFileConsensusOp.PIECE, node.getOp()); + Assert.assertEquals(1, node.getPieceRefs().size()); + final LoadTsFileConsensusNode.PieceRef actual = node.getPieceRefs().get(0); + Assert.assertEquals("task-1/partition.tsfile", actual.getRelativePath()); + Assert.assertEquals(4096L, actual.getOffset()); + Assert.assertEquals(3L, actual.getSize()); + Assert.assertTrue(Arrays.equals(new byte[] {1, 2, 3}, actual.getContent())); + } + + @Test + public void testPieceMarkerSerializeRoundTrip() { + final LoadTsFileConsensusNode marker = + LoadTsFileConsensusNode.pieceMarker( + new PlanNodeId("n4"), "load-4", "file-4", 7L, 123456L, 1024L); + final ByteBuffer buffer = marker.serializeToByteBuffer(); + final PlanNode deserialized = PlanNodeType.deserialize(buffer); + Assert.assertTrue(deserialized instanceof LoadTsFileConsensusNode); + final LoadTsFileConsensusNode node = (LoadTsFileConsensusNode) deserialized; + Assert.assertEquals(LoadTsFileConsensusOp.PIECE, node.getOp()); + Assert.assertEquals("load-4", node.getLoadId()); + Assert.assertEquals("file-4", node.getTsFileId()); + Assert.assertEquals(7L, node.getPieceIndex()); + Assert.assertEquals(123456L, node.getChecksum()); + Assert.assertEquals(1024L, node.getDataSize()); + Assert.assertFalse(node.hasChunkData()); + Assert.assertTrue(node.getPieceRefs().isEmpty()); + } + + @Test + public void testPullSerializeRoundTrip() { + final LoadTsFileConsensusNode pull = + LoadTsFileConsensusNode.pull( + new PlanNodeId("n5"), "load-5", "file-5", 3L, 999L, "192.168.1.10:6667"); + final ByteBuffer buffer = pull.serializeToByteBuffer(); + final PlanNode deserialized = PlanNodeType.deserialize(buffer); + Assert.assertTrue(deserialized instanceof LoadTsFileConsensusNode); + final LoadTsFileConsensusNode node = (LoadTsFileConsensusNode) deserialized; + Assert.assertEquals(LoadTsFileConsensusOp.PULL, node.getOp()); + Assert.assertEquals(3L, node.getPieceIndex()); + Assert.assertEquals(999L, node.getChecksum()); + Assert.assertEquals("192.168.1.10:6667", node.getPullSourceEndPoint()); + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java index 1b738f6496d49..1dd8a45c9b61b 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileSchedulerTest.java @@ -39,8 +39,9 @@ import java.io.File; import java.lang.reflect.Constructor; -import java.lang.reflect.Field; import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.Collections; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; @@ -97,16 +98,15 @@ public void testGetPartitionQueryDatabaseForTableModelLoad() { @Test public void testBuildRetryTreeLoadStatementUpdatesDatabaseLevel() throws Exception { - final LoadTsFileScheduler scheduler = - new LoadTsFileScheduler( - distributedQueryPlan, + final LoadFallbackHandler fallbackHandler = + new LoadFallbackHandler( mock(MPPQueryContext.class), - mock(QueryStateMachine.class), - mock(IClientManager.class), - mock(IPartitionFetcher.class), - true); + true, + Collections.emptyList(), + new ArrayList<>(), + mock(QueryStateMachine.class)); final Method method = - LoadTsFileScheduler.class.getDeclaredMethod( + LoadFallbackHandler.class.getDeclaredMethod( "buildRetryTreeLoadStatement", String.class, boolean.class, String.class); method.setAccessible(true); @@ -115,7 +115,7 @@ public void testBuildRetryTreeLoadStatementUpdatesDatabaseLevel() throws Excepti final LoadTsFileStatement statement = (LoadTsFileStatement) - method.invoke(scheduler, tsFile.getAbsolutePath(), true, "root.test.sg_0"); + method.invoke(fallbackHandler, tsFile.getAbsolutePath(), true, "root.test.sg_0"); Assert.assertEquals("root.test.sg_0", statement.getDatabase()); Assert.assertEquals(2, statement.getDatabaseLevel()); @@ -123,41 +123,24 @@ public void testBuildRetryTreeLoadStatementUpdatesDatabaseLevel() throws Excepti } @Test - public void testTsFileDataManagerClearReleasesCachedMemory() throws Exception { + public void testMemoryBoundedBufferClearReleasesCachedMemory() throws Exception { final Constructor memoryBlockConstructor = LoadTsFileDataCacheMemoryBlock.class.getDeclaredConstructor(long.class); memoryBlockConstructor.setAccessible(true); final LoadTsFileDataCacheMemoryBlock memoryBlock = memoryBlockConstructor.newInstance(1024 * 1024L); - final Class dataManagerClass = - Class.forName(LoadTsFileScheduler.class.getName() + "$TsFileDataManager"); - final Constructor dataManagerConstructor = - dataManagerClass.getDeclaredConstructor( - LoadTsFileScheduler.class, - LoadSingleTsFileNode.class, - LoadTsFileDataCacheMemoryBlock.class); - dataManagerConstructor.setAccessible(true); - final Object dataManager = - dataManagerConstructor.newInstance( - mock(LoadTsFileScheduler.class), mock(LoadSingleTsFileNode.class), memoryBlock); - // Simulate data buffered before split or routing aborts. clear() is the last chance to return // this accounting to the shared LOAD memory block. final long cachedMemorySize = 128L; - memoryBlock.addMemoryUsage(cachedMemorySize); - final Field dataSizeField = dataManagerClass.getDeclaredField("dataSize"); - dataSizeField.setAccessible(true); - dataSizeField.setLong(dataManager, cachedMemorySize); - - final Method clearMethod = dataManagerClass.getDeclaredMethod("clear"); - clearMethod.setAccessible(true); - clearMethod.invoke(dataManager); + final MemoryBoundedBuffer memoryBoundedBuffer = new MemoryBoundedBuffer(memoryBlock); + memoryBoundedBuffer.add(cachedMemorySize); + memoryBoundedBuffer.clear(); final Method getMemoryUsageMethod = LoadTsFileDataCacheMemoryBlock.class.getDeclaredMethod("getMemoryUsageInBytes"); getMemoryUsageMethod.setAccessible(true); Assert.assertEquals(0L, getMemoryUsageMethod.invoke(memoryBlock)); - Assert.assertEquals(0L, dataSizeField.getLong(dataManager)); + Assert.assertEquals(0L, memoryBoundedBuffer.getDataSize()); } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileSnapshotMetaTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileSnapshotMetaTest.java new file mode 100644 index 0000000000000..be49ebfa34286 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFileSnapshotMetaTest.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.storageengine.load; + +import org.junit.Assert; +import org.junit.Test; + +import java.io.File; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.List; + +public class LoadTsFileSnapshotMetaTest { + + @Test + public void testSnapshotMetaRoundTrip() throws IOException { + final File metaFile = File.createTempFile("load-snapshot-meta", ".meta"); + try { + final List snapshots = new ArrayList<>(); + snapshots.add( + new LoadSnapshotManager.StagedFileSnapshot( + "root.sg_1_0.tsfile", "root.sg", "1", 0L, false)); + snapshots.add( + new LoadSnapshotManager.StagedFileSnapshot( + "root.sg_1_1000.tsfile", "root.sg", "1", 1000L, true)); + + final LoadSnapshotManager.TaskSnapshot taskSnapshot = + new LoadSnapshotManager.TaskSnapshot(snapshots, "0=1,1=2,"); + LoadSnapshotManager.writeSnapshotMeta(metaFile, taskSnapshot); + + final LoadSnapshotManager.TaskSnapshot parsedTaskSnapshot = + LoadSnapshotManager.parseSnapshotMeta(metaFile); + Assert.assertEquals(2, parsedTaskSnapshot.getStagedFiles().size()); + Assert.assertEquals("0=1,1=2,", parsedTaskSnapshot.getAppliedPieces()); + + final LoadSnapshotManager.StagedFileSnapshot first = + parsedTaskSnapshot.getStagedFiles().get(0); + Assert.assertEquals("root.sg_1_0.tsfile", first.getFileName()); + Assert.assertEquals("root.sg", first.getDatabase()); + Assert.assertEquals("1", first.getRegionId()); + Assert.assertEquals(0L, first.getTimePartitionStart()); + Assert.assertFalse(first.isFinalized()); + + final LoadSnapshotManager.StagedFileSnapshot second = + parsedTaskSnapshot.getStagedFiles().get(1); + Assert.assertEquals("root.sg_1_1000.tsfile", second.getFileName()); + Assert.assertEquals(1000L, second.getTimePartitionStart()); + Assert.assertTrue(second.isFinalized()); + } finally { + Files.deleteIfExists(metaFile.toPath()); + } + } + + @Test + public void testSnapshotMetaRejectsMalformedLine() throws IOException { + final File metaFile = File.createTempFile("load-snapshot-meta", ".meta"); + try { + Files.write(metaFile.toPath(), "only-one-field\n".getBytes(StandardCharsets.UTF_8)); + try { + LoadSnapshotManager.parseSnapshotMeta(metaFile); + Assert.fail("expected IOException for malformed snapshot meta"); + } catch (IOException expected) { + // expected + } + } finally { + Files.deleteIfExists(metaFile.toPath()); + } + } +} diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManagerTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManagerTest.java new file mode 100644 index 0000000000000..c94022d8c157e --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/TsFileWriterManagerTest.java @@ -0,0 +1,142 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.storageengine.load; + +import org.apache.iotdb.db.storageengine.load.splitter.TsFileData; +import org.apache.iotdb.db.storageengine.load.splitter.TsFileDataType; + +import org.junit.Assert; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.io.DataOutputStream; +import java.io.File; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; + +/** Focused unit tests for the per-task applied/cached/retained piece bookkeeping. */ +public class TsFileWriterManagerTest { + + @Rule public TemporaryFolder temporaryFolder = new TemporaryFolder(); + + @Test + public void testAppliedContiguityFence() throws IOException { + final TsFileWriterManager manager = + new TsFileWriterManager(temporaryFolder.newFolder("task-1")); + try { + Assert.assertTrue(manager.hasAppliedAllUpTo(-1)); + Assert.assertFalse(manager.hasAppliedAllUpTo(0)); + + manager.recordAppliedPiece(1, 111L); + // Piece 1 without piece 0 is not a contiguous prefix yet. + Assert.assertFalse(manager.hasAppliedAllUpTo(1)); + + manager.recordAppliedPiece(0, 100L); + Assert.assertTrue(manager.hasAppliedAllUpTo(0)); + Assert.assertTrue(manager.hasAppliedAllUpTo(1)); + Assert.assertFalse(manager.hasAppliedAllUpTo(2)); + + manager.recordAppliedPiece(3, 333L); + // A hole at 2 keeps the prefix at 2. + Assert.assertFalse(manager.hasAppliedAllUpTo(3)); + manager.recordAppliedPiece(2, 222L); + Assert.assertTrue(manager.hasAppliedAllUpTo(3)); + } finally { + manager.close(); + } + } + + @Test + public void testRestoredAppliedPiecesRebuildsContiguity() throws IOException { + final TsFileWriterManager manager = + new TsFileWriterManager(temporaryFolder.newFolder("task-2")); + try { + manager.restoreAppliedPieces("0:100,1:111,2:222,"); + Assert.assertTrue(manager.hasAppliedAllUpTo(2)); + Assert.assertFalse(manager.hasAppliedAllUpTo(3)); + Assert.assertTrue(manager.isPieceAlreadyApplied(1, 111L)); + Assert.assertTrue(manager.isPieceConflicting(1, 999L)); + } finally { + manager.close(); + } + } + + @Test + public void testCachedPieceChecksumConflict() throws IOException { + final TsFileWriterManager manager = + new TsFileWriterManager(temporaryFolder.newFolder("task-3")); + try { + final byte[] data = "chunk-data".getBytes(StandardCharsets.UTF_8); + Assert.assertTrue(manager.cachePiece(0, 1L, Arrays.asList(new TestTsFileData(data)))); + // Same index with the same checksum is idempotent. + Assert.assertTrue(manager.cachePiece(0, 1L, Arrays.asList(new TestTsFileData(data)))); + Assert.assertTrue(manager.hasCachedPiece(0, 1L)); + // Same index with a different checksum is a divergent delivery and must be rejected. + Assert.assertFalse(manager.cachePiece(0, 2L, Arrays.asList(new TestTsFileData(data)))); + } finally { + manager.close(); + } + } + + @Test + public void testRetainedPiecesSurviveRestart() throws IOException { + final File taskDir = temporaryFolder.newFolder("task-4"); + final byte[] pieceBytes = "serialized-piece".getBytes(StandardCharsets.UTF_8); + final TsFileWriterManager first = new TsFileWriterManager(taskDir); + // Simulate a crash: the first manager is abandoned without close() so the retained bytes stay + // on disk, exactly like a process that dies mid-load. + first.retainPiece(0, pieceBytes); + + // A new manager over the same (restored) task dir must reload the retained bytes from disk. + final TsFileWriterManager second = new TsFileWriterManager(taskDir, false); + try { + Assert.assertTrue(Arrays.equals(pieceBytes, second.getRetainedPiece(0).orElse(null))); + } finally { + second.close(); + } + } + + /** Minimal TsFileData placeholder for cache bookkeeping tests (never applied to a writer). */ + private static final class TestTsFileData implements TsFileData { + + private final byte[] data; + + private TestTsFileData(byte[] data) { + this.data = data; + } + + @Override + public long getDataSize() { + return data.length; + } + + @Override + public TsFileDataType getType() { + return TsFileDataType.CHUNK; + } + + @Override + public void serialize(DataOutputStream stream) { + throw new UnsupportedOperationException("not used in this test"); + } + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java index e36796fca43c3..e036f91c056c7 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/queryengine/plan/planner/plan/node/PlanNodeType.java @@ -209,6 +209,7 @@ public enum PlanNodeType { RELATIONAL_INSERT_ROWS((short) 2002), RELATIONAL_DELETE_DATA((short) 2003), OBJECT_FILE_NODE((short) 2004), + LOAD_TSFILE_CONSENSUS((short) 2010), ; private static final IPlanNodeDeserializer DESERIALIZER;